SessionManager: Allow backend storage to be configurable

lib/inc/drogon/HttpAppFramework.h

@@ -142,6 +142,28 @@ class HttpAppFramework : public trantor::NonCopyable
     virtual const std::function<HttpResponsePtr(HttpStatusCode)>
         &getCustomErrorHandler() const = 0;
 
+    /// Set custom session storage creator
+    /**
+     * @param resp_generator is invoked when the system starts running and is
+     * used to provide backing data storage for the built-in session handler.
+     * The two parameters are the event loop of the system and the timeout for
+     * keys.
+     */
+    virtual HttpAppFramework &setCustomSessionStorageCreator(
+        std::function<
+            std::unique_ptr<SessionStorageProvider>(trantor::EventLoop *,
+                                                    size_t)> &&creator) = 0;
+
+    /// Get custom session storage creator
+    /**
+     * @return A const-reference to the session storage creator set using
+     * setCustomSessionStorageCreator. If none was provided, the default creator
+     * is returned.
+     */
+    virtual const std::function<
+        std::unique_ptr<SessionStorageProvider>(trantor::EventLoop *, size_t)>
+        &getCustomSessionStorageCreator() const = 0;
+
     /// Get the plugin object registered in the framework
     /**
      * @note

lib/inc/drogon/Session.h

@@ -166,4 +166,30 @@ class Session
 
 using SessionPtr = std::shared_ptr<Session>;
 
+class SessionStorageProvider
+{
+  public:
+    virtual ~SessionStorageProvider();
+
+    /// Atomically find and get the value of a keyword
+    /**
+     * Return true when the value is found, and the value
+     * is assigned to the value argument.
+     */
+    virtual bool findAndFetch(const std::string &key, SessionPtr &value) = 0;
+
+    /**
+     * @brief Insert a key-value pair into the storage provider.
+     *
+     * @param key The key
+     * @param value The value
+     * @param timeout The timeout in seconds, if timeout > 0, the value will be
+     * erased within the 'timeout' seconds after the last access. If the timeout
+     * is zero, the value exists until being removed explicitly.
+     */
+    virtual void insert(const std::string &key,
+                        const SessionPtr &value,
+                        size_t timeout) = 0;
+};
+
 }  // namespace drogon

lib/src/HttpAppFrameworkImpl.cc

@@ -498,8 +498,9 @@ void HttpAppFrameworkImpl::run()
 
     if (useSession_)
     {
-        sessionManagerPtr_ = std::unique_ptr<SessionManager>(
-            new SessionManager(getLoop(), sessionTimeout_));
+        sessionManagerPtr_ = std::make_unique<SessionManager>(
+            sessionStorageCreator_(getLoop(), sessionTimeout_),
+            sessionTimeout_);
     }
 
     // Initialize plugins
@@ -982,3 +983,18 @@ const std::function<HttpResponsePtr(HttpStatusCode)>
 {
     return customErrorHandler_;
 }
+
+HttpAppFramework &HttpAppFrameworkImpl::setCustomSessionStorageCreator(
+    std::function<std::unique_ptr<SessionStorageProvider>(trantor::EventLoop *,
+                                                          size_t)> &&creator)
+{
+    sessionStorageCreator_ = std::move(creator);
+    return *this;
+}
+
+const std::function<
+    std::unique_ptr<SessionStorageProvider>(trantor::EventLoop *, size_t)>
+    &HttpAppFrameworkImpl::getCustomSessionStorageCreator() const
+{
+    return sessionStorageCreator_;
+}

lib/src/HttpAppFrameworkImpl.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include "impl_forwards.h"
+#include "SessionManager.h"
 #include <drogon/HttpAppFramework.h>
 #include <drogon/config.h>
 #include <memory>
@@ -197,6 +198,15 @@ class HttpAppFrameworkImpl : public HttpAppFramework
         useSession_ = false;
         return *this;
     }
+
+    HttpAppFramework &setCustomSessionStorageCreator(
+        std::function<std::unique_ptr<SessionStorageProvider>(
+            trantor::EventLoop *,
+            size_t)> &&creator) override;
+    const std::function<
+        std::unique_ptr<SessionStorageProvider>(trantor::EventLoop *, size_t)>
+        &getCustomSessionStorageCreator() const override;
+
     virtual const std::string &getDocumentRoot() const override
     {
         return rootPath_;
@@ -547,6 +557,7 @@ class HttpAppFrameworkImpl : public HttpAppFramework
     size_t clientMaxWebSocketMessageSize_{128 * 1024};
     std::string homePageFile_{"index.html"};
     std::function<void()> termSignalHandler_{[]() { app().quit(); }};
+    std::function<std::unique_ptr<SessionStorageProvider>(trantor::EventLoop*, size_t)> sessionStorageCreator_{createCacheMapProvider};
     std::unique_ptr<SessionManager> sessionManagerPtr_;
     Json::Value jsonConfig_;
     HttpResponsePtr custom404_;

lib/src/SessionManager.cc

@@ -16,50 +16,92 @@
 
 using namespace drogon;
 
-SessionManager::SessionManager(trantor::EventLoop *loop, size_t timeout)
-    : loop_(loop), timeout_(timeout)
+class CacheMapSessionStorageProvider : public SessionStorageProvider
 {
-    assert(timeout_ >= 0);
-    if (timeout_ > 0)
+  public:
+    CacheMapSessionStorageProvider(trantor::EventLoop *loop,
+                                   float tickInterval,
+                                   size_t wheelsNum,
+                                   size_t bucketsNumPerWheel)
+        : data_(std::make_unique<CacheMap<std::string, SessionPtr>>(
+              loop,
+              tickInterval,
+              wheelsNum,
+              bucketsNumPerWheel))
+    {
+    }
+
+    bool findAndFetch(const std::string &key, SessionPtr &value) override
+    {
+        return data_->findAndFetch(key, value);
+    }
+
+    void insert(const std::string &key,
+                const SessionPtr &value,
+                size_t timeout) override
+    {
+        data_->insert(key, value, timeout);
+    }
+
+  private:
+    std::unique_ptr<CacheMap<std::string, SessionPtr>> data_;
+};
+
+SessionManager::SessionManager(std::unique_ptr<SessionStorageProvider> provider,
+                               size_t timeout)
+    : provider_(std::move(provider)), timeout_(timeout)
+{
+}
+
+SessionPtr SessionManager::getSession(const std::string &sessionID,
+                                      bool needToSet)
+{
+    assert(!sessionID.empty());
+    SessionPtr sessionPtr;
+    std::lock_guard<std::mutex> lock(mapMutex_);
+    if (!provider_->findAndFetch(sessionID, sessionPtr))
+    {
+        sessionPtr = std::make_shared<Session>(sessionID, needToSet);
+        provider_->insert(sessionID, sessionPtr, timeout_);
+        return sessionPtr;
+    }
+    return sessionPtr;
+}
+
+SessionStorageProvider::~SessionStorageProvider() = default;
+
+std::unique_ptr<SessionStorageProvider> drogon::createCacheMapProvider(
+    trantor::EventLoop *loop,
+    size_t timeout)
+{
+    assert(timeout >= 0);
+    if (timeout > 0)
     {
         size_t wheelNum = 1;
         size_t bucketNum = 0;
-        if (timeout_ < 500)
+        if (timeout < 500)
         {
-            bucketNum = timeout_ + 1;
+            bucketNum = timeout + 1;
         }
         else
         {
-            auto tmpTimeout = timeout_;
+            auto tmpTimeout = timeout;
             bucketNum = 100;
             while (tmpTimeout > 100)
             {
                 ++wheelNum;
                 tmpTimeout = tmpTimeout / 100;
             }
         }
-        sessionMapPtr_ = std::unique_ptr<CacheMap<std::string, SessionPtr>>(
-            new CacheMap<std::string, SessionPtr>(
-                loop_, 1.0, wheelNum, bucketNum));
+        return std::make_unique<CacheMapSessionStorageProvider>(loop,
+                                                                1.0,
+                                                                wheelNum,
+                                                                bucketNum);
     }
-    else if (timeout_ == 0)
+    else if (timeout == 0)
     {
-        sessionMapPtr_ = std::unique_ptr<CacheMap<std::string, SessionPtr>>(
-            new CacheMap<std::string, SessionPtr>(loop_, 0, 0, 0));
+        return std::make_unique<CacheMapSessionStorageProvider>(loop, 0, 0, 0);
     }
-}
 
-SessionPtr SessionManager::getSession(const std::string &sessionID,
-                                      bool needToSet)
-{
-    assert(!sessionID.empty());
-    SessionPtr sessionPtr;
-    std::lock_guard<std::mutex> lock(mapMutex_);
-    if (sessionMapPtr_->findAndFetch(sessionID, sessionPtr) == false)
-    {
-        sessionPtr = std::make_shared<Session>(sessionID, needToSet);
-        sessionMapPtr_->insert(sessionID, sessionPtr, timeout_);
-        return sessionPtr;
-    }
-    return sessionPtr;
-}
+    return nullptr;
+}

lib/src/SessionManager.h

@@ -24,20 +24,24 @@
 
 namespace drogon
 {
+std::unique_ptr<SessionStorageProvider> createCacheMapProvider(
+    trantor::EventLoop *loop,
+    size_t timeout);
+
 class SessionManager : public trantor::NonCopyable
 {
   public:
-    SessionManager(trantor::EventLoop *loop, size_t timeout);
+    SessionManager(std::unique_ptr<SessionStorageProvider> provider,
+                   size_t timeout);
     ~SessionManager()
     {
-        sessionMapPtr_.reset();
+        provider_.reset();
     }
     SessionPtr getSession(const std::string &sessionID, bool needToSet);
 
   private:
-    std::unique_ptr<CacheMap<std::string, SessionPtr>> sessionMapPtr_;
+    std::unique_ptr<SessionStorageProvider> provider_;
     std::mutex mapMutex_;
-    trantor::EventLoop *loop_;
     size_t timeout_;
 };
 }  // namespace drogon