add onSessionStart() and onSessionDestroy() events (#1412)

Co-authored-by: an-tao <antao2002@gmail.com>

lib/inc/drogon/CacheMap.h

@@ -86,12 +86,16 @@ class CacheMap
     CacheMap(trantor::EventLoop *loop,
              float tickInterval = TICK_INTERVAL,
              size_t wheelsNum = WHEELS_NUM,
-             size_t bucketsNumPerWheel = BUCKET_NUM_PER_WHEEL)
+             size_t bucketsNumPerWheel = BUCKET_NUM_PER_WHEEL,
+             std::function<void(const T1 &)> fnOnInsert = nullptr,
+             std::function<void(const T1 &)> fnOnErase = nullptr)
         : loop_(loop),
           tickInterval_(tickInterval),
           wheelsNumber_(wheelsNum),
           bucketsNumPerWheel_(bucketsNumPerWheel),
-          ctrlBlockPtr_(std::make_shared<ControlBlock>())
+          ctrlBlockPtr_(std::make_shared<ControlBlock>()),
+          fnOnInsert_(fnOnInsert),
+          fnOnErase_(fnOnErase)
     {
         wheels_.resize(wheelsNumber_);
         for (size_t i = 0; i < wheelsNumber_; ++i)
@@ -215,6 +219,8 @@ class CacheMap
             std::lock_guard<std::mutex> lock(mtx_);
             map_.insert(std::make_pair(key, std::move(v)));
         }
+        if (fnOnInsert_)
+            fnOnInsert_(key);
     }
     /**
      * @brief Insert a key-value pair into the cache.
@@ -244,6 +250,8 @@ class CacheMap
             std::lock_guard<std::mutex> lock(mtx_);
             map_.insert(std::make_pair(key, std::move(v)));
         }
+        if (fnOnInsert_)
+            fnOnInsert_(key);
     }
 
     /**
@@ -284,28 +292,35 @@ class CacheMap
      * the key doesn't exist, a new one is created and passed to the handler and
      * stored in the cache with the timeout parameter. The changing of the data
      * is protected by the mutex of the cache.
+     *
      */
     template <typename Callable>
     void modify(const T1 &key, Callable &&handler, size_t timeout = 0)
     {
-        std::lock_guard<std::mutex> lock(mtx_);
-        auto iter = map_.find(key);
-        if (iter != map_.end())
         {
-            timeout = iter->second.timeout_;
-            handler(iter->second.value_);
+            std::lock_guard<std::mutex> lock(mtx_);
+            auto iter = map_.find(key);
+            if (iter != map_.end())
+            {
+                timeout = iter->second.timeout_;
+                handler(iter->second.value_);
+                if (timeout > 0)
+                    eraseAfter(timeout, key);
+                return;
+            }
+
+            MapValue v{T2(), timeout};
+            handler(v.value_);
+            map_.insert(std::make_pair(key, std::move(v)));
             if (timeout > 0)
+            {
                 eraseAfter(timeout, key);
-            return;
-        }
-        MapValue v{T2(), timeout};
-        handler(v.value_);
-        map_.insert(std::make_pair(key, std::move(v)));
-        if (timeout > 0)
-        {
-            eraseAfter(timeout, key);
+            }
         }
+        if (fnOnInsert_)
+            fnOnInsert_(key);
     }
+
     /// Check if the value of the keyword exists
     bool find(const T1 &key)
     {
@@ -358,8 +373,12 @@ class CacheMap
     void erase(const T1 &key)
     {
         // in this case,we don't evoke the timeout callback;
-        std::lock_guard<std::mutex> lock(mtx_);
-        map_.erase(key);
+        {
+            std::lock_guard<std::mutex> lock(mtx_);
+            map_.erase(key);
+        }
+        if (fnOnErase_)
+            fnOnErase_(key);
     }
     /**
      * @brief Get the event loop object
@@ -425,6 +444,8 @@ class CacheMap
     size_t wheelsNumber_;
     size_t bucketsNumPerWheel_;
     std::shared_ptr<ControlBlock> ctrlBlockPtr_;
+    std::function<void(const T1 &)> fnOnInsert_;
+    std::function<void(const T1 &)> fnOnErase_;
 
     bool noWheels_{false};
 
@@ -486,21 +507,28 @@ class CacheMap
         else
         {
             std::function<void()> cb = [this, key]() {
-                std::lock_guard<std::mutex> lock(mtx_);
-                if (map_.find(key) != map_.end())
+                bool erased{false};
+                std::function<void()> timeoutCallback;
                 {
-                    auto &value = map_[key];
-                    auto entryPtr = value.weakEntryPtr_.lock();
-                    // entryPtr is used to avoid race conditions
-                    if (value.timeout_ > 0 && !entryPtr)
+                    std::lock_guard<std::mutex> lock(mtx_);
+                    auto iter = map_.find(key);
+                    if (iter != map_.end())
                     {
-                        if (value.timeoutCallback_)
+                        auto &value = iter->second;
+                        auto entryPtr = value.weakEntryPtr_.lock();
+                        // entryPtr is used to avoid race conditions
+                        if (value.timeout_ > 0 && !entryPtr)
                         {
-                            value.timeoutCallback_();
+                            erased = true;
+                            timeoutCallback = std::move(value.timeoutCallback_);
+                            map_.erase(key);
                         }
-                        map_.erase(key);
                     }
                 }
+                if (erased && fnOnErase_)
+                    fnOnErase_(key);
+                if (erased && timeoutCallback)
+                    timeoutCallback();
             };
             entryPtr = std::make_shared<CallbackEntry>(std::move(cb));
             map_[key].weakEntryPtr_ = entryPtr;

lib/inc/drogon/HttpAppFramework.h

@@ -745,6 +745,20 @@ class DROGON_EXPORT HttpAppFramework : public trantor::NonCopyable
         return enableSession((size_t)timeout.count(), sameSite);
     }
 
+    /// Register an advice called when starting a new session.
+    /**
+     * @param advice is called with the session id.
+     */
+    virtual HttpAppFramework &registerSessionStartAdvice(
+        const AdviceStartSessionCallback &advice) = 0;
+
+    /// Register an advice called when destroying a session.
+    /**
+     * @param advice is called with the session id.
+     */
+    virtual HttpAppFramework &registerSessionDestroyAdvice(
+        const AdviceDestroySessionCallback &advice) = 0;
+
     /// Disable sessions supporting.
     /**
      * @note

lib/inc/drogon/drogon_callbacks.h

@@ -26,6 +26,8 @@ class HttpRequest;
 using HttpRequestPtr = std::shared_ptr<HttpRequest>;
 using AdviceCallback = std::function<void(const HttpResponsePtr &)>;
 using AdviceChainCallback = std::function<void()>;
+using AdviceStartSessionCallback = std::function<void(const std::string &)>;
+using AdviceDestroySessionCallback = std::function<void(const std::string &)>;
 using FilterCallback = std::function<void(const HttpResponsePtr &)>;
 using FilterChainCallback = std::function<void()>;
 using HttpReqCallback = std::function<void(ReqResult, const HttpResponsePtr &)>;

lib/src/HttpAppFrameworkImpl.cc

@@ -571,7 +571,10 @@ void HttpAppFrameworkImpl::run()
     if (useSession_)
     {
         sessionManagerPtr_ =
-            std::make_unique<SessionManager>(getLoop(), sessionTimeout_);
+            std::make_unique<SessionManager>(getLoop(),
+                                             sessionTimeout_,
+                                             sessionStartAdvices_,
+                                             sessionDestroyAdvices_);
     }
     // now start running!!
     running_ = true;

lib/src/HttpAppFrameworkImpl.h

@@ -221,6 +221,21 @@ class HttpAppFrameworkImpl final : public HttpAppFramework
         useSession_ = false;
         return *this;
     }
+
+    HttpAppFramework &registerSessionStartAdvice(
+        const AdviceStartSessionCallback &advice) override
+    {
+        sessionStartAdvices_.emplace_back(advice);
+        return *this;
+    }
+
+    HttpAppFramework &registerSessionDestroyAdvice(
+        const AdviceDestroySessionCallback &advice) override
+    {
+        sessionDestroyAdvices_.emplace_back(advice);
+        return *this;
+    }
+
     const std::string &getDocumentRoot() const override
     {
         return rootPath_;
@@ -657,6 +672,8 @@ class HttpAppFrameworkImpl final : public HttpAppFramework
     std::function<void()> termSignalHandler_{[]() { app().quit(); }};
     std::function<void()> intSignalHandler_{[]() { app().quit(); }};
     std::unique_ptr<SessionManager> sessionManagerPtr_;
+    std::vector<AdviceStartSessionCallback> sessionStartAdvices_;
+    std::vector<AdviceDestroySessionCallback> sessionDestroyAdvices_;
     std::shared_ptr<trantor::AsyncFileLogger> asyncFileLoggerPtr_;
     Json::Value jsonConfig_;
     HttpResponsePtr custom404_;

lib/src/SessionManager.cc

@@ -17,8 +17,15 @@
 
 using namespace drogon;
 
-SessionManager::SessionManager(trantor::EventLoop *loop, size_t timeout)
-    : loop_(loop), timeout_(timeout)
+SessionManager::SessionManager(
+    trantor::EventLoop* loop,
+    size_t timeout,
+    const std::vector<AdviceStartSessionCallback>& startAdvices,
+    const std::vector<AdviceDestroySessionCallback>& destroyAdvices)
+    : loop_(loop),
+      timeout_(timeout),
+      sessionStartAdvices_(startAdvices),
+      sessionDestroyAdvices_(destroyAdvices)
 {
     if (timeout_ > 0)
     {
@@ -38,25 +45,57 @@ SessionManager::SessionManager(trantor::EventLoop *loop, size_t timeout)
                 tmpTimeout = tmpTimeout / 100;
             }
         }
+
         sessionMapPtr_ = std::unique_ptr<CacheMap<std::string, SessionPtr>>(
             new CacheMap<std::string, SessionPtr>(
-                loop_, 1.0, wheelNum, bucketNum));
+                loop_,
+                1.0,
+                wheelNum,
+                bucketNum,
+                [this](const std::string& key) {
+                    for (auto& advice : sessionStartAdvices_)
+                    {
+                        advice(key);
+                    }
+                },
+                [this](const std::string& key) {
+                    for (auto& advice : sessionDestroyAdvices_)
+                    {
+                        advice(key);
+                    }
+                }));
     }
     else if (timeout_ == 0)
     {
         sessionMapPtr_ = std::unique_ptr<CacheMap<std::string, SessionPtr>>(
-            new CacheMap<std::string, SessionPtr>(loop_, 0, 0, 0));
+            new CacheMap<std::string, SessionPtr>(
+                loop_,
+                0,
+                0,
+                0,
+                [this](const std::string& key) {
+                    for (auto& advice : sessionStartAdvices_)
+                    {
+                        advice(key);
+                    }
+                },
+                [this](const std::string& key) {
+                    for (auto& advice : sessionDestroyAdvices_)
+                    {
+                        advice(key);
+                    }
+                }));
     }
 }
 
-SessionPtr SessionManager::getSession(const std::string &sessionID,
+SessionPtr SessionManager::getSession(const std::string& sessionID,
                                       bool needToSet)
 {
     assert(!sessionID.empty());
     SessionPtr sessionPtr;
     sessionMapPtr_->modify(
         sessionID,
-        [&sessionPtr, &sessionID, needToSet](SessionPtr &sessionInCache) {
+        [&sessionPtr, &sessionID, needToSet](SessionPtr& sessionInCache) {
             if (sessionInCache)
             {
                 sessionPtr = sessionInCache;
@@ -69,10 +108,11 @@ SessionPtr SessionManager::getSession(const std::string &sessionID,
             }
         },
         timeout_);
+
     return sessionPtr;
 }
 
-void SessionManager::changeSessionId(const SessionPtr &sessionPtr)
+void SessionManager::changeSessionId(const SessionPtr& sessionPtr)
 {
     auto oldId = sessionPtr->sessionId();
     auto newId = utils::getUuid();

lib/src/SessionManager.h

@@ -15,19 +15,25 @@
 #pragma once
 
 #include <drogon/Session.h>
+#include <drogon/drogon_callbacks.h>
 #include <drogon/CacheMap.h>
 #include <trantor/utils/NonCopyable.h>
 #include <trantor/net/EventLoop.h>
 #include <memory>
 #include <string>
 #include <mutex>
+#include <vector>
 
 namespace drogon
 {
 class SessionManager : public trantor::NonCopyable
 {
   public:
-    SessionManager(trantor::EventLoop *loop, size_t timeout);
+    SessionManager(
+        trantor::EventLoop *loop,
+        size_t timeout,
+        const std::vector<AdviceStartSessionCallback> &startAdvices,
+        const std::vector<AdviceDestroySessionCallback> &destroyAdvices);
     ~SessionManager()
     {
         sessionMapPtr_.reset();
@@ -39,5 +45,7 @@ class SessionManager : public trantor::NonCopyable
     std::unique_ptr<CacheMap<std::string, SessionPtr>> sessionMapPtr_;
     trantor::EventLoop *loop_;
     size_t timeout_;
+    const std::vector<AdviceStartSessionCallback> &sessionStartAdvices_;
+    const std::vector<AdviceDestroySessionCallback> &sessionDestroyAdvices_;
 };
 }  // namespace drogon

lib/tests/integration_test/server/main.cc

@@ -319,6 +319,12 @@ int main()
         }
         return nullResp;
     });
+    app().registerSessionStartAdvice([](const std::string &sessionId) {
+        LOG_DEBUG << "session start:" << sessionId;
+    });
+    app().registerSessionDestroyAdvice([](const std::string &sessionId) {
+        LOG_DEBUG << "session destroy:" << sessionId;
+    });
     // Output information of all handlers
     auto handlerInfo = app().getHandlersInfo();
     for (auto &info : handlerInfo)