refactor parts of ShardsPool code

Author: sewenew

src/sw/redis++/connection_pool.cpp

@@ -149,14 +149,6 @@ bool ConnectionPool::_need_reconnect(const Connection &connection) {
     return false;
 }
 
-ConnectionPoolGuard::ConnectionPoolGuard(ConnectionPool &pool, Connection &connection) :
-                                            _pool(pool),
-                                            _connection(connection) {}
-
-ConnectionPoolGuard::~ConnectionPoolGuard() {
-    _pool.release(std::move(_connection));
-}
-
 }
 
 }

src/sw/redis++/connection_pool.h

@@ -82,16 +82,6 @@ class ConnectionPool {
     std::condition_variable _cv;
 };
 
-class ConnectionPoolGuard {
-public:
-    ConnectionPoolGuard(ConnectionPool &pool, Connection &connection);
-    ~ConnectionPoolGuard();
-
-private:
-    ConnectionPool &_pool;
-    Connection &_connection;
-};
-
 }
 
 }

src/sw/redis++/redis.h

@@ -975,6 +975,20 @@ class Redis {
     long long publish(const StringView &channel, const StringView &message);
 
 private:
+    class ConnectionPoolGuard {
+    public:
+        ConnectionPoolGuard(ConnectionPool &pool,
+                            Connection &connection) : _pool(pool), _connection(connection) {}
+
+        ~ConnectionPoolGuard() {
+            _pool.release(std::move(_connection));
+        }
+
+    private:
+        ConnectionPool &_pool;
+        Connection &_connection;
+    };
+
     template <typename Cmd, typename ...Args>
     ReplyUPtr _score_command(std::true_type, Cmd cmd, Args &&... args);
 

src/sw/redis++/redis_cluster.hpp

@@ -842,15 +842,9 @@ template <typename Cmd, typename ...Args>
 ReplyUPtr RedisCluster::_command(Cmd cmd, const StringView &key, Args &&...args) {
     for (auto idx = 0; idx < 2; ++idx) {
         try {
-            auto pool = _pool.fetch(key);
+            auto guarded_connection = _pool.fetch(key);
 
-            assert(bool(pool));
-
-            auto connection = pool->fetch();
-
-            ConnectionPoolGuard guard(*pool, connection);
-
-            return _command(cmd, connection, std::forward<Args>(args)...);
+            return _command(cmd, guarded_connection.connection(), std::forward<Args>(args)...);
         } catch (const ClosedError &err) {
             // Node might be removed.
             // 1. Get up-to-date slot mapping to check if the node still exists.
@@ -863,15 +857,8 @@ ReplyUPtr RedisCluster::_command(Cmd cmd, const StringView &key, Args &&...args)
             // Slot mapping has been changed, update it and try again.
             _pool.update();
         } catch (const AskError &err) {
-            auto pool = _pool.fetch(err.node());
-
-            assert(bool(pool));
-
-            auto connection = pool->fetch();
-
-            assert(!connection.broken());
-
-            ConnectionPoolGuard guard(*pool, connection);
+            auto guarded_connection = _pool.fetch(err.node());
+            auto &connection = guarded_connection.connection();
 
             // 1. send ASKING command.
             _asking(connection);

src/sw/redis++/shards_pool.cpp

@@ -57,50 +57,41 @@ ShardsPool& ShardsPool::operator=(ShardsPool &&that) {
     return *this;
 }
 
-ConnectionPoolSPtr ShardsPool::fetch(const StringView &key) {
+GuardedConnection ShardsPool::fetch(const StringView &key) {
     auto slot = _slot(key);
 
     return _fetch(slot);
 }
 
-ConnectionPoolSPtr ShardsPool::fetch() {
+GuardedConnection ShardsPool::fetch() {
     auto slot = _slot();
 
     return _fetch(slot);
 }
 
-ConnectionPoolSPtr ShardsPool::fetch(const Node &node) {
+GuardedConnection ShardsPool::fetch(const Node &node) {
     std::lock_guard<std::mutex> lock(_mutex);
 
-    auto iter = _pool.find(node);
-    if (iter == _pool.end()) {
+    auto iter = _pools.find(node);
+    if (iter == _pools.end()) {
         // Node doesn't exist, and it should be a newly created node.
         // So add a new connection pool.
         iter = _add_node(node);
     }
 
-    assert(iter != _pool.end());
+    assert(iter != _pools.end());
 
-    return iter->second;
+    return GuardedConnection(iter->second);
 }
 
 void ShardsPool::update() {
     // My might send command to a removed node.
     // Try at most 3 times.
     for (auto idx = 0; idx < 3; ++idx) {
         try {
-            // Randomly pick a connection pool.
-            auto pool = fetch();
-
-            assert(bool(pool));
-
-            auto connection = pool->fetch();
-
-            assert(!connection.broken());
-
-            ConnectionPoolGuard guard(*pool, connection);
-
-            auto shards = _cluster_slots(connection);
+            // Randomly pick a connection.
+            auto guarded_connection = fetch();
+            auto shards = _cluster_slots(guarded_connection.connection());
 
             std::unordered_set<Node, NodeHash> nodes;
             for (const auto &shard : shards) {
@@ -114,10 +105,10 @@ void ShardsPool::update() {
             _shards = std::move(shards);
 
             // Remove non-existent nodes.
-            for (auto iter = _pool.begin(); iter != _pool.end(); ) {
+            for (auto iter = _pools.begin(); iter != _pools.end(); ) {
                 if (nodes.find(iter->first) == nodes.end()) {
                     // Node has been removed.
-                    _pool.erase(iter++);
+                    _pools.erase(iter++);
                 } else {
                     ++iter;
                 }
@@ -126,7 +117,7 @@ void ShardsPool::update() {
             // Add connection pool for new nodes.
             // In fact, connections will be created lazily.
             for (const auto &node : nodes) {
-                if (_pool.find(node) == _pool.end()) {
+                if (_pools.find(node) == _pools.end()) {
                     _add_node(node);
                 }
             }
@@ -145,7 +136,7 @@ void ShardsPool::_move(ShardsPool &&that) {
     _pool_opts = that._pool_opts;
     _connection_opts = that._connection_opts;
     _shards = std::move(that._shards);
-    _pool = std::move(that._pool);
+    _pools = std::move(that._pools);
 }
 
 void ShardsPool::_init_pool(const Shards &shards) {
@@ -268,7 +259,7 @@ Slot ShardsPool::_slot() const {
     return uniform_dist(engine);
 }
 
-ConnectionPoolSPtr ShardsPool::_fetch(Slot slot) {
+GuardedConnection ShardsPool::_fetch(Slot slot) {
     std::lock_guard<std::mutex> lock(_mutex);
 
     auto shards_iter = _shards.lower_bound(SlotRange{slot, slot});
@@ -278,20 +269,20 @@ ConnectionPoolSPtr ShardsPool::_fetch(Slot slot) {
 
     const auto &node = shards_iter->second;
 
-    auto node_iter = _pool.find(node);
-    if (node_iter == _pool.end()) {
+    auto node_iter = _pools.find(node);
+    if (node_iter == _pools.end()) {
         throw Error("Slot is NOT covered: " + std::to_string(slot));
     }
 
-    return node_iter->second;
+    return GuardedConnection(node_iter->second);
 }
 
 auto ShardsPool::_add_node(const Node &node) -> NodeMap::iterator {
     auto opts = _connection_opts;
     opts.host = node.host;
     opts.port = node.port;
 
-    return _pool.emplace(node, std::make_shared<ConnectionPool>(_pool_opts, opts)).first;
+    return _pools.emplace(node, std::make_shared<ConnectionPool>(_pool_opts, opts)).first;
 }
 
 }

src/sw/redis++/shards_pool.h

@@ -17,6 +17,7 @@
 #ifndef SEWENEW_REDISPLUSPLUS_SHARDS_POOL_H
 #define SEWENEW_REDISPLUSPLUS_SHARDS_POOL_H
 
+#include <cassert>
 #include <unordered_map>
 #include <string>
 #include <random>
@@ -31,6 +32,32 @@ namespace redis {
 
 using ConnectionPoolSPtr = std::shared_ptr<ConnectionPool>;
 
+class GuardedConnection {
+public:
+    GuardedConnection(const ConnectionPoolSPtr &pool) : _pool(pool),
+                                                        _connection(_pool->fetch()) {
+        assert(!_connection.broken());
+    }
+
+    GuardedConnection(const GuardedConnection &) = delete;
+    GuardedConnection& operator=(const GuardedConnection &) = delete;
+
+    GuardedConnection(GuardedConnection &&) = default;
+    GuardedConnection& operator=(GuardedConnection &&) = default;
+
+    ~GuardedConnection() {
+        _pool->release(std::move(_connection));
+    }
+
+    Connection& connection() {
+        return _connection;
+    }
+
+private:
+    ConnectionPoolSPtr _pool;
+    Connection _connection;
+};
+
 class ShardsPool {
 public:
     ShardsPool() = default;
@@ -46,14 +73,14 @@ class ShardsPool {
     ShardsPool(const ConnectionPoolOptions &pool_opts,
                 const ConnectionOptions &connection_opts);
 
-    // Fetch a connection pool by key.
-    ConnectionPoolSPtr fetch(const StringView &key);
+    // Fetch a connection by key.
+    GuardedConnection fetch(const StringView &key);
 
-    // Randomly pick a connection pool.
-    ConnectionPoolSPtr fetch();
+    // Randomly pick a connection.
+    GuardedConnection fetch();
 
-    // Fetch a connection pool by node.
-    ConnectionPoolSPtr fetch(const Node &node);
+    // Fetch a connection by node.
+    GuardedConnection fetch(const Node &node);
 
     void update();
 
@@ -76,7 +103,7 @@ class ShardsPool {
     // Randomly pick a slot.
     std::size_t _slot() const;
 
-    ConnectionPoolSPtr _fetch(Slot slot);
+    GuardedConnection _fetch(Slot slot);
 
     using NodeMap = std::unordered_map<Node, ConnectionPoolSPtr, NodeHash>;
 
@@ -88,7 +115,7 @@ class ShardsPool {
 
     Shards _shards;
 
-    NodeMap _pool;
+    NodeMap _pools;
 
     std::mutex _mutex;