Core/Database: Close the databases correctly when the DBUpdater fails

* Also fixes a memory leak spotted by Aokromes

3 files changed, 39 insertions, 45 deletions

diff --git a/src/server/database/Database/DatabaseLoader.cpp b/src/server/database/Database/DatabaseLoader.cpp
index 92d8730cd12..6a8e86ffca9 100644
--- a/src/server/database/Database/DatabaseLoader.cpp
+++ b/src/server/database/Database/DatabaseLoader.cpp
@@ -32,7 +32,7 @@ DatabaseLoader& DatabaseLoader::AddDatabase(DatabaseWorkerPool<T>& pool, std::st
 {
     bool const updatesEnabledForThis = DBUpdater<T>::IsEnabled(_updateFlags);
 
-    _open.push(std::make_pair([this, name, updatesEnabledForThis, &pool]() -> bool
+    _open.push([this, name, updatesEnabledForThis, &pool]() -> bool
     {
         std::string const dbString = sConfigMgr->GetStringDefault(name + "DatabaseInfo", "");
         if (dbString.empty())
@@ -71,12 +71,13 @@ DatabaseLoader& DatabaseLoader::AddDatabase(DatabaseWorkerPool<T>& pool, std::st
                 return false;
             }
         }
+        // Add the close operation
+        _close.push([&pool]
+        {
+            pool.Close();
+        });
         return true;
-    },
-    [&pool]()
-    {
-        pool.Close();
-    }));
+    });
 
     // Populate and update only if updates are enabled for this pool
     if (updatesEnabledForThis)
@@ -137,38 +138,7 @@ bool DatabaseLoader::Load()
 
 bool DatabaseLoader::OpenDatabases()
 {
-    while (!_open.empty())
-    {
-        std::pair<Predicate, std::function<void()>> const load = _open.top();
-        if (load.first())
-            _close.push(load.second);
-        else
-        {
-            // Close all loaded databases
-            while (!_close.empty())
-            {
-                _close.top()();
-                _close.pop();
-            }
-            return false;
-        }
-
-        _open.pop();
-    }
-    return true;
-}
-
-// Processes the elements of the given stack until a predicate returned false.
-bool DatabaseLoader::Process(std::stack<Predicate>& stack)
-{
-    while (!stack.empty())
-    {
-        if (!stack.top()())
-            return false;
-
-        stack.pop();
-    }
-    return true;
+    return Process(_open);
 }
 
 bool DatabaseLoader::PopulateDatabases()
@@ -186,6 +156,27 @@ bool DatabaseLoader::PrepareStatements()
     return Process(_prepare);
 }
 
+bool DatabaseLoader::Process(std::queue<Predicate>& queue)
+{
+    while (!queue.empty())
+    {
+        if (!queue.front()())
+        {
+            // Close all open databases which have a registered close operation
+            while (!_close.empty())
+            {
+                _close.top()();
+                _close.pop();
+            }
+
+            return false;
+        }
+
+        queue.pop();
+    }
+    return true;
+}
+
 template
 DatabaseLoader& DatabaseLoader::AddDatabase<LoginDatabaseConnection>(DatabaseWorkerPool<LoginDatabaseConnection>& pool, std::string const& name);
 template
diff --git a/src/server/database/Database/DatabaseLoader.h b/src/server/database/Database/DatabaseLoader.h
index da92cf85a9f..ec390a427ad 100644
--- a/src/server/database/Database/DatabaseLoader.h
+++ b/src/server/database/Database/DatabaseLoader.h
@@ -21,8 +21,9 @@
 #include "DatabaseWorkerPool.h"
 #include "DatabaseEnv.h"
 
-#include <stack>
 #include <functional>
+#include <stack>
+#include <queue>
 
 // A helper class to initiate all database worker pools,
 // handles updating, delays preparing of statements and cleans up on failure.
@@ -56,16 +57,18 @@ private:
     bool PrepareStatements();
 
     using Predicate = std::function<bool()>;
+    using Closer = std::function<void()>;
 
-    static bool Process(std::stack<Predicate>& stack);
+    // Invokes all functions in the given queue and closes the databases on errors.
+    // Returns false when there was an error.
+    bool Process(std::queue<Predicate>& queue);
 
     std::string const _logger;
     bool const _autoSetup;
     uint32 const _updateFlags;
 
-    std::stack<std::pair<Predicate, std::function<void()>>> _open;
-    std::stack<std::function<void()>> _close;
-    std::stack<Predicate> _populate, _update, _prepare;
+    std::queue<Predicate> _open, _populate, _update, _prepare;
+    std::stack<Closer> _close;
 };
 
 #endif // DatabaseLoader_h__
diff --git a/src/server/worldserver/Main.cpp b/src/server/worldserver/Main.cpp
index 53c5f250851..58ddce532c0 100644
--- a/src/server/worldserver/Main.cpp
+++ b/src/server/worldserver/Main.cpp
@@ -449,9 +449,9 @@ bool StartDB()
     // Load databases
     DatabaseLoader loader("server.worldserver", DatabaseLoader::DATABASE_NONE);
     loader
-        .AddDatabase(WorldDatabase, "World")
+        .AddDatabase(LoginDatabase, "Login")
         .AddDatabase(CharacterDatabase, "Character")
-        .AddDatabase(LoginDatabase, "Login");
+        .AddDatabase(WorldDatabase, "World");
 
     if (!loader.Load())
         return false;