From 59a04315122919a658ccb5a10b036934a7ba304e Mon Sep 17 00:00:00 2001
From: antao <antao2002@gmail.com>
Date: Thu, 10 Jan 2019 17:54:51 +0800
Subject: [PATCH] Add a mutex

---
 orm_lib/src/DbClientImpl.cc                 | 21 +++++++++++++++------
 orm_lib/src/mysql_impl/MysqlConnection.cc   |  4 +++-
 orm_lib/src/postgresql_impl/PgConnection.cc |  4 +++-
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/orm_lib/src/DbClientImpl.cc b/orm_lib/src/DbClientImpl.cc
index 0a03e093..8bbcba1c 100644
--- a/orm_lib/src/DbClientImpl.cc
+++ b/orm_lib/src/DbClientImpl.cc
@@ -44,7 +44,7 @@ using namespace drogon::orm;
 DbClientImpl::DbClientImpl(const std::string &connInfo, const size_t connNum, ClientType type)
     : _connInfo(connInfo),
       _connectNum(connNum),
-      _loops((connNum / 100 > 0 ? connNum / 100 : 1),"DbLoop")
+      _loops((connNum / 100 > 0 ? connNum / 100 : 1), "DbLoop")
 {
     _type = type;
     LOG_TRACE << "type=" << (int)type;
@@ -56,10 +56,12 @@ DbClientImpl::DbClientImpl(const std::string &connInfo, const size_t connNum, Cl
         {
             auto loop = _loops.getNextLoop();
             loop->runInLoop([this, loop]() {
+                std::lock_guard<std::mutex> lock(_connectionsMutex);
                 _connections.insert(newConnection(loop));
             });
         }
-    }).detach();
+    })
+        .detach();
 }
 
 DbClientImpl::~DbClientImpl() noexcept
@@ -288,16 +290,23 @@ DbConnectionPtr DbClientImpl::newConnection(trantor::EventLoop *loop)
 
     std::weak_ptr<DbClientImpl> weakPtr = shared_from_this();
     connPtr->setCloseCallback([weakPtr, loop](const DbConnectionPtr &closeConnPtr) {
+        //Erase the connection
+        auto thisPtr = weakPtr.lock();
+        if (!thisPtr)
+            return;
+        {
+            std::lock_guard<std::mutex> guard(thisPtr->_connectionsMutex);
+            thisPtr->_readyConnections.erase(closeConnPtr);
+            thisPtr->_busyConnections.erase(closeConnPtr);
+            assert(thisPtr->_connections.find(closeConnPtr) != thisPtr->_connections.end());
+            thisPtr->_connections.erase(closeConnPtr);
+        }
         //Reconnect after 1 second
         loop->runAfter(1, [weakPtr, closeConnPtr, loop] {
             auto thisPtr = weakPtr.lock();
             if (!thisPtr)
                 return;
             std::lock_guard<std::mutex> guard(thisPtr->_connectionsMutex);
-            thisPtr->_readyConnections.erase(closeConnPtr);
-            thisPtr->_busyConnections.erase(closeConnPtr);
-            assert(thisPtr->_connections.find(closeConnPtr) != thisPtr->_connections.end());
-            thisPtr->_connections.erase(closeConnPtr);
             thisPtr->_connections.insert(thisPtr->newConnection(loop));
         });
     });
diff --git a/orm_lib/src/mysql_impl/MysqlConnection.cc b/orm_lib/src/mysql_impl/MysqlConnection.cc
index fd62a8c4..05e6976e 100644
--- a/orm_lib/src/mysql_impl/MysqlConnection.cc
+++ b/orm_lib/src/mysql_impl/MysqlConnection.cc
@@ -133,8 +133,10 @@ void MysqlConnection::setChannel()
 
 void MysqlConnection::handleClosed()
 {
-    _status = ConnectStatus_Bad;
     _loop->assertInLoopThread();
+    if (_status == ConnectStatus_Bad)
+        return;
+    _status = ConnectStatus_Bad;
     _channelPtr->disableAll();
     _channelPtr->remove();
     assert(_closeCb);
diff --git a/orm_lib/src/postgresql_impl/PgConnection.cc b/orm_lib/src/postgresql_impl/PgConnection.cc
index d8b833af..01206fb9 100644
--- a/orm_lib/src/postgresql_impl/PgConnection.cc
+++ b/orm_lib/src/postgresql_impl/PgConnection.cc
@@ -78,8 +78,10 @@ PgConnection::PgConnection(trantor::EventLoop *loop, const std::string &connInfo
 // }
 void PgConnection::handleClosed()
 {
-    _status = ConnectStatus_Bad;
     _loop->assertInLoopThread();
+    if (_status == ConnectStatus_Bad)
+        return;
+    _status = ConnectStatus_Bad;
     _channel.disableAll();
     _channel.remove();
     assert(_closeCb);