One pool per node test

src/Database/Redis/Cluster.hs

@@ -25,6 +25,7 @@ module Database.Redis.Cluster
   , nodes
   , Pipeline
   , newPipelineVar
+  , NodeResource(..)
 ) where
 
 import qualified Data.ByteString as B
@@ -71,16 +72,18 @@ type IsReadOnly = Bool
 data Connection = Connection (MVar NodeConnectionMap) (MVar ShardMap) CMD.InfoMap IsReadOnly TcpInfo
 
 -- | A connection to a single node in the cluster, similar to 'ProtocolPipelining.Connection'
-data NodeConnection = NodeConnection (Pool CC.ConnectionContext) (IOR.IORef (Maybe B.ByteString)) NodeID
+data NodeConnection = NodeConnection (Pool NodeResource) NodeID
+
+data NodeResource = NodeResource CC.ConnectionContext (IOR.IORef (Maybe B.ByteString))
 
 instance Show NodeConnection where
-    show (NodeConnection _ _ id1) = "nodeId: " <> show id1
+    show (NodeConnection _ id1) = "nodeId: " <> show id1
 
 instance Eq NodeConnection where
-    (NodeConnection _ _ id1) == (NodeConnection _ _ id2) = id1 == id2
+    (NodeConnection _ id1) == (NodeConnection _ id2) = id1 == id2
 
 instance Ord NodeConnection where
-    compare (NodeConnection _ _ id1) (NodeConnection _ _ id2) = compare id1 id2
+    compare (NodeConnection _ id1) (NodeConnection _ id2) = compare id1 id2
 
 data PipelineState =
       -- Nothing in the pipeline has been evaluated yet so nothing has been
@@ -185,10 +188,16 @@ connect withAuth commandInfos shardMapVar isReadOnly refreshShardMap (tcpInfo@Tc
                     Left (_ :: SomeException) -> (acc, True)
            ) (mempty, False) info
     connectNode :: Node -> IO (NodeID, NodeConnection)
-    connectNode (Node n _ host port) = do
-        ctx <- createPool (withAuth host (CC.PortNumber $ toEnum port) timeoutOpt) CC.disconnect 1 idleTime maxResources
+    connectNode node@(Node n _ _ _) = do
+        pool <- createPool (createNodeResource node) destroyNodeResource 1 idleTime maxResources
+        return (n, NodeConnection pool n)
+    createNodeResource :: Node -> IO NodeResource
+    createNodeResource (Node _ _ host port) = do
+        ctx <- withAuth host (CC.PortNumber $ toEnum port) timeoutOpt
         ref <- IOR.newIORef Nothing
-        return (n, NodeConnection ctx ref n)
+        return $ NodeResource ctx ref
+    destroyNodeResource :: NodeResource -> IO ()
+    destroyNodeResource (NodeResource ctx _) = CC.disconnect ctx
     refreshShardMapVar :: ShardMap -> IO ()
     refreshShardMapVar shardMap = hasLocked $ modifyMVar_ shardMapVar (const (pure shardMap))
 
@@ -199,7 +208,7 @@ newPipelineVar = do
 
 destroyNodeResources :: Connection -> IO ()
 destroyNodeResources (Connection nodeConnMapVar _ _ _ _) = (readMVar nodeConnMapVar) >>= (mapM_ disconnectNode . HM.elems) where
-    disconnectNode (NodeConnection nodePool _ _) = destroyAllResources nodePool
+    disconnectNode (NodeConnection nodePool _) = destroyAllResources nodePool
 
 -- Add a request to the current pipeline for this connection. The pipeline will
 -- be executed implicitly as soon as any result returned from this function is
@@ -494,22 +503,25 @@ allMasterNodes (Connection nodeConnsVar _ _ _ _) (ShardMap shardMap) = do
     onlyMasterNodes = (\(Shard master _) -> master) <$> nub (IntMap.elems shardMap)
 
 requestNode :: NodeConnection -> [[B.ByteString]] -> IO [Reply]
-requestNode (NodeConnection pool lastRecvRef _) requests = withResource pool $ \ctx -> do
+requestNode (NodeConnection pool _) requests = withResource pool (`requestNodeResource` requests)
+
+requestNodeResource :: NodeResource -> [[B.ByteString]] -> IO [Reply]
+requestNodeResource (NodeResource ctx lastRecvRef) requests = do
     envTimeout <- round . (\x -> (x :: Time.NominalDiffTime) * 1000000) . realToFrac . fromMaybe (0.5 :: Double) . (>>= readMaybe) <$> lookupEnv "REDIS_REQUEST_NODE_TIMEOUT"
-    eresp <- race (requestNodeImpl ctx) (threadDelay envTimeout)
+    eresp <- race requestNodeImpl (threadDelay envTimeout)
     case eresp of
       Left e -> return e
       Right _ -> putStrLn "timeout happened" *> throwIO (RequestTimingOut requests)
     where
-    requestNodeImpl :: CC.ConnectionContext -> IO [Reply]
-    requestNodeImpl ctx = do
-        mapM_ (sendNode ctx . renderRequest) requests
+    requestNodeImpl :: IO [Reply]
+    requestNodeImpl = do
+        mapM_ (sendNode . renderRequest) requests
         _ <- CC.flush ctx
-        replicateM (length requests) $ recvNode ctx
-    sendNode :: CC.ConnectionContext -> B.ByteString -> IO ()
-    sendNode = CC.send
-    recvNode :: CC.ConnectionContext -> IO Reply
-    recvNode ctx = do
+        replicateM (length requests) $ recvNode
+    sendNode :: B.ByteString -> IO ()
+    sendNode = CC.send ctx
+    recvNode :: IO Reply
+    recvNode = do
         maybeLastRecv <- IOR.readIORef lastRecvRef
         scanResult <- case maybeLastRecv of
             Just lastRecv -> Scanner.scanWith (CC.recv ctx) reply lastRecv

src/Database/Redis/Connection.hs

@@ -265,8 +265,8 @@ connectCluster bootstrapConnInfo = do
           return clusterConn
           where
           ctxToConn :: Cluster.NodeConnection -> IO (Maybe PP.Connection)
-          ctxToConn (Cluster.NodeConnection pool _ nid) = do
-            maybeConn <- try $ withResource pool PP.fromCtx
+          ctxToConn (Cluster.NodeConnection pool nid) = do
+            maybeConn <- try $ withResource pool (\(Cluster.NodeResource ctx _) -> PP.fromCtx ctx)
             case maybeConn of
                 Right ppConn -> return $ Just ppConn
                 Left (_ :: SomeException) -> do
@@ -337,18 +337,24 @@ refreshShardMap (Cluster.Connection nodeConnsVar shardMapVar _ _ Cluster.TcpInfo
                             Left (_ :: SomeException) -> acc
                 ) mempty info
         connectNode :: Cluster.Node -> IO (Cluster.NodeID, Cluster.NodeConnection)
-        connectNode (Cluster.Node n _ host port) = do
-            ctx <- createPool (withAuth host (CC.PortNumber $ toEnum port) timeoutOpt) CC.disconnect 1 idleTime maxResources
+        connectNode node@(Cluster.Node n _ _ _) = do
+            pool <- createPool (createNodeResource node) destroyNodeResource 1 idleTime maxResources
+            return (n, Cluster.NodeConnection pool n)
+        createNodeResource :: Cluster.Node -> IO Cluster.NodeResource
+        createNodeResource (Cluster.Node _ _ host port) = do
+            ctx <- withAuth host (CC.PortNumber $ toEnum port) timeoutOpt
             ref <- IOR.newIORef Nothing
-            return (n, Cluster.NodeConnection ctx ref n)
+            return $ Cluster.NodeResource ctx ref
+        destroyNodeResource :: Cluster.NodeResource -> IO ()
+        destroyNodeResource (Cluster.NodeResource ctx _) = CC.disconnect ctx
 
 refreshShardMapWithNodeConn :: [Cluster.NodeConnection] -> IO ShardMap
 refreshShardMapWithNodeConn [] = throwIO $ ClusterConnectError (Error "Couldn't refresh shardMap due to connection error")
 refreshShardMapWithNodeConn nodeConnsList = do
     selectedIdx <- randomRIO (0, (length nodeConnsList) - 1)
     print selectedIdx
-    let (Cluster.NodeConnection pool _ _) = nodeConnsList !! selectedIdx
-    withResource pool $ \ctx -> do
+    let (Cluster.NodeConnection pool _) = nodeConnsList !! selectedIdx
+    withResource pool $ \(Cluster.NodeResource ctx _) -> do
         pipelineConn <- PP.fromCtx ctx
         envTimeout <- fromMaybe (10 ^ (3 :: Int)) . (>>= readMaybe) <$> lookupEnv "REDIS_CLUSTER_SLOTS_TIMEOUT"
         raceResult <- race (threadDelay envTimeout) (try $ refreshShardMapWithConn pipelineConn True) -- racing with delay of default 1 ms