Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SQLite: Don't bill for internal queries. #3605

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

SQLite: Don't bill for internal queries. #3605

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/workerd/api/sql.c++
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,11 @@ bool SqlStorage::allowTransactions() const {
"write coalescing.");
}

bool SqlStorage::shouldAddQueryStats() const {
// Bill for queries executed from JavaScript.
return true;
}

SqlStorage::StatementCache::~StatementCache() noexcept(false) {
for (auto& entry: lru) {
lru.remove(entry);
Expand Down
1 change: 1 addition & 0 deletions src/workerd/api/sql.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ class SqlStorage final: public jsg::Object, private SqliteDatabase::Regulator {
bool isAllowedTrigger(kj::StringPtr name) const override;
void onError(kj::Maybe<int> sqliteErrorCode, kj::StringPtr message) const override;
bool allowTransactions() const override;
bool shouldAddQueryStats() const override;

SqliteDatabase& getDb(jsg::Lock& js) {
return storage->getSqliteDb(js);
Expand Down
27 changes: 26 additions & 1 deletion src/workerd/util/sqlite-kv-test.c++
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,32 @@ namespace workerd {
namespace {

KJ_TEST("SQLite-KV") {
class TestSqliteObserver: public SqliteObserver {
public:
void addQueryStats(uint64_t read, uint64_t written) override {
rowsRead += read;
rowsWritten += written;
}

uint64_t rowsRead = 0;
uint64_t rowsWritten = 0;
};

auto dir = kj::newInMemoryDirectory(kj::nullClock());
SqliteDatabase::Vfs vfs(*dir);
SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY);
TestSqliteObserver sqliteObserver;
SqliteDatabase db(
vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, sqliteObserver);
SqliteKv kv(db);

kv.put("foo", "abc"_kj.asBytes());
kv.put("bar", "def"_kj.asBytes());
kv.put("baz", "123"_kj.asBytes());
kv.put("qux", "321"_kj.asBytes());

KJ_EXPECT(sqliteObserver.rowsWritten == 4);
KJ_EXPECT(sqliteObserver.rowsRead == 0);

{
bool called = false;
KJ_EXPECT(kv.get("foo", [&](kj::ArrayPtr<const byte> value) {
Expand All @@ -29,6 +45,9 @@ KJ_TEST("SQLite-KV") {
KJ_EXPECT(called);
}

KJ_EXPECT(sqliteObserver.rowsWritten == 4);
KJ_EXPECT(sqliteObserver.rowsRead == 1);

{
bool called = false;
KJ_EXPECT(kv.get("bar", [&](kj::ArrayPtr<const byte> value) {
Expand All @@ -38,10 +57,16 @@ KJ_TEST("SQLite-KV") {
KJ_EXPECT(called);
}

KJ_EXPECT(sqliteObserver.rowsWritten == 4);
KJ_EXPECT(sqliteObserver.rowsRead == 2);

KJ_EXPECT(!kv.get("corge", [&](kj::ArrayPtr<const byte> value) {
KJ_FAIL_EXPECT("should not call callback when no match", value.asChars());
}));

KJ_EXPECT(sqliteObserver.rowsWritten == 4);
KJ_EXPECT(sqliteObserver.rowsRead == 2);

auto list = [&](auto&&... params) {
kj::Vector<kj::String> results;
auto callback = [&](kj::StringPtr key, kj::ArrayPtr<const byte> value) {
Expand Down
5 changes: 5 additions & 0 deletions src/workerd/util/sqlite-kv.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,11 @@ namespace workerd {
// of KJ exceptions which become internal errors.
class SqliteKvRegulator: public SqliteDatabase::Regulator {
void onError(kj::Maybe<int> sqliteErrorCode, kj::StringPtr message) const override;

// We bill for KV operations as rows read/written.
virtual bool shouldAddQueryStats() const override {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(I guess the stmtCountKeys prepared statement doesn't use this regulator yet, so would be excluded from stats, but not sure if that query represents notable overhead.)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah good catch, probably should fix that.

return true;
}
};

// Class which implements KV storage on top of SQLite. This is intended to be used for Durable
Expand Down
33 changes: 27 additions & 6 deletions src/workerd/util/sqlite-test.c++
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -834,9 +834,17 @@ KJ_TEST("SQLite observer addQueryStats") {
uint64_t rowsWritten = 0;
};

class TestQueryStatsRegulator: public SqliteDatabase::Regulator {
public:
bool shouldAddQueryStats() const override {
return true;
}
};

TempDirOnDisk dir;
SqliteDatabase::Vfs vfs(*dir);
TestSqliteObserver sqliteObserver = TestSqliteObserver();
TestQueryStatsRegulator regulator;
SqliteDatabase db(
vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, sqliteObserver);

Expand All @@ -851,17 +859,17 @@ KJ_TEST("SQLite observer addQueryStats") {
int rowsWrittenBefore = sqliteObserver.rowsWritten;
constexpr int dbRowCount = 3;
{
db.run("INSERT INTO things (id) VALUES (10)");
db.run("INSERT INTO things (id) VALUES (11)");
db.run("INSERT INTO things (id) VALUES (12)");
db.run(regulator, "INSERT INTO things (id) VALUES (10)");
db.run(regulator, "INSERT INTO things (id) VALUES (11)");
db.run(regulator, "INSERT INTO things (id) VALUES (12)");
}
KJ_EXPECT(sqliteObserver.rowsRead - rowsReadBefore == dbRowCount);
KJ_EXPECT(sqliteObserver.rowsWritten - rowsWrittenBefore == dbRowCount);

rowsReadBefore = sqliteObserver.rowsRead;
rowsWrittenBefore = sqliteObserver.rowsWritten;
{
auto getCount = db.prepare("SELECT COUNT(*) FROM things");
auto getCount = db.prepare(regulator, "SELECT COUNT(*) FROM things");
KJ_EXPECT(getCount.run().getInt(0) == dbRowCount);
}
KJ_EXPECT(sqliteObserver.rowsRead - rowsReadBefore == dbRowCount);
Expand All @@ -871,7 +879,7 @@ KJ_TEST("SQLite observer addQueryStats") {
rowsReadBefore = sqliteObserver.rowsRead;
rowsWrittenBefore = sqliteObserver.rowsWritten;
{
auto stmt = db.prepare("SELECT * FROM things");
auto stmt = db.prepare(regulator, "SELECT * FROM things");
auto query = stmt.run();
KJ_ASSERT(!query.isDone());
while (!query.isDone()) {
Expand All @@ -881,11 +889,24 @@ KJ_TEST("SQLite observer addQueryStats") {
KJ_EXPECT(sqliteObserver.rowsRead - rowsReadBefore == dbRowCount);
KJ_EXPECT(sqliteObserver.rowsWritten - rowsWrittenBefore == 0);

// Verify system queries don't affect stats.
rowsReadBefore = sqliteObserver.rowsRead;
rowsWrittenBefore = sqliteObserver.rowsWritten;
db.run("INSERT INTO things (id) VALUES (13)");
{
auto query = db.run("SELECT * FROM things");
while (!query.isDone()) {
query.nextRow();
}
}
KJ_EXPECT(sqliteObserver.rowsRead == rowsReadBefore);
KJ_EXPECT(sqliteObserver.rowsWritten == rowsWrittenBefore);

// Verify addQueryStats works correctly when db is reset
rowsReadBefore = sqliteObserver.rowsRead;
rowsWrittenBefore = sqliteObserver.rowsWritten;
{
auto query = db.run("INSERT INTO things (id) VALUES (100)");
auto query = db.run(regulator, "INSERT INTO things (id) VALUES (100)");
db.reset();
}
KJ_EXPECT(sqliteObserver.rowsRead - rowsReadBefore == 1);
Expand Down
6 changes: 4 additions & 2 deletions src/workerd/util/sqlite.c++
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1291,8 +1291,10 @@ SqliteDatabase::Query::~Query() noexcept(false) {
}

void SqliteDatabase::Query::destroy() {
//Update the db stats that we have collected for the query
db.sqliteObserver.addQueryStats(rowsRead, rowsWritten);
if (regulator.shouldAddQueryStats()) {
//Update the db stats that we have collected for the query
db.sqliteObserver.addQueryStats(rowsRead, rowsWritten);
}

// We only need to reset the statement if we don't own it. If we own it, it's about to be
// destroyed anyway.
Expand Down
8 changes: 8 additions & 0 deletions src/workerd/util/sqlite.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,14 @@ class SqliteDatabase {
virtual bool allowTransactions() const {
return true;
}

// Whether or not this query's rows read and written should be recorded to the SqliteObserver
// when the query is done. (In other words, determines whether this query is billed.)
//
// We don't bill for TRUSTED queries since they are used internally by the system.
virtual bool shouldAddQueryStats() const {
return false;
}
};

// Use as the `Regulator&` for queries that are fully trusted. As a general rule, this should
Expand Down
Loading