Merge pull request #3117 from cloudflare/kenton/fix-sql-cache

Fix SQL statement caching of computed query strings.

Author: kentonv

src/workerd/api/sql-test.js

@@ -1012,6 +1012,25 @@ async function test(state) {
   await state.blockConcurrencyWhile(async () => {
     sql.exec(`PRAGMA foreign_keys = ON;`);
   });
+
+  // Verify caching.
+  {
+    let isCached = (q) => {
+      let cursor = sql.exec(q);
+      cursor.toArray();
+      return cursor.reusedCachedQueryForTest;
+    };
+
+    // Query based on literal string is cached.
+    assert.equal(false, isCached('SELECT 179321'));
+    assert.equal(true, isCached('SELECT 179321'));
+    assert.equal(true, isCached('SELECT 179321'));
+
+    // Qeury based on computed string is cached.
+    assert.equal(false, isCached('SELECT "' + 'x'.repeat(4) + '"'));
+    assert.equal(true, isCached('SELECT "' + 'x'.repeat(4) + '"'));
+    assert.equal(true, isCached('SELECT "' + 'x'.repeat(4) + '"'));
+  }
 }
 
 async function testIoStats(storage) {

src/workerd/api/sql.c++

@@ -25,6 +25,11 @@ SqlStorage::~SqlStorage() {}
 
 jsg::Ref<SqlStorage::Cursor> SqlStorage::exec(
     jsg::Lock& js, jsg::JsString querySql, jsg::Arguments<BindingValue> bindings) {
+  // Internalize the string, so that the cache can be keyed by string identity rather than content.
+  // Any string we put into the cache is expected to live there for a while anyway, so even if it
+  // is a one-off, internalizing it (which moves it to the old generation) shouldn't hurt.
+  querySql = querySql.internalize(js);
+
   auto& db = getDb(js);
   auto& statementCache = *this->statementCache;
 
@@ -159,6 +164,8 @@ void SqlStorage::Cursor::initColumnNames(jsg::Lock& js, State& stateRef) {
   // TODO(cleanup): Make `js.withinHandleScope` understand `jsg::JsValue` types in addition to
   //   `v8::Local`.
   KJ_IF_SOME(cached, stateRef.cachedStatement) {
+    reusedCachedQuery = cached->useCount++ > 0;
+
     KJ_IF_SOME(names, cached->cachedColumnNames) {
       // Use cached copy.
       columnNames = names.addRef(js);

src/workerd/api/sql.h

@@ -87,6 +87,7 @@ class SqlStorage final: public jsg::Object, private SqliteDatabase::Regulator {
     SqliteDatabase::Statement statement;
     kj::Maybe<jsg::JsRef<jsg::JsArray>> cachedColumnNames;  // initialized on first exec
     kj::ListLink<CachedStatement> lruLink;
+    uint useCount = 0;
 
     CachedStatement(jsg::Lock& js,
         SqlStorage& sqlStorage,
@@ -190,7 +191,7 @@ class SqlStorage::Cursor final: public jsg::Object {
   double getRowsWritten();
 
   jsg::JsArray getColumnNames(jsg::Lock& js);
-  JSG_RESOURCE_TYPE(Cursor) {
+  JSG_RESOURCE_TYPE(Cursor, CompatibilityFlags::Reader flags) {
     JSG_METHOD(next);
     JSG_METHOD(toArray);
     JSG_METHOD(one);
@@ -210,6 +211,10 @@ class SqlStorage::Cursor final: public jsg::Object {
       one(): T;
       columnNames: string[];
     });
+
+    if (flags.getWorkerdExperimental()) {
+      JSG_READONLY_PROTOTYPE_PROPERTY(reusedCachedQueryForTest, getReusedCachedQueryForTest);
+    }
   }
 
   JSG_ITERATOR(RowIterator, rows, jsg::JsObject, jsg::Ref<Cursor>, rowIteratorNext);
@@ -226,6 +231,10 @@ class SqlStorage::Cursor final: public jsg::Object {
     tracker.trackField("columnNames", columnNames);
   }
 
+  bool getReusedCachedQueryForTest() {
+    return reusedCachedQuery;
+  }
+
 private:
   struct State {
     kj::Maybe<kj::Rc<CachedStatement>> cachedStatement;
@@ -251,6 +260,9 @@ class SqlStorage::Cursor final: public jsg::Object {
   // flag an error if the application tries to reuse the cursor.
   bool canceled = false;
 
+  // Did we reuse a query from the query cache? Tracked for testing purposes.
+  bool reusedCachedQuery = false;
+
   // Reference to a weak reference that might point back to this object. If so, null it out at
   // destruction. Used by Statement to invalidate past cursors when the statement is
   // executed again.
@@ -295,7 +307,9 @@ class SqlStorage::Statement final: public jsg::Object {
 public:
   Statement(jsg::Lock& js, jsg::Ref<SqlStorage> sqlStorage, jsg::JsString query)
       : sqlStorage(kj::mv(sqlStorage)),
-        query(js.v8Isolate, query) {}
+        // Internalize the string before constructing the statement so that it doesn't have to
+        // re-lookup the internalized string for every invocation.
+        query(js.v8Isolate, query.internalize(js)) {}
 
   jsg::Ref<Cursor> run(jsg::Lock& js, jsg::Arguments<BindingValue> bindings);
 

src/workerd/jsg/jsvalue.c++

@@ -278,6 +278,10 @@ bool JsString::operator==(const JsString& other) const {
   return inner->StringEquals(other.inner);
 }
 
+JsString JsString::internalize(Lock& js) const {
+  return JsString(inner->InternalizeString(js.v8Isolate));
+}
+
 JsString::WriteIntoStatus JsString::writeInto(
     Lock& js, kj::ArrayPtr<char> buffer, WriteOptions options) const {
   WriteIntoStatus result = {0, 0};

src/workerd/jsg/jsvalue.h

@@ -240,6 +240,12 @@ class JsString final: public JsBase<v8::String, JsString> {
 
   bool operator==(const JsString& other) const;
 
+  // "Internalize" the string. Returns a string with the same content but which is identity-equal
+  // to all other internalized strings with the same content. If the string is already
+  // internalized, this returns the same value. Note that strings originating from literals in the
+  // code are always internalized.
+  JsString internalize(Lock& js) const;
+
   static JsString concat(Lock& js, const JsString& one, const JsString& two) KJ_WARN_UNUSED_RESULT;
 
   enum WriteOptions {