perfetto: extract out lifecycle management from the window operator

This CL extracts out the lifecycle management of vtabs from the window
operator. This prepares the ground for the use of these methods from other module implementations.

Change-Id: Iab8e9321ef28e110ffcd54da1b0dad1c200fa3e7
diff --git a/BUILD b/BUILD
index 7f7ff58..b5ce58d 100644
--- a/BUILD
+++ b/BUILD
@@ -2693,6 +2693,7 @@
     srcs = [
         "src/trace_processor/sqlite/db_sqlite_table.cc",
         "src/trace_processor/sqlite/db_sqlite_table.h",
+        "src/trace_processor/sqlite/module_lifecycle_manager.h",
         "src/trace_processor/sqlite/query_cache.h",
         "src/trace_processor/sqlite/scoped_db.h",
         "src/trace_processor/sqlite/sql_source.cc",
diff --git a/src/trace_processor/perfetto_sql/intrinsics/operators/window_operator.cc b/src/trace_processor/perfetto_sql/intrinsics/operators/window_operator.cc
index eb31bb3..3092af1 100644
--- a/src/trace_processor/perfetto_sql/intrinsics/operators/window_operator.cc
+++ b/src/trace_processor/perfetto_sql/intrinsics/operators/window_operator.cc
@@ -22,6 +22,7 @@
 
 #include "perfetto/base/logging.h"
 #include "src/trace_processor/sqlite/bindings/sqlite_result.h"
+#include "src/trace_processor/sqlite/module_lifecycle_manager.h"
 #include "src/trace_processor/sqlite/sqlite_utils.h"
 
 namespace perfetto::trace_processor {
@@ -52,24 +53,15 @@
     return ret;
   }
   auto* ctx = GetContext(raw_ctx);
-  auto it_and_inserted = ctx->state_by_name.Insert(argv[2], nullptr);
-  PERFETTO_CHECK(
-      it_and_inserted.second ||
-      (it_and_inserted.first && it_and_inserted.first->get()->disconnected));
-  *it_and_inserted.first = std::make_unique<State>();
-
   std::unique_ptr<Vtab> res = std::make_unique<Vtab>();
-  res->context = ctx;
-  res->name = argv[2];
-  res->state = it_and_inserted.first->get();
+  res->state = ctx->manager.OnCreate(argv, std::make_unique<State>());
   *vtab = res.release();
   return SQLITE_OK;
 }
 
 int WindowOperatorModule::Destroy(sqlite3_vtab* vtab) {
-  auto* tab = GetVtab(vtab);
-  PERFETTO_CHECK(tab->context->state_by_name.Erase(tab->name));
-  delete tab;
+  std::unique_ptr<Vtab> tab(GetVtab(vtab));
+  sqlite::ModuleStateManager<WindowOperatorModule>::OnDestroy(tab->state);
   return SQLITE_OK;
 }
 
@@ -84,24 +76,15 @@
     return ret;
   }
   auto* ctx = GetContext(raw_ctx);
-  auto* ptr = ctx->state_by_name.Find(argv[2]);
-  PERFETTO_CHECK(ptr);
-  ptr->get()->disconnected = false;
-
   std::unique_ptr<Vtab> res = std::make_unique<Vtab>();
-  res->context = ctx;
-  res->name = argv[2];
-  res->state = ptr->get();
+  res->state = ctx->manager.OnConnect(argv);
   *vtab = res.release();
   return SQLITE_OK;
 }
 
 int WindowOperatorModule::Disconnect(sqlite3_vtab* vtab) {
-  auto* tab = GetVtab(vtab);
-  auto* ptr = tab->context->state_by_name.Find(tab->name);
-  PERFETTO_CHECK(ptr);
-  ptr->get()->disconnected = true;
-  delete tab;
+  std::unique_ptr<Vtab> tab(GetVtab(vtab));
+  sqlite::ModuleStateManager<WindowOperatorModule>::OnDisconnect(tab->state);
   return SQLITE_OK;
 }
 
@@ -143,13 +126,12 @@
                                  sqlite3_value** argv) {
   auto* t = GetVtab(cursor->pVtab);
   auto* c = GetCursor(cursor);
+  auto* s =
+      sqlite::ModuleStateManager<WindowOperatorModule>::GetState(t->state);
 
-  c->window_start = t->state->window_start;
-  c->window_end = t->state->window_start + t->state->window_dur;
-  c->step_size =
-      t->state->quantum == 0 ? t->state->window_dur : t->state->quantum;
-
-  c->current_ts = c->window_start;
+  c->window_end = s->window_start + s->window_dur;
+  c->step_size = s->quantum == 0 ? s->window_dur : s->quantum;
+  c->current_ts = s->window_start;
 
   if (is_row_id_constraint) {
     PERFETTO_CHECK(argc == 1);
@@ -186,18 +168,19 @@
                                  int N) {
   auto* t = GetVtab(cursor->pVtab);
   auto* c = GetCursor(cursor);
+  auto* s =
+      sqlite::ModuleStateManager<WindowOperatorModule>::GetState(t->state);
   switch (N) {
     case Column::kQuantum: {
-      sqlite::result::Long(ctx, static_cast<sqlite_int64>(t->state->quantum));
+      sqlite::result::Long(ctx, static_cast<sqlite_int64>(s->quantum));
       break;
     }
     case Column::kWindowStart: {
-      sqlite::result::Long(ctx,
-                           static_cast<sqlite_int64>(t->state->window_start));
+      sqlite::result::Long(ctx, static_cast<sqlite_int64>(s->window_start));
       break;
     }
     case Column::kWindowDur: {
-      sqlite::result::Long(ctx, static_cast<int>(t->state->window_dur));
+      sqlite::result::Long(ctx, static_cast<int>(s->window_dur));
       break;
     }
     case Column::kTs: {
@@ -233,6 +216,8 @@
                                  sqlite3_value** argv,
                                  sqlite_int64*) {
   auto* t = GetVtab(tab);
+  auto* s =
+      sqlite::ModuleStateManager<WindowOperatorModule>::GetState(t->state);
 
   // We only support updates to ts and dur. Disallow deletes (argc == 1) and
   // inserts (argv[0] == null).
@@ -249,9 +234,9 @@
         tab, "Cannot set duration of window table to zero.");
   }
 
-  t->state->quantum = new_quantum;
-  t->state->window_start = new_start;
-  t->state->window_dur = new_dur;
+  s->quantum = new_quantum;
+  s->window_start = new_start;
+  s->window_dur = new_dur;
 
   return SQLITE_OK;
 }
diff --git a/src/trace_processor/perfetto_sql/intrinsics/operators/window_operator.h b/src/trace_processor/perfetto_sql/intrinsics/operators/window_operator.h
index d023042..f7cf805 100644
--- a/src/trace_processor/perfetto_sql/intrinsics/operators/window_operator.h
+++ b/src/trace_processor/perfetto_sql/intrinsics/operators/window_operator.h
@@ -24,6 +24,7 @@
 
 #include "perfetto/ext/base/flat_hash_map.h"
 #include "src/trace_processor/sqlite/bindings/sqlite_module.h"
+#include "src/trace_processor/sqlite/module_lifecycle_manager.h"
 
 namespace perfetto::trace_processor {
 
@@ -39,7 +40,6 @@
     kReturnFirst = 1,
   };
   struct State {
-    bool disconnected = false;
     int64_t quantum = 0;
     int64_t window_start = 0;
 
@@ -48,15 +48,12 @@
     int64_t window_dur = std::numeric_limits<int64_t>::max();
   };
   struct Context {
-    base::FlatHashMap<std::string, std::unique_ptr<State>> state_by_name;
+    sqlite::ModuleStateManager<WindowOperatorModule> manager;
   };
   struct Vtab : sqlite::Module<WindowOperatorModule>::Vtab {
-    Context* context;
-    std::string name;
-    State* state = nullptr;
+    sqlite::ModuleStateManager<WindowOperatorModule>::PerVtabState* state;
   };
   struct Cursor : sqlite::Module<WindowOperatorModule>::Cursor {
-    int64_t window_start = 0;
     int64_t window_end = 0;
     int64_t step_size = 0;
 
diff --git a/src/trace_processor/sqlite/BUILD.gn b/src/trace_processor/sqlite/BUILD.gn
index 0a3beef..927b4a5 100644
--- a/src/trace_processor/sqlite/BUILD.gn
+++ b/src/trace_processor/sqlite/BUILD.gn
@@ -20,6 +20,7 @@
   sources = [
     "db_sqlite_table.cc",
     "db_sqlite_table.h",
+    "module_lifecycle_manager.h",
     "query_cache.h",
     "scoped_db.h",
     "sql_source.cc",
diff --git a/src/trace_processor/sqlite/module_lifecycle_manager.h b/src/trace_processor/sqlite/module_lifecycle_manager.h
new file mode 100644
index 0000000..0e44a6d
--- /dev/null
+++ b/src/trace_processor/sqlite/module_lifecycle_manager.h
@@ -0,0 +1,126 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SRC_TRACE_PROCESSOR_SQLITE_MODULE_LIFECYCLE_MANAGER_H_
+#define SRC_TRACE_PROCESSOR_SQLITE_MODULE_LIFECYCLE_MANAGER_H_
+
+#include <memory>
+#include <string>
+
+#include "perfetto/base/logging.h"
+#include "perfetto/ext/base/flat_hash_map.h"
+
+namespace perfetto::trace_processor::sqlite {
+
+// Helper class which abstracts away management of per-vtab state of an SQLite
+// virtual table module.
+//
+// SQLite has some subtle semantics around lifecycle of vtabs which makes state
+// management complex. This class attempts to encapsulate some of that
+// complexity as a central place where we can document the quirks.
+//
+// Usage of this class:
+// struct MyModule : sqlite::Module<MyModule> {
+//   struct Context {
+//     // Store the manager in the context object.
+//     ModuleStateManager<MyModule> manager.
+//     ... (other fields)
+//   }
+//   struct Vtab : sqlite::Module<MyModule>::Vtab {
+//     // Store the per-vtab-state pointer in the vtab object.
+//     ModuleStateManager<MyModule>::PerVtabState* state;
+//     ... (other fields)
+//   }
+//   static void OnCreate(...) {
+//     ...
+//     // Call OnCreate on the manager object and store the returned pointer.
+//     tab->state = ctx->manager.OnCreate(argv);
+//     ...
+//   }
+//   static void OnDestroy(...) {
+//     ...
+//     // Call OnDestroy with the stored state pointer.
+//     sqlite::ModuleStateManager<MyModule>::OnDestroy(tab->state);
+//     ...
+//   }
+//   // Do the same in OnConnect and OnDisconnect as in OnCreate and OnDestroy
+//   // respectively.
+//   static void OnConnect(...)
+//   static void OnDisconnect(...)
+// }
+template <typename Module>
+class ModuleStateManager {
+ public:
+  // Per-vtab state. The pointer to this class should be stored in the Vtab.
+  struct PerVtabState {
+    ModuleStateManager* manager;
+    bool disconnected = false;
+    std::string table_name;
+    std::unique_ptr<typename Module::State> state;
+  };
+
+  // Lifecycle method to be called from Module::Create.
+  [[nodiscard]] PerVtabState* OnCreate(
+      const char* const* argv,
+      std::unique_ptr<typename Module::State> state) {
+    auto it_and_inserted = state_by_name_.Insert(argv[2], nullptr);
+    PERFETTO_CHECK(
+        it_and_inserted.second ||
+        (it_and_inserted.first && it_and_inserted.first->get()->disconnected));
+
+    auto s = std::make_unique<PerVtabState>();
+    auto* s_ptr = s.get();
+    *it_and_inserted.first = std::move(s);
+
+    s_ptr->manager = this;
+    s_ptr->table_name = argv[2];
+    s_ptr->state = std::move(state);
+    return it_and_inserted.first->get();
+  }
+
+  // Lifecycle method to be called from Module::Connect.
+  [[nodiscard]] PerVtabState* OnConnect(const char* const* argv) {
+    auto* ptr = state_by_name_.Find(argv[2]);
+    PERFETTO_CHECK(ptr);
+    ptr->get()->disconnected = false;
+    return ptr->get();
+  }
+
+  // Lifecycle method to be called from Module::Disconnect.
+  static void OnDisconnect(PerVtabState* state) {
+    auto* ptr = state->manager->state_by_name_.Find(state->table_name);
+    PERFETTO_CHECK(ptr);
+    ptr->get()->disconnected = true;
+  }
+
+  // Lifecycle method to be called from Module::Destroy.
+  static void OnDestroy(PerVtabState* state) {
+    PERFETTO_CHECK(state->manager->state_by_name_.Erase(state->table_name));
+  }
+
+  // Method to be called from module callbacks to extract the module state
+  // from the manager state.
+  static typename Module::State* GetState(PerVtabState* s) {
+    return s->state.get();
+  }
+
+ private:
+  base::FlatHashMap<std::string, std::unique_ptr<PerVtabState>> state_by_name_;
+};
+
+}  // namespace perfetto::trace_processor::sqlite
+
+#endif  // SRC_TRACE_PROCESSOR_SQLITE_MODULE_LIFECYCLE_MANAGER_H_