Fix timeouts not worknig when not leading or follwoing

BedrockCore.cpp

@@ -52,12 +52,12 @@ uint64_t BedrockCore::_getRemainingTime(const unique_ptr<BedrockCommand>& comman
     return isProcessing ? min(processTimeout, adjustedTimeout) : adjustedTimeout;
 }
 
-bool BedrockCore::isTimedOut(unique_ptr<BedrockCommand>& command) {
+bool BedrockCore::isTimedOut(unique_ptr<BedrockCommand>& command, SQLite* db, const BedrockServer* server) {
     try {
         _getRemainingTime(command, false);
     } catch (const SException& e) {
         // Yep, timed out.
-        _handleCommandException(command, e);
+        _handleCommandException(command, e, db, server);
         command->complete = true;
         return true;
     }
@@ -104,7 +104,7 @@ void BedrockCore::prePeekCommand(unique_ptr<BedrockCommand>& command, bool isBlo
             STHROW("555 Timeout prePeeking command");
         }
     } catch (const SException& e) {
-        _handleCommandException(command, e);
+        _handleCommandException(command, e, &_db, &_server);
         command->complete = true;
     } catch (...) {
         SALERT("Unhandled exception typename: " << SGetCurrentExceptionName() << ", command: " << request.methodLine);
@@ -186,7 +186,7 @@ BedrockCore::RESULT BedrockCore::peekCommand(unique_ptr<BedrockCommand>& command
         }
     } catch (const SException& e) {
         command->repeek = false;
-        _handleCommandException(command, e);
+        _handleCommandException(command, e, &_db, &_server);
     } catch (const SHTTPSManager::NotLeading& e) {
         command->repeek = false;
         returnValue = RESULT::SHOULD_PROCESS;
@@ -284,7 +284,7 @@ BedrockCore::RESULT BedrockCore::processCommand(unique_ptr<BedrockCommand>& comm
             }
         }
     } catch (const SException& e) {
-        _handleCommandException(command, e);
+        _handleCommandException(command, e, &_db, &_server);
         _db.rollback();
         needsCommit = false;
     } catch (const SQLite::constraint_error& e) {
@@ -353,7 +353,7 @@ void BedrockCore::postProcessCommand(unique_ptr<BedrockCommand>& command, bool i
             STHROW("555 Timeout postProcessing command");
         }
     } catch (const SException& e) {
-        _handleCommandException(command, e);
+        _handleCommandException(command, e, &_db, &_server);
     } catch (...) {
         SALERT("Unhandled exception typename: " << SGetCurrentExceptionName() << ", command: " << request.methodLine);
         command->response.methodLine = "500 Unhandled Exception";
@@ -367,7 +367,7 @@ void BedrockCore::postProcessCommand(unique_ptr<BedrockCommand>& command, bool i
     _db.setQueryOnly(false);
 }
 
-void BedrockCore::_handleCommandException(unique_ptr<BedrockCommand>& command, const SException& e) {
+void BedrockCore::_handleCommandException(unique_ptr<BedrockCommand>& command, const SException& e, SQLite* db, const BedrockServer* server) {
     string msg = "Error processing command '" + command->request.methodLine + "' (" + e.what() + "), ignoring.";
     if (!e.body.empty()) {
         msg = msg + " Request body: " + e.body;
@@ -396,9 +396,11 @@ void BedrockCore::_handleCommandException(unique_ptr<BedrockCommand>& command, c
     }
 
     // Add the commitCount header to the response.
-    command->response["commitCount"] = to_string(_db.getCommitCount());
+    if (db) {
+        command->response["commitCount"] = to_string(db->getCommitCount());
+    }
 
-    if (_server.args.isSet("-extraExceptionLogging")) {
+    if (server && server->args.isSet("-extraExceptionLogging")) {
         auto stack = e.details();
         command->response["exceptionSource"] = stack.back();
     }

BedrockCore.h

@@ -34,7 +34,7 @@ class BedrockCore : public SQLiteCore {
     // Checks if a command has already timed out. Like `peekCommand` without doing any work. Returns `true` and sets
     // the same command state as `peekCommand` would if the command has timed out. Returns `false` and does nothing if
     // the command hasn't timed out.
-    bool isTimedOut(unique_ptr<BedrockCommand>& command);
+    static bool isTimedOut(unique_ptr<BedrockCommand>& command, SQLite* db = nullptr, const BedrockServer* server = nullptr);
 
     void prePeekCommand(unique_ptr<BedrockCommand>& command, bool isBlockingCommitThread);
 
@@ -71,8 +71,8 @@ class BedrockCore : public SQLiteCore {
     // Gets the amount of time remaining until this command times out. This is the difference between the command's
     // 'timeout' value (or the default timeout, if not set) and the time the command was initially scheduled to run. If
     // this time is already expired, this throws `555 Timeout`
-    uint64_t _getRemainingTime(const unique_ptr<BedrockCommand>& command, bool isProcessing);
+    static uint64_t _getRemainingTime(const unique_ptr<BedrockCommand>& command, bool isProcessing);
 
-    void _handleCommandException(unique_ptr<BedrockCommand>& command, const SException& e);
+    static void _handleCommandException(unique_ptr<BedrockCommand>& command, const SException& e, SQLite* db = nullptr, const BedrockServer* server = nullptr);
     const BedrockServer& _server;
 };

BedrockServer.cpp

@@ -801,6 +801,14 @@ void BedrockServer::runCommand(unique_ptr<BedrockCommand>&& _command, bool isBlo
         // We just spin until the node looks ready to go. Typically, this doesn't happen expect briefly at startup.
         size_t waitCount = 0;
         while (_upgradeInProgress || (getState() != SQLiteNodeState::LEADING && getState() != SQLiteNodeState::FOLLOWING)) {
+
+            // It's feasible that our command times out in this loop. In this case, we do not have a DB object to pass.
+            // The only implication of this is the response does not get the commitCount attached to it.
+            if (BedrockCore::isTimedOut(command, nullptr, this)) {
+                _reply(command);
+                return;
+            }
+
             // This sleep call is pretty ugly, but it should almost never happen. We're accepting the potential
             // looping sleep call for the general case where we just check some bools and continue, instead of
             // avoiding the sleep call but having every thread lock a mutex here on every loop.
@@ -880,7 +888,7 @@ void BedrockServer::runCommand(unique_ptr<BedrockCommand>&& _command, bool isBlo
             // to be returned to the main queue, where they would have timed out in `peek`, but it was never called
             // because the commands already had a HTTPS request attached, and then they were immediately re-sent to the
             // sync queue, because of the QUORUM consistency requirement, resulting in an endless loop.
-            if (core.isTimedOut(command)) {
+            if (core.isTimedOut(command, &db, this)) {
                 _reply(command);
                 return;
             }

sqlitecluster/SQLiteCore.h

@@ -2,6 +2,11 @@
 class SQLite;
 class SQLiteNode;
 
+#include <cstdint>
+#include <string>
+
+using namespace std;
+
 class SQLiteCore {
   public:
     // Constructor that stores the database object we'll be working on.

test/lib/BedrockTester.cpp

@@ -546,10 +546,10 @@ void BedrockTester::freeDB() {
     _db = nullptr;
 }
 
-string BedrockTester::readDB(const string& query, bool online)
+string BedrockTester::readDB(const string& query, bool online, int64_t timeoutMS)
 {
     SQResult result;
-    bool success = readDB(query, result, online);
+    bool success = readDB(query, result, online, timeoutMS);
     if (!success) {
         return "";
     }
@@ -565,7 +565,7 @@ string BedrockTester::readDB(const string& query, bool online)
     return result.rows[0][0];
 }
 
-bool BedrockTester::readDB(const string& query, SQResult& result, bool online)
+bool BedrockTester::readDB(const string& query, SQResult& result, bool online, int64_t timeoutMS)
 {
     if (ENABLE_HCTREE && online) {
         string fixedQuery = query;
@@ -576,6 +576,9 @@ bool BedrockTester::readDB(const string& query, SQResult& result, bool online)
         SData command("Query");
         command["Query"] = fixedQuery;
         command["Format"] = "JSON";
+        if (timeoutMS) {
+            command["timeout"] = to_string(timeoutMS);
+        }
         auto commandResult = executeWaitMultipleData({command}, 1);
         auto row0 = SParseJSONObject(commandResult[0].content)["rows"];
         auto headerString = SParseJSONObject(commandResult[0].content)["headers"];

test/lib/BedrockTester.h

@@ -79,8 +79,9 @@ class BedrockTester {
 
     // Read from the DB file, without going through the bedrock server. Two interfaces are provided to maintain
     // compatibility with the `SQLite` class.
-    string readDB(const string& query, bool online = true);
-    bool readDB(const string& query, SQResult& result, bool online = true);
+    // Note that timeoutMS only applies in HC-Tree mode. It is ignored in WAL2 mode.
+    string readDB(const string& query, bool online = true, int64_t timeoutMS = 0);
+    bool readDB(const string& query, SQResult& result, bool online = true, int64_t timeoutMS = 0);
 
     // Closes and releases any existing DB file.
     void freeDB();