From f6594978040d01ea4dbee69d20bb82c187ee1083 Mon Sep 17 00:00:00 2001
From: FredyH <frederikhaselmeier@gmail.com>
Date: Mon, 4 Sep 2023 03:22:34 +0200
Subject: [PATCH] Added options to set timeout values on database object

---
 src/lua/LuaDatabase.cpp | 39 +++++++++++++++++++++++++++++++++++++++
 src/mysql/Database.cpp  | 30 ++++++++++++++++++++++++++++--
 src/mysql/Database.h    | 11 +++++++++++
 3 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/src/lua/LuaDatabase.cpp b/src/lua/LuaDatabase.cpp
index a5ed0d9..518426e 100644
--- a/src/lua/LuaDatabase.cpp
+++ b/src/lua/LuaDatabase.cpp
@@ -148,6 +148,36 @@ MYSQLOO_LUA_FUNCTION(setSSLSettings) {
     return 0;
 }
 
+MYSQLOO_LUA_FUNCTION(setReadTimeout) {
+    auto database = LuaObject::getLuaObject<LuaDatabase>(LUA);
+    unsigned int timeout = (int) LUA->GetNumber(2);
+    if (timeout == 0) {
+        LUA->ThrowError("Timeout must be at least 1");
+    }
+    database->m_database->setReadTimeout(timeout);
+    return 0;
+}
+
+MYSQLOO_LUA_FUNCTION(setWriteTimeout) {
+    auto database = LuaObject::getLuaObject<LuaDatabase>(LUA);
+    unsigned int timeout = (int) LUA->GetNumber(2);
+    if (timeout == 0) {
+        LUA->ThrowError("Timeout must be at least 1");
+    }
+    database->m_database->setWriteTimeout(timeout);
+    return 0;
+}
+
+MYSQLOO_LUA_FUNCTION(setConnectTimeout) {
+    auto database = LuaObject::getLuaObject<LuaDatabase>(LUA);
+    unsigned int timeout = (int) LUA->GetNumber(2);
+    if (timeout == 0) {
+        LUA->ThrowError("Timeout must be at least 1");
+    }
+    database->m_database->setConnectTimeout(timeout);
+    return 0;
+}
+
 MYSQLOO_LUA_FUNCTION(disconnect) {
     auto database = LuaObject::getLuaObject<LuaDatabase>(LUA);
     bool wait = false;
@@ -252,6 +282,15 @@ void LuaDatabase::createMetaTable(ILuaBase *LUA) {
     LUA->PushCFunction(setSSLSettings);
     LUA->SetField(-2, "setSSLSettings");
 
+    LUA->PushCFunction(setReadTimeout);
+    LUA->SetField(-2, "setReadTimeout");
+
+    LUA->PushCFunction(setWriteTimeout);
+    LUA->SetField(-2, "setWriteTimeout");
+
+    LUA->PushCFunction(setConnectTimeout);
+    LUA->SetField(-2, "setConnectTimeout");
+
     LUA->PushCFunction(disconnect);
     LUA->SetField(-2, "disconnect");
 
diff --git a/src/mysql/Database.cpp b/src/mysql/Database.cpp
index d6c5a1f..fcfaa62 100644
--- a/src/mysql/Database.cpp
+++ b/src/mysql/Database.cpp
@@ -1,7 +1,6 @@
 #include "Database.h"
 #include "MySQLOOException.h"
 #include <string>
-#include <cstring>
 #include <iostream>
 #include <utility>
 #include "mysqld_error.h"
@@ -374,6 +373,7 @@ void Database::connectRun() {
             m_status = DATABASE_CONNECTION_FAILED;
             return;
         }
+        this->applyTimeoutSettings();
         this->customSSLSettings.applySSLSettings(this->m_sql);
         const char *socketStr = (this->socket.length() == 0) ? nullptr : this->socket.c_str();
         unsigned long clientFlag = (this->useMultiStatements) ? CLIENT_MULTI_STATEMENTS : 0;
@@ -406,6 +406,8 @@ void Database::connectRun() {
     }
 }
 
+#pragma clang diagnostic push
+#pragma ide diagnostic ignored "misc-no-recursion"
 void Database::runQuery(const std::shared_ptr<IQuery>& query, const std::shared_ptr<IQueryData>& data, bool retry) {
     try {
         query->executeStatement(*this, this->m_sql, data);
@@ -427,6 +429,7 @@ void Database::runQuery(const std::shared_ptr<IQuery>& query, const std::shared_
         }
     }
 }
+#pragma clang diagnostic pop
 
 /* The run method of the thread of the database instance.
  */
@@ -469,11 +472,34 @@ void Database::run() {
 bool Database::attemptReconnect() {
     bool success;
     my_bool reconnect = '1';
-    mysql_optionsv(this->m_sql, MYSQL_OPT_RECONNECT, &reconnect);
     success = mariadb_reconnect(this->m_sql) == 0;
     reconnect = '0';
     mysql_optionsv(this->m_sql, MYSQL_OPT_RECONNECT, &reconnect);
     return success;
 }
 
+void Database::setConnectTimeout(unsigned int timeout) {
+    this->connectTimeout = timeout;
+}
+
+void Database::setReadTimeout(unsigned int timeout) {
+    this->readTimeout = timeout;
+}
+
+void Database::setWriteTimeout(unsigned int timeout) {
+    this->writeTimeout = timeout;
+}
+
+void Database::applyTimeoutSettings() {
+    if (this->connectTimeout > 0) {
+        mysql_optionsv(this->m_sql, MYSQL_OPT_CONNECT_TIMEOUT, &this->connectTimeout);
+    }
+    if (this->readTimeout > 0) {
+        mysql_optionsv(this->m_sql, MYSQL_OPT_READ_TIMEOUT, &this->readTimeout);
+    }
+    if (this->writeTimeout > 0) {
+        mysql_optionsv(this->m_sql, MYSQL_OPT_WRITE_TIMEOUT, &this->writeTimeout);
+    }
+}
+
 StatementHandle::StatementHandle(MYSQL_STMT *stmt, bool valid) : stmt(stmt), valid(valid) {}
diff --git a/src/mysql/Database.h b/src/mysql/Database.h
index da6ea5e..ff0b3e2 100644
--- a/src/mysql/Database.h
+++ b/src/mysql/Database.h
@@ -95,6 +95,12 @@ class Database : public std::enable_shared_from_this<Database> {
 
     void disconnect(bool wait);
 
+    void setConnectTimeout(unsigned int timeout);
+
+    void setReadTimeout(unsigned int timeout);
+
+    void setWriteTimeout(unsigned int timeout);
+
     void setSSLSettings(const SSLSettings &settings);
 
     bool isConnectionDone() { return m_connectionDone; }
@@ -132,6 +138,8 @@ class Database : public std::enable_shared_from_this<Database> {
 
     void waitForQuery(const std::shared_ptr<IQuery> &query, const std::shared_ptr<IQueryData> &data);
 
+    void applyTimeoutSettings();
+
     BlockingQueue<std::pair<std::shared_ptr<IQuery>, std::shared_ptr<IQueryData>>> finishedQueries{};
     BlockingQueue<std::pair<std::shared_ptr<IQuery>, std::shared_ptr<IQueryData>>> queryQueue{};
     std::unordered_set<std::shared_ptr<StatementHandle>> cachedStatements{};
@@ -165,6 +173,9 @@ class Database : public std::enable_shared_from_this<Database> {
     std::string socket;
     unsigned int port;
     SSLSettings customSSLSettings{};
+    unsigned int readTimeout = 0;
+    unsigned int writeTimeout = 0;
+    unsigned int connectTimeout = 0;
     std::atomic<DatabaseStatus> m_status{DATABASE_NOT_CONNECTED};
 };