Core/DBLayer: Add nicer api for SQLQueryHolders

8 files changed, 50 insertions, 14 deletions

diff --git a/src/server/database/Database/DatabaseEnvFwd.h b/src/server/database/Database/DatabaseEnvFwd.h
index 8d09e502f72..ba41655e213 100644
--- a/src/server/database/Database/DatabaseEnvFwd.h
+++ b/src/server/database/Database/DatabaseEnvFwd.h
@@ -82,6 +82,8 @@ using CharacterDatabaseQueryHolder = SQLQueryHolder<CharacterDatabaseConnection>
 using LoginDatabaseQueryHolder = SQLQueryHolder<LoginDatabaseConnection>;
 using WorldDatabaseQueryHolder = SQLQueryHolder<WorldDatabaseConnection>;
 
+class SQLQueryHolderCallback;
+
 // mysql
 struct MySQLHandle;
 struct MySQLResult;
diff --git a/src/server/database/Database/DatabaseWorkerPool.cpp b/src/server/database/Database/DatabaseWorkerPool.cpp
index c0a24c6ed11..8ef1e28d70f 100644
--- a/src/server/database/Database/DatabaseWorkerPool.cpp
+++ b/src/server/database/Database/DatabaseWorkerPool.cpp
@@ -226,7 +226,7 @@ QueryCallback DatabaseWorkerPool<T>::AsyncQuery(PreparedStatement<T>* stmt)
 }
 
 template <class T>
-QueryResultHolderFuture DatabaseWorkerPool<T>::DelayQueryHolder(SQLQueryHolder<T>* holder)
+SQLQueryHolderCallback DatabaseWorkerPool<T>::DelayQueryHolder(SQLQueryHolder<T>* holder)
 {
     SQLQueryHolderTask* task = new SQLQueryHolderTask(holder);
     // Store future result before enqueueing - task might get already processed and deleted before returning from this method
diff --git a/src/server/database/Database/DatabaseWorkerPool.h b/src/server/database/Database/DatabaseWorkerPool.h
index 13d52f99539..f46c0e29c04 100644
--- a/src/server/database/Database/DatabaseWorkerPool.h
+++ b/src/server/database/Database/DatabaseWorkerPool.h
@@ -160,7 +160,7 @@ class DatabaseWorkerPool
         //! return object as soon as the query is executed.
         //! The return value is then processed in ProcessQueryCallback methods.
         //! Any prepared statements added to this holder need to be prepared with the CONNECTION_ASYNC flag.
-        QueryResultHolderFuture DelayQueryHolder(SQLQueryHolder<T>* holder);
+        SQLQueryHolderCallback DelayQueryHolder(SQLQueryHolder<T>* holder);
 
         /**
             Transaction context methods.
diff --git a/src/server/database/Database/QueryHolder.cpp b/src/server/database/Database/QueryHolder.cpp
index dacd12913da..eb9dbdcca22 100644
--- a/src/server/database/Database/QueryHolder.cpp
+++ b/src/server/database/Database/QueryHolder.cpp
@@ -92,3 +92,14 @@ bool SQLQueryHolderTask::Execute()
     m_result.set_value(m_holder);
     return true;
 }
+
+bool SQLQueryHolderCallback::InvokeIfReady()
+{
+    if (m_future.valid() && m_future.wait_for(std::chrono::seconds(0)) == std::future_status::ready)
+    {
+        m_callback(m_future.get());
+        return true;
+    }
+
+    return false;
+}
diff --git a/src/server/database/Database/QueryHolder.h b/src/server/database/Database/QueryHolder.h
index fe4a0e78331..f94eff49d2c 100644
--- a/src/server/database/Database/QueryHolder.h
+++ b/src/server/database/Database/QueryHolder.h
@@ -63,4 +63,23 @@ class TC_DATABASE_API SQLQueryHolderTask : public SQLOperation
         QueryResultHolderFuture GetFuture() { return m_result.get_future(); }
 };
 
+class TC_DATABASE_API SQLQueryHolderCallback
+{
+public:
+    SQLQueryHolderCallback(QueryResultHolderFuture&& future) : m_future(std::move(future)) { }
+    SQLQueryHolderCallback(SQLQueryHolderCallback&&) = default;
+
+    SQLQueryHolderCallback& operator=(SQLQueryHolderCallback&&) = default;
+
+    void AfterComplete(std::function<void(SQLQueryHolderBase*)> callback) &
+    {
+        m_callback = std::move(callback);
+    }
+
+    bool InvokeIfReady();
+
+    QueryResultHolderFuture m_future;
+    std::function<void(SQLQueryHolderBase*)> m_callback;
+};
+
 #endif
diff --git a/src/server/game/Handlers/CharacterHandler.cpp b/src/server/game/Handlers/CharacterHandler.cpp
index fbb7c4213cb..e5d8ae3e7ca 100644
--- a/src/server/game/Handlers/CharacterHandler.cpp
+++ b/src/server/game/Handlers/CharacterHandler.cpp
@@ -728,7 +728,10 @@ void WorldSession::HandlePlayerLoginOpcode(WorldPacket& recvData)
         return;
     }
 
-    _charLoginCallback = CharacterDatabase.DelayQueryHolder(holder);
+    AddQueryHolderCallback(CharacterDatabase.DelayQueryHolder(holder)).AfterComplete([this](SQLQueryHolderBase* holder)
+    {
+        HandlePlayerLogin(static_cast<LoginQueryHolder*>(holder));
+    });
 }
 
 void WorldSession::HandlePlayerLogin(LoginQueryHolder* holder)
diff --git a/src/server/game/Server/WorldSession.cpp b/src/server/game/Server/WorldSession.cpp
index 4ae28b0bf12..2f91160e308 100644
--- a/src/server/game/Server/WorldSession.cpp
+++ b/src/server/game/Server/WorldSession.cpp
@@ -1198,13 +1198,7 @@ void WorldSession::ProcessQueryCallbacks()
 {
     _queryProcessor.ProcessReadyCallbacks();
     _transactionCallbacks.ProcessReadyCallbacks();
-
-    if (_realmAccountLoginCallback.valid() && _realmAccountLoginCallback.wait_for(std::chrono::seconds(0)) == std::future_status::ready)
-        InitializeSessionCallback(static_cast<CharacterDatabaseQueryHolder*>(_realmAccountLoginCallback.get()));
-
-    //! HandlePlayerLoginOpcode
-    if (_charLoginCallback.valid() && _charLoginCallback.wait_for(std::chrono::seconds(0)) == std::future_status::ready)
-        HandlePlayerLogin(reinterpret_cast<LoginQueryHolder*>(_charLoginCallback.get()));
+    _queryHolderProcessor.ProcessReadyCallbacks();
 }
 
 TransactionCallback& WorldSession::AddTransactionCallback(TransactionCallback&& callback)
@@ -1212,6 +1206,11 @@ TransactionCallback& WorldSession::AddTransactionCallback(TransactionCallback&&
     return _transactionCallbacks.AddCallback(std::move(callback));
 }
 
+SQLQueryHolderCallback& WorldSession::AddQueryHolderCallback(SQLQueryHolderCallback&& callback)
+{
+    return _queryHolderProcessor.AddCallback(std::move(callback));
+}
+
 void WorldSession::InitWarden(SessionKey const& k, std::string const& os)
 {
     if (os == "Win")
@@ -1287,7 +1286,10 @@ void WorldSession::InitializeSession()
         return;
     }
 
-    _realmAccountLoginCallback = CharacterDatabase.DelayQueryHolder(realmHolder);
+    AddQueryHolderCallback(CharacterDatabase.DelayQueryHolder(realmHolder)).AfterComplete([this](SQLQueryHolderBase* holder)
+    {
+        InitializeSessionCallback(static_cast<AccountInfoQueryHolderPerRealm*>(holder));
+    });
 }
 
 void WorldSession::InitializeSessionCallback(CharacterDatabaseQueryHolder* realmHolder)
diff --git a/src/server/game/Server/WorldSession.h b/src/server/game/Server/WorldSession.h
index e2e0369eb5a..f929c013030 100644
--- a/src/server/game/Server/WorldSession.h
+++ b/src/server/game/Server/WorldSession.h
@@ -1095,15 +1095,14 @@ class TC_GAME_API WorldSession
     public:
         QueryCallbackProcessor& GetQueryProcessor() { return _queryProcessor; }
         TransactionCallback& AddTransactionCallback(TransactionCallback&& callback);
+        SQLQueryHolderCallback& AddQueryHolderCallback(SQLQueryHolderCallback&& callback);
 
     private:
         void ProcessQueryCallbacks();
 
-        QueryResultHolderFuture _realmAccountLoginCallback;
-        QueryResultHolderFuture _charLoginCallback;
-
         QueryCallbackProcessor _queryProcessor;
         AsyncCallbackProcessor<TransactionCallback> _transactionCallbacks;
+        AsyncCallbackProcessor<SQLQueryHolderCallback> _queryHolderProcessor;
 
     friend class World;
     protected: