sqlite: create first/last subscriber callback

Author: araujogui

benchmark/sqlite/sqlite-trace.js

@@ -0,0 +1,42 @@
+'use strict';
+const common = require('../common.js');
+const sqlite = require('node:sqlite');
+const dc = require('diagnostics_channel');
+const assert = require('assert');
+
+const bench = common.createBenchmark(main, {
+  n: [1e5],
+  mode: ['none', 'subscribed', 'unsubscribed'],
+});
+
+function main(conf) {
+  const { n, mode } = conf;
+
+  const db = new sqlite.DatabaseSync(':memory:');
+  db.exec('CREATE TABLE t (x INTEGER)');
+  const insert = db.prepare('INSERT INTO t VALUES (?)');
+
+  let subscriber;
+  if (mode === 'subscribed') {
+    subscriber = () => {};
+    dc.subscribe('sqlite.db.query', subscriber);
+  } else if (mode === 'unsubscribed') {
+    subscriber = () => {};
+    dc.subscribe('sqlite.db.query', subscriber);
+    dc.unsubscribe('sqlite.db.query', subscriber);
+  }
+  // mode === 'none': no subscription ever made
+
+  let result;
+  bench.start();
+  for (let i = 0; i < n; i++) {
+    result = insert.run(i);
+  }
+  bench.end(n);
+
+  if (mode === 'subscribed') {
+    dc.unsubscribe('sqlite.db.query', subscriber);
+  }
+
+  assert.ok(result !== undefined);
+}

lib/diagnostics_channel.js

@@ -72,11 +72,19 @@ function markActive(channel) {
   ObjectSetPrototypeOf(channel, ActiveChannel.prototype);
   channel._subscribers = [];
   channel._stores = new SafeMap();
+
+  // Notify native modules that this channel just got its first subscriber.
+  if (channel._index !== undefined)
+    dc_binding.notifyChannelActive(channel._index);
 }
 
 function maybeMarkInactive(channel) {
   // When there are no more active subscribers or bound, restore to fast prototype.
   if (!channel._subscribers.length && !channel._stores.size) {
+    // Notify native modules that this channel just lost its last subscriber.
+    if (channel._index !== undefined)
+      dc_binding.notifyChannelInactive(channel._index);
+
     // eslint-disable-next-line no-use-before-define
     ObjectSetPrototypeOf(channel, Channel.prototype);
     channel._subscribers = undefined;

src/base_object_types.h

@@ -24,7 +24,8 @@ namespace node {
 #define UNSERIALIZABLE_BINDING_TYPES(V)                                        \
   V(http2_binding_data, http2::BindingData)                                    \
   V(http_parser_binding_data, http_parser::BindingData)                        \
-  V(quic_binding_data, quic::BindingData)
+  V(quic_binding_data, quic::BindingData)                                      \
+  V(sqlite_binding_data, sqlite::BindingData)
 
 // List of (non-binding) BaseObjects that are serializable in the snapshot.
 // The first argument should match what the type passes to

src/node_diagnostics_channel.cc

@@ -127,12 +127,38 @@ void BindingData::Deserialize(Local<Context> context,
   CHECK_NOT_NULL(binding);
 }
 
+void BindingData::SetChannelStatusCallback(uint32_t index,
+                                           ChannelStatusCallback cb) {
+  channel_status_callbacks_[index] = std::move(cb);
+}
+
+void BindingData::NotifyChannelActive(const FunctionCallbackInfo<Value>& args) {
+  Realm* realm = Realm::GetCurrent(args);
+  BindingData* binding = realm->GetBindingData<BindingData>();
+  if (binding == nullptr) return;
+  uint32_t index = args[0].As<v8::Uint32>()->Value();
+  auto it = binding->channel_status_callbacks_.find(index);
+  if (it != binding->channel_status_callbacks_.end()) it->second(true);
+}
+
+void BindingData::NotifyChannelInactive(
+    const FunctionCallbackInfo<Value>& args) {
+  Realm* realm = Realm::GetCurrent(args);
+  BindingData* binding = realm->GetBindingData<BindingData>();
+  if (binding == nullptr) return;
+  uint32_t index = args[0].As<v8::Uint32>()->Value();
+  auto it = binding->channel_status_callbacks_.find(index);
+  if (it != binding->channel_status_callbacks_.end()) it->second(false);
+}
+
 void BindingData::CreatePerIsolateProperties(IsolateData* isolate_data,
                                              Local<ObjectTemplate> target) {
   Isolate* isolate = isolate_data->isolate();
   SetMethod(
       isolate, target, "getOrCreateChannelIndex", GetOrCreateChannelIndex);
   SetMethod(isolate, target, "linkNativeChannel", LinkNativeChannel);
+  SetMethod(isolate, target, "notifyChannelActive", NotifyChannelActive);
+  SetMethod(isolate, target, "notifyChannelInactive", NotifyChannelInactive);
 }
 
 void BindingData::CreatePerContextProperties(Local<Object> target,
@@ -148,6 +174,8 @@ void BindingData::RegisterExternalReferences(
     ExternalReferenceRegistry* registry) {
   registry->Register(GetOrCreateChannelIndex);
   registry->Register(LinkNativeChannel);
+  registry->Register(NotifyChannelActive);
+  registry->Register(NotifyChannelInactive);
 }
 
 Channel::Channel(Environment* env,

src/node_diagnostics_channel.h

@@ -4,6 +4,7 @@
 #if defined(NODE_WANT_INTERNALS) && NODE_WANT_INTERNALS
 
 #include <cinttypes>
+#include <functional>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -53,6 +54,14 @@ class BindingData : public SnapshotableObject {
   static void LinkNativeChannel(
       const v8::FunctionCallbackInfo<v8::Value>& args);
 
+  using ChannelStatusCallback = std::function<void(bool is_active)>;
+  void SetChannelStatusCallback(uint32_t index, ChannelStatusCallback cb);
+
+  static void NotifyChannelActive(
+      const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void NotifyChannelInactive(
+      const v8::FunctionCallbackInfo<v8::Value>& args);
+
   static void CreatePerIsolateProperties(IsolateData* isolate_data,
                                          v8::Local<v8::ObjectTemplate> target);
   static void CreatePerContextProperties(v8::Local<v8::Object> target,
@@ -63,6 +72,7 @@ class BindingData : public SnapshotableObject {
 
  private:
   InternalFieldInfo* internal_field_info_ = nullptr;
+  std::unordered_map<uint32_t, ChannelStatusCallback> channel_status_callbacks_;
 };
 
 class Channel : public BaseObject {

src/node_sqlite.cc

@@ -7,6 +7,7 @@
 #include "node.h"
 #include "node_diagnostics_channel.h"
 #include "node_errors.h"
+#include "node_external_reference.h"
 #include "node_mem-inl.h"
 #include "node_url.h"
 #include "simdutf.h"
@@ -80,6 +81,28 @@ inline MaybeLocal<String> Utf8StringMaybeOneByte(Isolate* isolate,
       isolate, input.data(), NewStringType::kNormal, len);
 }
 
+BindingData::BindingData(Realm* realm, Local<Object> wrap)
+    : BaseObject(realm, wrap) {
+  MakeWeak();
+}
+
+void BindingData::MemoryInfo(MemoryTracker* tracker) const {
+  tracker->TrackFieldWithSize("open_databases",
+                              open_databases.size() * sizeof(DatabaseSync*),
+                              "open_databases");
+}
+
+void BindingData::CreatePerContextProperties(Local<Object> target,
+                                             Local<Value> unused,
+                                             Local<Context> context,
+                                             void* priv) {
+  Realm* realm = Realm::GetCurrent(context);
+  realm->AddBindingData<BindingData>(target);
+}
+
+void BindingData::RegisterExternalReferences(
+    ExternalReferenceRegistry* registry) {}
+
 #define CHECK_ERROR_OR_THROW(isolate, db, expr, expected, ret)                 \
   do {                                                                         \
     int r_ = (expr);                                                           \
@@ -887,6 +910,9 @@ DatabaseSync::DatabaseSync(Environment* env,
   enable_load_extension_ = allow_load_extension;
   ignore_next_sqlite_error_ = false;
 
+  BindingData* binding = env->principal_realm()->GetBindingData<BindingData>();
+  if (binding != nullptr) binding->open_databases.insert(this);
+
   if (open) {
     Open();
   }
@@ -910,6 +936,10 @@ void DatabaseSync::DeleteSessions() {
 }
 
 DatabaseSync::~DatabaseSync() {
+  BindingData* binding =
+      env()->principal_realm()->GetBindingData<BindingData>();
+  if (binding != nullptr) binding->open_databases.erase(this);
+
   FinalizeBackups();
 
   if (IsOpen()) {
@@ -994,11 +1024,25 @@ bool DatabaseSync::Open() {
         env()->isolate(), this, load_extension_ret, SQLITE_OK, false);
   }
 
-  sqlite3_trace_v2(connection_, SQLITE_TRACE_STMT, TraceCallback, this);
+  diagnostics_channel::Channel* ch =
+      diagnostics_channel::Channel::Get(env(), "sqlite.db.query");
+  if (ch != nullptr && ch->HasSubscribers()) {
+    sqlite3_trace_v2(connection_, SQLITE_TRACE_STMT, TraceCallback, this);
+  }
 
   return true;
 }
 
+void DatabaseSync::EnableTracing() {
+  if (!IsOpen()) return;
+  sqlite3_trace_v2(connection_, SQLITE_TRACE_STMT, TraceCallback, this);
+}
+
+void DatabaseSync::DisableTracing() {
+  if (!IsOpen()) return;
+  sqlite3_trace_v2(connection_, 0, nullptr, nullptr);
+}
+
 void DatabaseSync::FinalizeBackups() {
   for (auto backup : backups_) {
     backup->Cleanup();
@@ -3966,7 +4010,31 @@ static void Initialize(Local<Object> target,
                        Local<Context> context,
                        void* priv) {
   Environment* env = Environment::GetCurrent(context);
+  Realm* realm = env->principal_realm();
   Isolate* isolate = env->isolate();
+
+  // Set up the per-Environment database registry.
+  BindingData::CreatePerContextProperties(target, unused, context, priv);
+
+  // Register a native callback on the sqlite.db.query diagnostic channel so
+  // that SQLite tracing is enabled/disabled as subscribers come and go.
+  auto* diag_binding =
+      realm->GetBindingData<diagnostics_channel::BindingData>();
+  auto* sqlite_bd = realm->GetBindingData<BindingData>();
+  if (diag_binding != nullptr && sqlite_bd != nullptr) {
+    uint32_t idx = diag_binding->GetOrCreateChannelIndex("sqlite.db.query");
+    BaseObjectPtr<BindingData> bd_ptr(sqlite_bd);
+    diag_binding->SetChannelStatusCallback(idx, [bd_ptr](bool is_active) {
+      BindingData* bd = bd_ptr.get();
+      if (bd == nullptr) return;
+      for (DatabaseSync* db : bd->open_databases) {
+        if (is_active)
+          db->EnableTracing();
+        else
+          db->DisableTracing();
+      }
+    });
+  }
   Local<FunctionTemplate> db_tmpl =
       NewFunctionTemplate(isolate, DatabaseSync::New);
   db_tmpl->InstanceTemplate()->SetInternalFieldCount(
@@ -4047,3 +4115,5 @@ static void Initialize(Local<Object> target,
 }  // namespace node
 
 NODE_BINDING_CONTEXT_AWARE_INTERNAL(sqlite, node::sqlite::Initialize)
+NODE_BINDING_EXTERNAL_REFERENCE(
+    sqlite, node::sqlite::BindingData::RegisterExternalReferences)

src/node_sqlite.h

@@ -18,6 +18,9 @@
 #include <vector>
 
 namespace node {
+
+class ExternalReferenceRegistry;
+
 namespace sqlite {
 
 // Mapping from JavaScript property names to SQLite limit constants
@@ -161,6 +164,27 @@ class StatementExecutionHelper {
                                        bool use_big_ints);
 };
 
+class DatabaseSync;
+
+class BindingData : public BaseObject {
+ public:
+  SET_BINDING_ID(sqlite_binding_data)
+
+  BindingData(Realm* realm, v8::Local<v8::Object> wrap);
+
+  void MemoryInfo(MemoryTracker* tracker) const override;
+  SET_MEMORY_INFO_NAME(BindingData)
+  SET_SELF_SIZE(BindingData)
+
+  std::unordered_set<DatabaseSync*> open_databases;
+
+  static void CreatePerContextProperties(v8::Local<v8::Object> target,
+                                         v8::Local<v8::Value> unused,
+                                         v8::Local<v8::Context> context,
+                                         void* priv);
+  static void RegisterExternalReferences(ExternalReferenceRegistry* registry);
+};
+
 class DatabaseSync : public BaseObject {
  public:
   enum InternalFields {
@@ -231,6 +255,8 @@ class DatabaseSync : public BaseObject {
   // enable that use case.
   void SetIgnoreNextSQLiteError(bool ignore);
   bool ShouldIgnoreSQLiteError();
+  void EnableTracing();
+  void DisableTracing();
 
   SET_MEMORY_INFO_NAME(DatabaseSync)
   SET_SELF_SIZE(DatabaseSync)