Merge pull request #1543 from Expensify/main

Update expensify_prod branch
Expensify · Jul 24, 2023 · cc77451 · cc77451
2 parents e8a59d2 + fa4ac12
commit cc77451
Show file tree

Hide file tree

Showing 6 changed files with 54 additions and 23 deletions.
diff --git a/BedrockServer.cpp b/BedrockServer.cpp
@@ -278,6 +278,17 @@ void BedrockServer::sync()
         // Ok, let the sync node to it's updating for as many iterations as it requires. We'll update the replication
         // state when it's finished.
         SQLiteNodeState preUpdateState = _syncNode->getState();
+        if(command && committingCommand) {
+            void (*onPrepareHandler)(SQLite& db, int64_t tableID) = nullptr;
+            bool enabled = command->shouldEnableOnPrepareNotification(db, &onPrepareHandler);
+            if (enabled) {
+                _syncNode->onPrepareHandlerEnabled = enabled;
+                _syncNode->onPrepareHandler = onPrepareHandler;
+            }
+        } else {
+            _syncNode->onPrepareHandlerEnabled = false;
+            _syncNode->onPrepareHandler = nullptr;
+        }
         while (_syncNode->update()) {}
         SQLiteNodeState nodeState = _syncNode->getState();
         _replicationState.store(nodeState);

diff --git a/libstuff/AutoScopeOnPrepare.cpp b/libstuff/AutoScopeOnPrepare.cpp
@@ -0,0 +1,16 @@
+#include "AutoScopeOnPrepare.h"
+
+AutoScopeOnPrepare::AutoScopeOnPrepare(bool enable, SQLite& db, void (*handler)(SQLite& _db, int64_t tableID))
+    : _enable(enable), _db(db), _handler(handler) {
+    if (_enable) {
+        _db.setOnPrepareHandler(_handler);
+        _db.enablePrepareNotifications(true);
+    }
+}
+
+AutoScopeOnPrepare ::~AutoScopeOnPrepare() {
+    if (_enable) {
+        _db.setOnPrepareHandler(nullptr);
+        _db.enablePrepareNotifications(false);
+    }
+}
diff --git a/libstuff/AutoScopeOnPrepare.h b/libstuff/AutoScopeOnPrepare.h
@@ -0,0 +1,15 @@
+#pragma once
+#include <sqlitecluster/SQLite.h>
+using namespace std;
+
+// RAII-style mechanism for automatically setting and unsetting an on prepare handler
+class AutoScopeOnPrepare {
+  public:
+    AutoScopeOnPrepare(bool enable, SQLite& db, void (*handler)(SQLite& _db, int64_t tableID));
+    ~AutoScopeOnPrepare();
+
+  private:
+    bool _enable;
+    SQLite& _db;
+    void (*_handler)(SQLite& _db, int64_t tableID);
+};
diff --git a/sqlitecluster/SQLiteCore.cpp b/sqlitecluster/SQLiteCore.cpp
@@ -1,3 +1,4 @@
+#include <libstuff/AutoScopeOnPrepare.h>
 #include <libstuff/libstuff.h>
 #include "SQLiteCore.h"
 #include "SQLite.h"
@@ -6,28 +7,6 @@
 SQLiteCore::SQLiteCore(SQLite& db) : _db(db)
 { }
 
-// RAII-style mechanism for automatically setting and unsetting an on prepare handler
-class AutoScopeOnPrepare {
-  public:
-    AutoScopeOnPrepare(bool enable, SQLite& db, void (*handler)(SQLite& _db, int64_t tableID)) : _enable(enable), _db(db), _handler(handler) {
-        if (_enable) {
-            _db.setOnPrepareHandler(_handler);
-            _db.enablePrepareNotifications(true);
-        }
-    }
-    ~AutoScopeOnPrepare() {
-        if (_enable) {
-            _db.setOnPrepareHandler(nullptr);
-            _db.enablePrepareNotifications(false);
-        }
-    }
-
-  private:
-    bool _enable;
-    SQLite& _db;
-    void (*_handler)(SQLite& _db, int64_t tableID);
-};
-
 bool SQLiteCore::commit(const string& description, uint64_t& commitID, string& transactionHash,
                         bool needsPluginNotification, void (*notificationHandler)(SQLite& _db, int64_t tableID)) {
 

diff --git a/sqlitecluster/SQLiteNode.cpp b/sqlitecluster/SQLiteNode.cpp
@@ -2,6 +2,7 @@
 
 #include <unistd.h>
 
+#include <libstuff/AutoScopeOnPrepare.h>
 #include <libstuff/libstuff.h>
 #include <libstuff/SRandom.h>
 #include <libstuff/SQResult.h>
@@ -127,6 +128,7 @@ SQLiteNode::SQLiteNode(SQLiteServer& server, shared_ptr<SQLitePool> dbPool, cons
       _syncPeer(nullptr)
 {
     SASSERT(_originalPriority >= 0);
+    onPrepareHandlerEnabled = false;
     SINFO("[NOTIFY] setting commit count to: " << _db.getCommitCount());
     _localCommitNotifier.notifyThrough(_db.getCommitCount());
 
@@ -999,7 +1001,10 @@ bool SQLiteNode::update() {
 
             // There's no handling for a failed prepare. This should only happen if the DB has been corrupted or
             // something catastrophic like that.
-            SASSERT(_db.prepare());
+            {
+                AutoScopeOnPrepare onPrepare(onPrepareHandlerEnabled, _db, onPrepareHandler);
+                SASSERT(_db.prepare());
+            }
 
             // Begin the distributed transaction
             SData transaction("BEGIN_TRANSACTION");

diff --git a/sqlitecluster/SQLiteNode.h b/sqlitecluster/SQLiteNode.h
@@ -164,6 +164,11 @@ class SQLiteNode : public STCPManager {
     // access _peerList and peer->name, both of which are const. So it is safe
     // to call from other public functions.
     SQLitePeer* getPeerByName(const string& name) const;
+
+    // Pointer to the current on prepare handler to be passed on commit to SQLite
+    void (*onPrepareHandler)(SQLite& _db, int64_t tableID);
+    bool onPrepareHandlerEnabled;
+
   private:
     // Utility class that can decrement _replicationThreadCount when objects go out of scope.
     template <typename CounterType>