diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,84 @@ +# Revision history for redis-client + +## 0.6.0.0 -- 2026-02-13 + +* **Multiplexing Now Default** + * Cluster mode (`ClusterConfig`) now uses multiplexing by default (`clusterUseMultiplexing = True`). + Set `clusterUseMultiplexing = False` to opt out. +* **New: Standalone Multiplexed Client** + * Added `StandaloneClient` module for pipelined throughput on a single (non-cluster) Redis server. + * Provides `createStandaloneClient`, `createStandaloneClientFromConfig`, `closeStandaloneClient`, + and `runStandaloneClient` lifecycle functions. + * `StandaloneConfig` allows tuning multiplexer count and toggling multiplexing on/off. + * Implements `RedisCommands` typeclass — all existing commands work transparently. + * Re-exported from the top-level `Redis` module. +* **Testing** + * Added `MultiplexerSpec` — 14 unit tests for multiplexer internals (slot pool, response slot, lifecycle). + * Added `MultiplexPoolSpec` — 8 unit tests for per-node multiplexer management and round-robin routing. + * Added standalone E2E tests for all command families (string, hash, list, set, sorted set). + * Added concurrency stress tests (50+ threads, submit-after-destroy, multi-multiplexer distribution). + +## 0.5.0.0 -- 2026-02-05 + +* **Major Feature: Redis Cluster Support** + * Implemented full cluster support including topology management, connection pooling, and slot routing. + * Added smart routing with automatic `MOVED` and `ASK` redirection handling. + * Updated `fill`, `cli`, and `tunnel` modes to work seamlessly with Redis Cluster. + * Added "Pinned Proxy" mode to support specific cluster proxy configurations. + * Fixed hostname rewriting for cluster nodes in tunnel mode. +* **Infrastructure & Testing** + * Added comprehensive End-to-End (E2E) tests for Cluster scenarios (Basic, Fill, Tunnel, CLI). + * Updated build system to prefer `nix-build` and `make test`. + * Added GitHub Actions workflows for automated testing and bootstrapping. +* **Fill Operation Enhancements** + * Added configurable `--key-size` and `--value-size` support for realistic workload simulation. + * Added `--pipeline` configuration for explicit batch size control. + * Removed `REDIS_CLIENT_FILL_CHUNK_KB` environment variable in favor of explicit flags. + * Fixed fill calculation logic to ensure exact data generation matching requested sizes. +* **Improvements** + * Fixed various race conditions and unsafe head usage. + * Improved error handling for `CROSSSLOT` operations. + * Refactored codebase for better maintainability (Cluster separation). + +## 0.4.0.0 -- 2026-02-02 + +* **Azure Integration** + * Added `azure-redis-connect` python script for Azure Redis Cache discovery and connectivity. + * Added support for Entra authentication (username/OID retrieval from JWT). +* **Performance Optimization** + * Optimized data filling with parallel noise generation strategies. + * Optimized connection handling (parallel connections adjustment). +* **DevOps** + * Added agent testing infrastructure. + +## 0.3.0.0 -- 2025-10-15 + +* **Command Expansion** + * Added support for Geo commands. + * Expanded standard Redis command set support. +* **Build System** + * Added `flake.nix` for modern Nix support. + * Fixed inefficient store usage in build. +* **Testing** + * Added tunneling tests. + +## 0.2.0.0 -- 2025-03-28 + +* **New Architectures** + * Added "Tunnel mode" for proxying Redis connections. + * Added TLS client support. + * Added Pipelining support for improved throughput. +* **CLI & Usability** + * Added argument parsing for the CLI application. + * Added support for arrow key navigation in CLI. + * Added DNS resolution support. +* **Core Improvements** + * Refactored RESP parsing logic. + * Improved Incremental network reading (State monad refactor). + * Added basic E2E testing framework. + +## 0.1.0.0 -- 2025-02-11 + +* First version. +* Basic RESP data types. +* Simple plaintext client implementation. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2025 Seth Speaks + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,99 @@ +# hask-redis-mux + +[![Hackage](https://img.shields.io/hackage/v/hask-redis-mux.svg)](https://hackage.haskell.org/package/hask-redis-mux) +[![CI](https://github.com/sspeaks/redis-client/actions/workflows/ci.yml/badge.svg)](https://github.com/sspeaks/redis-client/actions) + +A multiplexed Redis client library for Haskell with full RESP protocol support, +Redis Cluster topology discovery, connection pooling, and TLS. + +## Features + +- **Standalone & Cluster** — works with single-node Redis and Redis Cluster +- **Multiplexed pipelining** — concurrent commands share a single TCP connection +- **Typed returns** via `FromResp` — parse responses as `ByteString`, `Integer`, `Text`, `Bool`, or custom types +- **TLS support** — connect over TLS with `crypton` +- **Bracket-style resource management** — `withStandaloneClient` / `withClusterClient` for exception-safe cleanup +- **Connection pooling** — automatic pool management for cluster nodes + +## Installation + +Add to your `.cabal` file: + +```cabal +build-depends: hask-redis-mux >= 0.1 && < 0.2 +``` + +## Quick Start + +```haskell +{-# LANGUAGE OverloadedStrings #-} +import Database.Redis + +main :: IO () +main = do + -- Connect to localhost:6379, run commands, auto-close + result <- runRedis defaultStandaloneConfig $ do + set "greeting" "hello" + (val :: ByteString) <- get "greeting" + return val + print result -- "hello" +``` + +## Typed Returns with FromResp + +Commands return polymorphic types via the `FromResp` typeclass. Just add a +type annotation and the response is parsed automatically: + +```haskell +runRedis defaultStandaloneConfig $ do + set "counter" "42" + + (n :: Integer) <- get "counter" -- 42 + (bs :: ByteString) <- get "counter" -- "42" + (mt :: Maybe Text) <- get "missing" -- Nothing + (ok :: Bool) <- set "k" "v" -- True (from +OK) +``` + +## Bracket Pattern (Recommended) + +Use bracket-style functions for exception-safe resource management: + +```haskell +-- Standalone +withStandaloneClient config $ \client -> + runStandaloneClient client $ do + set "key" "value" + get "key" + +-- Cluster +withClusterClient clusterConfig connector $ \client -> + runClusterCommandClient client $ do + set "key" "value" + get "key" +``` + +## Custom Configuration + +```haskell +import Database.Redis + +main :: IO () +main = do + let config = StandaloneConfig + { standaloneNodeAddress = NodeAddress "redis.example.com" 6379 + , standaloneConnector = clusterPlaintextConnector + , standaloneMultiplexerCount = 4 -- 4 multiplexed connections + } + withStandaloneClient config $ \client -> + runStandaloneClient client $ do + set "key" "value" +``` + +## Documentation + +- [Haddock API docs](https://hackage.haskell.org/package/hask-redis-mux) +- [GitHub repository](https://github.com/sspeaks/redis-client) + +## License + +MIT — see [LICENSE](LICENSE) for details. diff --git a/hask-redis-mux.cabal b/hask-redis-mux.cabal new file mode 100644 --- /dev/null +++ b/hask-redis-mux.cabal @@ -0,0 +1,208 @@ +cabal-version: 3.0 +name: hask-redis-mux +version: 0.1.0.0 +synopsis: A RESP protocol implementation and multiplexed Redis client library +description: + hask-redis-mux provides a full RESP protocol parser, Redis client with TLS support, + CRC16 hash slot computation, and a multiplexed Redis Cluster client with connection + pooling and automatic topology discovery. +license: MIT +license-file: LICENSE +author: Seth Speaks +maintainer: sspeaks610@gmail.com +category: Network +build-type: Simple +extra-doc-files: CHANGELOG.md + , README.md +extra-source-files: lib/crc16/crc16.c +homepage: https://github.com/sspeaks/redis-client +bug-reports: https://github.com/sspeaks/redis-client/issues +tested-with: GHC == 9.10.3 + +source-repository head + type: git + location: https://github.com/sspeaks/redis-client.git + +-- --------------------------------------------------------------------------- +-- Common stanzas +-- --------------------------------------------------------------------------- + +common defaults + default-language: GHC2021 + ghc-options: -Wall -O2 + build-depends: base >=4.19 && <4.22 + +common test-defaults + import: defaults + ghc-options: -threaded -rtsopts "-with-rtsopts=-N -H1024M -A128m -n8m -qb" + build-depends: hspec >=2.11 && <2.12 + hs-source-dirs: test + +-- --------------------------------------------------------------------------- +-- Libraries +-- --------------------------------------------------------------------------- + +library + import: defaults + reexported-modules: + Database.Redis.Resp + , Database.Redis.FromResp + , Database.Redis.RedisError + , Database.Redis.Client + , Database.Redis.Crc16 + , Database.Redis.Command + , Database.Redis.Cluster + , Database.Redis.Cluster.Client + , Database.Redis.Cluster.Commands + , Database.Redis.Cluster.SlotMapping + , Database.Redis.Cluster.ConnectionPool + , Database.Redis.Connector + , Database.Redis.Internal.MultiplexPool + , Database.Redis.Internal.Multiplexer + , Database.Redis.Standalone + , Database.Redis + build-depends: resp + , client + , crc16 + , redis-command-client + , cluster + , redis + +library resp + import: defaults + hs-source-dirs: lib/resp + exposed-modules: Database.Redis.Resp + build-depends: attoparsec >=0.14 && <0.15 + , bytestring >=0.12 && <0.13 + , containers >=0.6 && <0.8 + +library client + import: defaults + hs-source-dirs: lib/client + exposed-modules: Database.Redis.Client + build-depends: bytestring >=0.12 && <0.13 + , crypton-x509-system >=1.6 && <1.7 + , data-default-class >=0.1 && <0.3 + , dns >=4.2 && <4.3 + , iproute >=1.7 && <1.8 + , network >=3.2 && <3.3 + , tls >=2.1 && <2.2 + +library crc16 + import: defaults + hs-source-dirs: lib/crc16 + exposed-modules: Database.Redis.Crc16 + c-sources: lib/crc16/crc16.c + build-depends: bytestring >=0.12 && <0.13 + +library redis-command-client + import: defaults + hs-source-dirs: lib/redis-command-client + exposed-modules: Database.Redis.Command + , Database.Redis.FromResp + , Database.Redis.RedisError + build-depends: attoparsec >=0.14 && <0.15 + , bytestring >=0.12 && <0.13 + , client + , mtl >=2.3 && <2.4 + , resp + , text >=2.1 && <2.2 + +library cluster + import: defaults + hs-source-dirs: lib/cluster + exposed-modules: Database.Redis.Cluster + , Database.Redis.Cluster.Client + , Database.Redis.Cluster.Commands + , Database.Redis.Cluster.SlotMapping + , Database.Redis.Cluster.ConnectionPool + , Database.Redis.Connector + , Database.Redis.Internal.MultiplexPool + , Database.Redis.Internal.Multiplexer + , Database.Redis.Standalone + build-depends: attoparsec >=0.14 && <0.15 + , bytestring >=0.12 && <0.13 + , client + , containers >=0.6 && <0.8 + , crc16 + , file-embed >=0.0.16 && <0.1 + , mtl >=2.3 && <2.4 + , redis-command-client + , resp + , stm >=2.5 && <2.6 + , time >=1.12 && <1.13 + , vector >=0.13 && <0.14 + +library redis + import: defaults + visibility: public + hs-source-dirs: lib/redis + exposed-modules: Database.Redis + build-depends: bytestring >=0.12 && <0.13 + , client + , cluster + , redis-command-client + , resp + +-- --------------------------------------------------------------------------- +-- Test suites +-- --------------------------------------------------------------------------- + +test-suite RespSpec + import: test-defaults + type: exitcode-stdio-1.0 + main-is: Spec.hs + build-depends: attoparsec >=0.14 && <0.15 + , bytestring >=0.12 && <0.13 + , containers >=0.6 && <0.8 + , redis-command-client + , resp + +test-suite ClusterSpec + import: test-defaults + type: exitcode-stdio-1.0 + main-is: ClusterSpec.hs + build-depends: bytestring >=0.12 && <0.13 + , cluster + , resp + , time >=1.12 && <1.13 + +test-suite ClusterCommandSpec + import: test-defaults + type: exitcode-stdio-1.0 + main-is: ClusterCommandSpec.hs + build-depends: bytestring >=0.12 && <0.13 + , client + , cluster + , containers >=0.6 && <0.8 + , resp + , stm >=2.5 && <2.6 + , time >=1.12 && <1.13 + , vector >=0.13 && <0.14 + +test-suite MultiplexerSpec + import: test-defaults + type: exitcode-stdio-1.0 + main-is: MultiplexerSpec.hs + build-depends: bytestring >=0.12 && <0.13 + , client + , cluster + , resp + +test-suite FromRespSpec + import: test-defaults + type: exitcode-stdio-1.0 + main-is: FromRespSpec.hs + build-depends: bytestring >=0.12 && <0.13 + , redis-command-client + , resp + , text >=2.1 && <2.2 + +test-suite MultiplexPoolSpec + import: test-defaults + type: exitcode-stdio-1.0 + main-is: MultiplexPoolSpec.hs + build-depends: bytestring >=0.12 && <0.13 + , client + , cluster + , resp diff --git a/lib/client/Database/Redis/Client.hs b/lib/client/Database/Redis/Client.hs new file mode 100644 --- /dev/null +++ b/lib/client/Database/Redis/Client.hs @@ -0,0 +1,221 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | TCP and TLS transport layer for Redis connections. +-- +-- The 'Client' typeclass uses a type-level 'ConnectionStatus' parameter to enforce +-- connection state at compile time: you can only 'send' and 'receive' on a +-- @client \'Connected@, and only 'connect' a @client \'NotConnected@. This prevents +-- use-after-close and send-before-connect bugs statically. +-- +-- @since 0.1.0.0 +module Database.Redis.Client + ( Client (..) + , serve + , PlainTextClient (NotConnectedPlainTextClient) + , TLSClient (NotConnectedTLSClient, NotConnectedTLSClientWithHostname, TLSTunnel) + , ConnectionStatus (..) + ) where + +import Control.Exception (IOException, bracket, catch, + finally, throwIO) +import Control.Monad (void) +import Control.Monad.IO.Class +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BS8 +import qualified Data.ByteString.Lazy as LBS +import Data.Default.Class (def) +import Data.IP (IPv4, toHostAddress) +import Data.Kind (Type) +import Data.Word (Word32) +import Network.DNS (defaultResolvConf, lookupA, + makeResolvSeed, withResolver) +import Network.Socket (Family (AF_INET), HostAddress, + SockAddr (SockAddrInet), + Socket, SocketOption (..), + SocketType (Stream), + defaultProtocol, + setSocketOption, socket, + tupleToHostAddress) +import qualified Network.Socket as S +import Network.Socket.ByteString (recv, sendMany) +import Network.Socket.ByteString.Lazy (sendAll) +import Network.TLS (ClientHooks (..), + ClientParams (..), Context, + Shared (..), Supported (..), + Version (..), bye, contextNew, + defaultParamsClient, handshake, + recvData, sendData) +import Network.TLS.Extra (ciphersuite_strong) +import Prelude hiding (getContents) +import System.Environment (lookupEnv) +import System.IO (BufferMode (LineBuffering), + hFlush, hSetBuffering, stdout) +import System.Timeout (timeout) +import System.X509.Unix (getSystemCertificateStore) +import Text.Printf (printf) + +-- | Connection lifecycle phase, used as a DataKinds-promoted type parameter +-- to statically track whether a client is connected. +data ConnectionStatus = Connected | NotConnected | Server + +-- | Transport abstraction indexed by 'ConnectionStatus'. The type parameter +-- ensures that 'send' and 'receive' can only be called on a connected client, +-- and 'connect' can only be called on a not-yet-connected client. +class Client (client :: ConnectionStatus -> Type) where + connect :: (MonadIO m) => client 'NotConnected -> m (client 'Connected) + close :: (MonadIO m) => client 'Connected -> m () + send :: (MonadIO m) => client 'Connected -> LBS.ByteString -> m () + -- | Send multiple strict ByteString chunks via vectored I/O (writev). + -- Default implementation falls back to 'send' with lazy concatenation. + -- PlainTextClient overrides with sendMany for zero-copy vectored I/O. + sendChunks :: (MonadIO m) => client 'Connected -> [BS.ByteString] -> m () + sendChunks conn chunks = send conn (LBS.fromChunks chunks) + receive :: (MonadIO m, MonadFail m) => client 'Connected -> m BS.ByteString + +-- | Plain TCP client. Construct with 'NotConnectedPlainTextClient' providing +-- a hostname and optional port (defaults to 6379). +data PlainTextClient (a :: ConnectionStatus) where + NotConnectedPlainTextClient :: String -> Maybe Int -> PlainTextClient 'NotConnected + ConnectedPlainTextClient :: String -> Word32 -> Socket -> PlainTextClient 'Connected + +instance Client PlainTextClient where + connect :: (MonadIO m) => PlainTextClient 'NotConnected -> m (PlainTextClient 'Connected) + connect (NotConnectedPlainTextClient hostname port) = liftIO $ do + (sock, ipCorrectEndian) <- createSocket hostname (maybe 6379 fromIntegral port) + S.connect sock (SockAddrInet (maybe 6379 fromIntegral port) ipCorrectEndian) `catch` \(e :: IOException) -> do + printf "Wasn't able to connect to the server: %s...\n" (show e) + putStrLn "Tried to use a plain text socket on port 6379. Did you mean to use TLS on port 6380?" + throwIO e + return $ ConnectedPlainTextClient hostname ipCorrectEndian sock + + close :: (MonadIO m) => PlainTextClient 'Connected -> m () + close (ConnectedPlainTextClient _ _ sock) = liftIO $ S.close sock + + send :: (MonadIO m) => PlainTextClient 'Connected -> LBS.ByteString -> m () + send (ConnectedPlainTextClient _ _ sock) dat = liftIO $ sendAll sock dat + + sendChunks :: (MonadIO m) => PlainTextClient 'Connected -> [BS.ByteString] -> m () + sendChunks (ConnectedPlainTextClient _ _ sock) chunks = liftIO $ sendMany sock chunks + + receive :: (MonadIO m, MonadFail m) => PlainTextClient 'Connected -> m BS.ByteString + receive (ConnectedPlainTextClient _ _ sock) = do + -- Timeout increased to 300s (5 minutes) to handle massive backlogs during fill operations + val <- liftIO $ timeout (300 * 1000000) $ recv sock 16384 + case val of + Nothing -> fail "recv socket timeout (plaintext)" + Just v -> return v + +-- | TLS-encrypted client. Construct with 'NotConnectedTLSClient' (hostname + optional port, +-- defaults to 6380) or 'NotConnectedTLSClientWithHostname' when the TLS certificate +-- hostname differs from the connection address (common in cluster mode). +-- Set @REDIS_CLIENT_TLS_INSECURE@ to skip certificate validation. +data TLSClient (a :: ConnectionStatus) where + NotConnectedTLSClient :: String -> Maybe Int -> TLSClient 'NotConnected + -- | TLS client with separate hostname for certificate validation + -- This is useful for cluster mode where CLUSTER SLOTS returns IP addresses + -- but we need to use the original hostname for TLS certificate validation + NotConnectedTLSClientWithHostname :: String -> String -> Maybe Int -> TLSClient 'NotConnected + ConnectedTLSClient :: String -> Word32 -> Socket -> Context -> TLSClient 'Connected + TLSTunnel :: TLSClient 'Connected -> TLSClient 'Server + +instance Client TLSClient where + connect :: (MonadIO m) => TLSClient 'NotConnected -> m (TLSClient 'Connected) + connect (NotConnectedTLSClient hostname port) = + connectTLS hostname hostname port + + connect (NotConnectedTLSClientWithHostname certHostname targetAddress port) = + connectTLS certHostname targetAddress port + + close :: (MonadIO m) => TLSClient 'Connected -> m () + close (ConnectedTLSClient _ _ sock ctx) = liftIO $ bye ctx `finally` S.close sock + + send :: (MonadIO m) => TLSClient 'Connected -> LBS.ByteString -> m () + send (ConnectedTLSClient _ _ _ ctx) dat = liftIO $ sendData ctx dat + + receive :: (MonadIO m, MonadFail m) => TLSClient 'Connected -> m BS.ByteString + receive (ConnectedTLSClient _ _ _ ctx) = do + -- Timeout increased to 300s (5 minutes) to handle massive backlogs + val <- liftIO $ timeout (300 * 1000000) $ recvData ctx + case val of + Nothing -> fail "recv socket timeout (TLS)" + Just v -> return v + +-- | Connect to a TLS server, using certHostname for certificate validation +-- and targetAddress for the actual network connection. +connectTLS :: (MonadIO m) => String -> String -> Maybe Int -> m (TLSClient 'Connected) +connectTLS certHostname targetAddress port = liftIO $ do + (sock, ipCorrectEndian) <- createSocket targetAddress (maybe 6380 fromIntegral port) + S.connect sock (SockAddrInet (maybe 6380 fromIntegral port) ipCorrectEndian) + store <- getSystemCertificateStore + insecureFlag <- lookupEnv "REDIS_CLIENT_TLS_INSECURE" + let allowInsecure = maybe False (not . null) insecureFlag + baseParams = + (defaultParamsClient certHostname "redis-server") + { clientSupported = + def + { supportedVersions = [TLS13, TLS12], + supportedCiphers = ciphersuite_strong + }, + clientShared = + def + { sharedCAStore = store + } + } + clientParams = + if allowInsecure + then baseParams {clientHooks = def {onServerCertificate = \_ _ _ _ -> pure []}} + else baseParams + context <- contextNew sock clientParams + handshake context + return $ ConnectedTLSClient certHostname ipCorrectEndian sock context + +-- | Start a local TCP proxy on @localhost:6379@ that forwards traffic through an +-- existing TLS connection. Useful for tunneling plain-text Redis tools over TLS. +serve :: (MonadIO m) => TLSClient 'Server -> m () +serve (TLSTunnel redisClient) = liftIO $ do + hSetBuffering stdout LineBuffering + bracket (socket AF_INET Stream defaultProtocol) S.close $ \sock -> do + setSocketOption sock ReuseAddr 1 + S.bind sock (SockAddrInet 6379 (tupleToHostAddress (127, 0, 0, 1))) + S.listen sock 1024 + putStrLn "Listening on localhost:6379" + hFlush stdout + (clientSock, _) <- S.accept sock + putStrLn "Accepted connection" + hFlush stdout + void $ + finally + (loop clientSock redisClient) + (S.close clientSock) + where + loop client redis = do + dat <- recv client 4096 + send redisClient (LBS.fromStrict dat) + receivedData <- receive redis + sendAll client (LBS.fromStrict receivedData) + loop client redis + +-- | Create a TCP socket with standard options (NoDelay, KeepAlive) and resolve the hostname. +createSocket :: String -> S.PortNumber -> IO (Socket, HostAddress) +createSocket hostname _port = do + ipAddr <- resolve hostname + sock <- socket AF_INET Stream defaultProtocol + setSocketOption sock NoDelay 1 + setSocketOption sock KeepAlive 1 + return (sock, ipAddr) + +-- | Resolve a hostname or IP address string to a 'HostAddress'. +-- Handles @\"localhost\"@, dotted-quad IPv4 literals, and DNS A-record lookups. +resolve :: String -> IO HostAddress +resolve "localhost" = return (tupleToHostAddress (127, 0, 0, 1)) +resolve address = + case reads address :: [(IPv4, String)] of + [(ip, "")] -> return (toHostAddress ip) + _ -> do + rs <- makeResolvSeed defaultResolvConf + result <- withResolver rs $ \resolver -> lookupA resolver (BS8.pack address) + case result of + Right (a : _) -> return (toHostAddress a) + _ -> error $ "no address found for: " ++ address diff --git a/lib/cluster/Database/Redis/Cluster.hs b/lib/cluster/Database/Redis/Cluster.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Cluster.hs @@ -0,0 +1,190 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Redis cluster topology model and slot routing. +-- +-- Provides types for representing cluster nodes, slot ranges, and the full topology, +-- plus functions to compute hash slots, parse @CLUSTER SLOTS@ responses, and look up +-- which node owns a given slot. +-- +-- @since 0.1.0.0 +module Database.Redis.Cluster + ( ClusterNode (..), + SlotRange (..), + ClusterTopology (..), + NodeRole (..), + NodeAddress (..), + calculateSlot, + extractHashTag, + parseClusterSlots, + findNodeForSlot, + findNodeAddressForSlot, + ) +where + +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BS8 +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Time (UTCTime) +import Data.Vector (Vector) +import qualified Data.Vector as V +import Data.Word (Word16) +import Database.Redis.Crc16 (crc16) +import Database.Redis.Resp (RespData (..)) + +-- | Node role in the cluster +data NodeRole = Master | Replica + deriving (Show, Eq) + +-- | Network address for connecting to a cluster node. +data NodeAddress = NodeAddress + { nodeHost :: String, + nodePort :: Int + } + deriving (Show, Eq, Ord) + +-- | A cluster node with its identity, address, role, and slot assignments. +data ClusterNode = ClusterNode + { nodeId :: ByteString, + nodeAddress :: NodeAddress, + nodeRole :: NodeRole, + nodeSlotsServed :: [SlotRange], + nodeReplicas :: [ByteString] -- Node IDs of replicas + } + deriving (Show, Eq) + +-- | A contiguous range of hash slots and the nodes responsible for them. +data SlotRange = SlotRange + { slotStart :: Word16, -- 0-16383 + slotEnd :: Word16, + slotMaster :: ByteString, -- Node ID reference + slotReplicas :: [ByteString] -- Node ID references + } + deriving (Show, Eq) + +-- | Full snapshot of the cluster topology: a fast O(1) slot-to-node vector, +-- a map of all known nodes, and the time the snapshot was taken. +data ClusterTopology = ClusterTopology + { topologySlots :: Vector ByteString, -- 16384 slots, each mapped to node ID + topologyAddresses :: Vector NodeAddress, -- 16384 slots, each mapped directly to NodeAddress (hot path) + topologyNodes :: Map ByteString ClusterNode, -- Node ID -> full node details + topologyUpdateTime :: UTCTime + } + deriving (Show) + +-- | Calculate the hash slot (0–16383) for a Redis key. +-- Respects hash tags: if the key starts with @{tag}@, only @tag@ is hashed. +calculateSlot :: ByteString -> Word16 +calculateSlot key = + let !hashKey = extractHashTag key + in crc16 hashKey +{-# INLINE calculateSlot #-} + +-- | Extract hash tag from a key if present +-- Pattern: {tag} - returns the content within the first valid {} pair +-- Examples: +-- "{user}:profile" -> "user" +-- "key" -> "key" +-- "{}" -> "{}" +-- "{user" -> "{user" +-- "key{tag}" -> "key{tag}" (no tag at end) +extractHashTag :: ByteString -> ByteString +extractHashTag key = + case BS.breakSubstring "{" key of + (before, rest) + | not (BS.null rest) && not (BS.null before) -> + -- Found { but there's content before it - no valid tag + key + | not (BS.null rest) -> + -- Found { at start, check for closing } + case BS.breakSubstring "}" (BS.tail rest) of + (tag, after) + | not (BS.null after) && not (BS.null tag) -> tag + | otherwise -> key + | otherwise -> key + +-- | Parse the RESP response from @CLUSTER SLOTS@ into a 'ClusterTopology'. +-- Returns 'Left' with an error message if the response format is unexpected. +parseClusterSlots :: RespData -> UTCTime -> Either String ClusterTopology +parseClusterSlots (RespArray slots) currentTime = do + ranges <- mapM parseSlotRange slots + let (slotMap, nodeMap) = buildTopology ranges + addrMap = buildAddressVector slotMap nodeMap + return $ ClusterTopology slotMap addrMap nodeMap currentTime + where + parseSlotRange :: RespData -> Either String (SlotRange, [(ByteString, NodeAddress)]) + parseSlotRange (RespArray (RespInteger start : RespInteger end : masterInfo : replicaInfos)) = do + master <- parseNodeInfo masterInfo + replicas <- mapM parseNodeInfo replicaInfos + let range = SlotRange + { slotStart = fromIntegral start, + slotEnd = fromIntegral end, + slotMaster = fst master, + slotReplicas = map fst replicas + } + return (range, master : replicas) + parseSlotRange other = Left $ "Invalid slot range format: " ++ show other + + parseNodeInfo :: RespData -> Either String (ByteString, NodeAddress) + parseNodeInfo (RespArray (RespBulkString host : RespInteger port : RespBulkString nodeIdBS : _)) = + Right (nodeIdBS, NodeAddress (BS8.unpack host) (fromIntegral port)) + parseNodeInfo other = Left $ "Invalid node info format: " ++ show other + + buildTopology :: [(SlotRange, [(ByteString, NodeAddress)])] -> (Vector ByteString, Map ByteString ClusterNode) + buildTopology rangesWithNodes = + let slotVector = V.replicate 16384 "" + slotMap = foldl (\v (range, _) -> assignSlots v range) slotVector rangesWithNodes + nodeMap = foldl buildNodeMap Map.empty rangesWithNodes + in (slotMap, nodeMap) + + assignSlots :: Vector ByteString -> SlotRange -> Vector ByteString + assignSlots vec range = + foldl + (\v slot -> v V.// [(fromIntegral slot, slotMaster range)]) + vec + [slotStart range .. slotEnd range] + + buildNodeMap :: Map ByteString ClusterNode -> (SlotRange, [(ByteString, NodeAddress)]) -> Map ByteString ClusterNode + buildNodeMap nodeMap (range, nodeInfos) = + case nodeInfos of + [] -> nodeMap -- No nodes in this range (shouldn't happen) + (masterId, masterAddr) : replicaInfos -> + -- Insert master node + let nodeMapWithMaster = insertNode nodeMap (masterId, masterAddr) Master + -- Insert replica nodes + nodeMapWithReplicas = foldl (\nm (replicaId, replicaAddr) -> + insertNode nm (replicaId, replicaAddr) Replica) + nodeMapWithMaster replicaInfos + in nodeMapWithReplicas + where + insertNode :: Map ByteString ClusterNode -> (ByteString, NodeAddress) -> NodeRole -> Map ByteString ClusterNode + insertNode nm (nodeId, addr) role = + if Map.member nodeId nm + then nm -- Node already exists, don't overwrite + else Map.insert nodeId (ClusterNode nodeId addr role [range] []) nm +parseClusterSlots other _ = Left $ "Expected array of slot ranges, got: " ++ show other + +-- | Build a slot-to-NodeAddress vector from the slot-to-nodeId vector and node map. +-- Used for O(1) hot path lookups that skip the Map entirely. +buildAddressVector :: Vector ByteString -> Map ByteString ClusterNode -> Vector NodeAddress +buildAddressVector slotVec nodeMap = + V.map (\nid -> case Map.lookup nid nodeMap of + Just node -> nodeAddress node + Nothing -> NodeAddress "" 0 -- placeholder for unassigned slots + ) slotVec + +-- | Look up the master node ID responsible for a given slot. +-- Returns 'Nothing' if the slot is out of range (≥ 16384). +findNodeForSlot :: ClusterTopology -> Word16 -> Maybe ByteString +findNodeForSlot topology slot + | slot < 16384 = Just $ topologySlots topology V.! fromIntegral slot + | otherwise = Nothing + +-- | Look up the 'NodeAddress' responsible for a given slot directly (O(1), no Map lookup). +-- Returns 'Nothing' if the slot is out of range (≥ 16384). +findNodeAddressForSlot :: ClusterTopology -> Word16 -> Maybe NodeAddress +findNodeAddressForSlot topology slot + | slot < 16384 = Just $! topologyAddresses topology V.! fromIntegral slot + | otherwise = Nothing +{-# INLINE findNodeAddressForSlot #-} diff --git a/lib/cluster/Database/Redis/Cluster/Client.hs b/lib/cluster/Database/Redis/Cluster/Client.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Cluster/Client.hs @@ -0,0 +1,631 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} + +-- | Cluster-aware Redis command client with automatic slot routing, MOVED\/ASK +-- redirection handling, and connection pooling. +-- +-- == Quick Start +-- +-- @ +-- import Redis +-- +-- client <- 'createClusterClient' config connector +-- +-- -- Use the monadic interface (implements 'RedisCommands'): +-- 'runClusterCommandClient' client $ do +-- set \"key1\" \"val1\" +-- set \"key2\" \"val2\" +-- get \"key1\" +-- +-- -- One-shot from IO: +-- result <- 'runClusterCommandClient' client (get \"mykey\") +-- +-- 'closeClusterClient' client +-- @ +-- +-- The 'ClusterCommandClient' monad implements 'RedisCommands', providing +-- the same @get@\/@set@\/@del@\/… API as single-node Redis with transparent +-- cluster slot routing, MOVED\/ASK handling, and connection pooling. +-- +-- For advanced use (e.g.\ forwarding raw RESP commands), the low-level +-- 'executeKeyedClusterCommand' and 'executeKeylessClusterCommand' are also +-- available but are not re-exported by the convenience "Redis" module. +-- +-- @since 0.1.0.0 +module Database.Redis.Cluster.Client + ( -- * Client Types + ClusterClient (..), + ClusterCommandClient, + ClusterError (..), + ClusterConfig (..), + -- * Client Lifecycle + createClusterClient, + closeClusterClient, + withClusterClient, + refreshTopology, + -- * Running Commands (monadic, recommended) + runClusterCommandClient, + -- * Low-Level Command Execution (advanced) + -- | These are intended for internal use or advanced scenarios like RESP + -- proxying. Prefer 'runClusterCommandClient' with 'RedisCommands' for + -- normal Redis operations. + executeKeyedClusterCommand, + executeKeylessClusterCommand, + -- * Re-export RedisCommands for convenience + module RedisCommandClient, + -- * Internal (exported for testing) + RedirectionInfo (..), + parseRedirectionError, + detectRedirection, + ) +where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.MVar (MVar, newMVar, putMVar, + tryTakeMVar) +import Control.Concurrent.STM (TVar, atomically, + newTVarIO, readTVarIO, + writeTVar) +import Control.Exception (SomeException, bracket, + finally, throwIO, try) +import Control.Monad (when) +import Control.Monad.IO.Class (MonadIO (..)) +import qualified Control.Monad.State as State +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Char8 as BS8 +import qualified Data.Map.Strict as Map +import Data.Time.Clock (NominalDiffTime, + diffUTCTime, + getCurrentTime) +import Data.Word (Word16) +import Database.Redis.Client (Client (..)) +import Database.Redis.Cluster (ClusterNode (..), + ClusterTopology (..), + NodeAddress (..), + NodeRole (..), + calculateSlot, + findNodeAddressForSlot, + parseClusterSlots) +import Database.Redis.Cluster.ConnectionPool (ConnectionPool, + PoolConfig (..), + closePool, createPool, + withConnection) +import Database.Redis.Command (ClientState (..), + RedisCommandClient (..), + RedisCommands (..), + convertResp, + encodeCommandBuilder, + geoRadiusFlagToList, + geoSearchByToList, + geoSearchFromToList, + geoSearchOptionToList, + geoUnitKeyword, + runRedisCommandClient, + showBS) +import qualified Database.Redis.Command as RedisCommandClient +import Database.Redis.Connector (Connector) +import Database.Redis.FromResp (FromResp (..)) +import Database.Redis.Internal.MultiplexPool (MultiplexPool, + closeMultiplexPool, + createMultiplexPool, + submitToNode, + submitToNodeWithAsking) +import Database.Redis.Resp (RespData (..)) + +-- | Error types specific to cluster operations. +data ClusterError + = MovedError Word16 NodeAddress -- ^ Permanent redirect: the slot has migrated to a different node. + | AskError Word16 NodeAddress -- ^ Temporary redirect during slot migration; retry at the given node. + | ClusterDownError String -- ^ The cluster is in a down or error state. + | TryAgainError String -- ^ Transient failure; the operation should be retried. + | CrossSlotError String -- ^ Multi-key command spans multiple hash slots. + | MaxRetriesExceeded String -- ^ All retry attempts exhausted. + | TopologyError String -- ^ Slot or node lookup failed (e.g., empty topology). + | ConnectionError String -- ^ Network-level failure connecting to a node. + deriving (Show, Eq) + +-- | Redirection information parsed from errors +data RedirectionInfo = RedirectionInfo + { redirSlot :: Word16, + redirHost :: String, + redirPort :: Int + } + deriving (Show, Eq) + +-- | Configuration for a cluster client. +data ClusterConfig = ClusterConfig + { clusterSeedNode :: NodeAddress -- ^ Initial node used to discover the cluster topology. + , clusterPoolConfig :: PoolConfig -- ^ Connection pool settings applied to every node. + , clusterMaxRetries :: Int -- ^ Maximum retry attempts on MOVED\/ASK\/transient errors (default: 3). + , clusterRetryDelay :: Int -- ^ Initial retry delay in microseconds; doubled on each retry (default: 100000 = 100ms). + , clusterTopologyRefreshInterval :: Int -- ^ Seconds between automatic background topology refreshes (default: 600 = 10 min). + } + deriving (Show) + +-- | A cluster client that manages topology discovery, a per-node connection pool +-- (for keyless commands and topology refresh), and a multiplexer pool for +-- pipelined keyed command execution. +-- Created via 'createClusterClient' and closed with 'closeClusterClient'. +data ClusterClient client = ClusterClient + { clusterTopology :: TVar ClusterTopology, + clusterConnectionPool :: ConnectionPool client, + clusterConfig :: ClusterConfig, + clusterConnector :: Connector client, -- ^ Connector factory used for all connections + clusterRefreshLock :: MVar () -- ^ Lock to prevent concurrent topology refreshes + , clusterMultiplexPool :: MultiplexPool client -- ^ Multiplexer pool for pipelined command execution + } + +-- | Monad for executing Redis commands on a cluster +-- Wraps StateT to abstract away the client state +data ClusterCommandClient client a where + ClusterCommandClient :: (Client client) => + State.StateT (ClusterClient client) IO a + -> ClusterCommandClient client a + +-- | Run Redis commands against the cluster. This is the primary API. +-- +-- The 'ClusterCommandClient' monad implements 'RedisCommands', so you can use +-- the familiar @get@\/@set@\/@del@\/… functions with transparent cluster routing. +-- Each command routes independently to the correct node. Works for both +-- single commands and multi-command sequences. +-- +-- @ +-- -- Single command: +-- result <- runClusterCommandClient client (get \"mykey\") +-- +-- -- Multi-command sequence: +-- runClusterCommandClient client $ do +-- set \"key1\" \"val1\" +-- set \"key2\" \"val2\" +-- get \"key1\" +-- @ +runClusterCommandClient :: + (Client client) => + ClusterClient client -> + ClusterCommandClient client a -> + IO a +runClusterCommandClient client (ClusterCommandClient action) = + State.evalStateT action client + +instance (Client client) => Functor (ClusterCommandClient client) where + fmap :: (a -> b) -> ClusterCommandClient client a -> ClusterCommandClient client b + fmap f (ClusterCommandClient s) = ClusterCommandClient (fmap f s) + +instance (Client client) => Applicative (ClusterCommandClient client) where + pure :: a -> ClusterCommandClient client a + pure = ClusterCommandClient . pure + (<*>) :: ClusterCommandClient client (a -> b) -> ClusterCommandClient client a -> ClusterCommandClient client b + ClusterCommandClient f <*> ClusterCommandClient s = ClusterCommandClient (f <*> s) + +instance (Client client) => Monad (ClusterCommandClient client) where + (>>=) :: ClusterCommandClient client a -> (a -> ClusterCommandClient client b) -> ClusterCommandClient client b + ClusterCommandClient s >>= f = ClusterCommandClient (s >>= \a -> let ClusterCommandClient s' = f a in s') + +instance (Client client) => MonadIO (ClusterCommandClient client) where + liftIO :: IO a -> ClusterCommandClient client a + liftIO = ClusterCommandClient . liftIO + +instance (Client client) => State.MonadState (ClusterClient client) (ClusterCommandClient client) where + get :: ClusterCommandClient client (ClusterClient client) + get = ClusterCommandClient State.get + put :: ClusterClient client -> ClusterCommandClient client () + put = ClusterCommandClient . State.put + +instance (Client client) => MonadFail (ClusterCommandClient client) where + fail :: String -> ClusterCommandClient client a + fail = ClusterCommandClient . liftIO . Prelude.fail + +-- | Connect to the seed node, issue @CLUSTER SLOTS@, and build the initial topology. +-- Throws on failure to connect or parse the topology response. +createClusterClient :: + (Client client) => + ClusterConfig -> + Connector client -> + IO (ClusterClient client) +createClusterClient config connector = do + pool <- createPool (clusterPoolConfig config) + + -- Discover initial topology before creating TVar + let seedNode = clusterSeedNode config + response <- withConnection pool seedNode connector $ \conn -> do + let clientState = ClientState conn BS8.empty + State.evalStateT (runRedisCommandClient clusterSlots) clientState + + currentTime <- getCurrentTime + case parseClusterSlots response currentTime of + Left err -> throwIO $ userError $ "Failed to parse cluster topology: " <> err + Right initialTopology -> do + topology <- newTVarIO initialTopology + refreshLock <- newMVar () + muxPool <- createMultiplexPool connector 1 + return $ ClusterClient topology pool config connector refreshLock muxPool + +-- | Close all pooled connections across every node. +-- +-- Consider using 'withClusterClient' instead for automatic cleanup. +closeClusterClient :: (Client client) => ClusterClient client -> IO () +closeClusterClient client = do + closePool (clusterConnectionPool client) + closeMultiplexPool (clusterMultiplexPool client) + +-- | Bracket-style resource management for cluster clients. +-- +-- Creates a client, runs the given action, and ensures the client is closed +-- even if an exception occurs. Prefer this over manual 'createClusterClient' +-- and 'closeClusterClient'. +-- +-- @ +-- withClusterClient config connector $ \\client -> +-- runClusterCommandClient client $ do +-- set \"key\" \"value\" +-- get \"key\" +-- @ +withClusterClient + :: (Client client) + => ClusterConfig + -> Connector client + -> (ClusterClient client -> IO a) + -> IO a +withClusterClient config connector = + bracket (createClusterClient config connector) closeClusterClient + +-- | Refresh cluster topology by querying CLUSTER SLOTS. +-- Uses a lock to prevent thundering herd: if another thread is already +-- refreshing, this call returns immediately (the other thread's refresh +-- will update the shared topology). +refreshTopology :: + (Client client) => + ClusterClient client -> + IO () +refreshTopology client = do + acquired <- tryTakeMVar (clusterRefreshLock client) + case acquired of + Nothing -> return () -- Another thread is already refreshing + Just _ -> finally doRefresh (putMVar (clusterRefreshLock client) ()) + where + connector = clusterConnector client + doRefresh = do + let seedNode = clusterSeedNode (clusterConfig client) + response <- withConnection (clusterConnectionPool client) seedNode connector $ \conn -> do + let clientState = ClientState conn BS8.empty + State.evalStateT (runRedisCommandClient clusterSlots) clientState + + currentTime <- getCurrentTime + case parseClusterSlots response currentTime of + Left err -> throwIO $ userError $ "Failed to parse cluster topology: " <> err + Right topology -> atomically $ writeTVar (clusterTopology client) topology + +-- | Check if topology is stale and refresh if needed +-- Called before every keyed command execution. +-- Performance: ~100-500ns (non-blocking read + time check) +-- Only triggers refresh when topology is older than clusterTopologyRefreshInterval. +refreshTopologyIfStale :: + (Client client) => + ClusterClient client -> + IO () +refreshTopologyIfStale client = do + topology <- readTVarIO (clusterTopology client) + currentTime <- getCurrentTime + let timeSinceUpdate = diffUTCTime currentTime (topologyUpdateTime topology) + refreshInterval = fromIntegral (clusterTopologyRefreshInterval (clusterConfig client)) :: NominalDiffTime + when (timeSinceUpdate >= refreshInterval) $ do + refreshTopology client + +-- | Detect MOVED or ASK errors from RespData. +-- Common case (no redirect) is a constructor check plus at most a single byte +-- comparison, with zero allocation. +{-# INLINE detectRedirection #-} +detectRedirection :: RespData -> Maybe (Either RedirectionInfo RedirectionInfo) +detectRedirection (RespError msg) + | BS.length msg >= 6 -- shortest redirect: "ASK x y" needs >= 7 chars + , let w = BS.index msg 0 + = if w == 0x4D -- 'M' + then if BS.isPrefixOf "MOVED " msg + then case parseMovedAsk (BS.drop 6 msg) of + Just redir -> Just (Left redir) + Nothing -> Nothing + else Nothing + else if w == 0x41 -- 'A' + then if BS.isPrefixOf "ASK " msg + then case parseMovedAsk (BS.drop 4 msg) of + Just redir -> Just (Right redir) + Nothing -> Nothing + else Nothing + else Nothing + | otherwise = Nothing +detectRedirection _ = Nothing + +-- | Execute a command on a specific node (used for keyless commands and topology refresh) +executeOnNode :: + (Client client) => + ClusterClient client -> + NodeAddress -> + RedisCommandClient client a -> + Connector client -> + IO (Either ClusterError a) +executeOnNode client nodeAddr action connector = do + result <- try $ withConnection (clusterConnectionPool client) nodeAddr connector $ \conn -> do + let clientState = ClientState conn BS8.empty + State.evalStateT (runRedisCommandClient action) clientState + + case result of + Left (e :: SomeException) -> return $ Left $ ConnectionError $ show e + Right value -> return $ Right value + +-- | Execute a command that does not target a specific key (e.g., PING, AUTH, FLUSHALL). +-- Routed to an arbitrary master node. +executeKeylessClusterCommand :: + (Client client) => + ClusterClient client -> + RedisCommandClient client a -> + IO (Either ClusterError a) +executeKeylessClusterCommand client action = do + let connector = clusterConnector client + topology <- readTVarIO (clusterTopology client) + let masterNodes = [node | node <- Map.elems (topologyNodes topology), nodeRole node == Master] + case masterNodes of + [] -> return $ Left $ TopologyError "No master nodes available" + (node:_) -> executeOnNode client (nodeAddress node) action connector + +-- | Retry logic with exponential backoff and topology refresh on MOVED errors +-- +-- MOVED errors trigger immediate topology refresh and retry. This ensures the client +-- quickly adapts to cluster topology changes (e.g., slot migrations, node failures). +-- +-- Performance considerations: +-- - During cluster reconfiguration, multiple MOVED errors may trigger concurrent refreshes +-- - Each refresh costs ~1-5ms (network + parsing) +-- - For clusters with high concurrency and frequent rebalancing, consider implementing +-- refresh rate limiting or deduplication to prevent refresh storms +-- +-- ASK errors follow the Redis protocol: retry at the target node with an ASKING prefix. +-- No topology refresh is needed since ASK indicates a temporary, in-progress migration. +withRetryAndRefresh :: + (Client client) => + ClusterClient client -> + Int -> -- Max retries + Int -> -- Initial delay (microseconds) + (Maybe NodeAddress -> IO (Either ClusterError a)) -> + IO (Either ClusterError a) +withRetryAndRefresh client maxRetries initialDelay action = go 0 initialDelay Nothing + where + go attempt delay askTarget + | attempt >= maxRetries = return $ Left $ MaxRetriesExceeded $ "Max retries (" ++ show maxRetries ++ ") exceeded" + | otherwise = do + result <- action askTarget + case result of + Left (TryAgainError _) -> do + threadDelay delay + go (attempt + 1) (delay * 2) Nothing + Left (MovedError _ _) -> do + refreshTopology client + go (attempt + 1) delay Nothing + Left (AskError _ addr) -> do + go (attempt + 1) delay (Just addr) + Left (ConnectionError _) -> do + refreshTopology client + go (attempt + 1) delay Nothing + Left err -> return $ Left err + Right value -> return $ Right value + +-- | Parse the payload after "MOVED " or "ASK " prefix. +-- Input format: "3999 127.0.0.1:6381" (slot, space, host:port) +-- Avoids BS8.words allocation by using break/drop directly. +{-# INLINE parseMovedAsk #-} +parseMovedAsk :: ByteString -> Maybe RedirectionInfo +parseMovedAsk rest = + case BS8.readInt rest of + Just (slot, afterSlot) + | not (BS8.null afterSlot) + , BS8.head afterSlot == ' ' + -> let hostPort = BS8.tail afterSlot + in case BS8.break (== ':') hostPort of + (host, portPart) + | not (BS8.null portPart) + -> case BS8.readInt (BS8.tail portPart) of + Just (port, rest') + | BS8.null rest' + -> Just $ RedirectionInfo (fromIntegral slot) (BS8.unpack host) port + _ -> Nothing + _ -> Nothing + _ -> Nothing + +-- | Parse redirection error messages (backward-compatible wrapper). +-- Format: "MOVED 3999 127.0.0.1:6381" or "ASK 3999 127.0.0.1:6381" +parseRedirectionError :: ByteString -> ByteString -> Maybe RedirectionInfo +parseRedirectionError errorType msg + | BS.isPrefixOf errorType msg + , BS.length msg > BS.length errorType + , BS.index msg (BS.length errorType) == 0x20 -- ' ' + = parseMovedAsk (BS.drop (BS.length errorType + 1) msg) + | otherwise = Nothing + +-- | Internal helper to execute a keyless command within ClusterCommandClient monad +executeKeylessCommand :: + (Client client) => + RedisCommandClient client a -> + ClusterCommandClient client (Either ClusterError a) +executeKeylessCommand action = do + client <- State.get + liftIO $ executeKeylessClusterCommand client action + +-- | Helper to unwrap Either ClusterError or fail +unwrapClusterResult :: (Client client) => Either ClusterError a -> ClusterCommandClient client a +unwrapClusterResult (Right a) = pure a +unwrapClusterResult (Left err) = Prelude.fail $ "Cluster error: " ++ show err + +-- | Execute a keyed command and unwrap the result. +-- Routes through the multiplexer pool for pipelined execution. +executeKeyed :: (Client client) => ByteString -> [ByteString] -> ClusterCommandClient client RespData +executeKeyed key cmdArgs = do + client <- State.get + result <- liftIO $ executeKeyedClusterCommand client key cmdArgs + unwrapClusterResult result + +-- | Execute a keyed command, unwrap, and convert via 'FromResp'. +executeKeyedAs :: (Client client, FromResp a) => ByteString -> [ByteString] -> ClusterCommandClient client a +executeKeyedAs key cmdArgs = executeKeyed key cmdArgs >>= convertResp + +-- | Execute a keyless command and unwrap the result +executeKeyless :: (Client client) => RedisCommandClient client a -> ClusterCommandClient client a +executeKeyless action = do + result <- executeKeylessCommand action + unwrapClusterResult result + +-- | Execute a keyed command via the multiplexer pool. +-- Pre-encodes the command to a Builder, routes by slot, and handles MOVED/ASK redirection. +-- +-- This is the low-level API for executing commands with explicit routing key. +-- For most operations, prefer 'runClusterCommandClient' with 'RedisCommands'. +executeKeyedClusterCommand :: + (Client client) => + ClusterClient client -> + ByteString -> -- key for routing + [ByteString] -> -- command args + IO (Either ClusterError RespData) +executeKeyedClusterCommand client key cmdArgs = do + refreshTopologyIfStale client + let muxPool = clusterMultiplexPool client + cmdBuilder = encodeCommandBuilder cmdArgs + !slot = calculateSlot key + withRetryAndRefresh client (clusterMaxRetries (clusterConfig client)) (clusterRetryDelay (clusterConfig client)) $ \askTarget -> + case askTarget of + Nothing -> executeOnSlotMux client muxPool slot cmdBuilder + Just addr -> executeOnNodeWithAsking client muxPool addr cmdBuilder + +-- | Execute a pre-encoded command via multiplexer on the node for a given slot. +-- Uses findNodeAddressForSlot for O(1) direct address lookup (no Map needed). +executeOnSlotMux :: + (Client client) => + ClusterClient client -> + MultiplexPool client -> + Word16 -> + Builder.Builder -> + IO (Either ClusterError RespData) +executeOnSlotMux client muxPool slot cmdBuilder = do + topology <- readTVarIO (clusterTopology client) + case findNodeAddressForSlot topology slot of + Nothing -> return $ Left $ TopologyError $ "No node found for slot " ++ show slot + Just addr -> do + result <- try $ submitToNode muxPool addr cmdBuilder + case result of + Left (e :: SomeException) -> return $ Left $ ConnectionError $ show e + Right respData -> case detectRedirection respData of + Just (Left (RedirectionInfo s host port)) -> + return $ Left $ MovedError s (NodeAddress host port) + Just (Right (RedirectionInfo s host port)) -> + return $ Left $ AskError s (NodeAddress host port) + Nothing -> return $ Right respData + +-- | Execute a command on a specific node with ASKING prefix (for ASK redirects). +-- Per Redis protocol, ASK requires sending ASKING before the actual command to the +-- target node. Both commands are submitted atomically so no other command can be +-- interleaved between them on the multiplexed connection. +executeOnNodeWithAsking :: + (Client client) => + ClusterClient client -> + MultiplexPool client -> + NodeAddress -> + Builder.Builder -> + IO (Either ClusterError RespData) +executeOnNodeWithAsking _client muxPool addr cmdBuilder = do + let askingBuilder = encodeCommandBuilder ["ASKING"] + result <- try $ submitToNodeWithAsking muxPool addr askingBuilder cmdBuilder + case result of + Left (e :: SomeException) -> return $ Left $ ConnectionError $ show e + Right respData -> case detectRedirection respData of + Just (Left (RedirectionInfo s host port)) -> + return $ Left $ MovedError s (NodeAddress host port) + Just (Right (RedirectionInfo s host port)) -> + return $ Left $ AskError s (NodeAddress host port) + Nothing -> return $ Right respData + +instance (Client client) => RedisCommands (ClusterCommandClient client) where + auth username password = executeKeyless (RedisCommandClient.auth username password) + ping = executeKeyless RedisCommandClient.ping + set k v = executeKeyedAs k ["SET", k, v] + get k = executeKeyedAs k ["GET", k] + mget keys = case keys of + [] -> executeKeyless (RedisCommandClient.mget []) + (k:_) -> executeKeyedAs k ("MGET" : keys) + setnx k v = executeKeyedAs k ["SETNX", k, v] + decr k = executeKeyedAs k ["DECR", k] + psetex k ms v = executeKeyedAs k ["PSETEX", k, showBS ms, v] + bulkSet kvs = case kvs of + [] -> executeKeyless (RedisCommandClient.bulkSet []) + ((k, _):_) -> executeKeyedAs k (["MSET"] <> concatMap (\(k', v') -> [k', v']) kvs) + flushAll = executeKeyless RedisCommandClient.flushAll + dbsize = executeKeyless RedisCommandClient.dbsize + del keys = case keys of + [] -> executeKeyless (RedisCommandClient.del []) + (k:_) -> executeKeyedAs k ("DEL" : keys) + exists keys = case keys of + [] -> executeKeyless (RedisCommandClient.exists []) + (k:_) -> executeKeyedAs k ("EXISTS" : keys) + incr k = executeKeyedAs k ["INCR", k] + hset k f v = executeKeyedAs k ["HSET", k, f, v] + hget k f = executeKeyedAs k ["HGET", k, f] + hmget k fs = executeKeyedAs k ("HMGET" : k : fs) + hexists k f = executeKeyedAs k ["HEXISTS", k, f] + lpush k vs = executeKeyedAs k ("LPUSH" : k : vs) + lrange k start stop = executeKeyedAs k ["LRANGE", k, showBS start, showBS stop] + expire k secs = executeKeyedAs k ["EXPIRE", k, showBS secs] + ttl k = executeKeyedAs k ["TTL", k] + rpush k vs = executeKeyedAs k ("RPUSH" : k : vs) + lpop k = executeKeyedAs k ["LPOP", k] + rpop k = executeKeyedAs k ["RPOP", k] + sadd k vs = executeKeyedAs k ("SADD" : k : vs) + smembers k = executeKeyedAs k ["SMEMBERS", k] + scard k = executeKeyedAs k ["SCARD", k] + sismember k v = executeKeyedAs k ["SISMEMBER", k, v] + hdel k fs = executeKeyedAs k ("HDEL" : k : fs) + hkeys k = executeKeyedAs k ["HKEYS", k] + hvals k = executeKeyedAs k ["HVALS", k] + llen k = executeKeyedAs k ["LLEN", k] + lindex k idx = executeKeyedAs k ["LINDEX", k, showBS idx] + clientSetInfo args = executeKeyless (RedisCommandClient.clientSetInfo args) + clientReply val = executeKeyless (RedisCommandClient.clientReply val) + zadd k members = + let payload = concatMap (\(score, member) -> [showBS score, member]) members + in executeKeyedAs k ("ZADD" : k : payload) + zrange k start stop withScores = + let base = ["ZRANGE", k, showBS start, showBS stop] + command = if withScores then base ++ ["WITHSCORES"] else base + in executeKeyedAs k command + geoadd k entries = + let payload = concatMap (\(lon, lat, member) -> [showBS lon, showBS lat, member]) entries + in executeKeyedAs k ("GEOADD" : k : payload) + geodist k m1 m2 unit = + let unitPart = maybe [] (\u -> [geoUnitKeyword u]) unit + in executeKeyedAs k (["GEODIST", k, m1, m2] ++ unitPart) + geohash k members = executeKeyedAs k ("GEOHASH" : k : members) + geopos k members = executeKeyedAs k ("GEOPOS" : k : members) + georadius k lon lat radius unit flags = + let base = ["GEORADIUS", k, showBS lon, showBS lat, showBS radius, geoUnitKeyword unit] + in executeKeyedAs k (base ++ concatMap geoRadiusFlagToList flags) + georadiusRo k lon lat radius unit flags = + let base = ["GEORADIUS_RO", k, showBS lon, showBS lat, showBS radius, geoUnitKeyword unit] + in executeKeyedAs k (base ++ concatMap geoRadiusFlagToList flags) + georadiusByMember k member radius unit flags = + let base = ["GEORADIUSBYMEMBER", k, member, showBS radius, geoUnitKeyword unit] + in executeKeyedAs k (base ++ concatMap geoRadiusFlagToList flags) + georadiusByMemberRo k member radius unit flags = + let base = ["GEORADIUSBYMEMBER_RO", k, member, showBS radius, geoUnitKeyword unit] + in executeKeyedAs k (base ++ concatMap geoRadiusFlagToList flags) + geosearch k fromSpec bySpec options = + executeKeyedAs k (["GEOSEARCH", k] + ++ geoSearchFromToList fromSpec + ++ geoSearchByToList bySpec + ++ concatMap geoSearchOptionToList options) + geosearchstore dest src fromSpec bySpec options storeDist = + let base = ["GEOSEARCHSTORE", dest, src] + ++ geoSearchFromToList fromSpec + ++ geoSearchByToList bySpec + ++ concatMap geoSearchOptionToList options + command = if storeDist then base ++ ["STOREDIST"] else base + in executeKeyedAs dest command + clusterSlots = executeKeyless RedisCommandClient.clusterSlots diff --git a/lib/cluster/Database/Redis/Cluster/Commands.hs b/lib/cluster/Database/Redis/Cluster/Commands.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Cluster/Commands.hs @@ -0,0 +1,119 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Shared command classification for cluster routing +-- This module provides lists of Redis commands categorized by their routing requirements +-- +-- @since 0.1.0.0 +module Database.Redis.Cluster.Commands + ( keylessCommands, + requiresKeyCommands, + CommandRouting (..), + classifyCommand, + ) +where + +import Data.ByteString (ByteString) +import qualified Data.ByteString.Char8 as BS8 +import Data.Char (toUpper) + +-- | Result of classifying a command for cluster routing +data CommandRouting + = KeylessRoute -- ^ Route to any master node + | KeyedRoute ByteString -- ^ Route by this key's hash slot + | CommandError String -- ^ Invalid command (e.g., missing required key) + +-- | Classify a Redis command for cluster routing. +-- Returns 'KeylessRoute' for commands like PING or AUTH that can go to any master, +-- 'KeyedRoute' with the routing key for commands that target a specific slot, +-- or 'CommandError' if a key-requiring command is missing its key argument. +classifyCommand :: ByteString -> [ByteString] -> CommandRouting +classifyCommand cmd args = + let cmdUpper = BS8.map toUpper cmd + in if cmdUpper `elem` keylessCommands + then KeylessRoute + else case args of + [] -> if cmdUpper `elem` requiresKeyCommands + then CommandError $ "Command " ++ BS8.unpack cmd ++ " requires a key argument" + else KeylessRoute -- Unknown command without args, try keyless + (key:_) -> KeyedRoute key + +-- | Commands that don't require a key argument (route to any master node) +-- These commands can be executed on any master node in the cluster +-- Note: This list is based on Redis 7.x commands and may need updates for newer versions +-- See CLUSTERING_IMPLEMENTATION_PLAN.md Phase 17 for future work on eliminating this hardcoded list +keylessCommands :: [ByteString] +keylessCommands = + [ "PING", + "AUTH", + "FLUSHALL", + "FLUSHDB", + "DBSIZE", + "CLUSTER", + "INFO", + "TIME", + "CLIENT", + "CONFIG", + "BGREWRITEAOF", + "BGSAVE", + "SAVE", + "LASTSAVE", + "SHUTDOWN", + "SLAVEOF", + "REPLICAOF", + "ROLE", + "ECHO", + "SELECT", + "QUIT", + "COMMAND" + ] + +-- | Commands that require a key argument (route by key's hash slot) +-- These commands must be routed to the node responsible for the key's hash slot +-- Note: This list is based on Redis 7.x commands and may need updates for newer versions +-- See CLUSTERING_IMPLEMENTATION_PLAN.md Phase 17 for future work on eliminating this hardcoded list +requiresKeyCommands :: [ByteString] +requiresKeyCommands = + [ "GET", + "SET", + "DEL", + "EXISTS", + "INCR", + "DECR", + "HGET", + "HSET", + "HDEL", + "HKEYS", + "HVALS", + "HGETALL", + "HEXISTS", + "LPUSH", + "RPUSH", + "LPOP", + "RPOP", + "LRANGE", + "LLEN", + "LINDEX", + "SADD", + "SREM", + "SMEMBERS", + "SCARD", + "SISMEMBER", + "ZADD", + "ZREM", + "ZRANGE", + "ZCARD", + "EXPIRE", + "TTL", + "PERSIST", + "MGET", + "MSET", + "SETNX", + "PSETEX", + "APPEND", + "GETRANGE", + "SETRANGE", + "STRLEN", + "GETEX", + "GETDEL", + "SETEX" + ] diff --git a/lib/cluster/Database/Redis/Cluster/ConnectionPool.hs b/lib/cluster/Database/Redis/Cluster/ConnectionPool.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Cluster/ConnectionPool.hs @@ -0,0 +1,216 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} + +-- | Thread-safe connection pool for managing Redis connections. +-- +-- Connections are created lazily and managed per-node. Each call to +-- 'withConnection' checks out an exclusive connection for the caller, +-- preventing RESP protocol interleaving between threads. Connections +-- are returned to the pool after use, or discarded if an error occurred. +-- +-- When the pool is at capacity, callers block until a connection becomes +-- available rather than creating unbounded overflow connections. +-- +-- @since 0.1.0.0 +module Database.Redis.Cluster.ConnectionPool + ( ConnectionPool (..), + PoolConfig (..), + createPool, + withConnection, + closePool, + ) +where + +import Control.Concurrent.MVar (MVar, modifyMVar, newEmptyMVar, + newMVar, putMVar, takeMVar) +import Control.Exception (SomeException, catch, throwIO, + toException, try) +import Control.Monad (forM_) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Database.Redis.Client (Client (..), ConnectionStatus (..)) +import Database.Redis.Cluster (NodeAddress (..)) +import Database.Redis.Connector (Connector) + +-- | Configuration for the connection pool. +data PoolConfig = PoolConfig + { maxConnectionsPerNode :: Int -- ^ Maximum number of connections kept per node. Callers block when all connections are in use. + , connectionTimeout :: Int -- ^ Connection timeout in seconds (reserved for future use). + , maxRetries :: Int -- ^ Maximum retry attempts for cluster operations. + , useTLS :: Bool -- ^ Whether to use TLS connections. + } + deriving (Show) + +-- | Per-node connection state: available connections, total count, and waiters +data NodePool client = NodePool + { availableConns :: [client 'Connected] -- ^ Idle connections ready for checkout + , totalConns :: !Int -- ^ Total connections created (available + in-use) + , waitQueue :: [MVar (Either SomeException (client 'Connected))] + -- ^ Threads waiting for a connection. Right = success, Left = pool error (retry). + } + +-- | Thread-safe connection pool using MVar for atomic access. +-- Each node has a pool of connections; callers check out exclusive +-- connections and return them after use. When no connections are +-- available and the pool is at capacity, callers block until one +-- is returned. +data ConnectionPool client = ConnectionPool + { poolConnections :: MVar (Map NodeAddress (NodePool client)) + , poolConfig :: PoolConfig + } + +-- | Create a new empty connection pool. +-- Connections are created lazily when first requested. +createPool :: PoolConfig -> IO (ConnectionPool client) +createPool config = do + connections <- newMVar Map.empty + return $ ConnectionPool connections config + +-- | What to do after acquiring the MVar lock +data CheckoutResult client + = UseExisting (client 'Connected) -- ^ Reuse an idle connection + | CreateNew -- ^ Create a new connection (slot reserved) + | Wait (MVar (Either SomeException (client 'Connected))) -- ^ Block until a connection is returned + +-- | Check out a connection, run an action, and return the connection to the pool. +-- If the action throws an exception, the connection is discarded (not returned) +-- since its RESP parse state may be corrupted. A fresh connection will be created +-- on the next checkout for that node. +withConnection :: + (Client client) => + ConnectionPool client -> + NodeAddress -> + Connector client -> + (client 'Connected -> IO a) -> + IO a +withConnection pool addr connector action = do + conn <- checkoutConnection pool addr connector + result <- try (action conn) + case result of + Right val -> do + returnConnection pool addr conn + return val + Left (e :: SomeException) -> do + discardConnection pool addr conn connector + throwIO e + +-- | Check out a connection from the pool. Creates a new one if none available +-- and the max hasn't been reached. Blocks if pool is at capacity. +checkoutConnection :: + (Client client) => + ConnectionPool client -> + NodeAddress -> + Connector client -> + IO (client 'Connected) +checkoutConnection pool addr connector = do + result <- modifyMVar (poolConnections pool) $ \m -> do + let nodePool = Map.findWithDefault (NodePool [] 0 []) addr m + case availableConns nodePool of + (conn : rest) -> do + let updated = nodePool { availableConns = rest } + return (Map.insert addr updated m, UseExisting conn) + [] -> + if totalConns nodePool < maxConnectionsPerNode (poolConfig pool) + then do + -- Reserve a slot, create connection outside the lock + let updated = nodePool { totalConns = totalConns nodePool + 1 } + return (Map.insert addr updated m, CreateNew) + else do + -- At capacity — enqueue a waiter + waiter <- newEmptyMVar + let updated = nodePool { waitQueue = waitQueue nodePool ++ [waiter] } + return (Map.insert addr updated m, Wait waiter) + case result of + UseExisting conn -> return conn + CreateNew -> do + -- Create connection outside the MVar lock + connResult <- try (connector addr) + case connResult of + Right conn -> return conn + Left (e :: SomeException) -> do + -- Creation failed — release the reserved slot + modifyMVar (poolConnections pool) $ \m -> do + let m' = Map.adjust (\np -> np { totalConns = totalConns np - 1 }) addr m + return (m', ()) + throwIO e + Wait waiter -> takeMVar waiter >>= either throwIO return + +-- | Return a connection to the pool for reuse. +-- If threads are waiting, hand the connection directly to the next waiter. +returnConnection :: + (Client client) => + ConnectionPool client -> + NodeAddress -> + client 'Connected -> + IO () +returnConnection pool addr conn = + modifyMVar (poolConnections pool) $ \m -> do + let nodePool = Map.findWithDefault (NodePool [] 0 []) addr m + case waitQueue nodePool of + (waiter : rest) -> do + -- Hand connection directly to a waiting thread + putMVar waiter (Right conn) + let updated = nodePool { waitQueue = rest } + return (Map.insert addr updated m, ()) + [] -> + if length (availableConns nodePool) < maxConnectionsPerNode (poolConfig pool) + then do + let updated = nodePool { availableConns = conn : availableConns nodePool } + return (Map.insert addr updated m, ()) + else do + -- Shouldn't happen, but close just in case + close conn `catch` \(_ :: SomeException) -> return () + let updated = nodePool { totalConns = totalConns nodePool - 1 } + return (Map.insert addr updated m, ()) + +-- | Discard a connection (on error) and wake a waiter or release the slot. +-- If threads are waiting, attempts to create a replacement connection. +-- If replacement creation fails, the waiter receives the error. +discardConnection :: + (Client client) => + ConnectionPool client -> + NodeAddress -> + client 'Connected -> + Connector client -> + IO () +discardConnection pool addr conn connector = do + close conn `catch` \(_ :: SomeException) -> return () + maybeWaiter <- modifyMVar (poolConnections pool) $ \m -> do + let nodePool = Map.findWithDefault (NodePool [] 0 []) addr m + case waitQueue nodePool of + (waiter : rest) -> do + -- Keep slot reserved for the waiter (don't decrement totalConns) + let updated = nodePool { waitQueue = rest } + return (Map.insert addr updated m, Just waiter) + [] -> do + -- No waiters, just release the slot + let updated = nodePool { totalConns = totalConns nodePool - 1 } + return (Map.insert addr updated m, Nothing) + case maybeWaiter of + Nothing -> return () + Just waiter -> do + -- Try to create a replacement connection for the waiter + connResult <- try (connector addr) + case connResult of + Right newConn -> putMVar waiter (Right newConn) + Left (e :: SomeException) -> do + -- Failed — release the reserved slot and notify waiter of the error + modifyMVar (poolConnections pool) $ \m -> do + let m' = Map.adjust (\np -> np { totalConns = totalConns np - 1 }) addr m + return (m', ()) + putMVar waiter (Left e) + +-- | Close all connections in the pool and wake any blocked waiters. +-- Exceptions during close are caught and ignored. +closePool :: (Client client) => ConnectionPool client -> IO () +closePool pool = + modifyMVar (poolConnections pool) $ \m -> do + let poolClosed = toException (userError "Connection pool closed") + forM_ (Map.elems m) $ \nodePool -> do + forM_ (availableConns nodePool) $ \conn -> + close conn `catch` \(_ :: SomeException) -> return () + -- Wake all blocked waiters with an error + forM_ (waitQueue nodePool) $ \waiter -> + putMVar waiter (Left poolClosed) + return (Map.empty, ()) diff --git a/lib/cluster/Database/Redis/Cluster/SlotMapping.hs b/lib/cluster/Database/Redis/Cluster/SlotMapping.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Cluster/SlotMapping.hs @@ -0,0 +1,42 @@ +{-# LANGUAGE TemplateHaskell #-} + +-- | Static mapping from Redis cluster hash slots to hash tags. +-- The mapping is fully resolved at compile time via Template Haskell, +-- so no parsing or file I/O occurs at runtime. +-- +-- @since 0.1.0.0 +module Database.Redis.Cluster.SlotMapping + ( SlotMapping + , slotMappings + ) where + +import qualified Data.ByteString.Char8 as BS8 +import Data.FileEmbed (embedFile) +import qualified Data.Map.Strict as Map +import Data.Vector (Vector) +import qualified Data.Vector as V +import Data.Word (Word16) + +-- | Type alias for the slot-to-hash-tag lookup vector. +-- Index by slot number (0–16383) to get the corresponding hash tag 'ByteString'. +type SlotMapping = Vector BS8.ByteString + +-- | Pre-computed mapping from slot number (0-16383) to hash tag. +-- The raw file is embedded at compile time; the Vector is constructed +-- once on first access and shared thereafter (top-level CAF). +{-# NOINLINE slotMappings #-} +slotMappings :: SlotMapping +slotMappings = + let raw = $(embedFile "data/cluster_slot_mapping.txt") + entries = map parseLine $ BS8.lines raw + validEntries = [(slot, tag) | Just (slot, tag) <- entries] + entryMap = Map.fromList validEntries + in V.generate 16384 (\i -> Map.findWithDefault BS8.empty (fromIntegral i) entryMap) + where + parseLine :: BS8.ByteString -> Maybe (Word16, BS8.ByteString) + parseLine line = + case BS8.words line of + [slotStr, tag] -> case reads (BS8.unpack slotStr) of + [(slot, "")] -> Just (slot, tag) + _ -> Nothing + _ -> Nothing diff --git a/lib/cluster/Database/Redis/Connector.hs b/lib/cluster/Database/Redis/Connector.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Connector.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE RankNTypes #-} + +-- | Connector factories for creating Redis connections. +-- +-- The 'Connector' type alias represents a function that creates a connected +-- client for a given 'NodeAddress'. Connector values are passed to +-- 'ClusterCommandClient.createClusterClient' and related functions. +-- +-- @ +-- import Redis +-- +-- main :: IO () +-- main = do +-- -- Standalone plaintext +-- conn <- connectPlaintext "localhost" 6379 +-- ... +-- +-- -- Cluster with TLS +-- let connector = clusterTLSConnector "redis.example.com" +-- client <- createClusterClient config connector +-- ... +-- @ +-- +-- @since 0.1.0.0 +module Database.Redis.Connector + ( -- * Connector type + Connector + -- * Standalone connections + , connectPlaintext + , connectTLS + -- * Cluster connector factories + , clusterPlaintextConnector + , clusterTLSConnector + ) where + +import Database.Redis.Client (Client (connect), + ConnectionStatus (..), + PlainTextClient (NotConnectedPlainTextClient), + TLSClient (NotConnectedTLSClient, NotConnectedTLSClientWithHostname)) +import Database.Redis.Cluster (NodeAddress (..)) + +-- | A function that creates a connected client for a given node address. +-- Used throughout the cluster layer to establish connections on demand. +type Connector client = NodeAddress -> IO (client 'Connected) + +-- | Connect a plaintext client to a specific host and port. +-- +-- @ +-- conn <- connectPlaintext "localhost" 6379 +-- @ +connectPlaintext :: String -> Int -> IO (PlainTextClient 'Connected) +connectPlaintext host port = + connect $ NotConnectedPlainTextClient host (Just port) + +-- | Connect a TLS client to a specific host and port. +-- +-- @ +-- conn <- connectTLS "redis.example.com" 6380 +-- @ +connectTLS :: String -> Int -> IO (TLSClient 'Connected) +connectTLS host port = + connect $ NotConnectedTLSClient host (Just port) + +-- | Create a cluster connector for plaintext connections. +-- Each cluster node will be connected to using its advertised address. +-- +-- @ +-- let connector = clusterPlaintextConnector +-- client <- createClusterClient config connector +-- @ +clusterPlaintextConnector :: Connector PlainTextClient +clusterPlaintextConnector addr = + connect $ NotConnectedPlainTextClient (nodeHost addr) (Just $ nodePort addr) + +-- | Create a cluster connector for TLS connections. +-- The @certHostname@ is used for TLS certificate validation, while each +-- node's advertised address is used for the network connection. This is +-- needed because @CLUSTER SLOTS@ often returns IP addresses that don't +-- match the TLS certificate's hostname. +-- +-- @ +-- let connector = clusterTLSConnector "redis.example.com" +-- client <- createClusterClient config connector +-- @ +clusterTLSConnector :: String -> Connector TLSClient +clusterTLSConnector certHostname addr = + connect $ NotConnectedTLSClientWithHostname certHostname (nodeHost addr) (Just $ nodePort addr) diff --git a/lib/cluster/Database/Redis/Internal/MultiplexPool.hs b/lib/cluster/Database/Redis/Internal/MultiplexPool.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Internal/MultiplexPool.hs @@ -0,0 +1,241 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} + +-- | Pool of 'Multiplexer's for cluster-mode usage. +-- +-- Manages one or more multiplexed connections per node, matching the +-- StackExchange.Redis architecture. Automatically reconnects dead multiplexers. +-- +-- @ +-- pool <- createMultiplexPool connector 1 +-- resp <- submitToNode pool nodeAddr cmdBytes +-- closeMultiplexPool pool +-- @ +-- +-- @since 0.1.0.0 +module Database.Redis.Internal.MultiplexPool + ( MultiplexPool + , createMultiplexPool + , submitToNode + , submitToNodeWithAsking + , submitToNodeAsync + , waitSlotResult + , closeMultiplexPool + ) where + +import Control.Concurrent.MVar (MVar, modifyMVar, newMVar) +import Control.Exception (SomeException, catch, + throwIO) +import qualified Data.ByteString.Builder as Builder +import Data.IORef (IORef, atomicModifyIORef', + atomicWriteIORef, + newIORef, readIORef) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Vector (Vector) +import qualified Data.Vector as V +import Database.Redis.Client (Client (..)) +import Database.Redis.Cluster (NodeAddress (..)) +import Database.Redis.Connector (Connector) +import Database.Redis.Internal.Multiplexer (Multiplexer, ResponseSlot, + SlotPool, + createMultiplexer, + createSlotPool, + destroyMultiplexer, + isMultiplexerAlive, + submitCommandAsync, + submitCommandPairPooled, + submitCommandPooled, + waitSlot) +import Database.Redis.Resp (RespData) + +-- | Per-node multiplexer group with its own round-robin counter. +-- Keeping the counter per-node eliminates cross-node CAS contention +-- on the shared counter that existed before. +data NodeMuxes = NodeMuxes + { nmMuxes :: !(Vector Multiplexer) + , nmCounter :: !(IORef Int) + } + +-- | A pool of multiplexers, N per node address (round-robin selected). +-- Uses IORef for fast lock-free reads on the hot path, +-- with MVar protecting creation/replacement (exclusive writes). +-- Includes a per-pool SlotPool for ResponseSlot reuse. +data MultiplexPool client = MultiplexPool + { poolNodesRef :: !(IORef (Map NodeAddress NodeMuxes)) -- fast reads + , poolNodesLock :: !(MVar ()) -- protects writes + , poolConnector :: !(Connector client) + , poolSlotPool :: !SlotPool -- reusable ResponseSlots + , poolMuxCount :: !Int -- multiplexers per node + } + +-- | Create a new empty multiplexer pool. +-- Multiplexers are created lazily when a node is first accessed. +-- @muxCount@ controls how many multiplexers are created per node. +createMultiplexPool + :: (Client client) + => Connector client + -> Int + -> IO (MultiplexPool client) +createMultiplexPool connector muxCnt = do + nodesRef <- newIORef Map.empty + nodesLock <- newMVar () + slotPool <- createSlotPool 256 + return $ MultiplexPool nodesRef nodesLock connector slotPool (max 1 muxCnt) + +-- | Submit a pre-encoded RESP command (as a Builder) to the multiplexer for a given node. +-- Creates the multiplexer on demand if the node hasn't been seen before. +-- On submission failure (dead multiplexer), replaces it and retries once. +submitToNode + :: (Client client) + => MultiplexPool client + -> NodeAddress + -> Builder.Builder + -> IO RespData +submitToNode pool addr cmdBuilder = do + mux <- getMultiplexer pool addr + submitCommandPooled (poolSlotPool pool) mux cmdBuilder + `catch` \(e :: SomeException) -> do + -- Multiplexer may be dead; try to replace and retry once + alive <- isMultiplexerAlive mux + if alive + then throwIO e -- mux is alive, error is something else + else do + newMux <- replaceMux pool addr mux + submitCommandPooled (poolSlotPool pool) newMux cmdBuilder +{-# INLINE submitToNode #-} + +-- | Submit an ASKING command followed by a real command atomically to a node. +-- Both commands are enqueued in a single atomic operation so no other command +-- can be interleaved between them on the same connection. The ASKING response +-- is discarded; only the real command's response is returned. +submitToNodeWithAsking + :: (Client client) + => MultiplexPool client + -> NodeAddress + -> Builder.Builder -- ^ ASKING command builder + -> Builder.Builder -- ^ The actual command builder + -> IO RespData +submitToNodeWithAsking pool addr askingBuilder cmdBuilder = do + mux <- getMultiplexer pool addr + submitCommandPairPooled (poolSlotPool pool) mux askingBuilder cmdBuilder + `catch` \(e :: SomeException) -> do + alive <- isMultiplexerAlive mux + if alive + then throwIO e + else do + newMux <- replaceMux pool addr mux + submitCommandPairPooled (poolSlotPool pool) newMux askingBuilder cmdBuilder +{-# INLINE submitToNodeWithAsking #-} + +-- | Async version of submitToNode: enqueue the command and return a ResponseSlot. +-- Caller must later call 'waitSlotResult' to get the response. +submitToNodeAsync + :: (Client client) + => MultiplexPool client + -> NodeAddress + -> Builder.Builder + -> IO ResponseSlot +submitToNodeAsync pool addr cmdBuilder = do + mux <- getMultiplexer pool addr + submitCommandAsync (poolSlotPool pool) mux cmdBuilder +{-# INLINE submitToNodeAsync #-} + +-- | Wait for an async submission's result and release the slot. +waitSlotResult :: MultiplexPool client -> ResponseSlot -> IO RespData +waitSlotResult pool slot = waitSlot (poolSlotPool pool) slot +{-# INLINE waitSlotResult #-} + +-- | Get or create a multiplexer for a node, round-robin among N muxes. +-- Uses readIORef for the common path (lock-free, no MVar overhead). +-- Per-node counter eliminates cross-node CAS contention. +-- When only 1 mux per node, skips the counter entirely. +getMultiplexer + :: (Client client) + => MultiplexPool client + -> NodeAddress + -> IO Multiplexer +getMultiplexer pool addr = do + m <- readIORef (poolNodesRef pool) + case Map.lookup addr m of + Just nm -> pickMux nm + Nothing -> modifyMVar (poolNodesLock pool) $ \() -> do + -- Double-check after acquiring lock + m' <- readIORef (poolNodesRef pool) + case Map.lookup addr m' of + Just nm -> do + mux <- pickMux nm + return ((), mux) + Nothing -> do + nm <- createNodeMuxes (poolConnector pool) addr (poolMuxCount pool) + atomicWriteIORef (poolNodesRef pool) (Map.insert addr nm m') + return ((), V.head (nmMuxes nm)) +{-# INLINE getMultiplexer #-} + +-- | Pick a multiplexer from a NodeMuxes using round-robin. +-- Fast path: single mux skips atomic counter entirely. +pickMux :: NodeMuxes -> IO Multiplexer +pickMux nm + | V.length (nmMuxes nm) == 1 = return $! V.unsafeHead (nmMuxes nm) + | otherwise = do + idx <- atomicModifyIORef' (nmCounter nm) (\n -> (n + 1, n)) + return $! nmMuxes nm `V.unsafeIndex` (idx `mod` V.length (nmMuxes nm)) +{-# INLINE pickMux #-} + +-- | Create N multiplexers for a node address, bundled with a per-node counter. +createNodeMuxes + :: (Client client) + => Connector client + -> NodeAddress + -> Int + -> IO NodeMuxes +createNodeMuxes connector addr count = do + muxes <- V.generateM count $ \_ -> do + conn <- connector addr + createMultiplexer conn (receive conn) + counter <- newIORef 1 + return $ NodeMuxes muxes counter + +-- | Replace a dead multiplexer for a node. +replaceMux + :: (Client client) + => MultiplexPool client + -> NodeAddress + -> Multiplexer + -> IO Multiplexer +replaceMux pool addr oldMux = do + destroyMultiplexer oldMux `catch` \(_ :: SomeException) -> return () + modifyMVar (poolNodesLock pool) $ \() -> do + m <- readIORef (poolNodesRef pool) + case Map.lookup addr m of + Just nm -> do + -- Find and replace the dead mux in the vector + newMuxes <- V.mapM (\mux -> do + alive <- isMultiplexerAlive mux + if alive + then return mux + else do + conn <- (poolConnector pool) addr + createMultiplexer conn (receive conn) + ) (nmMuxes nm) + let nm' = nm { nmMuxes = newMuxes } + atomicWriteIORef (poolNodesRef pool) (Map.insert addr nm' m) + -- Return the first alive one + mux <- pickMux nm' + return ((), mux) + Nothing -> do + nm <- createNodeMuxes (poolConnector pool) addr (poolMuxCount pool) + atomicWriteIORef (poolNodesRef pool) (Map.insert addr nm m) + return ((), V.head (nmMuxes nm)) + +-- | Tear down all multiplexers across all nodes. +closeMultiplexPool + :: MultiplexPool client + -> IO () +closeMultiplexPool pool = do + modifyMVar (poolNodesLock pool) $ \() -> do + m <- readIORef (poolNodesRef pool) + mapM_ (\nm -> V.mapM_ (\mux -> destroyMultiplexer mux `catch` \(_ :: SomeException) -> return ()) (nmMuxes nm)) + (Map.elems m) + atomicWriteIORef (poolNodesRef pool) Map.empty + return ((), ()) diff --git a/lib/cluster/Database/Redis/Internal/Multiplexer.hs b/lib/cluster/Database/Redis/Internal/Multiplexer.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Internal/Multiplexer.hs @@ -0,0 +1,552 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} + +-- | Multiplexed command pipelining over a single Redis connection. +-- +-- A 'Multiplexer' wraps a connected client with a writer thread (batches and +-- sends commands) and a reader thread (parses responses and dispatches them +-- to callers). Multiple threads submit commands concurrently via +-- 'submitCommand'; the FIFO ordering guarantee of Redis pipelining ensures +-- correct response demultiplexing without message IDs. +-- +-- @ +-- mux <- createMultiplexer conn recv +-- resp <- submitCommand mux (encode [\"GET\", \"key\"]) +-- destroyMultiplexer mux +-- @ +-- +-- @since 0.1.0.0 +module Database.Redis.Internal.Multiplexer + ( Multiplexer + , MultiplexerException (..) + , SlotPool + , ResponseSlot + , createSlotPool + , createMultiplexer + , submitCommand + , submitCommandPooled + , submitCommandPairPooled + , submitCommandAsync + , waitSlot + , destroyMultiplexer + , isMultiplexerAlive + ) where + +import Control.Concurrent (ThreadId, forkIO, killThread, + myThreadId) +import Control.Concurrent.MVar (MVar, newEmptyMVar, takeMVar, + tryPutMVar) +import Control.Exception (Exception, SomeException, + mask_, throwIO, toException, + try) +import Control.Monad (forM_, void) +import qualified Data.Attoparsec.ByteString.Char8 as StrictParse +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Builder.Extra as Builder (toLazyByteStringWith, + untrimmedStrategy) +import qualified Data.ByteString.Lazy as LBS +import Data.IORef (IORef, atomicModifyIORef', + atomicWriteIORef, newIORef, + readIORef, writeIORef) +import Data.List (foldl') +import Data.Sequence (Seq) +import qualified Data.Sequence as Seq +import Data.Typeable (Typeable) +import qualified Data.Vector as V +import Database.Redis.Client (Client (..), + ConnectionStatus (..)) +import Database.Redis.Resp (RespData, parseRespData) +import qualified GHC.Conc as GHC (threadCapability) + +-- | Exception thrown when submitting to a dead multiplexer. +data MultiplexerException + = MultiplexerDead String + | MultiplexerParseError String + | MultiplexerConnectionClosed + deriving (Show, Typeable) + +instance Exception MultiplexerException + +-- | Response slot: an IORef for the result and an MVar for signaling. +-- The reader writes the result to the IORef, then signals the MVar. +-- The caller waits on the MVar, then reads the IORef. +-- This avoids the heavier MVar write+wakeup pattern for the result itself. +data ResponseSlot = ResponseSlot + { slotResult :: !(IORef (Maybe (Either SomeException RespData))) + , slotSignal :: !(MVar ()) + } + +-- | Striped pool of pre-allocated ResponseSlots. +-- Uses multiple IORef-based stacks indexed by capability (core) to reduce +-- CAS contention between threads on different cores. +data SlotPool = SlotPool + { spStripes :: !(V.Vector (IORef [ResponseSlot])) + , spNumStripes :: !Int + } + +-- | Create a striped pool. Each stripe gets @n `div` numStripes@ pre-allocated slots. +createSlotPool :: Int -> IO SlotPool +createSlotPool n = do + let numStripes = 16 + perStripe = max 4 (n `div` numStripes) + stripes <- V.replicateM numStripes $ do + slots <- mapM (\_ -> do + r <- newIORef Nothing + s <- newEmptyMVar + return $ ResponseSlot r s + ) [1..perStripe] + newIORef slots + return $ SlotPool stripes numStripes + +-- | Pick a stripe based on the current thread's capability. +getStripe :: SlotPool -> IO (IORef [ResponseSlot]) +getStripe sp = do + tid <- myThreadId + (cap, _) <- GHC.threadCapability tid + let !idx = cap `mod` spNumStripes sp + return $! spStripes sp V.! idx +{-# INLINE getStripe #-} + +-- | Acquire a ResponseSlot from the pool, or allocate a fresh one if empty. +-- Resets the slot's IORef to Nothing before returning. +acquireSlot :: SlotPool -> IO ResponseSlot +acquireSlot sp = do + ref <- getStripe sp + mSlot <- atomicModifyIORef' ref $ \case + [] -> ([], Nothing) + (x:xs) -> (xs, Just x) + case mSlot of + Just slot -> do + writeIORef (slotResult slot) Nothing + return slot + Nothing -> do + r <- newIORef Nothing + s <- newEmptyMVar + return $ ResponseSlot r s +{-# INLINE acquireSlot #-} + +-- | Return a ResponseSlot to the pool for reuse. +releaseSlot :: SlotPool -> ResponseSlot -> IO () +releaseSlot sp slot = do + ref <- getStripe sp + atomicModifyIORef' ref $ \xs -> (slot : xs, ()) +{-# INLINE releaseSlot #-} + +-- | A command waiting to be sent, paired with a response slot. +data PendingCommand = PendingCommand + { pcBuilder :: !Builder.Builder + , pcSlot :: !ResponseSlot + } + +-- | SPSC queue for pending response slots. +-- Writer is sole producer, reader is sole consumer. +-- Uses IORef + MVar signaling instead of STM TQueue. +data PendingQueue = PendingQueue + { pqSlots :: !(IORef (Seq ResponseSlot)) + , pqSignal :: !(MVar ()) -- signaled when new items are available + } + +newPendingQueue :: IO PendingQueue +newPendingQueue = do + slots <- newIORef Seq.empty + signal <- newEmptyMVar + return $ PendingQueue slots signal + +-- | Enqueue a Seq of response slots directly (avoids Seq.fromList conversion). +pendingEnqueueSeq :: PendingQueue -> Seq ResponseSlot -> IO () +pendingEnqueueSeq pq newSlots = do + atomicModifyIORef' (pqSlots pq) $ \s -> + (s <> newSlots, ()) + void $ tryPutMVar (pqSignal pq) () +{-# INLINE pendingEnqueueSeq #-} + +-- | Dequeue one response slot (reader thread only — single consumer). +-- Blocks if empty. +pendingDequeue :: PendingQueue -> IO ResponseSlot +pendingDequeue pq = do + mSlot <- atomicModifyIORef' (pqSlots pq) $ \s -> + case Seq.viewl s of + Seq.EmptyL -> (s, Nothing) + x Seq.:< xs -> (xs, Just x) + case mSlot of + Just slot -> return slot + Nothing -> do + takeMVar (pqSignal pq) + pendingDequeue pq +{-# INLINE pendingDequeue #-} + +-- | Non-blocking dequeue of up to N response slots. +-- Returns empty Seq if none available. +pendingDequeueUpTo :: PendingQueue -> Int -> IO (Seq ResponseSlot) +pendingDequeueUpTo pq n = do + atomicModifyIORef' (pqSlots pq) $ \s -> + let (taken, rest) = Seq.splitAt n s + in (rest, taken) +{-# INLINE pendingDequeueUpTo #-} + +-- | Drain all pending slots (for error propagation). +pendingDrainAll :: PendingQueue -> IO [ResponseSlot] +pendingDrainAll pq = do + slots <- atomicModifyIORef' (pqSlots pq) $ \s -> (Seq.empty, s) + return $ foldr (:) [] slots + +-- | Lock-free MPSC (multi-producer, single-consumer) command queue. +-- Producers use atomicModifyIORef' to cons onto the list (single CAS). +-- The consumer reverses once per drain. MVar signals new item availability. +data CommandQueue = CommandQueue + { cqItems :: !(IORef [PendingCommand]) -- reverse order (newest first) + , cqSignal :: !(MVar ()) -- wake writer when items available + } + +newCommandQueue :: IO CommandQueue +newCommandQueue = do + items <- newIORef [] + signal <- newEmptyMVar + return $ CommandQueue items signal + +-- | Enqueue a command (caller thread — multi-producer safe). +commandEnqueue :: CommandQueue -> PendingCommand -> IO () +commandEnqueue cq pc = do + atomicModifyIORef' (cqItems cq) $ \xs -> (pc : xs, ()) + void $ tryPutMVar (cqSignal cq) () +{-# INLINE commandEnqueue #-} + +-- | Enqueue two commands atomically (caller thread — multi-producer safe). +-- Both commands are added in a single CAS so no other command can be +-- interleaved between them. The first command will appear before the second +-- in the pipeline. +commandEnqueuePair :: CommandQueue -> PendingCommand -> PendingCommand -> IO () +commandEnqueuePair cq pc1 pc2 = do + atomicModifyIORef' (cqItems cq) $ \xs -> (pc2 : pc1 : xs, ()) + void $ tryPutMVar (cqSignal cq) () +{-# INLINE commandEnqueuePair #-} + +-- | Drain all commands (writer thread only — single consumer). +-- Blocks if empty. Returns commands in submission order. +commandDrain :: CommandQueue -> IO [PendingCommand] +commandDrain cq = do + takeMVar (cqSignal cq) + batch <- atomicModifyIORef' (cqItems cq) $ \xs -> ([], xs) + return (reverse batch) + +-- | Non-blocking drain of any additional commands that have arrived. +-- Returns commands in submission order. Returns [] if none available. +commandTryDrain :: CommandQueue -> IO [PendingCommand] +commandTryDrain cq = do + batch <- atomicModifyIORef' (cqItems cq) $ \xs -> ([], xs) + case batch of + [] -> return [] + _ -> return (reverse batch) + +-- | Drain remaining commands without blocking (for cleanup). +commandFlush :: CommandQueue -> IO [PendingCommand] +commandFlush cq = do + batch <- atomicModifyIORef' (cqItems cq) $ \xs -> ([], xs) + return (reverse batch) + +-- | A multiplexer wrapping a single Redis connection. +data Multiplexer = Multiplexer + { muxCommandQueue :: !CommandQueue + , muxWriterThread :: !ThreadId + , muxReaderThread :: !ThreadId + , muxAlive :: !(IORef Bool) + } + +-- | Create a multiplexer over an already-connected client. +-- +-- Spawns a writer and reader green thread. The caller must eventually call +-- 'destroyMultiplexer' to clean up. +createMultiplexer + :: (Client client) + => client 'Connected + -> IO ByteString -- ^ Action to receive bytes from the connection + -> IO Multiplexer +createMultiplexer conn recv = do + cmdQueue <- newCommandQueue + pendingQueue <- newPendingQueue + alive <- newIORef True + + readerId <- forkIO $ readerLoop pendingQueue recv alive + writerId <- forkIO $ writerLoop cmdQueue pendingQueue conn alive + + return $ Multiplexer cmdQueue writerId readerId alive + +-- | Submit a pre-encoded RESP command as a Builder and block until the response arrives. +submitCommand :: Multiplexer -> Builder.Builder -> IO RespData +submitCommand mux cmdBuilder = do + isAlive <- readIORef (muxAlive mux) + if not isAlive + then throwIO $ MultiplexerDead "Multiplexer is not alive" + else do + resultRef <- newIORef Nothing + signal <- newEmptyMVar + let slot = ResponseSlot resultRef signal + pending = PendingCommand cmdBuilder slot + commandEnqueue (muxCommandQueue mux) pending + takeMVar signal + mResult <- readIORef resultRef + case mResult of + Just (Right resp) -> return resp + Just (Left e) -> throwIO e + Nothing -> throwIO $ MultiplexerDead "Response slot empty after signal" + +-- | Like 'submitCommand', but acquires a 'ResponseSlot' from the pool +-- instead of allocating a fresh IORef+MVar per call. +submitCommandPooled :: SlotPool -> Multiplexer -> Builder.Builder -> IO RespData +submitCommandPooled pool mux cmdBuilder = do + isAlive <- readIORef (muxAlive mux) + if not isAlive + then throwIO $ MultiplexerDead "Multiplexer is not alive" + else do + slot <- acquireSlot pool + let pending = PendingCommand cmdBuilder slot + commandEnqueue (muxCommandQueue mux) pending + takeMVar (slotSignal slot) + mResult <- readIORef (slotResult slot) + releaseSlot pool slot + case mResult of + Just (Right resp) -> return resp + Just (Left e) -> throwIO e + Nothing -> throwIO $ MultiplexerDead "Response slot empty after signal" +{-# INLINE submitCommandPooled #-} + +-- | Submit two commands atomically as a pair. Both are enqueued in a single +-- atomic operation so no other command can be interleaved between them. +-- Returns only the second command's response; the first response is discarded. +-- Used for ASKING + command sequences where ASKING must immediately precede +-- the target command on the same connection. +submitCommandPairPooled :: SlotPool -> Multiplexer -> Builder.Builder -> Builder.Builder -> IO RespData +submitCommandPairPooled pool mux firstBuilder secondBuilder = do + isAlive <- readIORef (muxAlive mux) + if not isAlive + then throwIO $ MultiplexerDead "Multiplexer is not alive" + else do + slot1 <- acquireSlot pool + slot2 <- acquireSlot pool + let pending1 = PendingCommand firstBuilder slot1 + pending2 = PendingCommand secondBuilder slot2 + commandEnqueuePair (muxCommandQueue mux) pending1 pending2 + -- Wait for and discard the first response (ASKING → +OK) + takeMVar (slotSignal slot1) + releaseSlot pool slot1 + -- Wait for the actual command response + takeMVar (slotSignal slot2) + mResult <- readIORef (slotResult slot2) + releaseSlot pool slot2 + case mResult of + Just (Right resp) -> return resp + Just (Left e) -> throwIO e + Nothing -> throwIO $ MultiplexerDead "Response slot empty after signal" +{-# INLINE submitCommandPairPooled #-} + +-- | Submit a command asynchronously: enqueue it and return the ResponseSlot. +-- The caller must later call 'waitSlot' to get the result, then 'releaseSlot'. +submitCommandAsync :: SlotPool -> Multiplexer -> Builder.Builder -> IO ResponseSlot +submitCommandAsync pool mux cmdBuilder = do + isAlive <- readIORef (muxAlive mux) + if not isAlive + then throwIO $ MultiplexerDead "Multiplexer is not alive" + else do + slot <- acquireSlot pool + let pending = PendingCommand cmdBuilder slot + commandEnqueue (muxCommandQueue mux) pending + return slot +{-# INLINE submitCommandAsync #-} + +-- | Wait for an async submission's result and release the slot back to the pool. +waitSlot :: SlotPool -> ResponseSlot -> IO RespData +waitSlot pool slot = do + takeMVar (slotSignal slot) + mResult <- readIORef (slotResult slot) + releaseSlot pool slot + case mResult of + Just (Right resp) -> return resp + Just (Left e) -> throwIO e + Nothing -> throwIO $ MultiplexerDead "Response slot empty after signal" +{-# INLINE waitSlot #-} + +-- | Tear down the multiplexer: kill both threads and fail all pending commands. +destroyMultiplexer :: Multiplexer -> IO () +destroyMultiplexer mux = mask_ $ do + atomicWriteIORef (muxAlive mux) False + killThread (muxWriterThread mux) + killThread (muxReaderThread mux) + remaining <- commandFlush (muxCommandQueue mux) + let err = MultiplexerDead "Multiplexer destroyed" + forM_ remaining $ \pc -> + failSlot (pcSlot pc) (toException err) + +-- | Check if the multiplexer's threads are still running. +isMultiplexerAlive :: Multiplexer -> IO Bool +isMultiplexerAlive = readIORef . muxAlive + +-- Writer thread: drains command queue, pushes response slots onto pending +-- queue (in IO, not STM), and sends batched bytes over the wire. +writerLoop + :: (Client client) + => CommandQueue + -> PendingQueue + -> client 'Connected + -> IORef Bool + -> IO () +writerLoop cmdQueue pendingQueue conn alive = go + where + go = do + isAlive <- readIORef alive + if not isAlive + then return () + else do + -- Drain command queue (lock-free MPSC, blocks if empty) + batch <- commandDrain cmdQueue + -- Non-blocking double-drain: pick up extra commands that arrived + extra <- commandTryDrain cmdQueue + let allCmds = batch ++ extra + + -- Single-pass: extract slots (as Seq) and build the combined Builder + let (!slots, !builder) = foldl' + (\(!sAcc, !bAcc) pc -> (sAcc Seq.|> pcSlot pc, bAcc <> pcBuilder pc)) + (Seq.empty, mempty) + allCmds + + -- Push response slots to pending queue (Seq avoids fromList conversion) + pendingEnqueueSeq pendingQueue slots + + -- Materialize with large buffer strategy and send via vectored I/O. + -- untrimmedStrategy avoids trimming/copying the final chunk. + -- 32KB initial / 64KB growth reduces chunk count vs default 4KB. + -- sendChunks uses writev(2) for zero-copy vectored I/O on plain sockets. + let !lbs = Builder.toLazyByteStringWith + (Builder.untrimmedStrategy 32768 65536) LBS.empty builder + !chunks = LBS.toChunks lbs + result <- try $ sendChunks conn chunks + case result of + Right () -> go + Left (e :: SomeException) -> do + atomicWriteIORef alive False + remaining <- pendingDrainAll pendingQueue + forM_ remaining $ \slot -> failSlot slot e + forM_ allCmds $ \pc -> failSlot (pcSlot pc) e + +-- Reader thread: pops response slots from the pending queue and fills +-- them with parsed RESP responses. When the buffer contains additional +-- data after parsing, batch-dequeues more slots and parses in a tight +-- inner loop to reduce per-response dequeue overhead. +-- Uses Attoparsec IResult directly to avoid Either allocation per response. +readerLoop + :: PendingQueue + -> IO ByteString + -> IORef Bool + -> IO () +readerLoop pendingQueue recv alive = go BS.empty + where + go !buffer = do + isAlive <- readIORef alive + if not isAlive + then return () + else do + slot <- pendingDequeue pendingQueue + feedParse slot (StrictParse.parse parseRespData buffer) + + -- Drive the incremental parser, feeding data until Done or Fail. + -- Avoids allocating Either/tuple wrappers on the hot path. + feedParse !slot (StrictParse.Done !remainder !resp) = do + writeIORef (slotResult slot) (Just (Right resp)) + void $ tryPutMVar (slotSignal slot) () + -- If there's remaining data, try to parse more in a tight loop + if BS.null remainder + then go remainder + else drainBuffer remainder + feedParse !slot (StrictParse.Fail _ _ err) = do + let !e = toException $ MultiplexerParseError err + writeIORef (slotResult slot) (Just (Left e)) + void $ tryPutMVar (slotSignal slot) () + atomicWriteIORef alive False + remaining <- pendingDrainAll pendingQueue + forM_ remaining $ \s -> failSlot s e + feedParse !slot (StrictParse.Partial cont) = do + moreResult <- try recv + case moreResult of + Left (e :: SomeException) -> do + writeIORef (slotResult slot) (Just (Left e)) + void $ tryPutMVar (slotSignal slot) () + atomicWriteIORef alive False + remaining <- pendingDrainAll pendingQueue + forM_ remaining $ \s -> failSlot s e + Right moreData + | BS.null moreData -> do + let !e = toException MultiplexerConnectionClosed + writeIORef (slotResult slot) (Just (Left e)) + void $ tryPutMVar (slotSignal slot) () + atomicWriteIORef alive False + remaining <- pendingDrainAll pendingQueue + forM_ remaining $ \s -> failSlot s e + | otherwise -> feedParse slot (cont moreData) + + -- Tight inner loop: buffer has data, grab available slots and parse + -- without blocking on empty queue. Falls back to outer loop when + -- no more slots are available or buffer is exhausted. + drainBuffer !buffer = do + isAlive <- readIORef alive + if not isAlive + then return () + else do + extraSlots <- pendingDequeueUpTo pendingQueue 128 + case Seq.viewl extraSlots of + Seq.EmptyL -> go buffer -- no slots ready, back to outer loop + firstSlot Seq.:< restSlots -> + fillSlots buffer firstSlot restSlots + + -- Parse and fill slots one at a time from the batch. + -- Uses Attoparsec IResult directly (no Either wrapper). + fillSlots !buffer !slot !remaining = + feedParseBatch slot remaining (StrictParse.parse parseRespData buffer) + + feedParseBatch !slot !remaining (StrictParse.Done !remainder !resp) = do + writeIORef (slotResult slot) (Just (Right resp)) + void $ tryPutMVar (slotSignal slot) () + case Seq.viewl remaining of + Seq.EmptyL -> + if BS.null remainder + then go remainder + else drainBuffer remainder + nextSlot Seq.:< restSlots -> + feedParseBatch nextSlot restSlots (StrictParse.parse parseRespData remainder) + feedParseBatch !slot !remaining (StrictParse.Fail _ _ err) = do + let !e = toException $ MultiplexerParseError err + writeIORef (slotResult slot) (Just (Left e)) + void $ tryPutMVar (slotSignal slot) () + atomicWriteIORef alive False + forM_ remaining $ \s -> failSlot s e + queued <- pendingDrainAll pendingQueue + forM_ queued $ \s -> failSlot s e + feedParseBatch !slot !remaining (StrictParse.Partial cont) = do + moreResult <- try recv + case moreResult of + Left (e :: SomeException) -> do + writeIORef (slotResult slot) (Just (Left e)) + void $ tryPutMVar (slotSignal slot) () + atomicWriteIORef alive False + forM_ remaining $ \s -> failSlot s e + queued <- pendingDrainAll pendingQueue + forM_ queued $ \s -> failSlot s e + Right moreData + | BS.null moreData -> do + let !e = toException MultiplexerConnectionClosed + writeIORef (slotResult slot) (Just (Left e)) + void $ tryPutMVar (slotSignal slot) () + atomicWriteIORef alive False + forM_ remaining $ \s -> failSlot s e + queued <- pendingDrainAll pendingQueue + forM_ queued $ \s -> failSlot s e + | otherwise -> feedParseBatch slot remaining (cont moreData) + +-- | Fail a response slot with an exception. +failSlot :: ResponseSlot -> SomeException -> IO () +failSlot slot e = do + writeIORef (slotResult slot) (Just (Left e)) + void $ tryPutMVar (slotSignal slot) () +{-# INLINE failSlot #-} + diff --git a/lib/cluster/Database/Redis/Standalone.hs b/lib/cluster/Database/Redis/Standalone.hs new file mode 100644 --- /dev/null +++ b/lib/cluster/Database/Redis/Standalone.hs @@ -0,0 +1,304 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Standalone multiplexed Redis client. +-- +-- Wraps a single 'Multiplexer' for standalone (non-cluster) Redis, providing +-- pipelined throughput without cluster mode. Implements 'RedisCommands' so all +-- existing commands work transparently. +-- +-- Use 'withStandaloneClient' for automatic resource management (recommended): +-- +-- @ +-- import Database.Redis +-- +-- main :: IO () +-- main = do +-- let config = StandaloneConfig +-- { standaloneNodeAddress = NodeAddress \"localhost\" 6379 +-- , standaloneConnector = clusterPlaintextConnector +-- , standaloneMultiplexerCount = 1 +-- } +-- withStandaloneClient config $ \\client -> +-- runStandaloneClient client $ do +-- set \"key\" \"value\" +-- result <- get \"key\" +-- ... +-- @ +-- +-- @since 0.1.0.0 +module Database.Redis.Standalone + ( -- * Configuration + StandaloneConfig (..) + , defaultStandaloneConfig + -- * Client type + , StandaloneClient + , StandaloneCommandClient + -- * Lifecycle + , createStandaloneClient + , createStandaloneClientFromConfig + , closeStandaloneClient + -- * Bracket-style lifecycle (recommended) + , withStandaloneClient + , runRedis + -- * Running commands + , runStandaloneClient + ) where + +import Control.Exception (SomeException, bracket, + catch) +import Control.Monad.IO.Class (MonadIO (..)) +import Control.Monad.Reader (ReaderT, ask, runReaderT) +import Data.ByteString (ByteString) +import Database.Redis.Client (Client (..), + PlainTextClient) +import Database.Redis.Cluster (NodeAddress (..)) +import Database.Redis.Command (ClientReplyValues (..), + RedisCommands (..), + convertResp, + encodeCommandBuilder, + geoRadiusFlagToList, + geoSearchByToList, + geoSearchFromToList, + geoSearchOptionToList, + geoUnitKeyword, showBS) +import Database.Redis.Connector (Connector, + clusterPlaintextConnector) +import Database.Redis.FromResp (FromResp (..)) +import Database.Redis.Internal.Multiplexer (Multiplexer, SlotPool, + createMultiplexer, + createSlotPool, + destroyMultiplexer, + submitCommandPooled) +import Database.Redis.Resp (RespData) + + +-- | Configuration for a standalone Redis client. +data StandaloneConfig client = StandaloneConfig + { standaloneNodeAddress :: !NodeAddress -- ^ Redis node to connect to. + , standaloneConnector :: !(Connector client) -- ^ Connection factory (plaintext or TLS). + , standaloneMultiplexerCount :: !Int -- ^ Number of multiplexers to create (default: 1). + } + +-- | Default configuration connecting to @localhost:6379@ over plaintext with 1 multiplexer. +-- +-- @ +-- runRedis defaultStandaloneConfig $ do +-- set \"key\" \"value\" +-- get \"key\" +-- @ +defaultStandaloneConfig :: StandaloneConfig PlainTextClient +defaultStandaloneConfig = StandaloneConfig + { standaloneNodeAddress = NodeAddress "localhost" 6379 + , standaloneConnector = clusterPlaintextConnector + , standaloneMultiplexerCount = 1 + } + +-- | A standalone Redis client backed by a multiplexer and slot pool. +data StandaloneClient = StandaloneClient + { standaloneMux :: !Multiplexer + , standalonePool :: !SlotPool + } + +-- | Create a standalone multiplexed client by connecting to a single Redis node. +-- This is the simple API; for more control, use 'createStandaloneClientFromConfig'. +-- +-- Consider using 'withStandaloneClient' instead for automatic cleanup. +createStandaloneClient + :: (Client client) + => Connector client + -> NodeAddress + -> IO StandaloneClient +createStandaloneClient connector addr = do + conn <- connector addr + mux <- createMultiplexer conn (receive conn) + pool <- createSlotPool 256 + return $ StandaloneClient mux pool + +-- | Create a standalone client from a 'StandaloneConfig'. +createStandaloneClientFromConfig + :: (Client client) + => StandaloneConfig client + -> IO StandaloneClient +createStandaloneClientFromConfig config = do + conn <- standaloneConnector config (standaloneNodeAddress config) + mux <- createMultiplexer conn (receive conn) + pool <- createSlotPool 256 + return $ StandaloneClient mux pool + +-- | Close the standalone client, destroying the underlying multiplexer. +-- +-- Consider using 'withStandaloneClient' instead for automatic cleanup. +closeStandaloneClient :: StandaloneClient -> IO () +closeStandaloneClient client = + destroyMultiplexer (standaloneMux client) + `catch` \(_ :: SomeException) -> return () + +-- | Bracket-style resource management for standalone clients. +-- +-- Creates a client, runs the given action, and ensures the client is closed +-- even if an exception occurs. Prefer this over manual 'createStandaloneClientFromConfig' +-- and 'closeStandaloneClient'. +-- +-- @ +-- withStandaloneClient config $ \\client -> +-- runStandaloneClient client $ do +-- set \"key\" \"value\" +-- get \"key\" +-- @ +withStandaloneClient + :: (Client client) + => StandaloneConfig client + -> (StandaloneClient -> IO a) + -> IO a +withStandaloneClient config = + bracket (createStandaloneClientFromConfig config) closeStandaloneClient + +-- | Convenience function that creates a client, runs commands, and closes +-- the client in one step. +-- +-- @ +-- result <- runRedis defaultStandaloneConfig $ do +-- set \"key\" \"value\" +-- get \"key\" +-- @ +runRedis + :: (Client client) + => StandaloneConfig client + -> StandaloneCommandClient a + -> IO a +runRedis config action = + withStandaloneClient config (`runStandaloneClient` action) + +-- | Monad for executing Redis commands on a standalone client. +newtype StandaloneCommandClient a = StandaloneCommandClient + { unStandaloneCommandClient :: ReaderT StandaloneClient IO a } + +-- | Run Redis commands against the standalone client. +runStandaloneClient :: StandaloneClient -> StandaloneCommandClient a -> IO a +runStandaloneClient client (StandaloneCommandClient action) = + runReaderT action client + +instance Functor StandaloneCommandClient where + fmap f (StandaloneCommandClient r) = StandaloneCommandClient (fmap f r) + +instance Applicative StandaloneCommandClient where + pure = StandaloneCommandClient . pure + StandaloneCommandClient f <*> StandaloneCommandClient r = StandaloneCommandClient (f <*> r) + +instance Monad StandaloneCommandClient where + StandaloneCommandClient r >>= f = StandaloneCommandClient (r >>= \a -> unStandaloneCommandClient (f a)) + +instance MonadIO StandaloneCommandClient where + liftIO = StandaloneCommandClient . liftIO + +instance MonadFail StandaloneCommandClient where + fail = StandaloneCommandClient . liftIO . Prelude.fail + +-- | Submit a command via the multiplexer backend. +submitMux :: [ByteString] -> StandaloneCommandClient RespData +submitMux args = do + client <- StandaloneCommandClient ask + let cmdBuilder = encodeCommandBuilder args + liftIO $ submitCommandPooled (standalonePool client) (standaloneMux client) cmdBuilder + +-- | Submit a command and convert the result via 'FromResp'. +submitMuxAs :: (FromResp a) => [ByteString] -> StandaloneCommandClient a +submitMuxAs args = submitMux args >>= convertResp + +instance RedisCommands StandaloneCommandClient where + auth username password = submitMuxAs ["HELLO", "3", "AUTH", username, password] + ping = submitMuxAs ["PING"] + set k v = submitMuxAs ["SET", k, v] + get k = submitMuxAs ["GET", k] + mget keys = submitMuxAs ("MGET" : keys) + setnx k v = submitMuxAs ["SETNX", k, v] + decr k = submitMuxAs ["DECR", k] + psetex k ms v = submitMuxAs ["PSETEX", k, showBS ms, v] + bulkSet kvs = submitMuxAs (["MSET"] <> concatMap (\(k, v) -> [k, v]) kvs) + flushAll = submitMuxAs ["FLUSHALL"] + dbsize = submitMuxAs ["DBSIZE"] + del keys = submitMuxAs ("DEL" : keys) + exists keys = submitMuxAs ("EXISTS" : keys) + incr k = submitMuxAs ["INCR", k] + hset k f v = submitMuxAs ["HSET", k, f, v] + hget k f = submitMuxAs ["HGET", k, f] + hmget k fs = submitMuxAs ("HMGET" : k : fs) + hexists k f = submitMuxAs ["HEXISTS", k, f] + lpush k vs = submitMuxAs ("LPUSH" : k : vs) + lrange k start stop = submitMuxAs ["LRANGE", k, showBS start, showBS stop] + expire k secs = submitMuxAs ["EXPIRE", k, showBS secs] + ttl k = submitMuxAs ["TTL", k] + rpush k vs = submitMuxAs ("RPUSH" : k : vs) + lpop k = submitMuxAs ["LPOP", k] + rpop k = submitMuxAs ["RPOP", k] + sadd k vs = submitMuxAs ("SADD" : k : vs) + smembers k = submitMuxAs ["SMEMBERS", k] + scard k = submitMuxAs ["SCARD", k] + sismember k v = submitMuxAs ["SISMEMBER", k, v] + hdel k fs = submitMuxAs ("HDEL" : k : fs) + hkeys k = submitMuxAs ["HKEYS", k] + hvals k = submitMuxAs ["HVALS", k] + llen k = submitMuxAs ["LLEN", k] + lindex k idx = submitMuxAs ["LINDEX", k, showBS idx] + clientSetInfo args = submitMuxAs (["CLIENT", "SETINFO"] ++ args) + clusterSlots = submitMuxAs ["CLUSTER", "SLOTS"] + + clientReply val = do + case val of + ON -> do + resp <- submitMux ["CLIENT", "REPLY", showBS val] + return (Just resp) + -- OFF/SKIP: Redis does not send a response, which would desync the + -- multiplexer. These are inherently incompatible with pipelined + -- multiplexing, so we silently ignore them. + _ -> return Nothing + + zadd k members = + let payload = concatMap (\(score, member) -> [showBS score, member]) members + in submitMuxAs ("ZADD" : k : payload) + + zrange k start stop withScores = + let base = ["ZRANGE", k, showBS start, showBS stop] + command = if withScores then base ++ ["WITHSCORES"] else base + in submitMuxAs command + + geoadd k entries = + let payload = concatMap (\(lon, lat, member) -> [showBS lon, showBS lat, member]) entries + in submitMuxAs ("GEOADD" : k : payload) + + geodist k m1 m2 unit = + let unitPart = maybe [] (\u -> [geoUnitKeyword u]) unit + in submitMuxAs (["GEODIST", k, m1, m2] ++ unitPart) + + geohash k members = submitMuxAs ("GEOHASH" : k : members) + geopos k members = submitMuxAs ("GEOPOS" : k : members) + + georadius k lon lat radius unit flags = + let base = ["GEORADIUS", k, showBS lon, showBS lat, showBS radius, geoUnitKeyword unit] + in submitMuxAs (base ++ concatMap geoRadiusFlagToList flags) + + georadiusRo k lon lat radius unit flags = + let base = ["GEORADIUS_RO", k, showBS lon, showBS lat, showBS radius, geoUnitKeyword unit] + in submitMuxAs (base ++ concatMap geoRadiusFlagToList flags) + + georadiusByMember k member radius unit flags = + let base = ["GEORADIUSBYMEMBER", k, member, showBS radius, geoUnitKeyword unit] + in submitMuxAs (base ++ concatMap geoRadiusFlagToList flags) + + georadiusByMemberRo k member radius unit flags = + let base = ["GEORADIUSBYMEMBER_RO", k, member, showBS radius, geoUnitKeyword unit] + in submitMuxAs (base ++ concatMap geoRadiusFlagToList flags) + + geosearch k fromSpec bySpec options = + submitMuxAs (["GEOSEARCH", k] + ++ geoSearchFromToList fromSpec + ++ geoSearchByToList bySpec + ++ concatMap geoSearchOptionToList options) + + geosearchstore dest src fromSpec bySpec options storeDist = + let base = ["GEOSEARCHSTORE", dest, src] + ++ geoSearchFromToList fromSpec + ++ geoSearchByToList bySpec + ++ concatMap geoSearchOptionToList options + command = if storeDist then base ++ ["STOREDIST"] else base + in submitMuxAs command diff --git a/lib/crc16/Database/Redis/Crc16.hs b/lib/crc16/Database/Redis/Crc16.hs new file mode 100644 --- /dev/null +++ b/lib/crc16/Database/Redis/Crc16.hs @@ -0,0 +1,27 @@ +-- | CRC16 hash computation for Redis cluster slot assignment. +-- Wraps a C implementation via FFI. The FFI call is safe to treat as pure +-- (deterministic, no side effects), so we use 'unsafeDupablePerformIO' to +-- avoid forcing callers into IO on the hot path. +-- +-- @since 0.1.0.0 +module Database.Redis.Crc16 ( crc16 ) where + +import Data.Bits ((.&.)) +import Data.ByteString (ByteString) +import qualified Data.ByteString.Char8 as BS8 +import Data.Word (Word16) +import Foreign.C.String (CString) +import Foreign.C.Types (CInt (..), CUShort (..)) +import System.IO.Unsafe (unsafeDupablePerformIO) + +-- Foreign function interface to the C function (marked unsafe for speed — no callbacks) +foreign import ccall unsafe "crc16" c_crc16 :: CString -> CInt -> CUShort + +-- | Compute the CRC16 hash of a key, reduced to the Redis cluster slot range (0–16383). +-- Pure: the underlying C function is deterministic with no side effects. +crc16 :: ByteString -> Word16 +crc16 bs = unsafeDupablePerformIO $ + BS8.useAsCStringLen bs $ \(cstr, len) -> do + let !result = c_crc16 cstr (fromIntegral len) + return $! fromIntegral result .&. (0x3FFF :: Word16) +{-# INLINE crc16 #-} diff --git a/lib/crc16/crc16.c b/lib/crc16/crc16.c new file mode 100644 --- /dev/null +++ b/lib/crc16/crc16.c @@ -0,0 +1,86 @@ +/* + * Copyright 2001-2010 Georges Menie (www.menie.org) + * Copyright 2010 Salvatore Sanfilippo (adapted to Redis coding style) + * All rights reserved. + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of the University of California, Berkeley nor the + * names of its contributors may be used to endorse or promote products + * derived from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE REGENTS AND CONTRIBUTORS BE LIABLE FOR ANY + * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +/* CRC16 implementation according to CCITT standards. + * + * Note by @antirez: this is actually the XMODEM CRC 16 algorithm, using the + * following parameters: + * + * Name : "XMODEM", also known as "ZMODEM", "CRC-16/ACORN" + * Width : 16 bit + * Poly : 1021 (That is actually x^16 + x^12 + x^5 + 1) + * Initialization : 0000 + * Reflect Input byte : False + * Reflect Output CRC : False + * Xor constant to output CRC : 0000 + * Output for "123456789" : 31C3 + */ +#include + +static const uint16_t crc16tab[256]= { + 0x0000,0x1021,0x2042,0x3063,0x4084,0x50a5,0x60c6,0x70e7, + 0x8108,0x9129,0xa14a,0xb16b,0xc18c,0xd1ad,0xe1ce,0xf1ef, + 0x1231,0x0210,0x3273,0x2252,0x52b5,0x4294,0x72f7,0x62d6, + 0x9339,0x8318,0xb37b,0xa35a,0xd3bd,0xc39c,0xf3ff,0xe3de, + 0x2462,0x3443,0x0420,0x1401,0x64e6,0x74c7,0x44a4,0x5485, + 0xa56a,0xb54b,0x8528,0x9509,0xe5ee,0xf5cf,0xc5ac,0xd58d, + 0x3653,0x2672,0x1611,0x0630,0x76d7,0x66f6,0x5695,0x46b4, + 0xb75b,0xa77a,0x9719,0x8738,0xf7df,0xe7fe,0xd79d,0xc7bc, + 0x48c4,0x58e5,0x6886,0x78a7,0x0840,0x1861,0x2802,0x3823, + 0xc9cc,0xd9ed,0xe98e,0xf9af,0x8948,0x9969,0xa90a,0xb92b, + 0x5af5,0x4ad4,0x7ab7,0x6a96,0x1a71,0x0a50,0x3a33,0x2a12, + 0xdbfd,0xcbdc,0xfbbf,0xeb9e,0x9b79,0x8b58,0xbb3b,0xab1a, + 0x6ca6,0x7c87,0x4ce4,0x5cc5,0x2c22,0x3c03,0x0c60,0x1c41, + 0xedae,0xfd8f,0xcdec,0xddcd,0xad2a,0xbd0b,0x8d68,0x9d49, + 0x7e97,0x6eb6,0x5ed5,0x4ef4,0x3e13,0x2e32,0x1e51,0x0e70, + 0xff9f,0xefbe,0xdfdd,0xcffc,0xbf1b,0xaf3a,0x9f59,0x8f78, + 0x9188,0x81a9,0xb1ca,0xa1eb,0xd10c,0xc12d,0xf14e,0xe16f, + 0x1080,0x00a1,0x30c2,0x20e3,0x5004,0x4025,0x7046,0x6067, + 0x83b9,0x9398,0xa3fb,0xb3da,0xc33d,0xd31c,0xe37f,0xf35e, + 0x02b1,0x1290,0x22f3,0x32d2,0x4235,0x5214,0x6277,0x7256, + 0xb5ea,0xa5cb,0x95a8,0x8589,0xf56e,0xe54f,0xd52c,0xc50d, + 0x34e2,0x24c3,0x14a0,0x0481,0x7466,0x6447,0x5424,0x4405, + 0xa7db,0xb7fa,0x8799,0x97b8,0xe75f,0xf77e,0xc71d,0xd73c, + 0x26d3,0x36f2,0x0691,0x16b0,0x6657,0x7676,0x4615,0x5634, + 0xd94c,0xc96d,0xf90e,0xe92f,0x99c8,0x89e9,0xb98a,0xa9ab, + 0x5844,0x4865,0x7806,0x6827,0x18c0,0x08e1,0x3882,0x28a3, + 0xcb7d,0xdb5c,0xeb3f,0xfb1e,0x8bf9,0x9bd8,0xabbb,0xbb9a, + 0x4a75,0x5a54,0x6a37,0x7a16,0x0af1,0x1ad0,0x2ab3,0x3a92, + 0xfd2e,0xed0f,0xdd6c,0xcd4d,0xbdaa,0xad8b,0x9de8,0x8dc9, + 0x7c26,0x6c07,0x5c64,0x4c45,0x3ca2,0x2c83,0x1ce0,0x0cc1, + 0xef1f,0xff3e,0xcf5d,0xdf7c,0xaf9b,0xbfba,0x8fd9,0x9ff8, + 0x6e17,0x7e36,0x4e55,0x5e74,0x2e93,0x3eb2,0x0ed1,0x1ef0 +}; + +uint16_t crc16(const char *buf, int len) { + int counter; + uint16_t crc = 0; + for (counter = 0; counter < len; counter++) + crc = (crc<<8) ^ crc16tab[((crc>>8) ^ *buf++)&0x00FF]; + return crc; +} diff --git a/lib/redis-command-client/Database/Redis/Command.hs b/lib/redis-command-client/Database/Redis/Command.hs new file mode 100644 --- /dev/null +++ b/lib/redis-command-client/Database/Redis/Command.hs @@ -0,0 +1,444 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | High-level Redis command interface built on top of 'Client'. +-- +-- Provides the 'RedisCommands' typeclass with methods for standard Redis commands +-- (strings, hashes, lists, sets, sorted sets, geo), a 'RedisCommandClient' monad +-- that manages connection state and incremental RESP parsing, and typed error handling +-- via 'RedisError'. +-- +-- @since 0.1.0.0 +module Database.Redis.Command + ( -- * Core types + ClientState (..) + , RedisCommandClient (..) + , RedisCommands (..) + , ClientReplyValues (..) + -- * Errors + , RedisError (..) + -- * Geo types + , GeoUnit (..) + , GeoRadiusFlag (..) + , GeoSearchFrom (..) + , GeoSearchBy (..) + , GeoSearchOption (..) + -- * Helpers + , wrapInRay + , encodeCommand + , encodeCommandBuilder + , encodeSetBuilder + , encodeGetBuilder + , encodeBulkArg + , showBS + , geoUnitKeyword + , geoRadiusFlagToList + , geoSearchFromToList + , geoSearchByToList + , geoSearchOptionToList + -- * Parsing + , parseWith + , parseManyWith + -- * FromResp conversion + , convertResp + ) where + +import Control.Exception (throwIO) +import Control.Monad.IO.Class (MonadIO (..)) +import Control.Monad.State as State (MonadState (get, put), + StateT) +import qualified Data.Attoparsec.ByteString.Char8 as StrictParse +import Data.ByteString (ByteString) +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Char8 as BS8 +import qualified Data.ByteString.Lazy as LBS +import Data.Kind (Type) +import Database.Redis.Client (Client (..), + ConnectionStatus (..)) +import Database.Redis.FromResp (FromResp (..)) +import Database.Redis.RedisError (RedisError (..)) +import Database.Redis.Resp (Encodable (encode), + RespData (..), parseRespData) + + +-- | Mutable state carried through a 'RedisCommandClient' session: the live connection +-- and an unparsed RESP byte buffer from previous receives. +data ClientState client = ClientState + { getClient :: client 'Connected, + getParseBuffer :: BS8.ByteString + } + +-- | A monad for sequencing Redis commands over a single connection. +-- Wraps 'StateT' over 'ClientState' to manage the connection handle and +-- an incremental parse buffer, so callers never deal with raw bytes. +data RedisCommandClient client (a :: Type) where + RedisCommandClient :: (Client client) => {runRedisCommandClient :: State.StateT (ClientState client) IO a} -> RedisCommandClient client a + +instance (Client client) => Functor (RedisCommandClient client) where + fmap :: (a -> b) -> RedisCommandClient client a -> RedisCommandClient client b + fmap f (RedisCommandClient s) = RedisCommandClient (fmap f s) + +instance (Client client) => Applicative (RedisCommandClient client) where + pure :: a -> RedisCommandClient client a + pure = RedisCommandClient . pure + (<*>) :: RedisCommandClient client (a -> b) -> RedisCommandClient client a -> RedisCommandClient client b + RedisCommandClient f <*> RedisCommandClient s = RedisCommandClient (f <*> s) + +instance (Client client) => Monad (RedisCommandClient client) where + (>>=) :: RedisCommandClient client a -> (a -> RedisCommandClient client b) -> RedisCommandClient client b + RedisCommandClient s >>= f = RedisCommandClient (s >>= \a -> let RedisCommandClient s' = f a in s') + +instance (Client client) => MonadIO (RedisCommandClient client) where + liftIO :: IO a -> RedisCommandClient client a + liftIO = RedisCommandClient . liftIO + +instance (Client client) => MonadState (ClientState client) (RedisCommandClient client) where + get :: RedisCommandClient client (ClientState client) + get = RedisCommandClient State.get + put :: ClientState client -> RedisCommandClient client () + put = RedisCommandClient . State.put + +instance (Client client) => MonadFail (RedisCommandClient client) where + fail :: String -> RedisCommandClient client a + fail = RedisCommandClient . liftIO . fail + +-- | The standard set of Redis commands. Implemented for both single-node +-- ('RedisCommandClient') and cluster ('ClusterCommandClient') monads. +-- Keys and values use strict 'ByteString' to avoid O(n) String conversions. +-- Command return types are polymorphic via 'FromResp', allowing typed results. +class (MonadIO m) => RedisCommands m where + auth :: (FromResp a) => ByteString -> ByteString -> m a + ping :: (FromResp a) => m a + set :: (FromResp a) => ByteString -> ByteString -> m a + get :: (FromResp a) => ByteString -> m a + mget :: (FromResp a) => [ByteString] -> m a + setnx :: (FromResp a) => ByteString -> ByteString -> m a + decr :: (FromResp a) => ByteString -> m a + psetex :: (FromResp a) => ByteString -> Int -> ByteString -> m a + bulkSet :: (FromResp a) => [(ByteString, ByteString)] -> m a + flushAll :: (FromResp a) => m a + dbsize :: (FromResp a) => m a + del :: (FromResp a) => [ByteString] -> m a + exists :: (FromResp a) => [ByteString] -> m a + incr :: (FromResp a) => ByteString -> m a + hset :: (FromResp a) => ByteString -> ByteString -> ByteString -> m a + hget :: (FromResp a) => ByteString -> ByteString -> m a + hmget :: (FromResp a) => ByteString -> [ByteString] -> m a + hexists :: (FromResp a) => ByteString -> ByteString -> m a + lpush :: (FromResp a) => ByteString -> [ByteString] -> m a + lrange :: (FromResp a) => ByteString -> Int -> Int -> m a + expire :: (FromResp a) => ByteString -> Int -> m a + ttl :: (FromResp a) => ByteString -> m a + rpush :: (FromResp a) => ByteString -> [ByteString] -> m a + lpop :: (FromResp a) => ByteString -> m a + rpop :: (FromResp a) => ByteString -> m a + sadd :: (FromResp a) => ByteString -> [ByteString] -> m a + smembers :: (FromResp a) => ByteString -> m a + scard :: (FromResp a) => ByteString -> m a + sismember :: (FromResp a) => ByteString -> ByteString -> m a + hdel :: (FromResp a) => ByteString -> [ByteString] -> m a + hkeys :: (FromResp a) => ByteString -> m a + hvals :: (FromResp a) => ByteString -> m a + llen :: (FromResp a) => ByteString -> m a + lindex :: (FromResp a) => ByteString -> Int -> m a + clientSetInfo :: (FromResp a) => [ByteString] -> m a + clientReply :: ClientReplyValues -> m (Maybe RespData) + zadd :: (FromResp a) => ByteString -> [(Int, ByteString)] -> m a + zrange :: (FromResp a) => ByteString -> Int -> Int -> Bool -> m a + geoadd :: (FromResp a) => ByteString -> [(Double, Double, ByteString)] -> m a + geodist :: (FromResp a) => ByteString -> ByteString -> ByteString -> Maybe GeoUnit -> m a + geohash :: (FromResp a) => ByteString -> [ByteString] -> m a + geopos :: (FromResp a) => ByteString -> [ByteString] -> m a + georadius :: (FromResp a) => ByteString -> Double -> Double -> Double -> GeoUnit -> [GeoRadiusFlag] -> m a + georadiusRo :: (FromResp a) => ByteString -> Double -> Double -> Double -> GeoUnit -> [GeoRadiusFlag] -> m a + georadiusByMember :: (FromResp a) => ByteString -> ByteString -> Double -> GeoUnit -> [GeoRadiusFlag] -> m a + georadiusByMemberRo :: (FromResp a) => ByteString -> ByteString -> Double -> GeoUnit -> [GeoRadiusFlag] -> m a + geosearch :: (FromResp a) => ByteString -> GeoSearchFrom -> GeoSearchBy -> [GeoSearchOption] -> m a + geosearchstore :: (FromResp a) => ByteString -> ByteString -> GeoSearchFrom -> GeoSearchBy -> [GeoSearchOption] -> Bool -> m a + clusterSlots :: (FromResp a) => m a + +-- | Helper to convert a showable value to ByteString for use in commands. +showBS :: (Show a) => a -> ByteString +showBS = BS8.pack . show + +-- | Wrap a list of strict ByteStrings into a RESP array of bulk strings. +wrapInRay :: [ByteString] -> RespData +wrapInRay inp = + let !res = RespArray . map RespBulkString $ inp + in res + +-- | Encode a Redis command (list of arguments) into a Builder for efficient batching. +-- Used by the multiplexer to defer materialization until the writer batches commands. +encodeCommandBuilder :: [ByteString] -> Builder.Builder +encodeCommandBuilder args = + Builder.char8 '*' <> Builder.intDec (length args) <> Builder.byteString "\r\n" <> + foldMap encodeBulkArg args + +-- | Encode a single bulk string argument: $LEN\r\nDATA\r\n +{-# INLINE encodeBulkArg #-} +encodeBulkArg :: ByteString -> Builder.Builder +encodeBulkArg a = Builder.char8 '$' <> Builder.intDec (BS8.length a) <> Builder.byteString "\r\n" + <> Builder.byteString a <> Builder.byteString "\r\n" + +-- | Pre-computed RESP preamble for SET: *3\r\n$3\r\nSET\r\n +setPreamble :: Builder.Builder +setPreamble = Builder.byteString "*3\r\n$3\r\nSET\r\n" +{-# NOINLINE setPreamble #-} + +-- | Pre-computed RESP preamble for GET: *2\r\n$3\r\nGET\r\n +getPreamble :: Builder.Builder +getPreamble = Builder.byteString "*2\r\n$3\r\nGET\r\n" +{-# NOINLINE getPreamble #-} + +-- | Specialized SET encoder: avoids list construction, length, and foldMap. +{-# INLINE encodeSetBuilder #-} +encodeSetBuilder :: ByteString -> ByteString -> Builder.Builder +encodeSetBuilder key val = setPreamble <> encodeBulkArg key <> encodeBulkArg val + +-- | Specialized GET encoder: avoids list construction, length, and foldMap. +{-# INLINE encodeGetBuilder #-} +encodeGetBuilder :: ByteString -> Builder.Builder +encodeGetBuilder key = getPreamble <> encodeBulkArg key + +-- | Encode a Redis command (list of arguments) into its RESP wire format as a strict ByteString. +-- Used by the multiplexer to pre-encode commands before queuing. +encodeCommand :: [ByteString] -> ByteString +encodeCommand args = LBS.toStrict $ Builder.toLazyByteString $ encodeCommandBuilder args + +-- | Send a command and parse the response. +executeCommand :: (Client client) => [ByteString] -> RedisCommandClient client RespData +executeCommand args = do + ClientState !client _ <- State.get + liftIO $ send client (Builder.toLazyByteString . encode $ wrapInRay args) + parseWith (receive client) + +-- | Convert a raw 'RespData' value using 'FromResp', throwing on failure. +convertResp :: (FromResp a, MonadIO m) => RespData -> m a +convertResp rd = case fromResp rd of + Right a -> return a + Left err -> liftIO $ throwIO err + +-- | Execute a command and convert the result via 'FromResp'. +executeCommandAs :: (Client client, FromResp a) => [ByteString] -> RedisCommandClient client a +executeCommandAs args = executeCommand args >>= convertResp + +-- | Distance unit for Redis GEO commands. +data GeoUnit + = Meters + | Kilometers + | Miles + | Feet + deriving (Eq, Show) + +-- | Convert a 'GeoUnit' to its Redis protocol keyword. +geoUnitKeyword :: GeoUnit -> ByteString +geoUnitKeyword unit = + case unit of + Meters -> "M" + Kilometers -> "KM" + Miles -> "MI" + Feet -> "FT" + +-- | Optional flags for GEORADIUS and GEORADIUSBYMEMBER commands. +data GeoRadiusFlag + = GeoWithCoord + | GeoWithDist + | GeoWithHash + | GeoRadiusCount Int Bool -- Bool indicates whether ANY is appended + | GeoRadiusAsc + | GeoRadiusDesc + | GeoRadiusStore ByteString + | GeoRadiusStoreDist ByteString + deriving (Eq, Show) + +-- | Convert a 'GeoRadiusFlag' to its Redis protocol argument list. +geoRadiusFlagToList :: GeoRadiusFlag -> [ByteString] +geoRadiusFlagToList flag = + case flag of + GeoWithCoord -> ["WITHCOORD"] + GeoWithDist -> ["WITHDIST"] + GeoWithHash -> ["WITHHASH"] + GeoRadiusCount n useAny -> ["COUNT", showBS n] <> ["ANY" | useAny] + GeoRadiusAsc -> ["ASC"] + GeoRadiusDesc -> ["DESC"] + GeoRadiusStore key -> ["STORE", key] + GeoRadiusStoreDist key -> ["STOREDIST", key] + +-- | Origin for a GEOSEARCH query: either a longitude\/latitude pair or an existing member. +data GeoSearchFrom + = GeoFromLonLat Double Double + | GeoFromMember ByteString + deriving (Eq, Show) + +-- | Convert a 'GeoSearchFrom' to its Redis protocol argument list. +geoSearchFromToList :: GeoSearchFrom -> [ByteString] +geoSearchFromToList fromSpec = + case fromSpec of + GeoFromLonLat lon lat -> ["FROMLONLAT", showBS lon, showBS lat] + GeoFromMember member -> ["FROMMEMBER", member] + +-- | Shape for a GEOSEARCH query: circular radius or rectangular box. +data GeoSearchBy + = GeoByRadius Double GeoUnit + | GeoByBox Double Double GeoUnit + deriving (Eq, Show) + +-- | Convert a 'GeoSearchBy' to its Redis protocol argument list. +geoSearchByToList :: GeoSearchBy -> [ByteString] +geoSearchByToList bySpec = + case bySpec of + GeoByRadius radius unit -> ["BYRADIUS", showBS radius, geoUnitKeyword unit] + GeoByBox width height unit -> ["BYBOX", showBS width, showBS height, geoUnitKeyword unit] + +-- | Optional modifiers for GEOSEARCH: include coordinates, distances, hashes, +-- limit count, or sort order. +data GeoSearchOption + = GeoSearchWithCoord + | GeoSearchWithDist + | GeoSearchWithHash + | GeoSearchCount Int Bool -- Bool indicates ANY + | GeoSearchAsc + | GeoSearchDesc + deriving (Eq, Show) + +-- | Convert a 'GeoSearchOption' to its Redis protocol argument list. +geoSearchOptionToList :: GeoSearchOption -> [ByteString] +geoSearchOptionToList opt = + case opt of + GeoSearchWithCoord -> ["WITHCOORD"] + GeoSearchWithDist -> ["WITHDIST"] + GeoSearchWithHash -> ["WITHHASH"] + GeoSearchCount n useAny -> ["COUNT", showBS n] <> ["ANY" | useAny] + GeoSearchAsc -> ["ASC"] + GeoSearchDesc -> ["DESC"] + +-- | Values for the CLIENT REPLY command. +data ClientReplyValues = OFF | ON | SKIP + deriving (Eq, Show) + +instance (Client client) => RedisCommands (RedisCommandClient client) where + ping = executeCommandAs ["PING"] + set k v = executeCommandAs ["SET", k, v] + get k = executeCommandAs ["GET", k] + mget keys = executeCommandAs ("MGET" : keys) + setnx key value = executeCommandAs ["SETNX", key, value] + decr key = executeCommandAs ["DECR", key] + psetex key milliseconds value = executeCommandAs ["PSETEX", key, showBS milliseconds, value] + auth username password = executeCommandAs ["HELLO", "3", "AUTH", username, password] + bulkSet kvs = executeCommandAs (["MSET"] <> concatMap (\(k, v) -> [k, v]) kvs) + flushAll = executeCommandAs ["FLUSHALL"] + dbsize = executeCommandAs ["DBSIZE"] + del keys = executeCommandAs ("DEL" : keys) + exists keys = executeCommandAs ("EXISTS" : keys) + incr key = executeCommandAs ["INCR", key] + hset key field value = executeCommandAs ["HSET", key, field, value] + hget key field = executeCommandAs ["HGET", key, field] + hmget key fields = executeCommandAs ("HMGET" : key : fields) + hexists key field = executeCommandAs ["HEXISTS", key, field] + lpush key values = executeCommandAs ("LPUSH" : key : values) + lrange key start stop = executeCommandAs ["LRANGE", key, showBS start, showBS stop] + expire key seconds = executeCommandAs ["EXPIRE", key, showBS seconds] + ttl key = executeCommandAs ["TTL", key] + rpush key values = executeCommandAs ("RPUSH" : key : values) + lpop key = executeCommandAs ["LPOP", key] + rpop key = executeCommandAs ["RPOP", key] + sadd key members = executeCommandAs ("SADD" : key : members) + smembers key = executeCommandAs ["SMEMBERS", key] + scard key = executeCommandAs ["SCARD", key] + sismember key member = executeCommandAs ["SISMEMBER", key, member] + hdel key fields = executeCommandAs ("HDEL" : key : fields) + hkeys key = executeCommandAs ["HKEYS", key] + hvals key = executeCommandAs ["HVALS", key] + llen key = executeCommandAs ["LLEN", key] + lindex key index = executeCommandAs ["LINDEX", key, showBS index] + clientSetInfo info = executeCommandAs (["CLIENT", "SETINFO"] ++ info) + clusterSlots = executeCommandAs ["CLUSTER", "SLOTS"] + + clientReply val = do + ClientState !client _ <- State.get + liftIO $ send client (Builder.toLazyByteString . encode $ wrapInRay ["CLIENT", "REPLY", showBS val]) + case val of + ON -> Just <$> parseWith (receive client) + _ -> return Nothing + + zadd key members = + let payload = concatMap (\(score, member) -> [showBS score, member]) members + in executeCommandAs ("ZADD" : key : payload) + + zrange key start stop withScores = + let base = ["ZRANGE", key, showBS start, showBS stop] + command = if withScores then base ++ ["WITHSCORES"] else base + in executeCommandAs command + + geoadd key entries = + let payload = concatMap (\(lon, lat, member) -> [showBS lon, showBS lat, member]) entries + in executeCommandAs ("GEOADD" : key : payload) + + geodist key member1 member2 unit = + let unitPart = maybe [] (\u -> [geoUnitKeyword u]) unit + in executeCommandAs (["GEODIST", key, member1, member2] ++ unitPart) + + geohash key members = executeCommandAs ("GEOHASH" : key : members) + geopos key members = executeCommandAs ("GEOPOS" : key : members) + + georadius key longitude latitude radius unit flags = + let base = ["GEORADIUS", key, showBS longitude, showBS latitude, showBS radius, geoUnitKeyword unit] + in executeCommandAs (base ++ concatMap geoRadiusFlagToList flags) + + georadiusRo key longitude latitude radius unit flags = + let base = ["GEORADIUS_RO", key, showBS longitude, showBS latitude, showBS radius, geoUnitKeyword unit] + in executeCommandAs (base ++ concatMap geoRadiusFlagToList flags) + + georadiusByMember key member radius unit flags = + let base = ["GEORADIUSBYMEMBER", key, member, showBS radius, geoUnitKeyword unit] + in executeCommandAs (base ++ concatMap geoRadiusFlagToList flags) + + georadiusByMemberRo key member radius unit flags = + let base = ["GEORADIUSBYMEMBER_RO", key, member, showBS radius, geoUnitKeyword unit] + in executeCommandAs (base ++ concatMap geoRadiusFlagToList flags) + + geosearch key fromSpec bySpec options = + executeCommandAs (["GEOSEARCH", key] + ++ geoSearchFromToList fromSpec + ++ geoSearchByToList bySpec + ++ concatMap geoSearchOptionToList options) + + geosearchstore dest source fromSpec bySpec options storeDist = + let base = ["GEOSEARCHSTORE", dest, source] + ++ geoSearchFromToList fromSpec + ++ geoSearchByToList bySpec + ++ concatMap geoSearchOptionToList options + command = if storeDist then base ++ ["STOREDIST"] else base + in executeCommandAs command + +-- | Receive exactly one RESP value, fetching more bytes from the connection as needed. +-- Throws 'ParseError' on malformed data and 'ConnectionClosed' if the remote end hangs up. +parseWith :: (Client client, MonadIO m, MonadState (ClientState client) m) => m BS8.ByteString -> m RespData +parseWith recv = do + result <- parseManyWith 1 recv + case result of + [x] -> return x + _ -> liftIO $ throwIO $ ParseError "parseWith: expected exactly one result" + +-- | Receive exactly @cnt@ RESP values from the connection, performing incremental +-- parsing against the internal buffer and fetching more bytes as needed. +parseManyWith :: (Client client, MonadIO m, MonadState (ClientState client) m) => Int -> m BS8.ByteString -> m [RespData] +parseManyWith cnt recv = do + (ClientState !client !input) <- State.get + case StrictParse.parse (StrictParse.count cnt parseRespData) input of + StrictParse.Fail _ _ err -> liftIO $ throwIO $ ParseError err + part@(StrictParse.Partial _) -> runUntilDone client part recv + StrictParse.Done remainder !r -> do + State.put (ClientState client remainder) + return r + where + runUntilDone :: (Client client, MonadIO m, MonadState (ClientState client) m) => client 'Connected -> StrictParse.IResult BS8.ByteString r -> m BS8.ByteString -> m r + runUntilDone _client (StrictParse.Fail _ _ err) _ = liftIO $ throwIO $ ParseError err + runUntilDone client (StrictParse.Partial f) getMore = do + moreData <- getMore + if BS8.null moreData + then liftIO $ throwIO ConnectionClosed + else runUntilDone client (f moreData) getMore + runUntilDone client (StrictParse.Done remainder !r) _ = do + State.put (ClientState client remainder) + return r diff --git a/lib/redis-command-client/Database/Redis/FromResp.hs b/lib/redis-command-client/Database/Redis/FromResp.hs new file mode 100644 --- /dev/null +++ b/lib/redis-command-client/Database/Redis/FromResp.hs @@ -0,0 +1,90 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Typeclass for converting 'RespData' values into user-chosen Haskell types. +-- +-- This module provides 'FromResp' together with a collection of instances for +-- common types so that Redis commands can return typed values instead of raw +-- 'RespData'. +-- +-- @since 0.1.0.0 +module Database.Redis.FromResp + ( FromResp (..) + ) where + +import Data.ByteString (ByteString) +import Data.Text (Text) +import qualified Data.Text.Encoding as TE +import Database.Redis.RedisError (RedisError (..)) +import Database.Redis.Resp (RespData (..)) + +-- | Convert a 'RespData' value to a Haskell type. +-- +-- Returns @Left (UnexpectedResp rd)@ when the conversion cannot be performed +-- for the given RESP value @rd@. +class FromResp a where + fromResp :: RespData -> Either RedisError a + +-- | Identity instance – always succeeds. +instance FromResp RespData where + fromResp :: RespData -> Either RedisError RespData + fromResp = Right + +-- | Extracts the payload of 'RespBulkString' or 'RespSimpleString'. +instance FromResp ByteString where + fromResp :: RespData -> Either RedisError ByteString + fromResp (RespBulkString bs) = Right bs + fromResp (RespSimpleString bs) = Right bs + fromResp rd = Left (UnexpectedResp rd) + +-- | 'Nothing' for 'RespNullBulkString', otherwise delegates to the 'ByteString' instance. +instance FromResp (Maybe ByteString) where + fromResp :: RespData -> Either RedisError (Maybe ByteString) + fromResp RespNullBulkString = Right Nothing + fromResp rd = Just <$> fromResp rd + +-- | Extracts 'RespInteger' payload. +instance FromResp Integer where + fromResp :: RespData -> Either RedisError Integer + fromResp (RespInteger i) = Right i + fromResp rd = Left (UnexpectedResp rd) + +-- | Interprets @RespSimpleString \"OK\"@ and @RespInteger 1@ as 'True', +-- @RespInteger 0@ as 'False'. +instance FromResp Bool where + fromResp :: RespData -> Either RedisError Bool + fromResp (RespSimpleString "OK") = Right True + fromResp (RespInteger 1) = Right True + fromResp (RespInteger 0) = Right False + fromResp rd = Left (UnexpectedResp rd) + +-- | Extracts each element of a 'RespArray' as a 'ByteString'. +instance FromResp [ByteString] where + fromResp :: RespData -> Either RedisError [ByteString] + fromResp (RespArray xs) = traverse fromResp xs + fromResp rd = Left (UnexpectedResp rd) + +-- | Extracts each element of a 'RespArray' as 'RespData'. +instance FromResp [RespData] where + fromResp :: RespData -> Either RedisError [RespData] + fromResp (RespArray xs) = Right xs + fromResp rd = Left (UnexpectedResp rd) + +-- | Always succeeds, discarding the response. +instance FromResp () where + fromResp :: RespData -> Either RedisError () + fromResp _ = Right () + +-- | UTF-8 decode a 'ByteString' payload into 'Text'. +instance FromResp Text where + fromResp :: RespData -> Either RedisError Text + fromResp rd = case fromResp rd of + Right (bs :: ByteString) -> case TE.decodeUtf8' bs of + Right t -> Right t + Left _ -> Left (UnexpectedResp rd) + Left e -> Left e + +-- | 'Nothing' for 'RespNullBulkString', otherwise delegates to the 'Text' instance. +instance FromResp (Maybe Text) where + fromResp :: RespData -> Either RedisError (Maybe Text) + fromResp RespNullBulkString = Right Nothing + fromResp rd = Just <$> fromResp rd diff --git a/lib/redis-command-client/Database/Redis/RedisError.hs b/lib/redis-command-client/Database/Redis/RedisError.hs new file mode 100644 --- /dev/null +++ b/lib/redis-command-client/Database/Redis/RedisError.hs @@ -0,0 +1,19 @@ +-- | Typed exceptions for Redis protocol and conversion errors. +-- +-- @since 0.1.0.0 +module Database.Redis.RedisError + ( RedisError (..) + ) where + +import Control.Exception (Exception) +import Data.Typeable (Typeable) +import Database.Redis.Resp (RespData) + +-- | Typed exceptions for Redis protocol errors. +data RedisError + = ParseError String -- ^ RESP parse failure + | ConnectionClosed -- ^ Remote end closed the connection + | UnexpectedResp RespData -- ^ Unexpected RESP value during 'FromResp' conversion + deriving (Eq, Show, Typeable) + +instance Exception RedisError diff --git a/lib/redis/Database/Redis.hs b/lib/redis/Database/Redis.hs new file mode 100644 --- /dev/null +++ b/lib/redis/Database/Redis.hs @@ -0,0 +1,125 @@ +-- | Convenience re-export module for the hask-redis-mux library. +-- +-- Import this single module for both standalone and cluster Redis usage. +-- +-- __Standalone usage with bracket pattern (recommended):__ +-- +-- @ +-- {-# LANGUAGE OverloadedStrings #-} +-- import Database.Redis +-- +-- main :: IO () +-- main = do +-- result <- runRedis defaultStandaloneConfig $ do +-- set \"mykey\" \"myvalue\" +-- (val :: ByteString) <- get \"mykey\" +-- return val +-- print result +-- @ +-- +-- __Typed returns via 'FromResp':__ +-- +-- @ +-- runRedis defaultStandaloneConfig $ do +-- set \"counter\" \"42\" +-- (n :: Integer) <- get \"counter\" -- automatically parsed +-- (bs :: ByteString) <- get \"counter\" -- raw bytes +-- (mt :: Maybe Text) <- get \"missing\" -- Nothing for missing keys +-- @ +-- +-- __Cluster usage with bracket pattern:__ +-- +-- @ +-- withClusterClient clusterConfig connector $ \\client -> +-- runClusterCommandClient client $ do +-- set \"mykey\" \"myvalue\" +-- get \"mykey\" +-- @ +-- +-- @since 0.1.0.0 +module Database.Redis + ( -- * RESP Protocol + module Database.Redis.Resp + -- * Transport + , module Database.Redis.Client + -- * Redis Commands + , module Database.Redis.Command + -- * FromResp conversion + , module Database.Redis.FromResp + -- * Cluster + , module Database.Redis.Cluster + , module Database.Redis.Cluster.Client + , module Database.Redis.Cluster.ConnectionPool + -- * Multiplexing + , module Database.Redis.Internal.Multiplexer + , module Database.Redis.Internal.MultiplexPool + -- * Standalone Multiplexed Client + , module Database.Redis.Standalone + -- * Connection Helpers + , module Database.Redis.Connector + -- * ByteString (re-exported for convenience) + , ByteString + ) where + +import Data.ByteString (ByteString) +import Database.Redis.Client (Client (..), + ConnectionStatus (..), + PlainTextClient (..), + TLSClient (..)) +import Database.Redis.Cluster (ClusterNode (..), + ClusterTopology (..), + NodeAddress (..), + NodeRole (..), + SlotRange (..)) +import Database.Redis.Cluster.Client (ClusterClient (..), + ClusterCommandClient, + ClusterConfig (..), + ClusterError (..), + closeClusterClient, + createClusterClient, + refreshTopology, + runClusterCommandClient) +import Database.Redis.Cluster.ConnectionPool (ConnectionPool (..), + PoolConfig (..), + closePool, createPool, + withConnection) +import Database.Redis.Command (ClientReplyValues (..), + ClientState (..), + RedisCommandClient (..), + RedisCommands (..), + RedisError (..), + convertResp, + encodeBulkArg, + encodeCommand, + encodeCommandBuilder, + encodeGetBuilder, + encodeSetBuilder, + parseManyWith, + parseWith, showBS) +import Database.Redis.Connector (Connector, + clusterPlaintextConnector, + clusterTLSConnector, + connectPlaintext, + connectTLS) +import Database.Redis.FromResp (FromResp (..)) +import Database.Redis.Internal.Multiplexer (Multiplexer, + MultiplexerException (..), + createMultiplexer, + destroyMultiplexer, + isMultiplexerAlive, + submitCommand) +import Database.Redis.Internal.MultiplexPool (MultiplexPool, + closeMultiplexPool, + createMultiplexPool, + submitToNode) +import Database.Redis.Resp (Encodable (..), + RespData (..), + parseRespData, + parseStrict) +import Database.Redis.Standalone (StandaloneClient, + StandaloneCommandClient, + StandaloneConfig (..), + closeStandaloneClient, + createStandaloneClient, + createStandaloneClientFromConfig, + runStandaloneClient) diff --git a/lib/resp/Database/Redis/Resp.hs b/lib/resp/Database/Redis/Resp.hs new file mode 100644 --- /dev/null +++ b/lib/resp/Database/Redis/Resp.hs @@ -0,0 +1,153 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Serialization and parsing for the Redis Serialization Protocol (RESP). +-- Supports RESP2 and RESP3 wire types including bulk strings, arrays, maps, and sets. +-- +-- @since 0.1.0.0 +module Database.Redis.Resp + ( RespData (..) + , Encodable (..) + , parseRespData + , parseStrict + ) where + +import Control.Applicative ((<|>)) +import qualified Data.Attoparsec.ByteString as StrictParse +import qualified Data.Attoparsec.ByteString.Char8 as Char8 +import Data.ByteString (ByteString) +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Char8 as BS8 +import qualified Data.Map as M +import qualified Data.Set as S + +-- | A value in the RESP wire protocol. Each constructor corresponds to a RESP type prefix. +data RespData where + RespSimpleString :: !ByteString -> RespData + RespError :: !ByteString -> RespData + RespInteger :: !Integer -> RespData + RespBulkString :: !ByteString -> RespData + RespNullBulkString :: RespData + RespArray :: ![RespData] -> RespData + RespMap :: !(M.Map RespData RespData) -> RespData + RespSet :: !(S.Set RespData) -> RespData + deriving (Eq, Ord) + +instance Show RespData where + show :: RespData -> String + show = BS8.unpack . byteShow + +byteShow :: RespData -> ByteString +byteShow (RespSimpleString !s) = s +byteShow (RespError !s) = "ERROR: " <> s +byteShow (RespInteger !i) = BS8.pack . show $ i +byteShow (RespBulkString !s) = s +byteShow RespNullBulkString = "NULL" +byteShow (RespArray !xs) = BS8.intercalate "\n" (map byteShow xs) +byteShow (RespSet !s) = BS8.intercalate "\n" (map byteShow (S.toList s)) +byteShow (RespMap !m) = BS8.intercalate "\n" (map showPair (M.toList m)) + where + showPair (k, v) = byteShow k <> ": " <> byteShow v + +-- | Types that can be serialized to the RESP wire format. +class Encodable a where + encode :: a -> Builder.Builder + +instance Encodable RespData where + encode :: RespData -> Builder.Builder + encode (RespSimpleString !s) = Builder.char8 '+' <> Builder.byteString s <> Builder.byteString "\r\n" + encode (RespError !s) = Builder.char8 '-' <> Builder.byteString s <> Builder.byteString "\r\n" + encode (RespInteger !i) = Builder.char8 ':' <> Builder.integerDec i <> Builder.byteString "\r\n" + encode (RespBulkString !s) = Builder.char8 '$' <> Builder.intDec (BS8.length s) <> Builder.byteString "\r\n" <> Builder.byteString s <> Builder.byteString "\r\n" + encode RespNullBulkString = Builder.byteString "$-1\r\n" + encode (RespArray !xs) = Builder.char8 '*' <> Builder.intDec (length xs) <> Builder.byteString "\r\n" <> foldMap encode xs + encode (RespSet !s) = Builder.char8 '~' <> Builder.intDec (S.size s) <> Builder.byteString "\r\n" <> foldMap encode (S.toList s) + encode (RespMap !m) = Builder.char8 '*' <> Builder.intDec (M.size m) <> Builder.byteString "\r\n" <> foldMap encodePair (M.toList m) + where + encodePair (k, v) = encode k <> encode v + +-- | Attoparsec parser for a single RESP value. Can be composed with other parsers +-- or used with 'parseStrict' for one-shot parsing. +parseRespData :: Char8.Parser RespData +parseRespData = + Char8.anyChar >>= \case + '+' -> parseSimpleString + '-' -> parseError + ':' -> parseInteger + '$' -> parseBulkString <|> parseNullBulkString + '*' -> parseArray + '~' -> parseSet + '%' -> parseMap + v -> do + rest <- Char8.takeByteString + let catted = BS8.cons v rest + fail ("Invalid RESP data type: Remaining string " <> BS8.unpack catted) +{-# INLINE parseRespData #-} + +parseSimpleString :: Char8.Parser RespData +parseSimpleString = do + !s <- Char8.takeTill (== '\r') + _ <- Char8.take 2 + return $ RespSimpleString s +{-# INLINE parseSimpleString #-} + +parseError :: Char8.Parser RespData +parseError = do + !s <- Char8.takeTill (== '\r') + _ <- Char8.take 2 + return $ RespError s +{-# INLINE parseError #-} + +parseInteger :: Char8.Parser RespData +parseInteger = do + !i <- Char8.signed Char8.decimal + _ <- Char8.endOfLine + return $ RespInteger i +{-# INLINE parseInteger #-} + +parseBulkString :: Char8.Parser RespData +parseBulkString = do + !len <- Char8.decimal + _ <- Char8.endOfLine + !s <- Char8.take len + _ <- Char8.endOfLine + return $ RespBulkString s +{-# INLINE parseBulkString #-} + +parseNullBulkString :: Char8.Parser RespData +parseNullBulkString = do + _ <- Char8.char '-' + _ <- Char8.char '1' + _ <- Char8.endOfLine + return RespNullBulkString + +parseArray :: Char8.Parser RespData +parseArray = do + !len <- Char8.decimal + _ <- Char8.endOfLine + !xs <- Char8.count len parseRespData + return $ RespArray xs + +parseSet :: Char8.Parser RespData +parseSet = do + !len <- Char8.decimal + _ <- Char8.endOfLine + !xs <- Char8.count len parseRespData + return $ RespSet (S.fromList xs) + +parseMap :: Char8.Parser RespData +parseMap = do + !len <- Char8.decimal + _ <- Char8.endOfLine + !pairs <- Char8.count len parsePair + return $ RespMap (M.fromList pairs) + where + parsePair = do + k <- parseRespData + v <- parseRespData + return (k, v) + +-- | Parse strict ByteString into RespData +parseStrict :: BS8.ByteString -> Either String RespData +parseStrict = StrictParse.parseOnly parseRespData diff --git a/test/ClusterCommandSpec.hs b/test/ClusterCommandSpec.hs new file mode 100644 --- /dev/null +++ b/test/ClusterCommandSpec.hs @@ -0,0 +1,430 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} + +module Main (main) where + +import Control.Concurrent (forkIO, threadDelay) +import Control.Concurrent.MVar (newEmptyMVar, newMVar, + putMVar, takeMVar) +import Control.Concurrent.STM (newTVarIO) +import Control.Monad.IO.Class (liftIO) +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Lazy as LBS +import Data.IORef (IORef, + atomicModifyIORef', + newIORef, readIORef) +import qualified Data.Map.Strict as Map +import Data.Time.Clock (getCurrentTime) +import qualified Data.Vector as V +import Database.Redis.Client (Client (..), + ConnectionStatus (..)) +import Database.Redis.Cluster (ClusterNode (..), + ClusterTopology (..), + NodeAddress (..), + NodeRole (..), + SlotRange (..)) +import Database.Redis.Cluster.Client +import Database.Redis.Cluster.ConnectionPool (PoolConfig (..), + createPool) +import Database.Redis.Internal.MultiplexPool (closeMultiplexPool, + createMultiplexPool) +import Database.Redis.Resp (Encodable (..), + RespData (..)) +import System.Timeout (timeout) +import Test.Hspec + +main :: IO () +main = hspec spec + +spec :: Spec +spec = do + describe "Redirection error parsing" $ do + describe "MOVED error parsing" $ do + it "parses valid MOVED error" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999 127.0.0.1:6381" + result `shouldBe` Just (RedirectionInfo 3999 "127.0.0.1" 6381) + + it "parses MOVED error with different slot" $ do + let result = parseRedirectionError "MOVED" "MOVED 12345 192.168.1.100:7000" + result `shouldBe` Just (RedirectionInfo 12345 "192.168.1.100" 7000) + + it "handles MOVED error with hostname containing colons" $ do + -- IPv6 addresses would need special handling + -- For now, test that we handle hostnames correctly + let result = parseRedirectionError "MOVED" "MOVED 3999 redis-node-1.example.com:6381" + result `shouldBe` Just (RedirectionInfo 3999 "redis-node-1.example.com" 6381) + + it "returns Nothing for malformed MOVED error" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999" + result `shouldBe` Nothing + + it "returns Nothing for MOVED with invalid slot" $ do + let result = parseRedirectionError "MOVED" "MOVED notanumber 127.0.0.1:6381" + result `shouldBe` Nothing + + it "returns Nothing for MOVED with invalid port" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999 127.0.0.1:notaport" + result `shouldBe` Nothing + + it "returns Nothing for MOVED with missing colon" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999 127.0.0.16381" + result `shouldBe` Nothing + + describe "ASK error parsing" $ do + it "parses valid ASK error" $ do + let result = parseRedirectionError "ASK" "ASK 3999 127.0.0.1:6381" + result `shouldBe` Just (RedirectionInfo 3999 "127.0.0.1" 6381) + + it "parses ASK error with different slot" $ do + let result = parseRedirectionError "ASK" "ASK 8765 10.0.0.1:6379" + result `shouldBe` Just (RedirectionInfo 8765 "10.0.0.1" 6379) + + it "returns Nothing for malformed ASK error" $ do + let result = parseRedirectionError "ASK" "ASK 3999" + result `shouldBe` Nothing + + it "returns Nothing for wrong error type prefix" $ do + let result = parseRedirectionError "ASK" "MOVED 3999 127.0.0.1:6381" + result `shouldBe` Nothing + + describe "Edge cases" $ do + it "handles slot 0" $ do + let result = parseRedirectionError "MOVED" "MOVED 0 127.0.0.1:6379" + result `shouldBe` Just (RedirectionInfo 0 "127.0.0.1" 6379) + + it "handles slot 16383 (max)" $ do + let result = parseRedirectionError "MOVED" "MOVED 16383 127.0.0.1:6379" + result `shouldBe` Just (RedirectionInfo 16383 "127.0.0.1" 6379) + + it "handles high port numbers" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999 127.0.0.1:65535" + result `shouldBe` Just (RedirectionInfo 3999 "127.0.0.1" 65535) + + it "handles hostname instead of IP" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999 redis-node-1:6379" + result `shouldBe` Just (RedirectionInfo 3999 "redis-node-1" 6379) + + it "returns Nothing for extra whitespace" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999 127.0.0.1:6381 " + -- Tighter parsing rejects non-standard formatting (Redis never produces this) + result `shouldBe` Nothing + + it "handles port 0" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999 127.0.0.1:0" + result `shouldBe` Just (RedirectionInfo 3999 "127.0.0.1" 0) + + it "returns Nothing for extra fields after host:port" $ do + let result = parseRedirectionError "MOVED" "MOVED 3999 127.0.0.1:6381 extra-data" + -- Tighter parsing rejects trailing data (Redis never produces this) + result `shouldBe` Nothing + + describe "detectRedirection (byte-level fast path)" $ do + it "returns Nothing for non-error responses (RespBulkString)" $ do + detectRedirection (RespBulkString "OK") `shouldBe` Nothing + + it "returns Nothing for non-error responses (RespSimpleString)" $ do + detectRedirection (RespSimpleString "OK") `shouldBe` Nothing + + it "returns Nothing for non-error responses (RespInteger)" $ do + detectRedirection (RespInteger 42) `shouldBe` Nothing + + it "returns Nothing for non-redirect errors" $ do + detectRedirection (RespError "ERR unknown command") `shouldBe` Nothing + + it "returns Nothing for short error messages" $ do + detectRedirection (RespError "ERR") `shouldBe` Nothing + + it "returns Nothing for empty error message" $ do + detectRedirection (RespError "") `shouldBe` Nothing + + it "detects MOVED redirect" $ do + detectRedirection (RespError "MOVED 3999 127.0.0.1:6381") + `shouldBe` Just (Left (RedirectionInfo 3999 "127.0.0.1" 6381)) + + it "detects ASK redirect" $ do + detectRedirection (RespError "ASK 3999 127.0.0.1:6381") + `shouldBe` Just (Right (RedirectionInfo 3999 "127.0.0.1" 6381)) + + it "returns Nothing for errors starting with M but not MOVED" $ do + detectRedirection (RespError "MASTERDOWN Link with MASTER is down") `shouldBe` Nothing + + it "returns Nothing for errors starting with A but not ASK" $ do + detectRedirection (RespError "AUTH required") `shouldBe` Nothing + + describe "ClusterError types" $ do + it "creates MovedError correctly" $ do + let err = MovedError 3999 (NodeAddress "127.0.0.1" 6381) + show err `shouldContain` "MovedError" + show err `shouldContain` "3999" + + it "creates AskError correctly" $ do + let err = AskError 3999 (NodeAddress "127.0.0.1" 6381) + show err `shouldContain` "AskError" + show err `shouldContain` "3999" + + it "creates ClusterDownError correctly" $ do + let err = ClusterDownError "Cluster is down" + show err `shouldContain` "ClusterDownError" + show err `shouldContain` "Cluster is down" + + it "creates TryAgainError correctly" $ do + let err = TryAgainError "Try again later" + show err `shouldContain` "TryAgainError" + + it "creates CrossSlotError correctly" $ do + let err = CrossSlotError "Keys in request don't hash to the same slot" + show err `shouldContain` "CrossSlotError" + + it "creates MaxRetriesExceeded correctly" $ do + let err = MaxRetriesExceeded "Max retries (3) exceeded" + show err `shouldContain` "MaxRetriesExceeded" + show err `shouldContain` "3" + + it "creates TopologyError correctly" $ do + let err = TopologyError "No node found for slot 3999" + show err `shouldContain` "TopologyError" + + it "creates ConnectionError correctly" $ do + let err = ConnectionError "Connection timeout" + show err `shouldContain` "ConnectionError" + + describe "ClusterConfig" $ do + it "creates valid cluster config" $ do + let poolConfig = PoolConfig + { maxConnectionsPerNode = 1, + connectionTimeout = 5000, + maxRetries = 3, + useTLS = False + } + config = ClusterConfig + { clusterSeedNode = NodeAddress "127.0.0.1" 7000, + clusterPoolConfig = poolConfig, + clusterMaxRetries = 3, + clusterRetryDelay = 100000, + clusterTopologyRefreshInterval = 600 + } + clusterMaxRetries config `shouldBe` 3 + clusterRetryDelay config `shouldBe` 100000 + clusterTopologyRefreshInterval config `shouldBe` 600 + + describe "RedirectionInfo" $ do + it "creates valid redirection info" $ do + let redir = RedirectionInfo 3999 "127.0.0.1" 6381 + redirSlot redir `shouldBe` 3999 + redirHost redir `shouldBe` "127.0.0.1" + redirPort redir `shouldBe` 6381 + + it "shows redirection info correctly" $ do + let redir = RedirectionInfo 3999 "127.0.0.1" 6381 + show redir `shouldContain` "3999" + show redir `shouldContain` "127.0.0.1" + show redir `shouldContain` "6381" + + it "compares redirection info for equality" $ do + let redir1 = RedirectionInfo 3999 "127.0.0.1" 6381 + redir2 = RedirectionInfo 3999 "127.0.0.1" 6381 + redir3 = RedirectionInfo 4000 "127.0.0.1" 6381 + redir1 `shouldBe` redir2 + redir1 `shouldNotBe` redir3 + + askRedirectSpec + +-- --------------------------------------------------------------------------- +-- Mock client (same pattern as MultiplexPoolSpec) +-- --------------------------------------------------------------------------- + +data MockClient (a :: ConnectionStatus) where + MockConnected :: !(IORef ByteString) -- sendBuf + -> !(IORef [ByteString]) -- recvQueue + -> MockClient 'Connected + +instance Client MockClient where + connect = error "MockClient: connect not supported" + close _ = return () + send (MockConnected sendBuf _) lbs = liftIO $ do + let !bs = LBS.toStrict lbs + atomicModifyIORef' sendBuf $ \old -> (old <> bs, ()) + receive (MockConnected sRef recvQueue) = liftIO $ recvLoop sRef recvQueue + +recvLoop :: IORef ByteString -> IORef [ByteString] -> IO ByteString +recvLoop sRef recvQueue = do + mChunk <- atomicModifyIORef' recvQueue $ \xs -> + case xs of + [] -> ([], Nothing) + (y:ys) -> (ys, Just y) + case mChunk of + Just chunk -> return chunk + Nothing -> do + threadDelay 1000 + recvLoop sRef recvQueue + +createMockClient :: IO (MockClient 'Connected, ByteString -> IO ()) +createMockClient = do + sendBuf <- newIORef BS.empty + recvQueue <- newIORef [] + let client = MockConnected sendBuf recvQueue + addRecv bs = atomicModifyIORef' recvQueue $ \xs -> (xs ++ [bs], ()) + return (client, addRecv) + +type AddRecvMap = IORef [(NodeAddress, ByteString -> IO ())] + +createMockConnector :: IO (NodeAddress -> IO (MockClient 'Connected), AddRecvMap, IORef Int) +createMockConnector = do + mapRef <- newIORef [] + connCount <- newIORef 0 + let connector addr = do + (client, addRecv) <- createMockClient + atomicModifyIORef' mapRef $ \xs -> (xs ++ [(addr, addRecv)], ()) + atomicModifyIORef' connCount $ \n -> (n + 1, ()) + return client + return (connector, mapRef, connCount) + +getAddRecvs :: AddRecvMap -> NodeAddress -> IO [ByteString -> IO ()] +getAddRecvs mapRef addr = do + xs <- readIORef mapRef + return [f | (a, f) <- xs, a == addr] + +firstOf :: [a] -> a +firstOf (x:_) = x +firstOf [] = error "firstOf: empty list" + +encodeResp :: RespData -> ByteString +encodeResp = LBS.toStrict . Builder.toLazyByteString . encode + +-- | Build a test topology where all 16384 slots map to a single node. +mkTopology :: NodeAddress -> IO ClusterTopology +mkTopology addr = do + now <- getCurrentTime + let nodeIdBS = "test-node-id-1" + node = ClusterNode nodeIdBS addr Master [SlotRange 0 16383 nodeIdBS []] [] + slotVec = V.replicate 16384 nodeIdBS + addrVec = V.replicate 16384 addr + nodeMap = Map.singleton nodeIdBS node + return $ ClusterTopology slotVec addrVec nodeMap now + +-- | Build a ClusterClient backed by a mock connector and a given topology. +mkClusterClient + :: (NodeAddress -> IO (MockClient 'Connected)) + -> ClusterTopology + -> IO (ClusterClient MockClient) +mkClusterClient connector topo = do + topoVar <- newTVarIO topo + pool <- createPool poolCfg + muxPool <- createMultiplexPool connector 1 + refreshLk <- newMVar () + return $ ClusterClient topoVar pool clusterCfg connector refreshLk muxPool + where + poolCfg = PoolConfig + { maxConnectionsPerNode = 1 + , connectionTimeout = 5000 + , maxRetries = 3 + , useTLS = False + } + clusterCfg = ClusterConfig + { clusterSeedNode = NodeAddress "127.0.0.1" 6379 + , clusterPoolConfig = poolCfg + , clusterMaxRetries = 5 + , clusterRetryDelay = 1000 + , clusterTopologyRefreshInterval = 600 + } + +-- --------------------------------------------------------------------------- +-- ASK redirect integration tests +-- --------------------------------------------------------------------------- + +node1, node2 :: NodeAddress +node1 = NodeAddress "127.0.0.1" 6379 +node2 = NodeAddress "127.0.0.2" 6380 + +askRedirectSpec :: Spec +askRedirectSpec = describe "ASK redirect integration (executeKeyedClusterCommand)" $ do + it "on ASK error, retries with ASKING prefix at the redirected node" $ do + result <- timeout 5000000 $ do + (connector, addRecvMap, _) <- createMockConnector + topo <- mkTopology node1 -- all slots → node1 + client <- mkClusterClient connector topo + + -- Run executeKeyedClusterCommand in a thread + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- executeKeyedClusterCommand client "mykey" ["GET", "mykey"] + putMVar resultMVar r + + -- Wait for node1 connection, then reply with ASK redirect to node2 + threadDelay 50000 + fns1 <- getAddRecvs addRecvMap node1 + length fns1 `shouldSatisfy` (>= 1) + -- node2 should have no connections yet (not in topology) + fns2Before <- getAddRecvs addRecvMap node2 + length fns2Before `shouldBe` 0 + (firstOf fns1) (encodeResp (RespError "ASK 3999 127.0.0.2:6380")) + + -- The retry should connect to node2 with ASKING + GET + threadDelay 50000 + fns2 <- getAddRecvs addRecvMap node2 + length fns2 `shouldSatisfy` (>= 1) + -- Feed +OK for ASKING, then the real response + (firstOf fns2) (encodeResp (RespSimpleString "OK")) + (firstOf fns2) (encodeResp (RespBulkString "myvalue")) + + r <- takeMVar resultMVar + r `shouldBe` Right (RespBulkString "myvalue") + + closeMultiplexPool (clusterMultiplexPool client) + result `shouldBe` Just () + + it "ASK redirect does NOT trigger topology refresh" $ do + result <- timeout 5000000 $ do + (connector, addRecvMap, connCount) <- createMockConnector + topo <- mkTopology node1 + client <- mkClusterClient connector topo + + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- executeKeyedClusterCommand client "testkey" ["SET", "testkey", "val"] + putMVar resultMVar r + + threadDelay 50000 + fns1 <- getAddRecvs addRecvMap node1 + (firstOf fns1) (encodeResp (RespError "ASK 100 127.0.0.2:6380")) + + threadDelay 50000 + fns2 <- getAddRecvs addRecvMap node2 + (firstOf fns2) (encodeResp (RespSimpleString "OK")) -- ASKING + (firstOf fns2) (encodeResp (RespSimpleString "OK")) -- SET + + r <- takeMVar resultMVar + r `shouldBe` Right (RespSimpleString "OK") + + -- Exactly 2 connections total: 1 for node1 (mux) + 1 for node2 (ASK redirect). + -- A topology refresh would have created a 3rd connection to the seed node, + -- and then hung waiting for CLUSTER SLOTS data (caught by the 5s timeout). + totalConns <- readIORef connCount + totalConns `shouldBe` 2 + + closeMultiplexPool (clusterMultiplexPool client) + result `shouldBe` Just () + + it "successful command without redirection returns directly" $ do + result <- timeout 5000000 $ do + (connector, addRecvMap, _) <- createMockConnector + topo <- mkTopology node1 + client <- mkClusterClient connector topo + + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- executeKeyedClusterCommand client "key1" ["GET", "key1"] + putMVar resultMVar r + + threadDelay 50000 + fns1 <- getAddRecvs addRecvMap node1 + (firstOf fns1) (encodeResp (RespBulkString "directvalue")) + + r <- takeMVar resultMVar + r `shouldBe` Right (RespBulkString "directvalue") + + closeMultiplexPool (clusterMultiplexPool client) + result `shouldBe` Just () diff --git a/test/ClusterSpec.hs b/test/ClusterSpec.hs new file mode 100644 --- /dev/null +++ b/test/ClusterSpec.hs @@ -0,0 +1,175 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Main (main) where + +import qualified Data.ByteString.Char8 as BS8 +import Data.Time.Clock (getCurrentTime) +import Database.Redis.Cluster +import Database.Redis.Resp (RespData (..)) +import Test.Hspec + +main :: IO () +main = hspec spec + +spec :: Spec +spec = do + describe "Hash tag extraction" $ do + it "extracts hash tag from valid key" $ do + extractHashTag "{user}:profile" `shouldBe` "user" + extractHashTag "{user}:settings" `shouldBe` "user" + + it "returns full key when no hash tag present" $ do + extractHashTag "simple-key" `shouldBe` "simple-key" + extractHashTag "key:with:colons" `shouldBe` "key:with:colons" + + it "handles edge cases correctly" $ do + -- Empty tag + extractHashTag "{}" `shouldBe` "{}" + -- No closing brace + extractHashTag "{user" `shouldBe` "{user" + -- No opening brace + extractHashTag "user}" `shouldBe` "user}" + -- Multiple braces - uses first valid pair + extractHashTag "{first}{second}" `shouldBe` "first" + -- Braces at end + extractHashTag "key{tag}" `shouldBe` "key{tag}" + + it "handles special characters in hash tags" $ do + extractHashTag "{user:1}" `shouldBe` "user:1" + extractHashTag "{a-b}" `shouldBe` "a-b" + extractHashTag "{tag.with.dots}" `shouldBe` "tag.with.dots" + + describe "Slot calculation" $ do + it "calculates slot within valid range" $ do + let slot = calculateSlot "test-key" + slot `shouldSatisfy` (< 16384) + + it "calculates same slot for keys with same hash tag" $ do + let slot1 = calculateSlot "{user}:profile" + slot2 = calculateSlot "{user}:settings" + slot1 `shouldBe` slot2 + + it "calculates different slots for different keys (usually)" $ do + let slot1 = calculateSlot "key1" + slot2 = calculateSlot "key2" + -- Note: This could theoretically fail if both keys hash to same slot, + -- but probability is very low + slot1 `shouldNotBe` slot2 + + it "handles empty key" $ do + let slot = calculateSlot "" + slot `shouldSatisfy` (< 16384) + + it "handles very long keys" $ do + let longKey = BS8.replicate 1000 'x' + slot = calculateSlot longKey + slot `shouldSatisfy` (< 16384) + + describe "Topology parsing" $ do + it "parses simple CLUSTER SLOTS response" $ do + currentTime <- getCurrentTime + let response = + RespArray + [ RespArray + [ RespInteger 0, + RespInteger 5460, + RespArray + [ RespBulkString "127.0.0.1", + RespInteger 7000, + RespBulkString "node-id-1" + ] + ] + ] + case parseClusterSlots response currentTime of + Left err -> expectationFailure $ "Parsing failed: " ++ err + Right topology -> do + -- Check that slots are assigned + case findNodeForSlot topology 0 of + Nothing -> expectationFailure "Slot 0 should be assigned" + Just nodeId -> nodeId `shouldNotBe` "" + + it "handles invalid responses" $ do + currentTime <- getCurrentTime + let invalidResponse = RespBulkString "invalid" + case parseClusterSlots invalidResponse currentTime of + Left _ -> return () -- Expected + Right _ -> expectationFailure "Should fail on invalid response" + + it "parses response with replicas" $ do + currentTime <- getCurrentTime + let response = + RespArray + [ RespArray + [ RespInteger 0, + RespInteger 5460, + RespArray + [ RespBulkString "127.0.0.1", + RespInteger 7000, + RespBulkString "master-1" + ], + RespArray + [ RespBulkString "127.0.0.1", + RespInteger 7003, + RespBulkString "replica-1" + ] + ] + ] + case parseClusterSlots response currentTime of + Left err -> expectationFailure $ "Parsing failed: " ++ err + Right topology -> do + -- Check that master is assigned + case findNodeForSlot topology 0 of + Nothing -> expectationFailure "Slot 0 should be assigned" + Just nodeId -> nodeId `shouldNotBe` "" + + describe "Node lookup" $ do + it "finds correct node for slot" $ do + currentTime <- getCurrentTime + let response = + RespArray + [ RespArray + [ RespInteger 0, + RespInteger 100, + RespArray + [ RespBulkString "127.0.0.1", + RespInteger 7000, + RespBulkString "node1" + ] + ], + RespArray + [ RespInteger 101, + RespInteger 200, + RespArray + [ RespBulkString "127.0.0.1", + RespInteger 7001, + RespBulkString "node2" + ] + ] + ] + case parseClusterSlots response currentTime of + Left err -> expectationFailure $ "Parsing failed: " ++ err + Right topology -> do + findNodeForSlot topology 50 `shouldSatisfy` (/= Nothing) + findNodeForSlot topology 150 `shouldSatisfy` (/= Nothing) + -- Out of range + findNodeForSlot topology 16384 `shouldBe` Nothing + + it "returns empty string for slots not covered by any node" $ do + currentTime <- getCurrentTime + let response = + RespArray + [ RespArray + [ RespInteger 0, + RespInteger 100, + RespArray + [ RespBulkString "127.0.0.1", + RespInteger 7000, + RespBulkString "node1" + ] + ] + ] + case parseClusterSlots response currentTime of + Left err -> expectationFailure $ "Parsing failed: " ++ err + Right topology -> do + findNodeForSlot topology 50 `shouldSatisfy` (/= Nothing) + findNodeForSlot topology 200 `shouldBe` Just "" diff --git a/test/FromRespSpec.hs b/test/FromRespSpec.hs new file mode 100644 --- /dev/null +++ b/test/FromRespSpec.hs @@ -0,0 +1,150 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Main where + +import Data.ByteString (ByteString) +import Data.Either (isLeft) +import Data.Text (Text) +import Database.Redis.Command (RedisError (..)) +import Database.Redis.FromResp (FromResp (..)) +import Database.Redis.Resp (RespData (..)) +import Test.Hspec + +main :: IO () +main = hspec $ do + describe "FromResp RespData (identity)" $ do + it "always succeeds for any RespData" $ do + fromResp (RespSimpleString "OK") `shouldBe` Right (RespSimpleString "OK") + fromResp (RespInteger 42) `shouldBe` Right (RespInteger 42) + fromResp RespNullBulkString `shouldBe` Right RespNullBulkString + fromResp (RespBulkString "hello") `shouldBe` Right (RespBulkString "hello") + fromResp (RespArray []) `shouldBe` Right (RespArray []) + + describe "FromResp ByteString" $ do + it "extracts RespBulkString payload" $ do + fromResp (RespBulkString "hello") `shouldBe` (Right "hello" :: Either RedisError ByteString) + + it "extracts RespSimpleString payload" $ do + fromResp (RespSimpleString "OK") `shouldBe` (Right "OK" :: Either RedisError ByteString) + + it "returns Left UnexpectedResp for wrong constructor" $ do + let result = fromResp (RespInteger 42) :: Either RedisError ByteString + result `shouldSatisfy` isLeft + result `shouldBe` Left (UnexpectedResp (RespInteger 42)) + + it "returns Left UnexpectedResp for RespNullBulkString" $ do + let result = fromResp RespNullBulkString :: Either RedisError ByteString + result `shouldBe` Left (UnexpectedResp RespNullBulkString) + + describe "FromResp (Maybe ByteString)" $ do + it "returns Nothing for RespNullBulkString" $ do + fromResp RespNullBulkString `shouldBe` (Right Nothing :: Either RedisError (Maybe ByteString)) + + it "returns Just payload for RespBulkString" $ do + fromResp (RespBulkString "val") `shouldBe` (Right (Just "val") :: Either RedisError (Maybe ByteString)) + + it "returns Just payload for RespSimpleString" $ do + fromResp (RespSimpleString "OK") `shouldBe` (Right (Just "OK") :: Either RedisError (Maybe ByteString)) + + it "returns Left UnexpectedResp for wrong constructor" $ do + let result = fromResp (RespInteger 1) :: Either RedisError (Maybe ByteString) + result `shouldSatisfy` isLeft + + describe "FromResp Integer" $ do + it "extracts RespInteger payload" $ do + fromResp (RespInteger 42) `shouldBe` (Right 42 :: Either RedisError Integer) + + it "extracts negative integers" $ do + fromResp (RespInteger (-1)) `shouldBe` (Right (-1) :: Either RedisError Integer) + + it "extracts zero" $ do + fromResp (RespInteger 0) `shouldBe` (Right 0 :: Either RedisError Integer) + + it "returns Left UnexpectedResp for wrong constructor" $ do + let result = fromResp (RespBulkString "42") :: Either RedisError Integer + result `shouldBe` Left (UnexpectedResp (RespBulkString "42")) + + describe "FromResp Bool" $ do + it "interprets RespSimpleString OK as True" $ do + fromResp (RespSimpleString "OK") `shouldBe` (Right True :: Either RedisError Bool) + + it "interprets RespInteger 1 as True" $ do + fromResp (RespInteger 1) `shouldBe` (Right True :: Either RedisError Bool) + + it "interprets RespInteger 0 as False" $ do + fromResp (RespInteger 0) `shouldBe` (Right False :: Either RedisError Bool) + + it "returns Left UnexpectedResp for other values" $ do + let result = fromResp (RespBulkString "true") :: Either RedisError Bool + result `shouldBe` Left (UnexpectedResp (RespBulkString "true")) + + it "returns Left UnexpectedResp for non-OK simple string" $ do + let result = fromResp (RespSimpleString "QUEUED") :: Either RedisError Bool + result `shouldBe` Left (UnexpectedResp (RespSimpleString "QUEUED")) + + describe "FromResp [ByteString]" $ do + it "extracts array of bulk strings" $ do + let input = RespArray [RespBulkString "a", RespBulkString "b"] + fromResp input `shouldBe` (Right ["a", "b"] :: Either RedisError [ByteString]) + + it "handles empty array" $ do + fromResp (RespArray []) `shouldBe` (Right [] :: Either RedisError [ByteString]) + + it "extracts array of simple strings" $ do + let input = RespArray [RespSimpleString "x", RespSimpleString "y"] + fromResp input `shouldBe` (Right ["x", "y"] :: Either RedisError [ByteString]) + + it "returns Left UnexpectedResp for non-array" $ do + let result = fromResp (RespBulkString "notarray") :: Either RedisError [ByteString] + result `shouldSatisfy` isLeft + + it "returns Left when element is wrong type" $ do + let input = RespArray [RespBulkString "a", RespInteger 1] + let result = fromResp input :: Either RedisError [ByteString] + result `shouldSatisfy` isLeft + + describe "FromResp [RespData]" $ do + it "extracts array elements as RespData" $ do + let input = RespArray [RespInteger 1, RespBulkString "two"] + fromResp input `shouldBe` Right [RespInteger 1, RespBulkString "two"] + + it "handles empty array" $ do + fromResp (RespArray []) `shouldBe` (Right [] :: Either RedisError [RespData]) + + it "returns Left UnexpectedResp for non-array" $ do + let result = fromResp (RespInteger 1) :: Either RedisError [RespData] + result `shouldBe` Left (UnexpectedResp (RespInteger 1)) + + describe "FromResp ()" $ do + it "always succeeds, discarding the response" $ do + fromResp (RespSimpleString "OK") `shouldBe` (Right () :: Either RedisError ()) + fromResp (RespInteger 42) `shouldBe` (Right () :: Either RedisError ()) + fromResp RespNullBulkString `shouldBe` (Right () :: Either RedisError ()) + fromResp (RespError "ERR") `shouldBe` (Right () :: Either RedisError ()) + fromResp (RespArray []) `shouldBe` (Right () :: Either RedisError ()) + + describe "FromResp Text" $ do + it "decodes UTF-8 bulk string to Text" $ do + fromResp (RespBulkString "hello") `shouldBe` (Right "hello" :: Either RedisError Text) + + it "decodes UTF-8 simple string to Text" $ do + fromResp (RespSimpleString "OK") `shouldBe` (Right "OK" :: Either RedisError Text) + + it "returns Left UnexpectedResp for invalid UTF-8" $ do + let result = fromResp (RespBulkString "\xff\xfe") :: Either RedisError Text + result `shouldSatisfy` isLeft + + it "returns Left UnexpectedResp for wrong constructor" $ do + let result = fromResp (RespInteger 42) :: Either RedisError Text + result `shouldBe` Left (UnexpectedResp (RespInteger 42)) + + describe "FromResp (Maybe Text)" $ do + it "returns Nothing for RespNullBulkString" $ do + fromResp RespNullBulkString `shouldBe` (Right Nothing :: Either RedisError (Maybe Text)) + + it "returns Just text for valid UTF-8 bulk string" $ do + fromResp (RespBulkString "hello") `shouldBe` (Right (Just "hello") :: Either RedisError (Maybe Text)) + + it "returns Left UnexpectedResp for wrong constructor" $ do + let result = fromResp (RespInteger 1) :: Either RedisError (Maybe Text) + result `shouldSatisfy` isLeft diff --git a/test/MultiplexPoolSpec.hs b/test/MultiplexPoolSpec.hs new file mode 100644 --- /dev/null +++ b/test/MultiplexPoolSpec.hs @@ -0,0 +1,465 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} + +module Main (main) where + +import Control.Concurrent (forkIO, threadDelay) +import Control.Concurrent.MVar (newEmptyMVar, putMVar, + takeMVar) +import Control.Monad.IO.Class (liftIO) +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Lazy as LBS +import Data.IORef (IORef, + atomicModifyIORef', + newIORef, readIORef) +import Database.Redis.Client (Client (..), + ConnectionStatus (..)) +import Database.Redis.Cluster (NodeAddress (..)) +import Database.Redis.Internal.MultiplexPool +import Database.Redis.Resp (Encodable (..), + RespData (..)) +import Test.Hspec + +-- --------------------------------------------------------------------------- +-- Mock client for testing without a real Redis connection +-- (Same pattern as MultiplexerSpec) +-- --------------------------------------------------------------------------- + +data MockClient (a :: ConnectionStatus) where + MockConnected :: !(IORef ByteString) -- sendBuf + -> !(IORef [ByteString]) -- recvQueue + -> MockClient 'Connected + +instance Client MockClient where + connect = error "MockClient: connect not supported" + close _ = return () + send (MockConnected sendBuf _) lbs = liftIO $ do + let !bs = LBS.toStrict lbs + atomicModifyIORef' sendBuf $ \old -> (old <> bs, ()) + receive (MockConnected sRef recvQueue) = liftIO $ recvLoop sRef recvQueue + +recvLoop :: IORef ByteString -> IORef [ByteString] -> IO ByteString +recvLoop sRef recvQueue = do + mChunk <- atomicModifyIORef' recvQueue $ \xs -> + case xs of + [] -> ([], Nothing) + (y:ys) -> (ys, Just y) + case mChunk of + Just chunk -> return chunk + Nothing -> do + threadDelay 1000 + recvLoop sRef recvQueue + +createMockClient :: IO (MockClient 'Connected, ByteString -> IO ()) +createMockClient = do + sendBuf <- newIORef BS.empty + recvQueue <- newIORef [] + let client = MockConnected sendBuf recvQueue + addRecv bs = atomicModifyIORef' recvQueue $ \xs -> (xs ++ [bs], ()) + return (client, addRecv) + +-- | Total version of 'head' that avoids the -Wx-partial warning in tests. +firstOf :: [a] -> a +firstOf (x:_) = x +firstOf [] = error "firstOf: empty list" + +encodeResp :: RespData -> ByteString +encodeResp = LBS.toStrict . Builder.toLazyByteString . encode + +encodeCmd :: [ByteString] -> Builder.Builder +encodeCmd args = + Builder.byteString ("*" <> bshow (length args) <> "\r\n") + <> foldMap (\a -> Builder.byteString ("$" <> bshow (BS.length a) <> "\r\n" <> a <> "\r\n")) args + where bshow x = LBS.toStrict (Builder.toLazyByteString (Builder.intDec x)) + +-- | A mock connector that creates MockClients and tracks per-node addRecv functions. +-- Returns (connector, getAddRecv) where getAddRecv retrieves the addRecv function +-- for a specific node after it has been connected. +type AddRecvMap = IORef [(NodeAddress, ByteString -> IO ())] + +createMockConnector :: IO (NodeAddress -> IO (MockClient 'Connected), AddRecvMap) +createMockConnector = do + mapRef <- newIORef [] + let connector addr = do + (client, addRecv) <- createMockClient + atomicModifyIORef' mapRef $ \xs -> (xs ++ [(addr, addRecv)], ()) + return client + return (connector, mapRef) + +-- | Get addRecv functions for a given node address. +getAddRecvs :: AddRecvMap -> NodeAddress -> IO [ByteString -> IO ()] +getAddRecvs mapRef addr = do + xs <- readIORef mapRef + return [f | (a, f) <- xs, a == addr] + + +-- --------------------------------------------------------------------------- +-- Test nodes +-- --------------------------------------------------------------------------- + +node1, node2, node3 :: NodeAddress +node1 = NodeAddress "127.0.0.1" 6379 +node2 = NodeAddress "127.0.0.2" 6380 +node3 = NodeAddress "127.0.0.3" 6381 + +-- --------------------------------------------------------------------------- +-- Tests +-- --------------------------------------------------------------------------- + +main :: IO () +main = hspec spec + +spec :: Spec +spec = do + roundRobinSpec + lazyCreationSpec + multiNodeRoutingSpec + poolClosureSpec + askingSpec + +roundRobinSpec :: Spec +roundRobinSpec = describe "Round-robin routing" $ do + it "sequential submits to same node rotate across multiplexers" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 3 -- 3 muxes per node + + -- First submit triggers lazy creation of 3 multiplexers for node1 + addRecvFns <- do + -- We need to feed responses for the first command to trigger creation, + -- but the connector creates clients lazily. Submit from a thread and + -- then feed responses once the connections are established. + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + + -- Wait for connections to be created + threadDelay 50000 + fns <- getAddRecvs addRecvMap node1 + -- Should have 3 connections (one per mux) + length fns `shouldBe` 3 + + -- Feed response to first mux (round-robin starts at 0) + (firstOf fns) (encodeResp (RespSimpleString "PONG1")) + r <- takeMVar resultMVar + r `shouldBe` RespSimpleString "PONG1" + return fns + + -- Second submit should go to mux index 1 + do + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + threadDelay 10000 + (addRecvFns !! 1) (encodeResp (RespSimpleString "PONG2")) + r <- takeMVar resultMVar + r `shouldBe` RespSimpleString "PONG2" + + -- Third submit should go to mux index 2 + do + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + threadDelay 10000 + (addRecvFns !! 2) (encodeResp (RespSimpleString "PONG3")) + r <- takeMVar resultMVar + r `shouldBe` RespSimpleString "PONG3" + + -- Fourth submit should wrap around to mux index 0 + do + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + threadDelay 10000 + (firstOf addRecvFns) (encodeResp (RespSimpleString "PONG4")) + r <- takeMVar resultMVar + r `shouldBe` RespSimpleString "PONG4" + + closeMultiplexPool pool + + it "single multiplexer skips round-robin counter" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 1 -- single mux per node + + -- Submit multiple commands — all go to the same (only) mux + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + threadDelay 50000 + fns <- getAddRecvs addRecvMap node1 + length fns `shouldBe` 1 + (firstOf fns) (encodeResp (RespSimpleString "OK")) + r1 <- takeMVar resultMVar + r1 `shouldBe` RespSimpleString "OK" + + -- Second command still goes to same mux + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + threadDelay 10000 + (firstOf fns) (encodeResp (RespSimpleString "OK2")) + r2 <- takeMVar resultMVar + r2 `shouldBe` RespSimpleString "OK2" + + closeMultiplexPool pool + +lazyCreationSpec :: Spec +lazyCreationSpec = describe "Lazy creation" $ do + it "node muxes created on first access, not at pool creation" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 2 + + -- At creation, no connections should exist + fns <- getAddRecvs addRecvMap node1 + length fns `shouldBe` 0 + + -- Access node1 — should trigger creation + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + threadDelay 50000 + + -- Now node1 should have 2 connections + fns1 <- getAddRecvs addRecvMap node1 + length fns1 `shouldBe` 2 + + -- node2 should still have 0 + fns2 <- getAddRecvs addRecvMap node2 + length fns2 `shouldBe` 0 + + -- Feed response and clean up + (firstOf fns1) (encodeResp (RespSimpleString "OK")) + _ <- takeMVar resultMVar + + closeMultiplexPool pool + + it "second access to same node reuses existing muxes" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 2 + + -- First access creates muxes + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + threadDelay 50000 + fns1 <- getAddRecvs addRecvMap node1 + (firstOf fns1) (encodeResp (RespSimpleString "OK")) + _ <- takeMVar resultMVar + + -- Second access should NOT create new connections + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar resultMVar r + threadDelay 10000 + fns2 <- getAddRecvs addRecvMap node1 + -- Still only 2 connections (not 4) + length fns2 `shouldBe` 2 + + -- Feed response via the second mux (round-robin moved to index 1) + (fns2 !! 1) (encodeResp (RespSimpleString "OK2")) + _ <- takeMVar resultMVar + + closeMultiplexPool pool + +multiNodeRoutingSpec :: Spec +multiNodeRoutingSpec = describe "Multi-node routing" $ do + it "commands to different nodes use different multiplexers" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 1 + + -- Submit to node1 + r1MVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["SET", "k1", "v1"]) + putMVar r1MVar r + threadDelay 50000 + + -- Submit to node2 + r2MVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node2 (encodeCmd ["SET", "k2", "v2"]) + putMVar r2MVar r + threadDelay 50000 + + -- Verify both nodes got separate connections + fns1 <- getAddRecvs addRecvMap node1 + fns2 <- getAddRecvs addRecvMap node2 + length fns1 `shouldBe` 1 + length fns2 `shouldBe` 1 + + -- Feed different responses to confirm independent routing + (firstOf fns1) (encodeResp (RespSimpleString "OK1")) + (firstOf fns2) (encodeResp (RespSimpleString "OK2")) + + r1 <- takeMVar r1MVar + r2 <- takeMVar r2MVar + r1 `shouldBe` RespSimpleString "OK1" + r2 `shouldBe` RespSimpleString "OK2" + + closeMultiplexPool pool + + it "three nodes each get independent multiplexer groups" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 2 + + -- Access all three nodes + mvars <- mapM (\node -> do + mv <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node (encodeCmd ["PING"]) + putMVar mv r + return mv + ) [node1, node2, node3] + threadDelay 50000 + + -- Each node should have exactly 2 connections + mapM_ (\node -> do + fns <- getAddRecvs addRecvMap node + length fns `shouldBe` 2 + ) [node1, node2, node3] + + -- Feed responses and verify + mapM_ (\(node, mv) -> do + fns <- getAddRecvs addRecvMap node + (firstOf fns) (encodeResp (RespSimpleString "PONG")) + r <- takeMVar mv + r `shouldBe` RespSimpleString "PONG" + ) (zip [node1, node2, node3] mvars) + + closeMultiplexPool pool + +poolClosureSpec :: Spec +poolClosureSpec = describe "Pool closure" $ do + it "closeMultiplexPool destroys all multiplexers" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 1 + + -- Create muxes for two nodes + r1MVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar r1MVar r + threadDelay 50000 + fns1 <- getAddRecvs addRecvMap node1 + (firstOf fns1) (encodeResp (RespSimpleString "OK")) + _ <- takeMVar r1MVar + + r2MVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node2 (encodeCmd ["PING"]) + putMVar r2MVar r + threadDelay 50000 + fns2 <- getAddRecvs addRecvMap node2 + (firstOf fns2) (encodeResp (RespSimpleString "OK")) + _ <- takeMVar r2MVar + + -- Close the pool + closeMultiplexPool pool + threadDelay 20000 + + -- After close, submitting should create new multiplexers (pool is empty) + -- The old ones should be destroyed — verify by submitting again and seeing new connections + r3MVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["PING"]) + putMVar r3MVar r + threadDelay 50000 + fns3 <- getAddRecvs addRecvMap node1 + -- Should now have 2 connections for node1 (1 old + 1 new) + length fns3 `shouldBe` 2 + -- Feed to the newest one (last in list) + (last fns3) (encodeResp (RespSimpleString "REOPENED")) + r3 <- takeMVar r3MVar + r3 `shouldBe` RespSimpleString "REOPENED" + + closeMultiplexPool pool + + it "closeMultiplexPool on empty pool does not crash" $ do + (connector, _) <- createMockConnector + pool <- createMultiplexPool connector 2 + -- Close without ever using it + closeMultiplexPool pool + -- Close again — should be idempotent + closeMultiplexPool pool + +askingSpec :: Spec +askingSpec = describe "ASKING support (submitToNodeWithAsking)" $ do + it "consumes the ASKING +OK and returns only the command response" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 1 + + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNodeWithAsking pool node1 + (encodeCmd ["ASKING"]) + (encodeCmd ["GET", "mykey"]) + putMVar resultMVar r + threadDelay 50000 + + -- Feed two responses: +OK for ASKING, then the real GET response + fns <- getAddRecvs addRecvMap node1 + length fns `shouldBe` 1 + (firstOf fns) (encodeResp (RespSimpleString "OK")) -- ASKING response + (firstOf fns) (encodeResp (RespBulkString "myvalue")) -- GET response + + r <- takeMVar resultMVar + -- Should return the GET response, not the ASKING +OK + r `shouldBe` RespBulkString "myvalue" + + closeMultiplexPool pool + + it "sends ASKING to the target node, not the original node" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 1 + + -- Submit ASKING+SET to node2 (the ASK target) + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNodeWithAsking pool node2 + (encodeCmd ["ASKING"]) + (encodeCmd ["SET", "migratingkey", "val"]) + putMVar resultMVar r + threadDelay 50000 + + -- node1 should have no connections + fns1 <- getAddRecvs addRecvMap node1 + length fns1 `shouldBe` 0 + + -- node2 should have the connection and receive both responses + fns2 <- getAddRecvs addRecvMap node2 + length fns2 `shouldBe` 1 + (firstOf fns2) (encodeResp (RespSimpleString "OK")) -- ASKING response + (firstOf fns2) (encodeResp (RespSimpleString "OK")) -- SET response + + r <- takeMVar resultMVar + r `shouldBe` RespSimpleString "OK" + + closeMultiplexPool pool + + it "regular submitToNode only consumes one response" $ do + (connector, addRecvMap) <- createMockConnector + pool <- createMultiplexPool connector 1 + + resultMVar <- newEmptyMVar + _ <- forkIO $ do + r <- submitToNode pool node1 (encodeCmd ["GET", "normalkey"]) + putMVar resultMVar r + threadDelay 50000 + + fns <- getAddRecvs addRecvMap node1 + -- Only feed one response — no ASKING involved + (firstOf fns) (encodeResp (RespBulkString "normalvalue")) + + r <- takeMVar resultMVar + r `shouldBe` RespBulkString "normalvalue" + + closeMultiplexPool pool diff --git a/test/MultiplexerSpec.hs b/test/MultiplexerSpec.hs new file mode 100644 --- /dev/null +++ b/test/MultiplexerSpec.hs @@ -0,0 +1,282 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +module Main (main) where + +import Control.Concurrent (forkIO, threadDelay) +import Control.Concurrent.MVar (newEmptyMVar, putMVar, + takeMVar) +import Control.Exception (SomeException, try) +import Control.Monad.IO.Class (liftIO) +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Lazy as LBS +import Data.IORef (IORef, atomicModifyIORef', + newIORef, readIORef) +import Database.Redis.Client (Client (..), + ConnectionStatus (..)) +import Database.Redis.Internal.Multiplexer +import Database.Redis.Resp (Encodable (..), + RespData (..)) +import Test.Hspec + +-- --------------------------------------------------------------------------- +-- Mock client for testing without a real Redis connection +-- --------------------------------------------------------------------------- + +-- | A mock client that uses IORef-based queues for send/receive. +-- Sent data is accumulated in sendBuf; recv reads from recvBuf. +data MockClient (a :: ConnectionStatus) where + MockConnected :: !(IORef ByteString) -- sendBuf (accumulates sent data) + -> !(IORef [ByteString]) -- recvQueue (list of chunks to return) + -> MockClient 'Connected + +instance Client MockClient where + connect = error "MockClient: connect not supported" + close _ = return () + send (MockConnected sendBuf _) lbs = liftIO $ do + let !bs = LBS.toStrict lbs + atomicModifyIORef' sendBuf $ \old -> (old <> bs, ()) + receive (MockConnected sRef recvQueue) = liftIO $ recvLoop sRef recvQueue + +-- | Polling recv loop — retries until data is available. +recvLoop :: IORef ByteString -> IORef [ByteString] -> IO ByteString +recvLoop sRef recvQueue = do + mChunk <- atomicModifyIORef' recvQueue $ \xs -> + case xs of + [] -> ([], Nothing) + (y:ys) -> (ys, Just y) + case mChunk of + Just chunk -> return chunk + Nothing -> do + threadDelay 1000 + recvLoop sRef recvQueue + +-- | Create a mock client and return (client, addRecvData). +-- addRecvData pushes response bytes that the reader thread will consume. +createMockClient :: IO (MockClient 'Connected, ByteString -> IO ()) +createMockClient = do + sendBuf <- newIORef BS.empty + recvQueue <- newIORef [] + let client = MockConnected sendBuf recvQueue + addRecv bs = atomicModifyIORef' recvQueue $ \xs -> (xs ++ [bs], ()) + return (client, addRecv) + +-- | Encode a RespData to strict ByteString (for feeding to mock recv). +encodeResp :: RespData -> ByteString +encodeResp = LBS.toStrict . Builder.toLazyByteString . encode + +-- | Encode a RESP command as a Builder (for submitting to multiplexer). +encodeCmd :: [ByteString] -> Builder.Builder +encodeCmd args = + Builder.byteString ("*" <> bshow (length args) <> "\r\n") + <> foldMap (\a -> Builder.byteString ("$" <> bshow (BS.length a) <> "\r\n" <> a <> "\r\n")) args + where bshow x = LBS.toStrict (Builder.toLazyByteString (Builder.intDec x)) + +-- --------------------------------------------------------------------------- +-- Tests +-- --------------------------------------------------------------------------- + +main :: IO () +main = hspec spec + +spec :: Spec +spec = do + slotPoolSpec + responseSlotSpec + commandQueueBatchingSpec + multiplexerLifecycleSpec + isMultiplexerAliveSpec + +slotPoolSpec :: Spec +slotPoolSpec = describe "SlotPool" $ do + it "allocation returns a valid slot" $ do + pool <- createSlotPool 64 + -- acquireSlot is not exported, but submitCommandPooled uses it internally. + -- We test via the public API: create multiplexer, submit, verify response. + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + addRecv (encodeResp (RespSimpleString "OK")) + resp <- submitCommandPooled pool mux (encodeCmd ["PING"]) + resp `shouldBe` RespSimpleString "OK" + destroyMultiplexer mux + + it "slots are reusable (return-and-reuse)" $ do + pool <- createSlotPool 4 + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + -- Submit multiple commands sequentially — slots should be reused from pool + let n = 20 + mapM_ (\i -> do + addRecv (encodeResp (RespInteger i)) + resp <- submitCommandPooled pool mux (encodeCmd ["GET", "key"]) + resp `shouldBe` RespInteger i + ) [1..n] + destroyMultiplexer mux + + it "striped distribution across cores does not crash" $ do + -- Verify that pool creation with various sizes works + pool1 <- createSlotPool 1 + pool2 <- createSlotPool 16 + pool3 <- createSlotPool 1024 + -- Use each pool to ensure stripes are functional + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + addRecv (encodeResp (RespSimpleString "OK")) + _ <- submitCommandPooled pool1 mux (encodeCmd ["PING"]) + addRecv (encodeResp (RespSimpleString "OK")) + _ <- submitCommandPooled pool2 mux (encodeCmd ["PING"]) + addRecv (encodeResp (RespSimpleString "OK")) + _ <- submitCommandPooled pool3 mux (encodeCmd ["PING"]) + destroyMultiplexer mux + +responseSlotSpec :: Spec +responseSlotSpec = describe "ResponseSlot" $ do + it "write-then-read returns correct value" $ do + pool <- createSlotPool 16 + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + -- Submit a command; the reader thread writes the response to the slot + addRecv (encodeResp (RespBulkString "hello")) + resp <- submitCommandPooled pool mux (encodeCmd ["GET", "key"]) + resp `shouldBe` RespBulkString "hello" + destroyMultiplexer mux + + it "waitSlot blocks until filled (async submit)" $ do + pool <- createSlotPool 16 + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + -- Submit async — slot should not be filled yet + slot <- submitCommandAsync pool mux (encodeCmd ["GET", "key"]) + -- Feed response after a short delay + _ <- forkIO $ do + threadDelay 50000 -- 50ms delay + addRecv (encodeResp (RespBulkString "delayed")) + resp <- waitSlot pool slot + resp `shouldBe` RespBulkString "delayed" + destroyMultiplexer mux + +commandQueueBatchingSpec :: Spec +commandQueueBatchingSpec = describe "Command queue batching" $ do + it "multiple enqueued commands are drained together" $ do + pool <- createSlotPool 64 + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + + -- Submit 5 commands concurrently, then feed 5 responses + let n = 5 + results <- newIORef ([] :: [RespData]) + barriers <- mapM (\_ -> newEmptyMVar) [1..n] + + -- Submit commands from separate threads + mapM_ (\(i, barrier) -> forkIO $ do + resp <- submitCommandPooled pool mux (encodeCmd ["GET", LBS.toStrict $ Builder.toLazyByteString $ Builder.intDec i]) + atomicModifyIORef' results $ \rs -> (rs ++ [resp], ()) + putMVar barrier () + ) (zip [1..n] barriers) + + -- Give commands time to be enqueued + threadDelay 10000 + + -- Feed all responses at once (they should be batched) + let allResponses = mconcat $ map (\i -> encodeResp (RespInteger (fromIntegral (i :: Int)))) [1..n] + addRecv allResponses + + -- Wait for all to complete + mapM_ takeMVar barriers + + -- Verify all responses received + rs <- readIORef results + length rs `shouldBe` n + destroyMultiplexer mux + +multiplexerLifecycleSpec :: Spec +multiplexerLifecycleSpec = describe "Multiplexer lifecycle" $ do + it "create and submit returns correct response" $ do + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + addRecv (encodeResp (RespSimpleString "PONG")) + resp <- submitCommand mux (encodeCmd ["PING"]) + resp `shouldBe` RespSimpleString "PONG" + destroyMultiplexer mux + + it "destroy then submit throws MultiplexerDead" $ do + (client, _) <- createMockClient + mux <- createMultiplexer client (receive client) + destroyMultiplexer mux + -- Small delay for destroy to take effect + threadDelay 10000 + result <- try $ submitCommand mux (encodeCmd ["PING"]) + case result of + Left (e :: SomeException) -> show e `shouldContain` "MultiplexerDead" + Right _ -> expectationFailure "Expected MultiplexerDead exception" + + it "submit-after-destroy with pooled also throws MultiplexerDead" $ do + pool <- createSlotPool 16 + (client, _) <- createMockClient + mux <- createMultiplexer client (receive client) + destroyMultiplexer mux + threadDelay 10000 + result <- try $ submitCommandPooled pool mux (encodeCmd ["GET", "key"]) + case result of + Left (e :: SomeException) -> show e `shouldContain` "MultiplexerDead" + Right _ -> expectationFailure "Expected MultiplexerDead exception" + + it "handles multiple sequential commands correctly" $ do + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + -- Send multiple commands and verify ordering + addRecv (encodeResp (RespSimpleString "OK")) + r1 <- submitCommand mux (encodeCmd ["SET", "k1", "v1"]) + r1 `shouldBe` RespSimpleString "OK" + + addRecv (encodeResp (RespBulkString "v1")) + r2 <- submitCommand mux (encodeCmd ["GET", "k1"]) + r2 `shouldBe` RespBulkString "v1" + + addRecv (encodeResp (RespInteger 1)) + r3 <- submitCommand mux (encodeCmd ["DEL", "k1"]) + r3 `shouldBe` RespInteger 1 + + addRecv (encodeResp RespNullBulkString) + r4 <- submitCommand mux (encodeCmd ["GET", "k1"]) + r4 `shouldBe` RespNullBulkString + + destroyMultiplexer mux + + it "handles RespArray responses" $ do + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + let arrResp = RespArray [RespBulkString "v1", RespBulkString "v2", RespNullBulkString] + addRecv (encodeResp arrResp) + resp <- submitCommand mux (encodeCmd ["MGET", "k1", "k2", "k3"]) + resp `shouldBe` arrResp + destroyMultiplexer mux + + it "handles RespError responses without crashing" $ do + (client, addRecv) <- createMockClient + mux <- createMultiplexer client (receive client) + addRecv (encodeResp (RespError "ERR wrong number of arguments")) + resp <- submitCommand mux (encodeCmd ["GET"]) + resp `shouldBe` RespError "ERR wrong number of arguments" + destroyMultiplexer mux + +isMultiplexerAliveSpec :: Spec +isMultiplexerAliveSpec = describe "isMultiplexerAlive" $ do + it "returns True for a live multiplexer" $ do + (client, _) <- createMockClient + mux <- createMultiplexer client (receive client) + alive <- isMultiplexerAlive mux + alive `shouldBe` True + destroyMultiplexer mux + + it "returns False after destroy" $ do + (client, _) <- createMockClient + mux <- createMultiplexer client (receive client) + destroyMultiplexer mux + threadDelay 10000 + alive <- isMultiplexerAlive mux + alive `shouldBe` False diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,208 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Main where + +import Data.Attoparsec.ByteString.Char8 (parseOnly) +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Char8 as BS8 +import qualified Data.ByteString.Lazy as LBS +import Data.Either (isLeft) +import qualified Data.Map as M +import qualified Data.Set as S +import Database.Redis.Command (wrapInRay) +import Database.Redis.Resp +import Test.Hspec + +main :: IO () +main = hspec $ do + describe "RESP Encoding" $ do + it "encodes simple strings correctly" $ do + Builder.toLazyByteString (encode (RespSimpleString "OK")) `shouldBe` "+OK\r\n" + Builder.toLazyByteString (encode (RespSimpleString "Hello")) `shouldBe` "+Hello\r\n" + + it "encodes errors correctly" $ do + Builder.toLazyByteString (encode (RespError "Error message")) `shouldBe` "-Error message\r\n" + Builder.toLazyByteString (encode (RespError "Another error")) `shouldBe` "-Another error\r\n" + + it "encodes integers correctly" $ do + Builder.toLazyByteString (encode (RespInteger 123)) `shouldBe` ":123\r\n" + Builder.toLazyByteString (encode (RespInteger 0)) `shouldBe` ":0\r\n" + + it "encodes bulk strings correctly" $ do + Builder.toLazyByteString (encode (RespBulkString "foobar")) `shouldBe` "$6\r\nfoobar\r\n" + Builder.toLazyByteString (encode (RespBulkString "")) `shouldBe` "$0\r\n\r\n" + + it "encodes arrays correctly" $ do + Builder.toLazyByteString (encode (RespArray [RespSimpleString "foo", RespInteger 42])) `shouldBe` "*2\r\n+foo\r\n:42\r\n" + Builder.toLazyByteString (encode (RespArray [])) `shouldBe` "*0\r\n" + + it "encodes maps correctly" $ do + let mapData = M.fromList [(RespSimpleString "first", RespInteger 1), (RespSimpleString "second", RespBulkString "bulkString")] + Builder.toLazyByteString (encode (RespMap mapData)) `shouldBe` "*2\r\n+first\r\n:1\r\n+second\r\n$10\r\nbulkString\r\n" + Builder.toLazyByteString (encode (RespMap M.empty)) `shouldBe` "*0\r\n" + + it "encodes sets correctly" $ do + let setData = S.fromList [RespSimpleString "Hello", RespInteger 5, RespBulkString "im very long"] + Builder.toLazyByteString (encode (RespSet setData)) `shouldBe` "~3\r\n+Hello\r\n:5\r\n$12\r\nim very long\r\n" + Builder.toLazyByteString (encode (RespSet S.empty)) `shouldBe` "~0\r\n" + + it "encodes lists of length greater than 10 correctly" $ do + let longList = RespArray (replicate 11 (RespSimpleString "item")) + Builder.toLazyByteString (encode longList) `shouldBe` "*11\r\n" <> mconcat (replicate 11 "+item\r\n") + + it "encodes negative integers correctly" $ do + Builder.toLazyByteString (encode (RespInteger (-1))) `shouldBe` ":-1\r\n" + Builder.toLazyByteString (encode (RespInteger (-999))) `shouldBe` ":-999\r\n" + + it "encodes null bulk strings correctly" $ do + Builder.toLazyByteString (encode RespNullBulkString) `shouldBe` "$-1\r\n" + + it "encodes nested arrays correctly" $ do + let nested = RespArray [RespArray [RespInteger 1, RespInteger 2], RespArray [RespInteger 3]] + Builder.toLazyByteString (encode nested) `shouldBe` "*2\r\n*2\r\n:1\r\n:2\r\n*1\r\n:3\r\n" + + it "encodes large bulk strings correctly" $ do + let bigStr = BS8.pack (replicate 1000 'A') + Builder.toLazyByteString (encode (RespBulkString bigStr)) `shouldBe` + "$1000\r\n" <> LBS.fromStrict bigStr <> "\r\n" + + describe "Redis Command Encoding" $ do + it "encodes MGET commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["MGET", "key1", "key2"])) + `shouldBe` "*3\r\n$4\r\nMGET\r\n$4\r\nkey1\r\n$4\r\nkey2\r\n" + + it "encodes ZADD commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["ZADD", "myzset", "1", "one", "2", "two"])) + `shouldBe` "*6\r\n$4\r\nZADD\r\n$6\r\nmyzset\r\n$1\r\n1\r\n$3\r\none\r\n$1\r\n2\r\n$3\r\ntwo\r\n" + + it "encodes ZRANGE WITHSCORES commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["ZRANGE", "myzset", "0", "-1", "WITHSCORES"])) + `shouldBe` "*5\r\n$6\r\nZRANGE\r\n$6\r\nmyzset\r\n$1\r\n0\r\n$2\r\n-1\r\n$10\r\nWITHSCORES\r\n" + + it "encodes DECR commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["DECR", "counter"])) + `shouldBe` "*2\r\n$4\r\nDECR\r\n$7\r\ncounter\r\n" + + it "encodes PSETEX commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["PSETEX", "temp", "500", "value"])) + `shouldBe` "*4\r\n$6\r\nPSETEX\r\n$4\r\ntemp\r\n$3\r\n500\r\n$5\r\nvalue\r\n" + + it "encodes HMGET commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["HMGET", "myhash", "field1", "field2"])) + `shouldBe` "*4\r\n$5\r\nHMGET\r\n$6\r\nmyhash\r\n$6\r\nfield1\r\n$6\r\nfield2\r\n" + + it "encodes HEXISTS commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["HEXISTS", "myhash", "field1"])) + `shouldBe` "*3\r\n$7\r\nHEXISTS\r\n$6\r\nmyhash\r\n$6\r\nfield1\r\n" + + it "encodes SCARD commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["SCARD", "myset"])) + `shouldBe` "*2\r\n$5\r\nSCARD\r\n$5\r\nmyset\r\n" + + it "encodes SISMEMBER commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["SISMEMBER", "myset", "one"])) + `shouldBe` "*3\r\n$9\r\nSISMEMBER\r\n$5\r\nmyset\r\n$3\r\none\r\n" + + it "encodes GEOADD commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEOADD", "locations", "13.361389", "38.115556", "Palermo"])) + `shouldBe` "*5\r\n$6\r\nGEOADD\r\n$9\r\nlocations\r\n$9\r\n13.361389\r\n$9\r\n38.115556\r\n$7\r\nPalermo\r\n" + + it "encodes GEOHASH commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEOHASH", "locations", "Palermo", "Catania"])) + `shouldBe` "*4\r\n$7\r\nGEOHASH\r\n$9\r\nlocations\r\n$7\r\nPalermo\r\n$7\r\nCatania\r\n" + + it "encodes GEOSEARCH commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEOSEARCH", "locations", "FROMLONLAT", "13", "38", "BYRADIUS", "50", "KM", "ASC"])) + `shouldBe` "*9\r\n$9\r\nGEOSEARCH\r\n$9\r\nlocations\r\n$10\r\nFROMLONLAT\r\n$2\r\n13\r\n$2\r\n38\r\n$8\r\nBYRADIUS\r\n$2\r\n50\r\n$2\r\nKM\r\n$3\r\nASC\r\n" + + it "encodes GEODIST commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEODIST", "geo", "Palermo", "Catania", "KM"])) + `shouldBe` "*5\r\n$7\r\nGEODIST\r\n$3\r\ngeo\r\n$7\r\nPalermo\r\n$7\r\nCatania\r\n$2\r\nKM\r\n" + + it "encodes GEOPOS commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEOPOS", "geo", "Palermo", "Catania"])) + `shouldBe` "*4\r\n$6\r\nGEOPOS\r\n$3\r\ngeo\r\n$7\r\nPalermo\r\n$7\r\nCatania\r\n" + + it "encodes GEORADIUS commands with flags correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEORADIUS", "geo", "15.0", "37.0", "200.0", "KM", "WITHDIST", "ASC"])) + `shouldBe` "*8\r\n$9\r\nGEORADIUS\r\n$3\r\ngeo\r\n$4\r\n15.0\r\n$4\r\n37.0\r\n$5\r\n200.0\r\n$2\r\nKM\r\n$8\r\nWITHDIST\r\n$3\r\nASC\r\n" + + it "encodes GEORADIUS_RO commands with flags correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEORADIUS_RO", "geo", "15.0", "37.0", "200.0", "KM", "WITHCOORD"])) + `shouldBe` "*7\r\n$12\r\nGEORADIUS_RO\r\n$3\r\ngeo\r\n$4\r\n15.0\r\n$4\r\n37.0\r\n$5\r\n200.0\r\n$2\r\nKM\r\n$9\r\nWITHCOORD\r\n" + + it "encodes GEORADIUSBYMEMBER commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEORADIUSBYMEMBER", "geo", "Palermo", "200.0", "KM", "WITHDIST", "DESC"])) + `shouldBe` "*7\r\n$17\r\nGEORADIUSBYMEMBER\r\n$3\r\ngeo\r\n$7\r\nPalermo\r\n$5\r\n200.0\r\n$2\r\nKM\r\n$8\r\nWITHDIST\r\n$4\r\nDESC\r\n" + + it "encodes GEORADIUSBYMEMBER_RO commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEORADIUSBYMEMBER_RO", "geo", "Palermo", "200.0", "KM", "COUNT", "5", "ANY"])) + `shouldBe` "*8\r\n$20\r\nGEORADIUSBYMEMBER_RO\r\n$3\r\ngeo\r\n$7\r\nPalermo\r\n$5\r\n200.0\r\n$2\r\nKM\r\n$5\r\nCOUNT\r\n$1\r\n5\r\n$3\r\nANY\r\n" + + it "encodes GEOSEARCHSTORE commands with STOREDIST correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["GEOSEARCHSTORE", "dest", "src", "FROMMEMBER", "Palermo", "BYBOX", "100.0", "200.0", "KM", "DESC", "STOREDIST"])) + `shouldBe` "*11\r\n$14\r\nGEOSEARCHSTORE\r\n$4\r\ndest\r\n$3\r\nsrc\r\n$10\r\nFROMMEMBER\r\n$7\r\nPalermo\r\n$5\r\nBYBOX\r\n$5\r\n100.0\r\n$5\r\n200.0\r\n$2\r\nKM\r\n$4\r\nDESC\r\n$9\r\nSTOREDIST\r\n" + + it "encodes CLIENT SETINFO commands correctly" $ do + Builder.toLazyByteString (encode (wrapInRay ["CLIENT", "SETINFO", "LIB-NAME", "redis-client"])) + `shouldBe` "*4\r\n$6\r\nCLIENT\r\n$7\r\nSETINFO\r\n$8\r\nLIB-NAME\r\n$12\r\nredis-client\r\n" + + describe "RESP Parsing" $ do + it "parses simple strings correctly" $ do + parseOnly parseRespData "+OK\r\n" `shouldBe` Right (RespSimpleString "OK") + parseOnly parseRespData "+Hello\r\n" `shouldBe` Right (RespSimpleString "Hello") + + it "parses errors correctly" $ do + parseOnly parseRespData "-Error message\r\n" `shouldBe` Right (RespError "Error message") + parseOnly parseRespData "-Another error\r\n" `shouldBe` Right (RespError "Another error") + + it "parses integers correctly" $ do + parseOnly parseRespData ":123\r\n" `shouldBe` Right (RespInteger 123) + parseOnly parseRespData ":0\r\n" `shouldBe` Right (RespInteger 0) + + it "parses bulk strings correctly" $ do + parseOnly parseRespData "$6\r\nfoobar\r\n" `shouldBe` Right (RespBulkString "foobar") + parseOnly parseRespData "$0\r\n\r\n" `shouldBe` Right (RespBulkString "") + + it "parses null bulk strings correctly" $ do + parseOnly parseRespData "$-1\r\n" `shouldBe` Right RespNullBulkString + + it "parses arrays correctly" $ do + parseOnly parseRespData "*2\r\n+foo\r\n:42\r\n" `shouldBe` Right (RespArray [RespSimpleString "foo", RespInteger 42]) + parseOnly parseRespData "*0\r\n" `shouldBe` Right (RespArray []) + + it "parses maps correctly" $ do + let mapData = M.fromList [(RespSimpleString "first", RespInteger 1), (RespSimpleString "second", RespBulkString "bulkString")] + parseOnly parseRespData "%2\r\n+first\r\n:1\r\n+second\r\n$10\r\nbulkString\r\n" `shouldBe` Right (RespMap mapData) + parseOnly parseRespData "%0\r\n" `shouldBe` Right (RespMap M.empty) + + it "parses sets correctly" $ do + let setData = S.fromList [RespSimpleString "Hello", RespInteger 5, RespBulkString "im very long"] + parseOnly parseRespData "~3\r\n+Hello\r\n:5\r\n$12\r\nim very long\r\n" `shouldBe` Right (RespSet setData) + parseOnly parseRespData "~0\r\n" `shouldBe` Right (RespSet S.empty) + + it "parses lists of length greater than 10 correctly" $ do + let longListEncoded = "*11\r\n" <> mconcat (replicate 11 "+item\r\n") + parseOnly parseRespData (LBS.toStrict longListEncoded) `shouldBe` Right (RespArray (replicate 11 (RespSimpleString "item"))) + + it "parses negative integers correctly" $ do + parseOnly parseRespData ":-1\r\n" `shouldBe` Right (RespInteger (-1)) + parseOnly parseRespData ":-999\r\n" `shouldBe` Right (RespInteger (-999)) + + it "parses null bulk strings and roundtrips correctly" $ do + let encoded = Builder.toLazyByteString (encode RespNullBulkString) + parseOnly parseRespData (LBS.toStrict encoded) `shouldBe` Right RespNullBulkString + + it "parses nested arrays correctly" $ do + let nested = RespArray [RespArray [RespInteger 1, RespInteger 2], RespArray [RespInteger 3]] + encoded = Builder.toLazyByteString (encode nested) + parseOnly parseRespData (LBS.toStrict encoded) `shouldBe` Right nested + + it "parses empty collections correctly" $ do + parseOnly parseRespData "*0\r\n" `shouldBe` Right (RespArray []) + parseOnly parseRespData "%0\r\n" `shouldBe` Right (RespMap M.empty) + parseOnly parseRespData "~0\r\n" `shouldBe` Right (RespSet S.empty) + + it "fails on incomplete RESP data" $ do + parseOnly parseRespData "+OK" `shouldSatisfy` isLeft + parseOnly parseRespData "$6\r\nfoo" `shouldSatisfy` isLeft