diff --git a/AUTHORS b/AUTHORS --- a/AUTHORS +++ b/AUTHORS @@ -1,6 +1,11 @@ Maintainers ----------- +- Kyle Butt + +Prior Maintainers +----------------- + - Toralf Wittner - Roman S. Borschel diff --git a/CHANGELOG b/CHANGELOG --- a/CHANGELOG +++ b/CHANGELOG @@ -1,3 +1,46 @@ +2.0.0 +----- +- Support multiple control connections with an upper limit and policy-controlled + priority. +- Cabal file improvements: bounds, version, etc. +- Switch from unmaintained cryptonite to crypton +- Add sufficient API tools for token-routing aware policy: + - Accessors to lookup the replication for a keyspace, and to get the replicas + for a given prepared query and its query values. +- Pass prepared query and values to policy when selecting host +- Track replication information for keyspaces and nodes +- Better handling of prepared queries: Prevent re-prepare if the user has + eagerly prepared a query. +- Better code sharing between default policies. +- Host events to pass Hosts, not addresses. This allows policies to better + handle ups and downs, as well as re-numbering. It makes it easier to write a + policy that prefers a given DC. +- Track peers more carefully, and track them via UUID, not address. +- Allow connection upgrading via min-max protocol version. Instead of setting a + fixed protocol version, allow setting a min and max, and upgrade when + possible. + +Breaking Changes: +- Policy has a new `priority` function +- Policy Events now have `Host`s instead of addresses, except for the new + `AddrDown` for the case of a `HostDown` CQL event where the `Host` in question + can't be found. +- Policy `select` function accepts an additional argument with token-based + routing information if it is available. +- Settings has `_minProtoVer` and `_maxProtoVer` instead of `_protoVersion` +- Settings has an additional field `_maxControlConnections` + +Porting Guide: + If using one of the included policies, the only necessary change should be to +set a minimum and maximum protocol versions instead of a single version. + If not using an included policy, an existing policy can be ported by ignoring +the additional argument to select, and by adjusting the event handler to handle +`Host`s rather than addresses. Hosts have an address field, so porting should be +straightforward. However, the changed policy API makes it possible to have a +token aware routing policy. It's worth taking a look at that while porting. The +improved `HostEvent` API should make it easier to handle host renumbering more +gracefully as well. + 1.1.1 ----- - Compatibility with `network >= 3`. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -33,6 +33,19 @@ Client to node communication can optionally use transport layer security (using HsOpenSSL). +Testing +======= + +The tests assume that a CQL cluster is available and running. The host should +either be available at localhost, or the environment variable CASSANDRA_HOST. To +spin up a single host for a simple test, you can do the following: + +``` +docker run --rm -ti --cpus 4 --memory 8G --name cql-io-cassandra \ + -p ::1:9042:9042 -v /tmp/cql-io-test-data:/var/lib/cassandra \ + --env MAX_HEAP_SIZE=1536M --env NEW_HEAP_SIZE=256M cassandra:5.0.5-jammy +``` + License ======= diff --git a/cql-io.cabal b/cql-io.cabal --- a/cql-io.cabal +++ b/cql-io.cabal @@ -1,18 +1,17 @@ name: cql-io -version: 1.1.1 +version: 2.0.0 synopsis: Cassandra CQL client. license: MPL-2.0 license-file: LICENSE author: Toralf Wittner -maintainer: Toralf Wittner , - Roman S. Borschel -copyright: (C) 2014-2016 Toralf Wittner -homepage: https://gitlab.com/twittner/cql-io/ -bug-reports: https://gitlab.com/twittner/cql-io/issues +maintainer: Kyle Butt +copyright: (C) 2014-2016 Toralf Wittner, 2025 Kyle Butt +homepage: https://github.com/iteratee/cql-io +bug-reports: https://github.com/iteratee/cql-io/issues category: Database build-type: Simple -cabal-version: >= 2.0 -extra-source-files: README.md +cabal-version: 2.0 +extra-doc-files: README.md CHANGELOG AUTHORS @@ -28,7 +27,8 @@ . * /Customisable load-balancing policies/. In addition to pre-built LB policies such as round-robin, users of this library can provide their - own policies if desired. + own policies if desired. The policy API is rich enough to support token + aware routing policies. . * /Support for connection streams/. Requests can be multiplexed over a few connections. @@ -45,12 +45,12 @@ source-repository head type: git - location: https://gitlab.com/twittner/cql-io + location: https://github.com/iteratee/cql-io library default-language: Haskell2010 hs-source-dirs: src/api - ghc-options: -Wall -O2 -fwarn-tabs + ghc-options: -Wall -fwarn-tabs exposed-modules: Database.CQL.IO @@ -66,7 +66,7 @@ library cql-io-lib default-language: Haskell2010 hs-source-dirs: src/lib - ghc-options: -Wall -O2 -fwarn-tabs + ghc-options: -Wall -fwarn-tabs exposed-modules: Database.CQL.IO.Batch @@ -84,6 +84,7 @@ Database.CQL.IO.Pool Database.CQL.IO.PrepQuery Database.CQL.IO.Protocol + Database.CQL.IO.Replication Database.CQL.IO.Settings Database.CQL.IO.Signal Database.CQL.IO.Sync @@ -91,58 +92,67 @@ Database.CQL.IO.Timeouts build-depends: - async >= 2.2 - , auto-update >= 0.1 - , base >= 4.9 && < 5.0 - , bytestring >= 0.10 - , containers >= 0.5 - , cql >= 4.0 - , cryptonite >= 0.13 - , data-default-class - , exceptions >= 0.4 - , hashable >= 1.2 - , iproute >= 1.3 - , HsOpenSSL >= 0.11 - , lens >= 4.4 - , mtl >= 2.1 - , mwc-random >= 0.13 - , retry >= 0.7 - , network >= 2.4 - , semigroups >= 0.15 - , stm >= 2.4 - , text >= 0.11 - , time >= 1.4 - , transformers >= 0.3 - , unliftio-core >= 0.1.1 - , unordered-containers >= 0.2 - , uuid >= 1.2.6 - , vector >= 0.10 + async >= 2.2 && < 2.3 + , attoparsec >= 0.14 && < 1 + , auto-update >= 0.1 && < 0.3 + , base >= 4.9 && < 5 + , bluefin >= 0.0.15 && < 0.18 + , bytestring >= 0.10 && < 0.13 + , containers >= 0.6.6 && < 0.9 + , cql >= 4.1 && < 4.2 + , crypton >= 1.0 && < 1.1 + , data-default-class >= 0.2 && < 0.3 + , exceptions >= 0.4 && < 0.11 + , hashable >= 1.2 && < 1.6 + , HsOpenSSL >= 0.11 && < 0.12 + , iproute >= 1.3 && < 1.8 + , lens >= 4.4 && < 5.4 + , mtl >= 2.1 && < 2.4 + , mwc-random >= 0.13 && < 0.16 + , network >= 2.4 && < 3.3 + , persist >= 1.0 && < 1.1 + , random-shuffle >= 0.0.4 && < 0.1 + , retry >= 0.7 && < 0.10 + , semigroups >= 0.15 && < 0.21 + , semver >= 0.4.0.1 && < 1 + , stm >= 2.5.1 && < 2.6 + , text >= 0.11 && < 2.2 + , time >= 1.4 && < 1.16 + , transformers >= 0.6.0.2 && < 0.7 + , unliftio-core >= 0.1.1 && < 0.3 + , unordered-containers >= 0.2 && < 0.3 + , uuid >= 1.2.6 && < 1.4 + , vector >= 0.10 && < 0.14 test-suite cql-io-tests type: exitcode-stdio-1.0 default-language: Haskell2010 main-is: Main.hs hs-source-dirs: src/test - ghc-options: -threaded -Wall -O2 -fwarn-tabs + ghc-options: -threaded -Wall -fwarn-tabs build-depends: base + , QuickCheck , async , containers , cql , cql-io , cql-io-lib + , formatting , Decimal , iproute , mtl - , tasty >= 0.11 - , tasty-hunit >= 0.9 - , text , primes , raw-strings-qq + , tasty >= 0.11 + , tasty-hunit >= 0.9 + , tasty-quickcheck >= 0.8 + , text , time , uuid other-modules: Test.Database.CQL.IO Test.Database.CQL.IO.Jobs + Test.Database.CQL.IO.Replication diff --git a/src/api/Database/CQL/IO.hs b/src/api/Database/CQL/IO.hs --- a/src/api/Database/CQL/IO.hs +++ b/src/api/Database/CQL/IO.hs @@ -46,7 +46,8 @@ , setPortNumber , PrepareStrategy (..) , setPrepareStrategy - , setProtocolVersion + , setMinProtocolVersion + , setMaxProtocolVersion , setResponseTimeout , setSendTimeout , setRetrySettings @@ -100,6 +101,8 @@ , HostEvent (..) , InetAddr (..) , hostAddr + , broadcastAddr + , hostId , dataCentre , rack @@ -157,6 +160,11 @@ , retry , once + -- ** Replication + , getReplication + , getSimpleReplicas + , getDCReplicas + -- ** Low-Level Queries -- $low-level-queries , RunQ (..) @@ -256,7 +264,7 @@ -> m (HostResponse k a b) instance RunQ QueryString where - runQ q p = request (RqQuery (Query q p)) + runQ q p = request (RqQuery (Query q p)) Nothing instance RunQ PrepQuery where runQ q = liftClient . execute q diff --git a/src/lib/Database/CQL/IO/Batch.hs b/src/lib/Database/CQL/IO/Batch.hs --- a/src/lib/Database/CQL/IO/Batch.hs +++ b/src/lib/Database/CQL/IO/Batch.hs @@ -5,6 +5,7 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} module Database.CQL.IO.Batch ( BatchM @@ -36,7 +37,7 @@ batch :: BatchM a -> Client () batch m = do b <- execStateT (unBatchM m) (Batch BatchLogged [] Quorum Nothing) - r <- executeWithPrepare Nothing (RqBatch b :: Raw Request) + r <- executeWithPrepare Nothing (RqBatch b :: Raw Request) Nothing getResult r >>= \case VoidResult -> return () _ -> unexpected r @@ -47,17 +48,20 @@ b { batchQuery = BatchQuery q p : batchQuery b } -- | Add a prepared query to this batch. -addPrepQuery :: (Show a, Tuple a, Tuple b) => PrepQuery W a b -> a -> BatchM () +addPrepQuery :: forall a b. (Show a, Tuple a, Tuple b) => PrepQuery W a b -> a -> BatchM () addPrepQuery q p = BatchM $ do pq <- lift preparedQueries - maybe (fresh pq) add =<< liftIO (atomically (lookupQueryId q pq)) + maybe (fresh pq) add =<< liftIO (atomically (lookupQueryLocal q pq)) where fresh pq = do - i <- snd <$> lift (prepare Nothing (queryString q)) - liftIO $ atomically (insert q i pq) - add i + spq <- snd <$> lift (prepare Nothing q) + liftIO $ atomically (insert q spq pq) + add spq - add i = modify' $ \b -> b { batchQuery = BatchPrepared i p : batchQuery b } + add spq = do + let qId :: QueryId W a b + qId = spqId spq + modify' $ \b -> b { batchQuery = BatchPrepared qId p : batchQuery b } -- | Set the type of this batch. setType :: BatchType -> BatchM () diff --git a/src/lib/Database/CQL/IO/Client.hs b/src/lib/Database/CQL/IO/Client.hs --- a/src/lib/Database/CQL/IO/Client.hs +++ b/src/lib/Database/CQL/IO/Client.hs @@ -3,804 +3,1754 @@ -- file, You can obtain one at http://mozilla.org/MPL/2.0/. {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE ViewPatterns #-} - -module Database.CQL.IO.Client - ( Client - , MonadClient (..) - , ClientState - , DebugInfo (..) - , ControlState (..) - , runClient - , init - , shutdown - , request - , requestN - , request1 - , execute - , executeWithPrepare - , prepare - , retry - , once - , debugInfo - , preparedQueries - , withPrepareStrategy - , getResult - , unexpected - , C.defQueryParams - ) where - -import Control.Applicative -import Control.Concurrent (threadDelay, forkIO) -import Control.Concurrent.Async (async, wait) -import Control.Concurrent.STM (STM, atomically) -import Control.Concurrent.STM.TVar -import Control.Exception (IOException, SomeAsyncException (..)) -import Control.Lens (makeLenses, (^.), set, over, view) -import Control.Monad (when, unless) -import Control.Monad.Catch -import Control.Monad.IO.Class -import Control.Monad.IO.Unlift -import Control.Monad.Reader (ReaderT (..), runReaderT, MonadReader, ask) -import Control.Monad.Trans.Class -import Control.Monad.Trans.Except -import Control.Retry (capDelay, exponentialBackoff, rsIterNumber) -import Control.Retry (recovering) -import Data.Foldable (for_, foldrM) -import Data.List (find) -import Data.Map.Strict (Map) -import Data.Maybe (fromMaybe, listToMaybe) -import Data.Semigroup -import Data.Text.Encoding (encodeUtf8) -import Data.Word -import Database.CQL.IO.Cluster.Host -import Database.CQL.IO.Cluster.Policies -import Database.CQL.IO.Connection (Connection, host, Raw) -import Database.CQL.IO.Connection.Settings -import Database.CQL.IO.Exception -import Database.CQL.IO.Jobs -import Database.CQL.IO.Log -import Database.CQL.IO.Pool (Pool) -import Database.CQL.IO.PrepQuery (PrepQuery, PreparedQueries) -import Database.CQL.IO.Settings -import Database.CQL.IO.Signal -import Database.CQL.IO.Timeouts (TimeoutManager) -import Database.CQL.Protocol hiding (Map) -import OpenSSL.Session (SomeSSLException) -import Prelude hiding (init) - -import qualified Control.Monad.Reader as Reader -import qualified Control.Monad.State.Strict as S -import qualified Control.Monad.State.Lazy as LS -import qualified Data.List.NonEmpty as NE -import qualified Data.Map.Strict as Map -import qualified Database.CQL.IO.Cluster.Discovery as Disco -import qualified Database.CQL.IO.Connection as C -import qualified Database.CQL.IO.Pool as Pool -import qualified Database.CQL.IO.PrepQuery as PQ -import qualified Database.CQL.IO.Timeouts as TM -import qualified Database.CQL.Protocol as Cql - -data ControlState - = Connected - | Reconnecting - | Disconnected - deriving (Eq, Ord, Show) - -data Control = Control - { _state :: !ControlState - , _connection :: !Connection - } - -data Context = Context - { _settings :: !Settings - , _timeouts :: !TimeoutManager - , _sigMonit :: !(Signal HostEvent) - } - --- | Opaque client state/environment. -data ClientState = ClientState - { _context :: !Context - , _policy :: !Policy - , _prepQueries :: !PreparedQueries - , _control :: !(TVar Control) - , _hostmap :: !(TVar (Map Host Pool)) - , _jobs :: !(Jobs InetAddr) - } - -makeLenses ''Control -makeLenses ''Context -makeLenses ''ClientState - --- | The Client monad. --- --- A simple reader monad on `IO` around some internal state. Prior to executing --- this monad via 'runClient', its state must be initialised through --- 'Database.CQL.IO.Client.init' and after finishing operation it should be --- terminated with 'shutdown'. --- --- To lift 'Client' actions into another monad, see 'MonadClient'. -newtype Client a = Client - { client :: ReaderT ClientState IO a - } deriving ( Functor - , Applicative - , Monad - , MonadIO - , MonadUnliftIO - , MonadThrow - , MonadCatch - , MonadMask - , MonadReader ClientState - ) - --- | Monads in which 'Client' actions may be embedded. -class (MonadIO m, MonadThrow m) => MonadClient m - where - -- | Lift a computation from the 'Client' monad. - liftClient :: Client a -> m a - -- | Execute an action with a modified 'ClientState'. - localState :: (ClientState -> ClientState) -> m a -> m a - -instance MonadClient Client where - liftClient = id - localState = Reader.local - -instance MonadClient m => MonadClient (ReaderT r m) where - liftClient = lift . liftClient - localState f m = ReaderT (localState f . runReaderT m) - -instance MonadClient m => MonadClient (S.StateT s m) where - liftClient = lift . liftClient - localState f m = S.StateT (localState f . S.runStateT m) - -instance MonadClient m => MonadClient (LS.StateT s m) where - liftClient = lift . liftClient - localState f m = LS.StateT (localState f . LS.runStateT m) - -instance MonadClient m => MonadClient (ExceptT e m) where - liftClient = lift . liftClient - localState f m = ExceptT $ localState f (runExceptT m) - ------------------------------------------------------------------------------ --- API - --- | Execute the client monad. -runClient :: MonadIO m => ClientState -> Client a -> m a -runClient p a = liftIO $ runReaderT (client a) p - --- | Use given 'RetrySettings' during execution of some client action. -retry :: MonadClient m => RetrySettings -> m a -> m a -retry r = localState (set (context.settings.retrySettings) r) - --- | Execute a client action once, without retries, i.e. --- --- @once action = retry noRetry action@. --- --- Primarily for use in applications where global 'RetrySettings' --- are configured and need to be selectively disabled for individual --- queries. -once :: MonadClient m => m a -> m a -once = retry noRetry - --- | Change the default 'PrepareStrategy' for the given client action. -withPrepareStrategy :: MonadClient m => PrepareStrategy -> m a -> m a -withPrepareStrategy s = localState (set (context.settings.prepStrategy) s) - --- | Send a 'Request' to the server and return a 'Response'. --- --- This function will first ask the clients load-balancing 'Policy' for --- some host and use its connection pool to acquire a connection for --- request transmission. --- --- If all available hosts are busy (i.e. their connection pools are fully --- utilised), the function will block until a connection becomes available --- or the maximum wait-queue length has been reached. --- --- The request is retried according to the configured 'RetrySettings'. -request :: (MonadClient m, Tuple a, Tuple b) => Request k a b -> m (HostResponse k a b) -request a = liftClient $ do - n <- liftIO . hostCount =<< view policy - withRetries (requestN n) a - --- | Send a request to a host chosen by the configured host policy. --- --- Tries up to @max(1,n)@ hosts. If no host can execute the request, --- a 'HostError' is thrown. Specifically: --- --- * If no host is available from the 'Policy', 'NoHostAvailable' is thrown. --- * If no host can execute the request, e.g. because all streams --- on all connections are occupied, 'HostsBusy' is thrown. -requestN :: (Tuple b, Tuple a) - => Word - -> Request k a b - -> ClientState - -> Client (HostResponse k a b) -requestN !n a s = liftIO (select (s^.policy)) >>= \case - Nothing -> replaceControl >> throwM NoHostAvailable - Just h -> tryRequest1 h a s >>= \case - Just hr -> return hr - Nothing -> if n > 1 - then requestN (n - 1) a s - else throwM HostsBusy - --- | Send a 'Request' to a specific 'Host'. --- --- If the request cannot be executed on the given host, e.g. --- because all connections are occupied, 'HostsBusy' is thrown. -request1 :: (Tuple a, Tuple b) - => Host - -> Request k a b - -> ClientState - -> Client (HostResponse k a b) -request1 h r s = do - rs <- tryRequest1 h r s - maybe (throwM HostsBusy) return rs - --- | Try to send a 'Request' to a specific 'Host'. --- --- If the request cannot be executed on the given host, e.g. --- because all connections are occupied, 'Nothing' is returned. -tryRequest1 :: (Tuple a, Tuple b) - => Host - -> Request k a b - -> ClientState - -> Client (Maybe (HostResponse k a b)) -tryRequest1 h a s = do - pool <- Map.lookup h <$> readTVarIO' (s^.hostmap) - case pool of - Just p -> do - result <- Pool.with p exec `catches` handlers - for_ result $ \(HostResponse _ r) -> - for_ (Cql.warnings r) $ \w -> - logWarn' $ "Server warning: " <> byteString (encodeUtf8 w) - return result - Nothing -> do - logError' $ "No pool for host: " <> string8 (show h) - p' <- mkPool (s^.context) h - atomically' $ modifyTVar' (s^.hostmap) (Map.alter (maybe (Just p') Just) h) - tryRequest1 h a s - where - exec c = do - r <- C.request c a - return $ HostResponse h r - - handlers = - [ Handler $ \(e :: ConnectionError) -> onConnectionError e - , Handler $ \(e :: IOException) -> onConnectionError e - , Handler $ \(e :: SomeSSLException) -> onConnectionError e - ] - - onConnectionError exc = do - e <- ask - logWarn' (string8 (show exc)) - -- Tell the policy that the host is down until monitoring confirms - -- it is still up, which will be signalled by a subsequent 'HostUp' - -- event. - liftIO $ ignore $ onEvent (e^.policy) (HostDown (h^.hostAddr)) - runJob_ (e^.jobs) (h^.hostAddr) $ - runClient e $ monitor (Ms 0) (Ms 30000) h - -- Any connection error may indicate a problem with the - -- control connection, if it uses the same host. - ch <- fmap (view (connection.host)) . readTVarIO' =<< view control - when (h == ch) $ do - ok <- checkControl - unless ok replaceControl - throwM exc - ------------------------------------------------------------------------------- --- Prepared queries - --- | Execute the given request. If an 'Unprepared' error is returned, this --- function will automatically try to re-prepare the query and re-execute --- the original request using the same host which was used for re-preparation. -executeWithPrepare :: (Tuple b, Tuple a) - => Maybe Host - -> Request k a b - -> Client (HostResponse k a b) -executeWithPrepare mh rq - | Just h <- mh = exec (request1 h) - | otherwise = do - p <- view policy - n <- liftIO $ hostCount p - exec (requestN n) - where - exec action = do - r <- withRetries action rq - case hrResponse r of - RsError _ _ (Unprepared _ i) -> do - pq <- preparedQueries - qs <- atomically' (PQ.lookupQueryString (QueryId i) pq) - case qs of - Nothing -> throwM $ UnexpectedQueryId (QueryId i) - Just s -> do - (h, _) <- prepare (Just LazyPrepare) (s :: Raw QueryString) - executeWithPrepare (Just h) rq - _ -> return r - --- | Prepare the given query according to the given 'PrepareStrategy', --- returning the resulting 'QueryId' and 'Host' which was used for --- preparation. -prepare :: (Tuple b, Tuple a) => Maybe PrepareStrategy -> QueryString k a b -> Client (Host, QueryId k a b) -prepare (Just LazyPrepare) qs = do - s <- ask - n <- liftIO $ hostCount (s^.policy) - r <- withRetries (requestN n) (RqPrepare (Prepare qs)) - getPreparedQueryId r - -prepare (Just EagerPrepare) qs = view policy - >>= liftIO . current - >>= mapM (action (RqPrepare (Prepare qs))) - >>= first - where - action rq h = withRetries (request1 h) rq >>= getPreparedQueryId - - first (x:_) = return x - first [] = replaceControl >> throwM NoHostAvailable - -prepare Nothing qs = do - ps <- view (context.settings.prepStrategy) - prepare (Just ps) qs - --- | Execute a prepared query (transparently re-preparing if necessary). -execute :: (Tuple b, Tuple a) => PrepQuery k a b -> QueryParams a -> Client (HostResponse k a b) -execute q p = do - pq <- view prepQueries - maybe (new pq) (exec Nothing) =<< atomically' (PQ.lookupQueryId q pq) - where - exec h i = executeWithPrepare h (RqExecute (Execute i p)) - new pq = do - (h, i) <- prepare (Just LazyPrepare) (PQ.queryString q) - atomically' (PQ.insert q i pq) - exec (Just h) i - -prepareAllQueries :: Host -> Client () -prepareAllQueries h = do - pq <- view prepQueries - qs <- atomically' $ PQ.queryStrings pq - for_ qs $ \q -> - let qry = QueryString q :: Raw QueryString in - withRetries (request1 h) (RqPrepare (Prepare qry)) - ------------------------------------------------------------------------------- --- Debug info - -data DebugInfo = DebugInfo - { policyInfo :: String - -- ^ Host 'Policy' string representation. - , jobInfo :: [InetAddr] - -- ^ Hosts with running background jobs (e.g. monitoring of hosts - -- currently considered down). - , hostInfo :: [Host] - -- ^ All known hosts. - , controlInfo :: (Host, ControlState) - -- ^ Control connection information. - } - -instance Show DebugInfo where - show dbg = showString "running jobs: " - . shows (jobInfo dbg) - . showString "\nknown hosts: " - . shows (hostInfo dbg) - . showString "\npolicy info: " - . shows (policyInfo dbg) - . showString "\ncontrol host: " - . shows (controlInfo dbg) - $ "" - -debugInfo :: MonadClient m => m DebugInfo -debugInfo = liftClient $ do - hosts <- Map.keys <$> (readTVarIO' =<< view hostmap) - pols <- liftIO . display =<< view policy - jbs <- listJobKeys =<< view jobs - ctrl <- (\(Control s c) -> (c^.host, s)) <$> (readTVarIO' =<< view control) - return $ DebugInfo pols jbs hosts ctrl - -preparedQueries :: Client PreparedQueries -preparedQueries = view prepQueries - ------------------------------------------------------------------------------ --- Initialisation - --- | Initialise client state with the given 'Settings' using the provided --- 'Logger' for all it's logging output. -init :: MonadIO m => Settings -> m ClientState -init s = liftIO $ do - tom <- TM.create (Ms 250) - ctx <- Context s tom <$> signal - bracketOnError (mkContact ctx) C.close $ \con -> do - pol <- s^.policyMaker - cst <- ClientState ctx - <$> pure pol - <*> PQ.new - <*> newTVarIO (Control Connected con) - <*> newTVarIO Map.empty - <*> newJobs - ctx^.sigMonit |-> onEvent pol - runClient cst (setupControl con) - return cst - --- | Try to establish a connection to one of the initial contacts. -mkContact :: Context -> IO Connection -mkContact (Context s t _) = tryAll (s^.contacts) mkConnection - where - mkConnection h = do - as <- C.resolve h (s^.portnumber) - NE.fromList as `tryAll` doConnect - - doConnect a = do - logDebug (s^.logger) $ "Connecting to " <> string8 (show a) - c <- C.connect (s^.connSettings) t (s^.protoVersion) (s^.logger) (Host a "" "") - return c - -discoverPeers :: MonadIO m => Context -> Connection -> m [Host] -discoverPeers ctx c = liftIO $ do - let p = ctx^.settings.portnumber - map (peer2Host p . asRecord) <$> C.query c One Disco.peers () - -mkPool :: MonadIO m => Context -> Host -> m Pool -mkPool ctx h = liftIO $ do - let s = ctx^.settings - let m = s^.connSettings.maxStreams - Pool.create (connOpen s) connClose (ctx^.settings.logger) (s^.poolSettings) m - where - lgr = ctx^.settings.logger - - connOpen s = do - c <- C.connect (s^.connSettings) (ctx^.timeouts) (s^.protoVersion) lgr h - logDebug lgr $ "Connection established: " <> string8 (show c) - return c - - connClose c = do - C.close c - logDebug lgr $ "Connection closed: " <> string8 (show c) - ------------------------------------------------------------------------------ --- Termination - --- | Terminate client state, i.e. end all running background checks and --- shutdown all connection pools. Once this is entered, the client --- will eventually be shut down, though an asynchronous exception can --- interrupt the wait for that to occur. -shutdown :: MonadIO m => ClientState -> m () -shutdown s = liftIO $ asyncShutdown >>= wait - where - asyncShutdown = async $ do - TM.destroy (s^.context.timeouts) True - cancelJobs (s^.jobs) - ignore $ C.close . view connection =<< readTVarIO (s^.control) - mapM_ Pool.destroy . Map.elems =<< readTVarIO (s^.hostmap) - ------------------------------------------------------------------------------ --- Monitoring - --- | @monitor initialDelay maxDelay host@ tries to establish a connection --- to @host@ after @initialDelay@. If the connection attempt fails, it is --- retried with exponentially increasing delays, up to a maximum delay of --- @maxDelay@. When a connection attempt suceeds, a 'HostUp' event is --- signalled. --- --- The function returns when one of the following conditions is met: --- --- 1. The connection attempt suceeds. --- 2. The host is no longer found to be in the client's known host map. --- --- I.e. as long as the host is still known to the client and is unreachable, the --- connection attempts continue. Both @initialDelay@ and @maxDelay@ are bounded --- by a limit of 5 minutes. -monitor :: Milliseconds -> Milliseconds -> Host -> Client () -monitor initial maxDelay h = do - liftIO $ threadDelay (toMicros initial) - logInfo' $ "Monitoring: " <> string8 (show h) - hostCheck 0 - where - hostCheck :: Int -> Client () - hostCheck !n = do - hosts <- liftIO . readTVarIO =<< view hostmap - when (Map.member h hosts) $ do - isUp <- C.canConnect h - if isUp then do - sig <- view (context.sigMonit) - liftIO $ sig $$ (HostUp (h^.hostAddr)) - logInfo' $ "Reachable: " <> string8 (show h) - else do - logInfo' $ "Unreachable: " <> string8 (show h) - liftIO $ threadDelay (2^n * minDelay) - hostCheck (min (n + 1) maxExp) - - -- Bounded to 5min - toMicros :: Milliseconds -> Int - toMicros (Ms s) = min (s * 1000) (5 * 60 * 1000000) - - minDelay :: Int - minDelay = 50000 -- 50ms - - maxExp :: Int - maxExp = let steps = fromIntegral (toMicros maxDelay `div` minDelay) :: Double - in floor (logBase 2 steps) - ------------------------------------------------------------------------------ --- Exception handling - --- [Note: Error responses] --- Cassandra error responses are locally thrown as 'ResponseError's to achieve --- a unified handling of retries in the context of a single retry policy, --- together with other recoverable (i.e. retryable) exceptions. However, this --- is just an internal technicality for handling retries - generally error --- responses must not escape this function as exceptions. Deciding if and when --- to actually throw a 'ResponseError' upon inspection of the 'HostResponse' --- must be left to the caller. -withRetries - :: (Tuple a, Tuple b) - => (Request k a b -> ClientState -> Client (HostResponse k a b)) - -> Request k a b - -> Client (HostResponse k a b) -withRetries fn a = do - s <- ask - let how = s^.context.settings.retrySettings.retryPolicy - let what = s^.context.settings.retrySettings.retryHandlers - r <- try $ recovering how what $ \i -> do - hr <- if rsIterNumber i == 0 - then fn a s - else fn (newRequest s) (adjust s) - -- [Note: Error responses] - maybe (return hr) throwM (toResponseError hr) - return $ either fromResponseError id r - where - adjust s = - let Ms x = s^.context.settings.retrySettings.sendTimeoutChange - Ms y = s^.context.settings.retrySettings.recvTimeoutChange - in over (context.settings.connSettings.sendTimeout) (Ms . (+ x) . ms) - . over (context.settings.connSettings.responseTimeout) (Ms . (+ y) . ms) - $ s - - newRequest s = - case s^.context.settings.retrySettings.reducedConsistency of - Nothing -> a - Just c -> - case a of - RqQuery (Query q p) -> RqQuery (Query q p { consistency = c }) - RqExecute (Execute q p) -> RqExecute (Execute q p { consistency = c }) - RqBatch b -> RqBatch b { batchConsistency = c } - _ -> a - ------------------------------------------------------------------------------- --- Control connection handling --- --- The control connection is dedicated to maintaining the client's --- view of the cluster topology. There is a single control connection in a --- client's 'ClientState' at any particular time. - --- | Setup and install the given connection as the new control --- connection, replacing the current one. -setupControl :: Connection -> Client () -setupControl c = do - env <- ask - pol <- view policy - ctx <- view context - l <- updateHost (c^.host) . listToMaybe <$> C.query c One Disco.local () - r <- discoverPeers ctx c - (up, down) <- mkHostMap ctx pol (l:r) - m <- view hostmap - let h = Map.union up down - atomically' $ writeTVar m h - liftIO $ setup pol (Map.keys up) (Map.keys down) - C.register c C.allEventTypes (runClient env . onCqlEvent) - logInfo' $ "Known hosts: " <> string8 (show (Map.keys h)) - j <- view jobs - for_ (Map.keys down) $ \d -> - runJob j (d^.hostAddr) $ - runClient env $ monitor (Ms 1000) (Ms 60000) d - ctl <- view control - let c' = set C.host l c - atomically' $ writeTVar ctl (Control Connected c') - logInfo' $ "New control connection: " <> string8 (show c') - --- | Initialise connection pools for the given hosts, checking for --- acceptability with the host policy and separating them by reachability. -mkHostMap :: Context -> Policy -> [Host] -> Client (Map Host Pool, Map Host Pool) -mkHostMap c p = liftIO . foldrM checkHost (Map.empty, Map.empty) - where - checkHost h (up, down) = do - okay <- acceptable p h - if okay then do - isUp <- C.canConnect h - if isUp then do - up' <- Map.insert h <$> mkPool c h <*> pure up - return (up', down) - else do - down' <- Map.insert h <$> mkPool c h <*> pure down - return (up, down') - else - return (up, down) - --- | Check if the control connection is healthy. -checkControl :: Client Bool -checkControl = do - cc <- view connection <$> (readTVarIO' =<< view control) - rs <- liftIO $ C.requestRaw cc (RqOptions Options) - return $ case rs of - RsSupported {} -> True - _ -> False - `recover` - False - --- | Asynchronously replace the control connection. --- --- Invariants: --- --- 1) When the control connection is in state 'Reconnecting' there --- is a thread running that attempts to establish a new control --- connection. --- --- 2) There is only one thread performing a reconnect at a time. --- --- To that end, the 'ControlState' acts as a mutex that is acquired --- in state 'Reconnecting' and must eventually be released by either --- a successful reconnect with state 'Connected' or a (fatal) failure --- with state 'Disconnected'. In the latter case, further failing --- requests may trigger another recovery attempt of the control --- connection. -replaceControl :: Client () -replaceControl = do - e <- ask - let l = e^.context.settings.logger - liftIO $ mask $ \restore -> do - cc <- setReconnecting e - for_ cc $ \c -> forkIO $ - restore $ do - ignore (C.close c) - reconnect e l - `catchAll` \ex -> do - logError l $ "Control connection reconnect aborted: " <> string8 (show ex) - atomically $ modifyTVar' (e^.control) (set state Disconnected) - where - setReconnecting e = atomically $ do - ctrl <- readTVar (e^.control) - if ctrl^.state /= Reconnecting - then do - writeTVar (e^.control) (set state Reconnecting ctrl) - return $ Just (ctrl^.connection) - else - return Nothing - - reconnect e l = recovering adInf (onExc l) $ \_ -> do - hosts <- NE.nonEmpty . Map.keys <$> readTVarIO (e^.hostmap) - case hosts of - Just hs -> hs `tryAll` (runClient e . renewControl) - `catch` \x -> case fromException x of - Just (SomeAsyncException _) -> throwM x - Nothing -> do - logError l "All known hosts unreachable." - runClient e rebootControl - Nothing -> do - logError l "No known hosts." - runClient e rebootControl - - adInf = capDelay 5000000 (exponentialBackoff 5000) - - onExc l = - [ const $ Handler $ \(_ :: SomeAsyncException) -> return False - , const $ Handler $ \(e :: SomeException) -> do - logError l $ "Replacement of control connection failed with: " - <> string8 (show e) - <> ". Retrying ..." - return True - ] - --- | Create a new connection to a known host and set it up --- as the new control connection. -renewControl :: Host -> Client () -renewControl h = do - ctx <- view context - logInfo' "Renewing control connection with known host ..." - let s = ctx^.settings - bracketOnError - (C.connect (s^.connSettings) (ctx^.timeouts) (s^.protoVersion) (s^.logger) h) - (liftIO . C.close) - setupControl - --- | Create a new connection to one of the initial contacts --- and set it up as the new control connection. -rebootControl :: Client () -rebootControl = do - e <- ask - logInfo' "Renewing control connection with initial contacts ..." - bracketOnError - (liftIO (mkContact (e^.context))) - (liftIO . C.close) - setupControl - ------------------------------------------------------------------------------ --- Event handling - -onCqlEvent :: Event -> Client () -onCqlEvent x = do - logInfo' $ "Event: " <> string8 (show x) - pol <- view policy - prt <- view (context.settings.portnumber) - case x of - StatusEvent Down (sock2inet prt -> a) -> - liftIO $ onEvent pol (HostDown a) - TopologyEvent RemovedNode (sock2inet prt -> a) -> do - hmap <- view hostmap - atomically' $ - modifyTVar' hmap (Map.filterWithKey (\h _ -> h^.hostAddr /= a)) - liftIO $ onEvent pol (HostGone a) - StatusEvent Up (sock2inet prt -> a) -> do - s <- ask - startMonitor s a - TopologyEvent NewNode (sock2inet prt -> a) -> do - s <- ask - let ctx = s^.context - let hmap = s^.hostmap - ctrl <- readTVarIO' (s^.control) - let c = ctrl^.connection - peers <- liftIO $ discoverPeers ctx c `recover` [] - let h = fromMaybe (Host a "" "") $ find ((a == ) . view hostAddr) peers - okay <- liftIO $ acceptable pol h - when okay $ do - p <- mkPool ctx h - atomically' $ modifyTVar' hmap (Map.alter (maybe (Just p) Just) h) - liftIO $ onEvent pol (HostNew h) - tryRunJob_ (s^.jobs) a $ runClient s (prepareAllQueries h) - SchemaEvent _ -> return () - where - startMonitor s a = do - hmp <- readTVarIO' (s^.hostmap) - case find ((a ==) . view hostAddr) (Map.keys hmp) of - Just h -> tryRunJob_ (s^.jobs) a $ runClient s $ do - monitor (Ms 3000) (Ms 60000) h - prepareAllQueries h - Nothing -> return () - ------------------------------------------------------------------------------ --- Utilities - --- | Get the 'Result' out of a 'HostResponse'. If the response is an 'RsError', --- a 'ResponseError' is thrown. If the response is neither --- 'RsResult' nor 'RsError', an 'UnexpectedResponse' is thrown. -getResult :: MonadThrow m => HostResponse k a b -> m (Result k a b) -getResult (HostResponse _ (RsResult _ _ r)) = return r -getResult (HostResponse h (RsError t w e)) = throwM (ResponseError h t w e) -getResult hr = unexpected hr -{-# INLINE getResult #-} - -getPreparedQueryId :: MonadThrow m => HostResponse k a b -> m (Host, QueryId k a b) -getPreparedQueryId hr = getResult hr >>= \case - PreparedResult i _ _ -> return (hrHost hr, i) - _ -> unexpected hr -{-# INLINE getPreparedQueryId #-} - -unexpected :: MonadThrow m => HostResponse k a b -> m c -unexpected (HostResponse h r) = throwM $ UnexpectedResponse h r - -atomically' :: STM a -> Client a -atomically' = liftIO . atomically - -readTVarIO' :: TVar a -> Client a -readTVarIO' = liftIO . readTVarIO - -logInfo' :: Builder -> Client () -logInfo' m = do - l <- view (context.settings.logger) - liftIO $ logInfo l m -{-# INLINE logInfo' #-} - -logWarn' :: Builder -> Client () -logWarn' m = do - l <- view (context.settings.logger) - liftIO $ logWarn l m -{-# INLINE logWarn' #-} - -logError' :: Builder -> Client () -logError' m = do - l <- view (context.settings.logger) - liftIO $ logError l m -{-# INLINE logError' #-} - +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE NumericUnderscores #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ViewPatterns #-} + +module Database.CQL.IO.Client + ( Client + , MonadClient (..) + , ClientState + , DebugInfo (..) + , runClient + , init + , shutdown + , request + , requestN + , request1 + , execute + , executeWithPrepare + , prepare + , retry + , once + , debugInfo + , preparedQueries + , withPrepareStrategy + , getResult + , unexpected + , C.defQueryParams + , getReplication + , getSimpleReplicas + , getDCReplicas + , buildTokenMap + ) where + +import Control.Applicative +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async (async, wait, Async, cancel, waitAnyCatchSTM, withAsync) +import Control.Concurrent.STM + ( STM + , TMVar + , atomically + , newEmptyTMVar + , putTMVar + , readTMVar + , tryReadTMVar + , tryTakeTMVar + ) +import qualified Control.Concurrent.STM as STM (retry) +import Control.Concurrent.STM.TVar +import Control.Exception (IOException) +import Control.Lens (makeLenses, (^.), set, over, view) +import Control.Monad (forM_, unless, void, when) +import Control.Monad.Catch +import Control.Monad.IO.Class +import Control.Monad.IO.Unlift +import Control.Monad.Reader (ReaderT (..), runReaderT, MonadReader, ask) +import Control.Monad.Trans.Class +import Control.Monad.Trans.Except +import Control.Monad.Trans.Maybe (hoistMaybe, MaybeT (..)) +import Control.Retry + ( RetryPolicyM + , fullJitterBackoff + , limitRetriesByCumulativeDelay + , recovering + , retrying + , rsIterNumber + ) +import Data.ByteString (ByteString) +import Data.Foldable (for_, foldrM, Foldable (..)) +import Data.Function ((&)) +import Data.Functor.Identity (Identity (..)) +import Data.Int (Int64) +import Data.IP (IP) +import Data.List (find) +import Data.Map.Strict (Map) +import Data.Maybe (fromMaybe, isNothing, listToMaybe, mapMaybe) +import Data.Semigroup +import Data.Sequence (Seq (..)) +import Data.Persist (runPut) +import Data.Text.Encoding (encodeUtf8) +import Data.Text (Text) +import Data.UUID (UUID) +import Data.Word +import Database.CQL.IO.Cluster.Host +import Database.CQL.IO.Cluster.Policies +import Database.CQL.IO.Connection (Connection, host, Raw) +import Database.CQL.IO.Connection.Settings +import Database.CQL.IO.Connection.Socket (PortNumber) +import Database.CQL.IO.Exception +import Database.CQL.IO.Jobs +import Database.CQL.IO.Log +import Database.CQL.IO.Pool (Pool) +import Database.CQL.IO.PrepQuery (PrepQuery (..), PreparedQueries, ServerPrepQuery (..)) +import Database.CQL.IO.Replication + ( NetworkTopologyReplicaMap + , ReplicationStrategy (..) + , SimpleReplicaMap + , buildMasterReplicaMaps + , parseReplicationStrategy + ) +import Database.CQL.IO.Settings +import Database.CQL.IO.Signal +import Database.CQL.IO.Timeouts (TimeoutManager) +import Database.CQL.Protocol hiding (Map) +import Network.Socket (SockAddr) +import OpenSSL.Session (SomeSSLException) +import System.Random.Shuffle (shuffleM) +import Prelude hiding (init) + +import qualified Data.Attoparsec.Text as AP (decimal, parseOnly, signed) +import qualified Control.Monad.Reader as Reader +import qualified Control.Monad.State.Strict as S +import qualified Control.Monad.State.Lazy as LS +import qualified Data.List.NonEmpty as NE +import qualified Data.Map.Strict as Map +import qualified Data.SemVer as SemVer (Version, fromText, version) +import qualified Data.Sequence as Seq +import qualified Data.Set as Set +import qualified Database.CQL.IO.Cluster.Discovery as Disco +import qualified Database.CQL.IO.Connection as C +import qualified Database.CQL.IO.Pool as Pool +import qualified Database.CQL.IO.PrepQuery as PQ +import qualified Database.CQL.IO.Timeouts as TM +import qualified Database.CQL.Protocol as Cql + +data Control = Control + -- These are TMVars so that can be initially empty. The presence of a + -- control structure in the map from host to control indicates that a + -- connection attempt is ongoing. + { _connection :: !(TMVar Connection) + , _sysLocal :: !(TMVar Disco.Local) + } + +-- Information used for tracking peers. Doesn't include the calculated replicas, +-- but includes the data used to do those calculations. +data PeerInfo = PeerInfo + { _hostsByRPC :: !(TVar (Map InetAddr Host)) + , _hostsByUUID :: !(TVar (Map UUID Host)) + , _dcHostCount :: !(TVar (Map Text (Int, Map Text Int))) + , _ownedTokens :: !(TVar (Map UUID (Set.Set Int64))) + -- Keyed by control connection hostId and address of new host + , _newHostJobs :: !(Jobs (UUID, InetAddr)) + } + +data ReplicationInfo = ReplicationInfo + { _keyspaces :: !(TVar (Map Keyspace ReplicationStrategy)) + , _tokenMap :: !(TVar (Map Int64 UUID)) + , _dcTokenMap :: !(TVar (Map Text (Map Int64 UUID))) + , _simpleReplicas :: !(TVar SimpleReplicaMap) + , _topoReplicas :: !(TVar NetworkTopologyReplicaMap) + , _addHostJobs :: !(Jobs UUID) + , _removeHostJobs :: !(Jobs UUID) + } + +minReleaseVersion :: SemVer.Version +minReleaseVersion = SemVer.version 3 0 0 [] [] + +pEERS_V2_MIN_VERSION :: SemVer.Version +pEERS_V2_MIN_VERSION = SemVer.version 4 0 0 [] [] + +data Context = Context + { _settings :: !Settings + , _timeouts :: !TimeoutManager + , _sigMonit :: !(Signal HostEvent) + } + +data VersionInfo = VersionInfo + { _releaseVer :: !SemVer.Version + , _cqlVer :: !SemVer.Version + } + +data MultiControlState = MultiControlState + { _multiControlMap :: !(TVar (Map UUID Control)) + , _multiControlQueue :: !(TVar (Map Int (Seq UUID))) + , _multiControlTask :: !(TVar (Maybe (Async ()))) + , _multiControlPriorityMap :: !(TVar (Map UUID Int)) + , _multiControlJobs :: !(Jobs InetAddr) + } + +-- | Opaque client state/environment. +data ClientState = ClientState + { _context :: !Context + , _policy :: !Policy + , _prepQueries :: !PreparedQueries + , _multiControl :: !MultiControlState + , _hostmap :: !(TVar (Map Host Pool)) + , _peerInfo :: !(TVar PeerInfo) + , _replicationInfo :: !(TVar ReplicationInfo) + , _versionInfo :: !(TVar VersionInfo) + , _jobs :: !(Jobs InetAddr) + , _protoVerInUse :: !(TVar Version) -- Use this instead of version from control connection. + } + +makeLenses ''Control +makeLenses ''PeerInfo +makeLenses ''ReplicationInfo +makeLenses ''Context +makeLenses ''VersionInfo +makeLenses ''MultiControlState +makeLenses ''ClientState + +-- | The Client monad. +-- +-- A simple reader monad on `IO` around some internal state. Prior to executing +-- this monad via 'runClient', its state must be initialised through +-- 'Database.CQL.IO.Client.init' and after finishing operation it should be +-- terminated with 'shutdown'. +-- +-- To lift 'Client' actions into another monad, see 'MonadClient'. +newtype Client a = Client + { client :: ReaderT ClientState IO a + } deriving ( Functor + , Applicative + , Monad + , MonadIO + , MonadUnliftIO + , MonadThrow + , MonadCatch + , MonadMask + , MonadReader ClientState + ) + +-- | Monads in which 'Client' actions may be embedded. +class (MonadIO m, MonadThrow m) => MonadClient m + where + -- | Lift a computation from the 'Client' monad. + liftClient :: Client a -> m a + -- | Execute an action with a modified 'ClientState'. + localState :: (ClientState -> ClientState) -> m a -> m a + +instance MonadClient Client where + liftClient = id + localState = Reader.local + +instance MonadClient m => MonadClient (ReaderT r m) where + liftClient = lift . liftClient + localState f m = ReaderT (localState f . runReaderT m) + +instance MonadClient m => MonadClient (S.StateT s m) where + liftClient = lift . liftClient + localState f m = S.StateT (localState f . S.runStateT m) + +instance MonadClient m => MonadClient (LS.StateT s m) where + liftClient = lift . liftClient + localState f m = LS.StateT (localState f . LS.runStateT m) + +instance MonadClient m => MonadClient (ExceptT e m) where + liftClient = lift . liftClient + localState f m = ExceptT $ localState f (runExceptT m) + +----------------------------------------------------------------------------- +-- API + +-- | Execute the client monad. +runClient :: MonadIO m => ClientState -> Client a -> m a +runClient p a = liftIO $ runReaderT (client a) p + +-- | Use given 'RetrySettings' during execution of some client action. +retry :: MonadClient m => RetrySettings -> m a -> m a +retry r = localState (set (context.settings.retrySettings) r) + +-- | Execute a client action once, without retries, i.e. +-- +-- @once action = retry noRetry action@. +-- +-- Primarily for use in applications where global 'RetrySettings' +-- are configured and need to be selectively disabled for individual +-- queries. +once :: MonadClient m => m a -> m a +once = retry noRetry + +-- | Change the default 'PrepareStrategy' for the given client action. +withPrepareStrategy :: MonadClient m => PrepareStrategy -> m a -> m a +withPrepareStrategy s = localState (set (context.settings.prepStrategy) s) + +-- | Send a 'Request' to the server and return a 'Response'. +-- +-- This function will first ask the clients load-balancing 'Policy' for +-- some host and use its connection pool to acquire a connection for +-- request transmission. +-- +-- If all available hosts are busy (i.e. their connection pools are fully +-- utilised), the function will block until a connection becomes available +-- or the maximum wait-queue length has been reached. +-- +-- The request is retried according to the configured 'RetrySettings'. +request :: (MonadClient m, Tuple a, Tuple b) => Request k a b -> Maybe (ServerPrepQuery k a b, a) -> m (HostResponse k a b) +request a mPq = liftClient $ do + n <- liftIO . hostCount =<< view policy + withRetries (requestN n mPq) a + +-- | Send a request to a host chosen by the configured host policy. +-- +-- Tries up to @max(1,n)@ hosts. If no host can execute the request, +-- a 'HostError' is thrown. Specifically: +-- +-- * If no host is available from the 'Policy', 'NoHostAvailable' is thrown. +-- * If no host can execute the request, e.g. because all streams +-- on all connections are occupied, 'HostsBusy' is thrown. +requestN :: (Tuple b, Tuple a) + => Word + -> Maybe (ServerPrepQuery k a b, a) + -> Request k a b + -> ClientState + -> Client (HostResponse k a b) +requestN !n mPq a s = do + liftIO (select (s^.policy) mPq) >>= \case + Nothing -> throwM NoHostAvailable + Just h -> do + tryRequest1 h a s >>= \case + Just hr -> return hr + Nothing -> if n > 1 + then requestN (n - 1) mPq a s + else throwM HostsBusy + +-- | Send a 'Request' to a specific 'Host'. +-- +-- If the request cannot be executed on the given host, e.g. +-- because all connections are occupied, 'HostsBusy' is thrown. +request1 :: (Tuple a, Tuple b) + => Host + -> Request k a b + -> ClientState + -> Client (HostResponse k a b) +request1 h r s = do + rs <- tryRequest1 h r s + maybe (throwM HostsBusy) return rs + +-- | Try to send a 'Request' to a specific 'Host'. +-- +-- If the request cannot be executed on the given host, e.g. +-- because all connections are occupied, 'Nothing' is returned. +tryRequest1 :: (Tuple a, Tuple b) + => Host + -> Request k a b + -> ClientState + -> Client (Maybe (HostResponse k a b)) +tryRequest1 h a s = do + pool <- Map.lookup h <$> readTVarIO' (s^.hostmap) + case pool of + Just p -> do + result <- Pool.with p exec `catches` handlers + for_ result $ \(HostResponse _ r) -> + for_ (Cql.warnings r) $ \w -> + logWarn' $ "Server warning: " <> byteString (encodeUtf8 w) + return result + Nothing -> do + logError' $ "No pool for host: " <> string8 (show h) + p' <- mkPool h + atomically' $ modifyTVar' (s^.hostmap) (Map.alter (maybe (Just p') Just) h) + tryRequest1 h a s + where + exec c = do + r <- C.request c a + return $ HostResponse h r + + handlers = + [ Handler $ \(e :: ConnectionError) -> onConnectionError e + , Handler $ \(e :: IOException) -> onConnectionError e + , Handler $ \(e :: SomeSSLException) -> onConnectionError e + ] + + onConnectionError exc = do + e <- ask + logWarn' (string8 (show exc)) + -- Tell the policy that the host is down until monitoring confirms + -- it is still up, which will be signalled by a subsequent 'HostUp' + -- event. + liftIO $ ignore $ onEvent (e^.policy) (HostDown h) + runJob_ (e^.jobs) (h^.hostAddr) $ + runClient e $ monitor (Ms 0) (Ms 30000) h + -- Any connection error may indicate a problem with the + -- control connection for the same host (if any) + controlMap <- readTVarIO' =<< view (multiControl.multiControlMap) + case Map.lookup (h ^. hostId) controlMap of + -- If there is no control connection to the host, do nothing right now. + Nothing -> pure () + -- If there is a connection, remove it unless it checks out as OK + Just ctrl -> do + mCC <- atomically' $ tryReadTMVar (ctrl^.connection) + forM_ mCC $ \cc -> do + ok <- checkControl cc + unless ok $ dropControl h + runJob_ (e^.jobs) (h^.hostAddr) $ do + runClient e $ monitor (Ms 0) (Ms 30_000) h + (controlMap', priorityMap) <- atomically $ do + cMap <- readTVar (s^.multiControl.multiControlMap) + pMap <- readTVar (s^.multiControl.multiControlPriorityMap) + pure (cMap, pMap) + case Map.lookup (h ^. hostId) controlMap' of + -- Need the priority here so that we can re-enqueue the host + -- Once the host is up, if there is no control connection, create one. + Nothing -> + for_ (Map.lookup (h ^. hostId) priorityMap) $ \prio -> runClient e $ enqueueControl h prio + -- Once the host is up, if there is a control connection, double check it. + Just ctrl -> runClient e $ do + mCC <- atomically' $ tryReadTMVar (ctrl^.connection) + forM_ mCC $ \cc -> do + ok <- checkControl cc + unless ok $ dropControl h + throwM exc + +------------------------------------------------------------------------------ +-- Prepared queries + +-- | Execute the given request. If an 'Unprepared' error is returned, this +-- function will automatically try to re-prepare the query and re-execute +-- the original request using the same host which was used for re-preparation. +executeWithPrepare :: + forall k a b. (Tuple a, Tuple b) + => Maybe Host + -> Request k a b + -> Maybe (ServerPrepQuery k a b, a) + -> Client (HostResponse k a b) +executeWithPrepare mh rq mPq + | Just h <- mh = exec (request1 h) + | otherwise = do + p <- view policy + n <- liftIO $ hostCount p + exec (requestN n mPq) + where + exec action = do + r <- withRetries action rq + case hrResponse r of + RsError _ _ (Unprepared _ i) -> do + let qId :: QueryId k a b + qId = QueryId i + pq <- preparedQueries + mSpq <- atomically' (PQ.lookupQueryServer qId pq) + case mSpq of + Nothing -> throwM $ UnexpectedQueryId qId + Just spq -> do + (h, _) <- prepare (Just LazyPrepare) (PQ.toLocal spq) + executeWithPrepare (Just h) rq mPq + _ -> return r + +-- | Prepare the given query according to the given 'PrepareStrategy', +-- returning the resulting 'QueryId' and 'Host' which was used for +-- preparation. +prepare :: (Tuple b, Tuple a) => Maybe PrepareStrategy -> PrepQuery k a b -> Client (Host, ServerPrepQuery k a b) +prepare (Just LazyPrepare) prepQuery = do + let queryStr = PQ.queryString prepQuery + s <- ask + n <- liftIO $ hostCount (s^.policy) + r <- withRetries (requestN n Nothing) (RqPrepare (Prepare queryStr)) + (h, serverPrepQuery) <- makePreparedQuery prepQuery r + pqStore <- view prepQueries + atomically' (PQ.insert prepQuery serverPrepQuery pqStore) + pure (h, serverPrepQuery) + +prepare (Just EagerPrepare) prepQuery = view policy + >>= liftIO . current + >>= mapM (action (RqPrepare (Prepare queryStr))) + >>= first + where + queryStr = PQ.queryString prepQuery + action rq h = withRetries (request1 h) rq >>= makePreparedQuery prepQuery + + first [] = throwM NoHostAvailable + first ((h, serverPrepQuery):_) = do + pqStore <- view prepQueries + atomically' (PQ.insert prepQuery serverPrepQuery pqStore) + pure (h, serverPrepQuery) + +prepare Nothing qs = do + ps <- view (context.settings.prepStrategy) + prepare (Just ps) qs + +-- | Execute a prepared query (transparently re-preparing if necessary). +execute :: (Tuple b, Tuple a) => PrepQuery k a b -> QueryParams a -> Client (HostResponse k a b) +execute q p = do + pq <- view prepQueries + maybe new (exec Nothing) =<< atomically' (PQ.lookupQueryLocal q pq) + where + qValues = values p + exec h spq = + executeWithPrepare h (RqExecute (Execute (spqId spq) p)) (Just (spq, qValues)) + new = do + (h, spq) <- prepare (Just LazyPrepare) q + exec (Just h) spq + +prepareAllQueries :: Host -> Client () +prepareAllQueries h = do + pq <- view prepQueries + qs <- atomically' $ PQ.queryStrings pq + for_ qs $ \q -> + let qry = QueryString q :: Raw QueryString in + withRetries (request1 h) (RqPrepare (Prepare qry)) + +------------------------------------------------------------------------------ +-- Debug info + +data DebugInfo = DebugInfo + { policyInfo :: String + -- ^ Host 'Policy' string representation. + , jobInfo :: [InetAddr] + -- ^ Hosts with running background jobs (e.g. monitoring of hosts + -- currently considered down). + , hostInfo :: [Host] + -- ^ All known hosts. + , controlInfo :: Map UUID Host + -- ^ Control connection information. + } + +instance Show DebugInfo where + show dbg = showString "running jobs: " + . shows (jobInfo dbg) + . showString "\nknown hosts: " + . shows (hostInfo dbg) + . showString "\npolicy info: " + . shows (policyInfo dbg) + . showString "\ncontrol host: " + . shows (controlInfo dbg) + $ "" + +debugInfo :: MonadClient m => m DebugInfo +debugInfo = liftClient $ do + hosts <- Map.keys <$> (readTVarIO' =<< view hostmap) + pols <- liftIO . display =<< view policy + jbs <- listJobKeys =<< view jobs + controlMap <- readTVarIO' =<< view (multiControl.multiControlMap) + ctrl <- liftClient $ mapM (\Control {_connection = cM} -> do + c <- atomically' $ readTMVar cM + pure (c^.host)) controlMap + return $ DebugInfo pols jbs hosts ctrl + +preparedQueries :: Client PreparedQueries +preparedQueries = view prepQueries + +----------------------------------------------------------------------------- +-- Initialisation + +-- | Initialise client state with the given 'Settings' using the provided +-- 'Logger' for all it's logging output. +init :: MonadIO m => Settings -> m ClientState +init s = liftIO $ do + tom <- TM.create (Ms 250) + let v = s^.minProtoVer + controlConCount = max 2 (s^.maxControlConnections) + s' = set maxControlConnections controlConCount s + ctx <- Context s' tom <$> signal + let processCon con = do + pol <- s^.policyMaker + pq <- PQ.new + multiControlMapTV <- newTVarIO Map.empty + multiControlQueueTV <- newTVarIO Map.empty + hostmapTV <- newTVarIO Map.empty + rInfo <- dummyReplicationInfo + pInfo <- dummyPeerInfo + replicationInfoTV <- newTVarIO rInfo + peerInfoTV <- newTVarIO pInfo + versionInfoTV <- newTVarIO dummyVersionInfo + protoVerTV <- newTVarIO v + cstJobs <- newJobs + multiControlPriorityMapTV <- newTVarIO Map.empty + multiControlTaskTV <- newTVarIO Nothing + mcJobs <- newJobs + let mcst = MultiControlState + { _multiControlMap = multiControlMapTV + , _multiControlQueue = multiControlQueueTV + , _multiControlTask = multiControlTaskTV + , _multiControlPriorityMap = multiControlPriorityMapTV + , _multiControlJobs = mcJobs + } + cst = ClientState + { _context = ctx + , _policy = pol + , _prepQueries = pq + , _multiControl = mcst + , _hostmap = hostmapTV + , _peerInfo = peerInfoTV + , _replicationInfo = replicationInfoTV + , _versionInfo = versionInfoTV + , _protoVerInUse = protoVerTV + , _jobs = cstJobs + } + ctx^.sigMonit |-> onEvent pol + runClient cst (setupInitControl con) + multiControlTaskHandle <- async (controlThread cst) + atomically $ writeTVar multiControlTaskTV (Just multiControlTaskHandle) + pure cst + mkContact v ctx processCon + where + dummyVersionInfo :: VersionInfo + dummyVersionInfo = VersionInfo + { _releaseVer = minReleaseVersion + , _cqlVer = minReleaseVersion + } + dummyReplicationInfo :: IO ReplicationInfo + dummyReplicationInfo = do + addJobs <- newJobs + removeJobs <- newJobs + ks <- newTVarIO Map.empty + sReps <- newTVarIO Map.empty + tReps <- newTVarIO Map.empty + tMap <- newTVarIO Map.empty + dcTMap <- newTVarIO Map.empty + pure ReplicationInfo + { _keyspaces = ks + , _tokenMap = tMap + , _dcTokenMap = dcTMap + , _simpleReplicas = sReps + , _topoReplicas = tReps + , _addHostJobs = addJobs + , _removeHostJobs = removeJobs + } + + dummyPeerInfo :: IO PeerInfo + dummyPeerInfo = do + nhJobs <- newJobs + hByRPC <- newTVarIO Map.empty + hByUUID <- newTVarIO Map.empty + hostCount <- newTVarIO Map.empty + owned <- newTVarIO Map.empty + pure PeerInfo + { _hostsByRPC = hByRPC + , _hostsByUUID = hByUUID + , _dcHostCount = hostCount + , _ownedTokens = owned + , _newHostJobs = nhJobs + } + +controlThread :: ClientState -> IO () +controlThread cst = do + waitSetTV <- newTVarIO Set.empty + withAsync (reapLoop waitSetTV) $ \_ -> + spawnLoop waitSetTV + where + -- The control thread is actually two threads. The spawn loop waits on the + -- semaphore for availability to open a new control connection. When the spawn + -- loop takes a token from the semaphore, it takes the highest priority host + -- and connects. If there are no hosts, it relies on STM to try again when a + -- new host is added to the queue. + spawnLoop :: TVar (Set.Set (Async ())) -> IO () + spawnLoop waitSetTV = do + let multiControlState = cst^.multiControl + mcMapTV = multiControlState^.multiControlMap + maxMapSize = cst^.context.settings.maxControlConnections + mNewControl <- atomically $ do + mcMap <- readTVar mcMapTV + when (Map.size mcMap >= maxMapSize) STM.retry + (nextHost, prio) <- getFirstHost + case Map.lookup (nextHost^.hostId) mcMap of + Just _ -> pure Nothing + Nothing -> do + connTMV <- newEmptyTMVar + localTMV <- newEmptyTMVar + let newControl = Control { _connection = connTMV, _sysLocal = localTMV } + writeTVar mcMapTV $ Map.insert (nextHost^.hostId) newControl mcMap + pure $ Just (newControl, nextHost, prio) + case mNewControl of + Nothing -> spawnLoop waitSetTV + Just (newControl, nextHost, prio) -> do + asyncControl <- async $ attemptConnect newControl nextHost prio + atomically $ do + waitSet <- readTVar waitSetTV + let waitSet' = Set.insert asyncControl waitSet + writeTVar waitSetTV waitSet' + spawnLoop waitSetTV + + -- The reap loop waits for connection attempts to succeed and removes them + -- from the wait set. If a connection fails, a job is created to monitor the + -- host and add them at the back of the queue when the host is up again. It is + -- important that the monitoring run as a job so that it can be canceled if + -- the host is renumbered. + reapLoop :: TVar (Set.Set (Async ())) -> IO () + reapLoop waitSetTV = do + atomically $ do + waitSet <- readTVar waitSetTV + when (Set.null waitSet) STM.retry + (toRemove, _) <- waitAnyCatchSTM (Set.toList waitSet) + let waitSet' = Set.delete toRemove waitSet + writeTVar waitSetTV waitSet' + reapLoop waitSetTV + + getFirstHost :: STM (Host, Int) + getFirstHost = do + controlSequenceMap <- readTVar (cst^.multiControl.multiControlQueue) + if Map.null controlSequenceMap + then STM.retry + else do + let (prio, minHostUUID, controlSequenceMap') = removeMinHost controlSequenceMap + pInfo <- readTVar (cst^.peerInfo) + uuidMap <- readTVar (pInfo^.hostsByUUID) + writeTVar (cst^.multiControl.multiControlQueue) controlSequenceMap' + case Map.lookup minHostUUID uuidMap of + Nothing -> STM.retry + Just h -> pure (h, prio) + + removeMinHost :: Map Int (Seq UUID) -> (Int, UUID, Map Int (Seq UUID)) + removeMinHost controlSequenceMap = case Map.minViewWithKey controlSequenceMap of + Nothing -> error "expecting nonEmpty map" + Just ((k, minSeq), remainingMap) -> case minSeq of + Empty -> error "expecting nonEmpty sequence" + firstUUID :<| Empty -> (k, firstUUID, remainingMap) + firstUUID :<| remainder -> (k, firstUUID, Map.insert k remainder remainingMap) + + attemptConnect :: Control -> Host -> Int -> IO () + attemptConnect ctrl h prio = do + bracketOnError + (runClient cst $ clientConnect h) + (onConnErr h prio) + (\c -> do + runClient cst $ setupControl ctrl c) + + onConnErr :: Host -> Int -> Connection -> IO () + onConnErr h prio c = do + runClient cst $ logInfo' $ "Control connection error for: " <> string8 (show h) + let jobQueue = cst^.multiControl.multiControlJobs + mcMapTV = cst^.multiControl.multiControlMap + atomically $ do + mcMap <- readTVar mcMapTV + writeTVar mcMapTV (Map.delete (h^.hostId) mcMap) + ignore $ C.close c + void $ runJob jobQueue (h^.hostAddr) $ monitorAndReconnect h prio + + monitorAndReconnect :: Host -> Int -> IO () + monitorAndReconnect h prio = runClient cst $ do + logInfo' $ "Control Monitoring job started for: " <> string8 (show h) + monitor (Ms 3000) (Ms 60_000) h + logInfo' $ "Control Monitor finished for: " <> string8 (show h) + <> ", setting up control." + enqueueControl h prio + +-- | Try to establish a connection to one of the initial contacts. +mkContact :: Version -> Context -> (Connection -> IO a) -> IO a +mkContact v (Context s t _) handler = tryAll (s^.contacts) mkConnection + where + mkConnection h = do + as <- C.resolve h (s^.portnumber) + tryAll (NE.fromList as) $ \a -> do + bracketOnError (doConnect a) C.close handler + + doConnect a = do + logDebug (s^.logger) $ "Connecting to " <> string8 (show a) + C.connect (s^.connSettings) t v (s^.logger) (dummyHost a) + +discoverPeers :: Connection -> Client [Disco.Peer] +discoverPeers c = liftIO $ + map asRecord <$> C.query c One Disco.peers () + +discoverPeersV2 :: Connection -> Client [Disco.PeerV2] +discoverPeersV2 c = liftIO $ + map asRecord <$> C.query c One Disco.peersV2 () + +discoverKeyspaces :: Connection -> Client [Disco.Keyspace] +discoverKeyspaces c = liftIO $ + map asRecord <$> C.query c One Disco.keyspaces () + +findKeyspace :: Keyspace -> Connection -> Client (Maybe Disco.Keyspace) +findKeyspace ks c = liftIO $ + listToMaybe . fmap asRecord <$> C.query c One Disco.keyspace (Identity ks) + +parseKeyspaces :: [Disco.Keyspace] -> Map Keyspace ReplicationStrategy +parseKeyspaces = foldMap' parseKeyspace + +parseKeyspace :: Disco.Keyspace -> Map Keyspace ReplicationStrategy +parseKeyspace ks = case parseReplicationStrategy (ks ^. Disco.replication) of + Nothing -> Map.empty + Just rs -> Map.singleton (ks ^. Disco.name) rs + +findPeer :: MonadIO m => IP -> Connection -> m (Maybe Disco.Peer) +findPeer peerAddr c = liftIO $ do + listToMaybe . fmap asRecord <$> C.query c One Disco.peer (Identity peerAddr) + +findPeerV2 :: MonadIO m => (IP, PortNumber) -> Connection -> m (Maybe Disco.PeerV2) +findPeerV2 (peerAddr, peerPort) c = liftIO $ do + find (\p2 -> Disco.peerV2NativePort p2 == fromIntegral peerPort) . fmap asRecord + <$> C.query c One Disco.peerV2 (Identity peerAddr) + +readLocal :: (MonadIO m, MonadThrow m) => Connection -> m Disco.Local +readLocal c = do + -- If we know we have the V4+ fields available in system.local, don't make two + -- queries. Jump straight to the query with all the data we want. + if c ^. C.protocol >= V4 + then readLocalV4 + else do + results <- C.query c One Disco.localV3 () + case results of + [sLocalTuple] -> do + let localV3 = asRecord sLocalTuple + case AP.parseOnly (AP.decimal @Int) (Disco.localV3NativeProtocolVer localV3) of + Left _ -> pure $ Disco.LocalV3_ localV3 + Right 3 -> pure $ Disco.LocalV3_ localV3 + Right x | x >= 4 -> readLocalV4 + _ -> pure $ Disco.LocalV3_ localV3 + [] -> throwM $ + SysLocalError (c^.host) "No results returned from system.local table" + _ -> throwM $ + SysLocalError (c^.host) "Too many results returned from system.local table" + where + readLocalV4 = do + results <- C.query c One Disco.localV4 () + case results of + [sLocalTuple] -> pure . Disco.LocalV4_ $ asRecord sLocalTuple + [] -> throwM $ + SysLocalError (c^.host) "No results returned from system.local table" + _ -> throwM $ + SysLocalError (c^.host) "Too many results returned from system.local table" + +mkPool :: Host -> Client Pool +mkPool h = do + cst <- ask + ctx <- view context + let s = ctx^.settings + m = s^.connSettings.maxStreams + lgr = s^.logger + liftIO $ Pool.create (connOpen cst) (connClose lgr) lgr (s^.poolSettings) m + where + connOpen cst = runClient cst $ do + c <- clientConnect h + logDebug' $ "Connection established: " <> string8 (show c) + pure c + + connClose lgr c = do + C.close c + logDebug lgr $ "Connection closed: " <> string8 (show c) + +----------------------------------------------------------------------------- +-- Termination + +-- | Terminate client state, i.e. end all running background checks and +-- shutdown all connection pools. Once this is entered, the client +-- will eventually be shut down, though an asynchronous exception can +-- interrupt the wait for that to occur. +shutdown :: MonadIO m => ClientState -> m () +shutdown s = liftIO $ asyncShutdown >>= wait + where + asyncShutdown = async $ do + TM.destroy (s^.context.timeouts) True + cancelJobs (s^.jobs) + mControlThread <- readTVarIO (s^.multiControl.multiControlTask) + forM_ mControlThread cancel + controlMap <- readTVarIO (s^.multiControl.multiControlMap) + forM_ controlMap $ \ctrl -> do + conn <- atomically $ readTMVar (ctrl^.connection) + ignore $ C.close conn + mapM_ Pool.destroy . Map.elems =<< readTVarIO (s^.hostmap) + +----------------------------------------------------------------------------- +-- Monitoring + +-- | @monitor initialDelay maxDelay host@ tries to establish a connection +-- to @host@ after @initialDelay@. If the connection attempt fails, it is +-- retried with exponentially increasing delays, up to a maximum delay of +-- @maxDelay@. When a connection attempt suceeds, a 'HostUp' event is +-- signalled. +-- +-- The function returns when one of the following conditions is met: +-- +-- 1. The connection attempt suceeds. +-- 2. The host is no longer found to be in the client's known host map. +-- +-- I.e. as long as the host is still known to the client and is unreachable, the +-- connection attempts continue. Both @initialDelay@ and @maxDelay@ are bounded +-- by a limit of 5 minutes. +monitor :: Milliseconds -> Milliseconds -> Host -> Client () +monitor initial maxDelay h = do + liftIO $ threadDelay (toMicros initial) + logInfo' $ "Monitoring: " <> string8 (show h) + hostCheck 0 + where + hostCheck :: Int -> Client () + hostCheck !n = do + cst <- ask + (uuidMap, poolMap) <- liftIO . atomically $ do + pInfo <- readTVar (cst^.peerInfo) + poolMap <- readTVar (cst^.hostmap) + uuidMap <- readTVar (pInfo^.hostsByUUID) + pure (uuidMap, poolMap) + let curAddress = (^.hostAddr) <$> Map.lookup (h^.hostId) uuidMap + when (Map.member h poolMap && curAddress == Just (h^.hostAddr)) $ do + isUp <- C.canConnect h + if isUp then do + sig <- view (context.sigMonit) + liftIO $ sig $$ HostUp h + logInfo' $ "Reachable: " <> string8 (show h) + else do + logInfo' $ "Unreachable: " <> string8 (show h) + liftIO $ threadDelay (2^n * minDelay) + hostCheck (min (n + 1) maxExp) + + -- Bounded to 5min + toMicros :: Milliseconds -> Int + toMicros (Ms s) = min (s * 1000) (5 * 60 * 1000000) + + minDelay :: Int + minDelay = 50000 -- 50ms + + maxExp :: Int + maxExp = let steps = fromIntegral (toMicros maxDelay `div` minDelay) :: Double + in floor (logBase 2 steps) + +----------------------------------------------------------------------------- +-- Exception handling + +-- [Note: Error responses] +-- Cassandra error responses are locally thrown as 'ResponseError's to achieve +-- a unified handling of retries in the context of a single retry policy, +-- together with other recoverable (i.e. retryable) exceptions. However, this +-- is just an internal technicality for handling retries - generally error +-- responses must not escape this function as exceptions. Deciding if and when +-- to actually throw a 'ResponseError' upon inspection of the 'HostResponse' +-- must be left to the caller. +withRetries + :: (Tuple a, Tuple b) + => (Request k a b -> ClientState -> Client (HostResponse k a b)) + -> Request k a b + -> Client (HostResponse k a b) +withRetries fn a = do + s <- ask + let how = s^.context.settings.retrySettings.retryPolicy + let what = s^.context.settings.retrySettings.retryHandlers + r <- try $ recovering how what $ \i -> do + hr <- if rsIterNumber i == 0 + then fn a s + else fn (newRequest s) (adjust s) + -- [Note: Error responses] + maybe (return hr) throwM (toResponseError hr) + return $ either fromResponseError id r + where + adjust s = + let Ms x = s^.context.settings.retrySettings.sendTimeoutChange + Ms y = s^.context.settings.retrySettings.recvTimeoutChange + in over (context.settings.connSettings.sendTimeout) (Ms . (+ x) . ms) + . over (context.settings.connSettings.responseTimeout) (Ms . (+ y) . ms) + $ s + + newRequest s = + case s^.context.settings.retrySettings.reducedConsistency of + Nothing -> a + Just c -> + case a of + RqQuery (Query q p) -> RqQuery (Query q p { consistency = c }) + RqExecute (Execute q p) -> RqExecute (Execute q p { consistency = c }) + RqBatch b -> RqBatch b { batchConsistency = c } + _ -> a + +------------------------------------------------------------------------------ +-- Control connection handling +-- +-- The control connection is dedicated to maintaining the client's +-- view of the cluster topology. There is a single control connection in a +-- client's 'ClientState' at any particular time. + +-- | Setup and install the given connection as the new control +-- connection, replacing the current one. If the protocol version needs to be +-- upgraded, the provided connection is closed, and a new one is created with +-- the correct protocol version. The protocol does not allow upgrades of an +-- established connection. +setupControl :: Control -> Connection -> Client () +setupControl hostControl c = do + env <- ask + newLocal <- readLocal c + newActiveVer <- computeActiveVer newLocal <$> view (context.settings) + if c ^. C.protocol /= newActiveVer + then replaceControlConnection newActiveVer + else do + let c' = setLocation newLocal c + atomically' $ do + putTMVar (hostControl^.connection) c' + putTMVar (hostControl^.sysLocal) newLocal + logInfo' $ "New control connection: " <> string8 (show c') + -- We register first, before reading anything from the peer tables. This + -- prevents missing anything in the gap between reading and registering. + C.register c' C.allEventTypes (runClient env . onCqlEvent hostControl) + pure () + where + replaceControlConnection :: Version -> Client () + replaceControlConnection v = do + logDebug' $ + "Reconnecting control connection to upgrade protocol version:" <> + string8 (show (c ^. C.protocol)) <> "->" <> string8 (show v) + liftIO $ C.close c + protoTV <- view protoVerInUse + atomically' $ writeTVar protoTV v + c' <- clientConnect (c ^. C.host) + setupControl hostControl c' + + computeActiveVer :: Disco.Local -> Settings -> Version + computeActiveVer l s = + case AP.parseOnly (AP.decimal @Int) (l ^. Disco.nativeProtocolVer) of + Left _ -> s^.minProtoVer + Right 3 -> min (s^.maxProtoVer) V3 + Right x | x >= 4 -> min (s^.maxProtoVer) V4 + _ -> s^.minProtoVer + setLocation :: Disco.Local -> Connection -> Connection + setLocation sLocal cLoc = + let (_, curPort) = inet2ipPort (cLoc^.host.hostAddr) + in set host (local2Host curPort sLocal) cLoc + +setupInitControl :: Connection -> Client () +setupInitControl c = do + env <- ask + pol <- view policy + ctrl <- newEmptyControl + mcMapTV <- view (multiControl.multiControlMap) + setupControl ctrl c + c' <- atomically' $ do + mcMap <- readTVar mcMapTV + c' <- readTMVar (ctrl^.connection) + let realHost = c' ^. C.host + writeTVar mcMapTV $ Map.insert (realHost^.hostId) ctrl mcMap + pure c' + curLocal <- atomically' $ readTMVar (ctrl^.sysLocal) + (rVer, cVer) <- parseVersions curLocal c' + versionInfoTV <- view versionInfo + atomically' $ do + writeTVar versionInfoTV VersionInfo + { _releaseVer = rVer + , _cqlVer = cVer + } + writeTVar (env^.protoVerInUse) (c' ^. C.protocol) + buildPeerInfo curLocal c' + buildReplicationInfo curLocal c' + buildPriorityMap + (up, down) <- mkHostMap + liftIO $ setup pol (Map.keys up) (Map.keys down) + for_ (Map.keys down) $ \d -> + runJob (env^.jobs) (d^.hostAddr) $ + runClient env $ monitor (Ms 1000) (Ms 60_000) d + keysShuffled <- liftIO $ shuffleM $ Map.keys up + for_ keysShuffled $ \u -> + when (u /= c'^.host) $ do + priorityMapTV <- view (multiControl.multiControlPriorityMap) + priorityMap <- liftIO $ readTVarIO priorityMapTV + for_ (Map.lookup (u^.hostId) priorityMap) $ \prio -> + runClient env $ enqueueControl u prio + where + newEmptyControl :: Client Control + newEmptyControl = atomically' $ do + connTMV <- newEmptyTMVar + localTMV <- newEmptyTMVar + pure Control { _connection = connTMV, _sysLocal = localTMV } + + buildPriorityMap :: Client () + buildPriorityMap = do + pInfoTV <- view peerInfo + uuidMap <- atomically' $ do + pInfo <- readTVar pInfoTV + readTVar (pInfo^.hostsByUUID) + let hosts = Map.elems uuidMap + getPriority :: Policy -> Host -> IO (UUID, Int) + getPriority pol h = do + prio <- priority pol h + pure (h^.hostId, prio) + positivePriority :: (UUID, Int) -> Bool + positivePriority = (>0) . snd + pol <- view policy + priorities <- mapM (liftIO . getPriority pol) hosts + let builtMap = Map.fromList $ filter positivePriority priorities + priorityMapTV <- view (multiControl.multiControlPriorityMap) + atomically' $ writeTVar priorityMapTV builtMap + + parseVersions + :: (MonadIO m, MonadThrow m) + => Disco.Local + -> Connection + -> m (SemVer.Version, SemVer.Version) + parseVersions l cParse = + case SemVer.fromText (l ^. Disco.releaseVer) of + Left parseError -> + throwM $ + SysLocalError + (cParse^.host) + ("Error while parsing release_version: " ++ parseError) + Right rVer -> case SemVer.fromText (l ^. Disco.cqlVer) of + Left parseError -> + throwM $ + SysLocalError + (cParse^.host) + ("Error while parsing cql_version: " ++ parseError) + Right cVer -> pure (rVer, cVer) + +type TokenMap = Map Int64 UUID +type DCTokenMap = Map Text TokenMap +type OwnedTokenMap = Map UUID (Set.Set Int64) +type DCHostCount = Map Text (Int, Map Text Int) + +buildOwnedTokenMapNoPort :: PortNumber -> [Disco.Peer] -> Disco.Local -> OwnedTokenMap +buildOwnedTokenMapNoPort p = buildOwnedTokenMap (peer2Host p) p +buildOwnedTokenMapV2 :: PortNumber -> [Disco.PeerV2] -> Disco.Local -> OwnedTokenMap +buildOwnedTokenMapV2 = buildOwnedTokenMap peer2HostV2 + +buildOwnedTokenMap :: + forall p. (Disco.HasTokens p (Cql.Set Text)) => + (p -> Host) -> + PortNumber -> + [p] -> + Disco.Local -> + OwnedTokenMap +buildOwnedTokenMap toHost portNumber peers curLocal = fullOwnedMap + where + fullOwnedMap :: OwnedTokenMap + fullOwnedMap = + peers & + foldMap' peerOwnedMap & + Set.insert localOwnedMap & + Map.fromArgSet + peerTokenList :: (Disco.HasTokens q (Cql.Set Text)) => q -> Set.Set Int64 + peerTokenList q = parseTokens (q ^. Disco.tokens) + peerOwnedMap :: p -> Set.Set (Arg UUID (Set.Set Int64)) + peerOwnedMap p = + let uuid = (^.hostId) . toHost $ p + in Set.singleton $ Arg uuid (peerTokenList p) + localOwnedMap :: Arg UUID (Set.Set Int64) + localOwnedMap = + let h = local2Host portNumber curLocal + in Arg (h^.hostId) (peerTokenList curLocal) + +parseTokens :: Cql.Set Text -> Set.Set Int64 +parseTokens = Set.fromList . mapMaybe maybeParse . Cql.fromSet + where + maybeParse :: Text -> Maybe Int64 + maybeParse t = case AP.parseOnly (AP.signed AP.decimal) t of + Left _ -> Nothing + Right x -> Just x + +buildTokenMap :: OwnedTokenMap -> Map UUID Host -> (TokenMap, DCTokenMap) +buildTokenMap ownedTokenMap uuidMap = (fullTokenMap, fullDCMap) + where + fullTokenMap :: TokenMap + fullTokenMap = Map.unions fullDCMap + fullDCMap :: DCTokenMap + dcSingleton :: (UUID, Set.Set Int64) -> DCTokenMap + dcSingleton (uuid, tSet) = + case Map.lookup uuid uuidMap of + Nothing -> Map.empty + Just h -> Map.singleton (h^.dataCentre) (Map.fromSet (const uuid) tSet) + fullDCMap = + ownedTokenMap & + Map.toList & + fmap dcSingleton & + Map.unionsWith Map.union + +buildDCHostCount :: forall p. (Disco.HasDC p Text, Disco.HasRack p Text) => [p] -> Disco.Local -> DCHostCount +buildDCHostCount peers curLocal = Map.unionWith mergeCounts peersMap localMap + where + mergeCounts :: (Int, Map Text Int) -> (Int, Map Text Int) -> (Int, Map Text Int) + mergeCounts (!count1, !racks1) (!count2, !racks2) = + let !count3 = count1 + count2 + !racks3 = Map.unionWith (+) racks1 racks2 + in (count3, racks3) + peersMap :: DCHostCount + peersMap = Map.unionsWith mergeCounts $ map peerMap peers + peerMap :: forall q. (Disco.HasDC q Text, Disco.HasRack q Text) => q -> DCHostCount + peerMap peer = + Map.singleton (peer ^. Disco.dC) (1, Map.singleton (peer ^. Disco.rack) 1) + localMap :: DCHostCount + localMap = peerMap curLocal + +buildPeerInfo :: Disco.Local -> Connection -> Client () +buildPeerInfo curLocal conn = do + vInfo <- readTVarIO' =<< view versionInfo + let rVer = vInfo^.releaseVer + pNum <- view $ context.settings.portnumber + pInfo <- readTVarIO' =<< view peerInfo + let withAmbient f = f pNum pInfo + if rVer >= pEERS_V2_MIN_VERSION + then withAmbient buildPeerInfoVersioned + discoverPeersV2 + peer2HostV2 + buildOwnedTokenMapV2 + else withAmbient buildPeerInfoVersioned + discoverPeers + (peer2Host pNum) + buildOwnedTokenMapNoPort + where + buildPeerInfoVersioned :: + (Disco.HasDC a Text, Disco.HasRack a Text, Disco.HasHostId a UUID) => + PortNumber -> + PeerInfo -> + (Connection -> Client [a]) -> + (a -> Host) -> + (PortNumber -> [a] -> Disco.Local -> OwnedTokenMap) -> + Client () + buildPeerInfoVersioned pNum pInfo getPeers toHost buildOwnedMap = do + rPeers <- getPeers conn + prt <- view (context.settings.portnumber) + let ownedMap = + if curLocal ^. Disco.partitioner == Disco.mURMUR3_PARTITIONER + then buildOwnedMap pNum rPeers curLocal + else Map.empty + hostCounts = buildDCHostCount rPeers curLocal + hostByHostAddr p = + let h = toHost p in (h^.hostAddr, h) + localByUUID = + let h = local2Host prt curLocal in Map.singleton (h^.hostId) h + localByHostAddr = + let h = local2Host prt curLocal in Map.singleton (h^.hostAddr) h + rpcMap = Map.fromList (map hostByHostAddr rPeers) <> localByHostAddr + uuidMap = Map.fromList (map (\p -> (p ^. Disco.hostId, toHost p)) rPeers) + <> localByUUID + atomically' $ do + writeTVar (pInfo^.hostsByRPC) rpcMap + writeTVar (pInfo^.hostsByUUID) uuidMap + writeTVar (pInfo^.dcHostCount) hostCounts + writeTVar (pInfo^.ownedTokens) ownedMap + +findHost :: InetAddr -> Client (Maybe Host) +findHost addr = do + pInfoTV <- view peerInfo + rpcMap <- atomically' $ do + pInfo <- readTVar pInfoTV + readTVar (pInfo^.hostsByRPC) + pure $ Map.lookup addr rpcMap + +buildReplicationInfo :: Disco.Local -> Connection -> Client () +buildReplicationInfo curLocal conn = do + cst <- ask + replicaKeyspaces <- parseKeyspaces <$> discoverKeyspaces conn + atomically' $ do + pInfo <- readTVar (cst^.peerInfo) + rInfo <- readTVar (cst^.replicationInfo) + when (curLocal ^. Disco.partitioner == Disco.mURMUR3_PARTITIONER) $ do + ownedMap <- readTVar (pInfo^.ownedTokens) + hostCounts <- readTVar (pInfo^.dcHostCount) + uuidMap <- readTVar (pInfo^.hostsByUUID) + let (tMap, dcTMap) = buildTokenMap ownedMap uuidMap + (simple, topo) = buildMasterReplicaMaps tMap dcTMap hostCounts uuidMap + writeTVar (rInfo^.keyspaces) replicaKeyspaces + writeTVar (rInfo^.tokenMap) tMap + writeTVar (rInfo^.dcTokenMap) dcTMap + writeTVar (rInfo^.simpleReplicas) simple + writeTVar (rInfo^.topoReplicas) topo + +-- | Initialise connection pools for the given hosts, checking for +-- acceptability with the host policy and separating them by reachability. +mkHostMap :: Client (Map Host Pool, Map Host Pool) +mkHostMap = do + env <- ask + pInfoTV <- view peerInfo + hostmapTV <- view hostmap + (uuidMap, rpcMap, hmap) <- atomically' $ do + pInfo <- readTVar pInfoTV + uuidMap <- readTVar (pInfo^.hostsByUUID) + rpcMap <- readTVar (pInfo^.hostsByRPC) + hmap <- readTVar hostmapTV + pure (uuidMap, rpcMap, hmap) + pol <- view policy + let hosts = Map.elems rpcMap + -- TODO: This isn't quite right. We need to verify that the hostmap didn't + -- change from where we read it above to where we are going to replace it. + -- Need to at least think about it. Maybe we need a mutex for the host map? + -- It shouldn't be replaced very often. + -- It's never replaced wholesale. But there could be events that came in from + -- the time we started until now that aren't reflected here. + (up, down) <- foldrM (checkHost uuidMap pol) (hmap, Map.empty) hosts + atomically' $ writeTVar (env^.hostmap) (Map.union up down) + pure (up, down) + where + checkHost :: Map UUID Host + -> Policy + -> Host + -> (Map Host Pool, Map Host Pool) + -> Client (Map Host Pool, Map Host Pool) + checkHost uuidMap pol h (up, down) = do + let mOldHost = Map.lookup (h^.hostId) uuidMap + case mOldHost of + Nothing -> pure (up, down) + Just oldHost -> do + let up' = + if oldHost^.hostAddr /= h^.hostAddr + then Map.delete h up + else up + okay <- liftIO $ acceptable pol h + if okay then do + isUp <- C.canConnect h + if isUp then do + up'' <- Map.insert h <$> mkPool h <*> pure up' + return (up'', down) + else do + down' <- Map.insert h <$> mkPool h <*> pure down + return (up', down') + else pure (up', down) + +-- | Check if the supplied control connection is healthy. +checkControl :: Connection -> Client Bool +checkControl cc = do + rs <- liftIO $ C.requestRaw cc (RqOptions Options) + return $ case rs of + RsSupported {} -> True + _ -> False + `recover` + False + +-- | Drop a control connection +dropControl :: Host -> Client () +dropControl h = do + logInfo' $ "Dropping control connection for: " <> string8 (show h) + cst <- ask + let mcst = cst^.multiControl + toClose <- atomically' $ do + controlMap <- readTVar (mcst^.multiControlMap) + case Map.lookup (h^.hostId) controlMap of + Nothing -> pure Nothing + Just ctrl -> do + let controlMap' = Map.delete (h^.hostId) controlMap + writeTVar (mcst^.multiControlMap) controlMap' + pure $ Just (ctrl^.connection) + forM_ toClose $ liftIO . closeIfPresent + +enqueueControl :: Host -> Int -> Client () +enqueueControl h prio = do + logInfo' $ "Enqueuing host as potential control connection: " <> string8 (show h) + mcst <- view multiControl + atomically' $ do + controlQueue <- readTVar (mcst^.multiControlQueue) + let updateSeq Nothing = Just (Seq.singleton (h^.hostId)) + updateSeq (Just curSeq) = Just (curSeq :|> h^.hostId) + controlQueue' = Map.alter updateSeq prio controlQueue + writeTVar (mcst^.multiControlQueue) controlQueue' + +----------------------------------------------------------------------------- +-- Event handling + +data NewPeerResult = + ExistingPeer | + RenumberedPeer !Host !Host | + NewPeer !(Set.Set Int64) !DCHostCount + deriving (Eq, Ord, Show) + +-- These events come from a given control connection. The source control +-- connection and its associated data is the first argument. +onCqlEvent :: Control -> Event -> Client () +onCqlEvent ctrl x = do + ctrlConn <- atomically' $ readTMVar (ctrl^.connection) + let uuid = ctrlConn^.host.hostId + logInfo' $ "Event: " <> string8 (show x) + <> " from control connection to host: " <> string8 (show uuid) + pol <- view policy + prt <- view (context.settings.portnumber) + relVer <- (^.releaseVer) <$> (readTVarIO' =<< view versionInfo) + let fromWire :: SockAddr -> InetAddr + fromWire = if relVer >= pEERS_V2_MIN_VERSION then InetAddr else sock2inet prt + case x of + StatusEvent Down (fromWire -> a) -> do + findHost a >>= \case + Just h -> do + s <- ask + liftIO $ onEvent pol (HostDown h) + liftIO $ checkHostControl s h + Nothing -> liftIO $ onEvent pol (AddrDown a) + TopologyEvent RemovedNode (fromWire -> a) -> do + (maybeH, tokenSet, dcRemains) <- removeHostPeerInfo a + for_ maybeH $ \h -> do + removeHostReplication h tokenSet dcRemains + liftIO $ onEvent pol (HostGone h) + StatusEvent Up (fromWire -> a) -> do + s <- ask + liftIO $ startMonitor s a + TopologyEvent NewNode (fromWire -> a) -> do + pInfoTV <- view peerInfo + rpcMap <- atomically' $ do + pInfo <- readTVar pInfoTV + readTVar (pInfo^.hostsByRPC) + let isNew = Map.notMember a rpcMap + when isNew $ startNewPeerJob a + unless isNew $ + logInfo' $ "Not starting new peer job for known peer: " <> string8 (show a) + SchemaEvent schemaChange -> + case schemaChange of + SchemaCreated (KeyspaceChange ks) -> upsertKeyspaceReplicationInfo ks + SchemaUpdated (KeyspaceChange ks) -> upsertKeyspaceReplicationInfo ks + SchemaDropped (KeyspaceChange ks) -> removeKeyspaceReplicationInfo ks + _ -> pure () + where + removeHostPeerInfo :: InetAddr -> Client (Maybe Host, Set.Set Int64, DCHostCount) + removeHostPeerInfo a = do + hmap <- view hostmap + pInfoTV <- view peerInfo + result@(removedHost, tokenSet, _) <- atomically' $ do + pInfo <- readTVar pInfoTV + rpcMap <- readTVar (pInfo^.hostsByRPC) + let deleteReturning :: (Ord k) => k -> Map k a -> (Maybe a, Map k a) + deleteReturning = Map.updateLookupWithKey (\_ _ -> Nothing) + (maybeHost, rpcMap') = deleteReturning a rpcMap + deleteReturningTv :: (Ord k) => TVar (Map k a) -> k -> STM (Maybe a) + deleteReturningTv tv k = do + mka <- readTVar tv + let (maybeVal, mka') = deleteReturning k mka + writeTVar tv mka' + pure maybeVal + case maybeHost of + Nothing -> do + hc <- readTVar (pInfo^.dcHostCount) + pure (maybeHost, Set.empty, hc) + Just h -> do + let dropIfRackEmpty n | n <= 1 = Nothing + | otherwise = Just (n - 1) + dropIfDCEmpty (n, racks) | n <= 1 = Nothing + | otherwise = + Just (n - 1, Map.update dropIfRackEmpty (h^.rack) racks) + dc = h^.dataCentre + mTokens <- deleteReturningTv (pInfo^.ownedTokens) (h^.hostId) + modifyTVar (pInfo^.dcHostCount) (Map.update dropIfDCEmpty dc) + countsWithout <- readTVar (pInfo^.dcHostCount) + modifyTVar (pInfo^.hostsByUUID) (Map.delete (h^.hostId)) + writeTVar (pInfo^.hostsByRPC) rpcMap' + modifyTVar hmap (Map.delete h) + let tokenSet = fromMaybe Set.empty mTokens + pure (maybeHost, tokenSet, countsWithout) + unless (Set.null tokenSet) $ do + for_ removedHost dropControl + pure result + + removeHostReplication :: Host -> Set.Set Int64 -> DCHostCount -> Client () + removeHostReplication h tokenSet counts = do + logInfo' $ "removing host: " <> string8 (show h) <> " and " <> + string8 (show $ Set.size tokenSet) <> " tokens." + pInfoTV <- view peerInfo + rInfoTV <- view replicationInfo + jobQueue <- (^.removeHostJobs) <$> readTVarIO' rInfoTV + runJob_ jobQueue (h^.hostId) $ do + atomically $ do + pInfo <- readTVar pInfoTV + rInfo <- readTVar rInfoTV + curLocal <- readTMVar (ctrl^.sysLocal) + -- sysLocal is particular to a given control connection + when (curLocal ^. Disco.partitioner == Disco.mURMUR3_PARTITIONER) $ do + tMap <- readTVar (rInfo^.tokenMap) + dcTMap <- readTVar (rInfo^.dcTokenMap) + uuidMap <- readTVar (pInfo^.hostsByUUID) + let tMap' = Map.withoutKeys tMap tokenSet + dc = h^.dataCentre + dcTMap' = if Map.member dc counts + then Map.update (Just . flip Map.withoutKeys tokenSet) dc dcTMap + else Map.delete dc dcTMap + (simple, topo) = buildMasterReplicaMaps tMap' dcTMap' counts uuidMap + writeTVar (rInfo^.tokenMap) tMap' + writeTVar (rInfo^.dcTokenMap) dcTMap' + writeTVar (rInfo^.simpleReplicas) simple + writeTVar (rInfo^.topoReplicas) topo + + addHostPeerInfo :: + (Disco.HasTokens p (Cql.Set Text)) => + ClientState -> + p -> + Host -> + IO () + addHostPeerInfo cst p h = do + newPeerResult <- atomically $ do + pInfo <- readTVar (cst^.peerInfo) + rpcMap <- readTVar (pInfo^.hostsByRPC) + uuidMap <- readTVar (pInfo^.hostsByUUID) + let existed = Map.member (h^.hostAddr) rpcMap + updateHostAddr curHost = do + writeTVar (pInfo^.hostsByUUID) (Map.insert (h^.hostId) h uuidMap) + modifyTVar (pInfo^.hostsByRPC) (Map.insert (h^.hostAddr) h) + modifyTVar (pInfo^.hostsByRPC) (Map.delete (curHost^.hostAddr)) + pure $ RenumberedPeer curHost h + if existed + then pure ExistingPeer + else do + let mHost = Map.lookup (h^.hostId) uuidMap + case mHost of + Just curHost -> updateHostAddr curHost + Nothing -> do + modifyTVar (pInfo^.hostsByRPC) (Map.insert (h^.hostAddr) h) + modifyTVar (pInfo^.hostsByUUID) (Map.insert (h^.hostId) h) + let addToCount (y1, rackMap1) (y2, rackMap2) = + let !y3 = y1 + y2 + !rackMap3 = Map.unionWith (+) rackMap1 rackMap2 + in (y3, rackMap3) + modifyTVar + (pInfo^.dcHostCount) + (Map.insertWith addToCount (h^.dataCentre) (1, Map.singleton (h^.rack) 1)) + counts <- readTVar (pInfo^.dcHostCount) + let tSet = parseTokens (p ^. Disco.tokens) + modifyTVar (pInfo^.ownedTokens) (Map.insert (h^.hostId) tSet) + pure $ NewPeer tSet counts + logInfo (cst^.context.settings.logger) $ "addHostPeerInfo: newPeerResult: " + <> string8 (show newPeerResult) + case newPeerResult of + NewPeer tSet counts -> do + let pol = cst^.policy + okay <- liftIO $ acceptable pol h + when okay $ do + newPool <- runClient cst $ mkPool h + atomically $ modifyTVar (cst^.hostmap) (Map.insert h newPool) + startMonitor cst (h^.hostAddr) + liftIO $ onEvent pol (HostNew h) + logInfo (cst^.context.settings.logger) $ + "New Peer: " <> string8 (show h) + rInfo <- readTVarIO (cst^.replicationInfo) + void $ tryRunJob (rInfo^.addHostJobs) (h^.hostId) $ addHostReplication cst h tSet counts + RenumberedPeer oldHost newHost -> do + let pol = cst^.policy + newOkay <- liftIO $ acceptable pol newHost + atomically $ modifyTVar (cst^.hostmap) (Map.delete oldHost) + liftIO $ onEvent pol (HostGone oldHost) + when newOkay $ do + newPool <- runClient cst $ mkPool newHost + atomically $ modifyTVar (cst^.hostmap) (Map.insert newHost newPool) + startMonitor cst (newHost^.hostAddr) + cancelJobRenumbered (cst^.jobs) (oldHost^.hostAddr) + cancelJobRenumbered (cst^.multiControl.multiControlJobs) (oldHost^.hostAddr) + logInfo (cst^.context.settings.logger) $ "Renumbered host: " <> + string8 (show oldHost) <> " -> " <> string8 (show newHost) + ExistingPeer -> do + logInfo (cst^.context.settings.logger) $ + "Host was already known: " <> string8 (show h) + + addHostReplication :: ClientState -> Host -> Set.Set Int64 -> DCHostCount -> IO () + addHostReplication cst h tSet counts = do + let pInfoTV = cst^.peerInfo + rInfoTV = cst^.replicationInfo + l = cst^.context.settings.logger + logInfo l $ "Host Replication Job for " <> string8 (show h) <> " started." + atomically $ do + pInfo <- readTVar pInfoTV + curLocal <- readTMVar (ctrl^.sysLocal) + when (curLocal ^. Disco.partitioner == Disco.mURMUR3_PARTITIONER) $ do + rInfo <- readTVar rInfoTV + tMap <- readTVar (rInfo^.tokenMap) + dcTMap <- readTVar (rInfo^.dcTokenMap) + uuidMap <- readTVar (pInfo^.hostsByUUID) + let singleMap = Map.fromArgSet $ Set.map (`Arg` (h^.hostId)) tSet + tMap' = Map.union tMap singleMap + dc = h^.dataCentre + dcTMap' = Map.insertWith Map.union dc singleMap dcTMap + (simple, topo) = buildMasterReplicaMaps tMap' dcTMap' counts uuidMap + writeTVar (rInfo^.tokenMap) tMap' + writeTVar (rInfo^.dcTokenMap) dcTMap' + writeTVar (rInfo^.simpleReplicas) simple + writeTVar (rInfo^.topoReplicas) topo + + startNewPeerJob :: InetAddr -> Client () + startNewPeerJob a = do + cst <- ask + (pInfo, verInfo, ctrlConn) <- atomically' $ do + pInfo <- readTVar (cst^.peerInfo) + verInfo <- readTVar (cst^.versionInfo) + ctrlConn <- readTMVar (ctrl^.connection) + pure (pInfo, verInfo, ctrlConn) + logInfo' $ "Starting new peer job for " <> string8 (show a) + <> " and control connection to: " <> string8 (show $ ctrlConn^.host.hostId) + if (verInfo^.releaseVer) >= pEERS_V2_MIN_VERSION + then runJob_ (pInfo^.newHostJobs) (ctrlConn^.host.hostId, a) (newPeerV2 cst) + else runJob_ (pInfo^.newHostJobs) (ctrlConn^.host.hostId, a) (newPeer cst) + where + tryTenMins :: MonadIO m => RetryPolicyM m + tryTenMins = limitRetriesByCumulativeDelay 600_000_000 $ fullJitterBackoff 5000 + newPeerV2 :: ClientState -> IO () + newPeerV2 cst = do + logInfo (cst^.context.settings.logger) $ "Starting NewPeer (V2) job: " <> string8 (show a) + conn <- atomically $ readTMVar (ctrl^.connection) + mPeer <- retrying tryTenMins (const $ pure . isNothing) (\_ -> findPeerV2 (inet2ipPort a) conn) + for_ mPeer $ \peer -> do + let h = peer2HostV2 peer + logInfo (cst^.context.settings.logger) $ "Found NewNode peer (V2): " <> string8 (show h) + addHostPeerInfo cst peer h + when (isNothing mPeer) $ logNotFound cst + newPeer :: ClientState -> IO () + newPeer cst = do + logInfo (cst^.context.settings.logger) $ "Starting NewPeer job: " <> string8 (show a) + conn <- atomically $ readTMVar (ctrl^.connection) + mPeer <- retrying tryTenMins (const $ pure . isNothing) (\_ -> findPeer (inet2ip a) conn) + for_ mPeer $ \peer -> do + let h = peer2Host (cst^.context.settings.portnumber) peer + logInfo (cst^.context.settings.logger) $ "Found NewNode peer: " <> string8 (show h) + addHostPeerInfo cst peer h + when (isNothing mPeer) $ logNotFound cst + logNotFound :: ClientState -> IO () + logNotFound cst = do + let l = cst^.context.settings.logger + logWarn l $ "Unable to find Peer for address: " <> string8 (show a) <> + " after ten minutes." + + upsertKeyspaceReplicationInfo :: Keyspace -> Client () + upsertKeyspaceReplicationInfo ksName = do + cst <- ask + conn <- atomically' $ readTMVar (ctrl^.connection) + mKs <- findKeyspace ksName conn + for_ mKs $ \ks -> do + let !replicationMapSingle = parseKeyspace ks + newKS = Map.lookup ksName replicationMapSingle + atomically' $ do + rInfo <- readTVar (cst^.replicationInfo) + oldKSMap <- readTVar (rInfo^.keyspaces) + unless (Map.lookup ksName oldKSMap == newKS) $ do + writeTVar (rInfo^.keyspaces) (Map.union replicationMapSingle oldKSMap) + + removeKeyspaceReplicationInfo :: Keyspace -> Client () + removeKeyspaceReplicationInfo ksName = do + cst <- ask + atomically' $ do + rInfo <- readTVar (cst^.replicationInfo) + modifyTVar (rInfo^.keyspaces) (Map.delete ksName) + + startMonitor :: ClientState -> InetAddr -> IO () + startMonitor s a = do + rpcMap <- atomically $ do + pInfo <- readTVar (s^.peerInfo) + readTVar (pInfo^.hostsByRPC) + for_ (Map.lookup a rpcMap) $ \h -> + tryRunJob_ (s^.jobs) a $ runClient s $ do + logInfo' $ "Monitoring job started for: " <> string8 (show a) + monitor (Ms 3000) (Ms 60_000) h + logInfo' $ "Monitor finished for: " <> string8 (show a) + <> ", preparing queries and setting up control." + prepareAllQueries h + (controlMap, priorityMap) <- atomically' $ do + cMap <- readTVar (s^.multiControl.multiControlMap) + pMap <- readTVar (s^.multiControl.multiControlPriorityMap) + pure (cMap, pMap) + case Map.lookup (h^.hostId) controlMap of + Nothing -> for_ (Map.lookup (h^.hostId) priorityMap) (enqueueControl h) + Just newCtrl -> do + cc <- atomically' $ readTMVar (newCtrl^.connection) + ok <- checkControl cc + unless ok $ dropControl h + + checkHostControl :: ClientState -> Host -> IO () + checkHostControl s h = do + controlMap <- readTVarIO (s^.multiControl.multiControlMap) + case Map.lookup (h ^. hostId) controlMap of + Nothing -> pure () + Just oldCtrl -> do + cc <- atomically $ readTMVar (oldCtrl^.connection) + runClient s $ do + ok <- checkControl cc + unless ok $ dropControl h + +----------------------------------------------------------------------------- +-- Replication Queries +getReplication :: ServerPrepQuery k a b -> Client (Maybe ReplicationStrategy) +getReplication ServerPrepQuery {spqMetadata} = runMaybeT $ do + colSpec <- hoistMaybe $ listToMaybe (Cql.columnSpecs spqMetadata) + let keyspace = Cql.keyspace colSpec + cst <- lift ask + strategyMap <- lift . atomically' $ do + rInfo <- readTVar (cst^.replicationInfo) + readTVar (rInfo^.keyspaces) + hoistMaybe $ Map.lookup keyspace strategyMap + +getSimpleReplicas :: Tuple a => Cql.Version -> Int -> ServerPrepQuery k a b -> a -> Client (Maybe [UUID]) +getSimpleReplicas v replicaCount spq a = do + computedKey <- computeHashKey v spq a + forMaybe computedKey $ \hashKey -> do + cst <- ask + repMap <- atomically' $ do + rInfo <- readTVar (cst^.replicationInfo) + readTVar (rInfo^.simpleReplicas) + case Map.lookupGE hashKey repMap of + Just (_k, replicas) -> pure . Just $ take replicaCount replicas + Nothing -> pure $ take replicaCount . snd <$> Map.lookupMin repMap + +getDCReplicas :: Tuple a => Cql.Version -> Text -> Int -> ServerPrepQuery k a b -> a -> Client (Maybe [UUID]) +getDCReplicas v dc replicaCount spq a = do + computedKey <- computeHashKey v spq a + forMaybe computedKey $ \hashKey -> do + cst <- ask + topoRepMap <- atomically' $ do + rInfo <- readTVar (cst^.replicationInfo) + readTVar (rInfo^.topoReplicas) + forMaybe (Map.lookup dc topoRepMap) $ \repMap -> do + case Map.lookupGE hashKey repMap of + Just (_k, replicas) -> pure . Just $ take replicaCount replicas + Nothing -> pure $ take replicaCount . snd <$> Map.lookupMin repMap + +computeHashKey :: Tuple a => Cql.Version -> ServerPrepQuery k a b -> a -> Client (Maybe Int64) +computeHashKey v ServerPrepQuery {..} a = + case Cql.primaryKeyIndices spqMetadata of + [] -> pure Nothing + ixs -> + let rawKey :: ByteString + !rawKey = runPut $ rowKey v ixs a + hashKey :: Int64 + !hashKey = murmur3HashKey rawKey + in pure $ Just hashKey + +forMaybe :: Applicative f => Maybe a -> (a -> f (Maybe b)) -> f (Maybe b) +forMaybe m f = maybe (pure Nothing) f m + +----------------------------------------------------------------------------- +-- Utilities + +clientConnect :: Host -> Client Connection +clientConnect h = do + ctx <- view context + protoTV <- view protoVerInUse + let s = ctx^.settings + v <- readTVarIO' protoTV + C.connect (s^.connSettings) (ctx^.timeouts) v (s^.logger) h + +-- | Get the 'Result' out of a 'HostResponse'. If the response is an 'RsError', +-- a 'ResponseError' is thrown. If the response is neither +-- 'RsResult' nor 'RsError', an 'UnexpectedResponse' is thrown. +getResult :: MonadThrow m => HostResponse k a b -> m (Result k a b) +getResult (HostResponse _ (RsResult _ _ r)) = return r +getResult (HostResponse h (RsError t w e)) = throwM (ResponseError h t w e) +getResult hr = unexpected hr +{-# INLINE getResult #-} + +makePreparedQuery :: MonadThrow m => PrepQuery k a b -> HostResponse k a b -> m (Host, ServerPrepQuery k a b) +makePreparedQuery PrepQuery {..} hr = getResult hr >>= \case + PreparedResult i prepMD _resultMD -> return (hrHost hr, ServerPrepQuery + { spqStr = pqStr + , spqId = i + , spqIdLocal = pqId + , spqMetadata = prepMD + }) + _ -> unexpected hr +{-# INLINE makePreparedQuery #-} + +unexpected :: MonadThrow m => HostResponse k a b -> m c +unexpected (HostResponse h r) = throwM $ UnexpectedResponse h r + +atomically' :: STM a -> Client a +atomically' = liftIO . atomically + +readTVarIO' :: TVar a -> Client a +readTVarIO' = liftIO . readTVarIO + +logInfo' :: Builder -> Client () +logInfo' m = do + l <- view (context.settings.logger) + liftIO $ logInfo l m +{-# INLINE logInfo' #-} + +logDebug' :: Builder -> Client () +logDebug' m = do + l <- view (context.settings.logger) + liftIO $ logDebug l m +{-# INLINE logDebug' #-} + +logWarn' :: Builder -> Client () +logWarn' m = do + l <- view (context.settings.logger) + liftIO $ logWarn l m +{-# INLINE logWarn' #-} + +logError' :: Builder -> Client () +logError' m = do + l <- view (context.settings.logger) + liftIO $ logError l m +{-# INLINE logError' #-} + +closeIfPresent :: TMVar Connection -> IO () +closeIfPresent connTMV = do + mConn <- atomically $ tryTakeTMVar connTMV + forM_ mConn $ \conn -> do + C.close conn diff --git a/src/lib/Database/CQL/IO/Cluster/Discovery.hs b/src/lib/Database/CQL/IO/Cluster/Discovery.hs --- a/src/lib/Database/CQL/IO/Cluster/Discovery.hs +++ b/src/lib/Database/CQL/IO/Cluster/Discovery.hs @@ -2,32 +2,234 @@ -- License, v. 2.0. If a copy of the MPL was not distributed with this -- file, You can obtain one at http://mozilla.org/MPL/2.0/. +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} module Database.CQL.IO.Cluster.Discovery where +import Control.Lens (makeFields) import Data.Functor.Identity (Identity) +import Data.Int (Int32) import Data.IP import Data.Text (Text) -import Database.CQL.Protocol +import Data.UUID (UUID, nil) +import Data.Word (Word32) +import Database.CQL.Protocol hiding (Keyspace) +import qualified Database.CQL.Protocol as CQL data Peer = Peer - { peerAddr :: !IP - , peerRPC :: !IP + { peerAddr :: !IP + , peerDC :: !Text + , peerHostId :: !UUID + , peerPreferredIp :: !(Maybe IP) + , peerRack :: !Text + , peerReleaseVer :: !Text + , peerRPC :: !IP -- ^ The address for the client to connect to. - , peerDC :: !Text - , peerRack :: !Text + , peerSchemaVer :: !UUID + , peerTokens :: !(CQL.Set Text) } deriving Show recordInstance ''Peer +makeFields ''Peer -peers :: QueryString R () (IP, IP, Text, Text) -peers = "SELECT peer, rpc_address, data_center, rack FROM system.peers" +data PeerV2 = PeerV2 + { peerV2Addr :: !IP + , peerV2Port :: !Word32 + , peerV2DC :: !Text + , peerV2HostId :: !UUID + , peerV2NativeAddr :: !IP + -- ^ The address for the client to connect to. + , peerV2NativePort :: !Word32 + -- ^ The port for the client to connect to. + , peerV2PreferredIp :: !(Maybe IP) + , peerV2PreferredPort :: !(Maybe Word32) + , peerV2Rack :: !Text + , peerV2ReleaseVer :: !Text + , peerV2SchemaVer :: !UUID + , peerV2Tokens :: !(CQL.Set Text) + } deriving Show -peer :: QueryString R (Identity IP) (IP, IP, Text, Text) -peer = "SELECT peer, rpc_address, data_center, rack FROM system.peers where peer = ?" +recordInstance ''PeerV2 +makeFields ''PeerV2 -local :: QueryString R () (Text, Text) -local = "SELECT data_center, rack FROM system.local WHERE key='local'" +data LocalV3 = LocalV3 + { localV3Bootstrapped :: !Text + , localV3BroadcastAddr :: !IP + , localV3ClusterName :: !Text + , localV3CqlVer :: !Text + , localV3DC :: !Text + , localV3GossipGeneration :: !Int32 + , localV3HostId :: !UUID + , localV3ListenAddr :: !IP + , localV3NativeProtocolVer :: !Text + , localV3Partitioner :: !Text + , localV3Rack :: !Text + , localV3ReleaseVer :: !Text + , localV3RpcAddr :: !IP + , localV3SchemaVer :: !UUID + , localV3Tokens :: !(CQL.Set Text) + } + deriving (Eq, Show) +recordInstance ''LocalV3 +makeFields ''LocalV3 + +data LocalV4 = LocalV4 + { localV4Bootstrapped :: !Text + , localV4BroadcastAddr :: !IP + , localV4BroadcastPort :: !Word32 + , localV4ClusterName :: !Text + , localV4CqlVer :: !Text + , localV4DC :: !Text + , localV4GossipGeneration :: !Int32 + , localV4HostId :: !UUID + , localV4ListenAddr :: !IP + , localV4ListenPort :: !Word32 + , localV4NativeProtocolVer :: !Text + , localV4Partitioner :: !Text + , localV4Rack :: !Text + , localV4ReleaseVer :: !Text + , localV4RpcAddr :: !IP + , localV4RpcPort :: !Word32 + , localV4SchemaVer :: !UUID + , localV4Tokens :: !(CQL.Set Text) + } + deriving (Eq, Show) + +recordInstance ''LocalV4 +makeFields ''LocalV4 + +data Local = LocalV3_ !LocalV3 | LocalV4_ !LocalV4 + deriving (Eq, Show) + +data Keyspace = Keyspace + { keyspaceName :: !CQL.Keyspace + , keyspaceDurableWrites :: !Bool + , keyspaceReplication :: !(CQL.Map Text Text) + } + deriving (Eq, Show) + +recordInstance ''Keyspace +makeFields ''Keyspace + +dummyLocal :: Local +dummyLocal = LocalV4_ LocalV4 + { localV4Bootstrapped = "COMPLETED" + , localV4BroadcastAddr = IPv4 $ toIPv4 [127,0,0,1] + , localV4BroadcastPort = 7000 + , localV4ClusterName = "Test Cluster" + , localV4CqlVer = "3.4.5" + , localV4GossipGeneration = 0 + , localV4DC = "DC1" + , localV4HostId = nil + , localV4ListenAddr = IPv4 $ toIPv4 [127,0,0,1] + , localV4ListenPort = 7000 + , localV4NativeProtocolVer = "4" + , localV4Partitioner = "org.apache.cassandra.dht.Murmur3Partitioner" + , localV4Rack = "RAC1" + , localV4ReleaseVer = "4.0.17" + , localV4RpcAddr = IPv4 $ toIPv4 [127,0,0,1] + , localV4RpcPort = 9042 + , localV4SchemaVer = nil + , localV4Tokens = CQL.Set [] + } + +instance HasBroadcastAddr Local IP where + broadcastAddr f l = case l of + LocalV3_ l3 -> LocalV3_ <$> broadcastAddr f l3 + LocalV4_ l4 -> LocalV4_ <$> broadcastAddr f l4 + +instance HasCqlVer Local Text where + cqlVer f l = case l of + LocalV3_ l3 -> LocalV3_ <$> cqlVer f l3 + LocalV4_ l4 -> LocalV4_ <$> cqlVer f l4 + +instance HasDC Local Text where + dC f l = case l of + LocalV3_ l3 -> LocalV3_ <$> dC f l3 + LocalV4_ l4 -> LocalV4_ <$> dC f l4 + +instance HasHostId Local UUID where + hostId f l = case l of + LocalV3_ l3 -> LocalV3_ <$> hostId f l3 + LocalV4_ l4 -> LocalV4_ <$> hostId f l4 + +instance HasNativeProtocolVer Local Text where + nativeProtocolVer f l = case l of + LocalV3_ l3 -> LocalV3_ <$> nativeProtocolVer f l3 + LocalV4_ l4 -> LocalV4_ <$> nativeProtocolVer f l4 + +instance HasPartitioner Local Text where + partitioner f l = case l of + LocalV3_ l3 -> LocalV3_ <$> partitioner f l3 + LocalV4_ l4 -> LocalV4_ <$> partitioner f l4 + +instance HasRack Local Text where + rack f l = case l of + LocalV3_ l3 -> LocalV3_ <$> rack f l3 + LocalV4_ l4 -> LocalV4_ <$> rack f l4 + + +instance HasReleaseVer Local Text where + releaseVer f l = case l of + LocalV3_ l3 -> LocalV3_ <$> releaseVer f l3 + LocalV4_ l4 -> LocalV4_ <$> releaseVer f l4 + +instance HasRpcAddr Local IP where + rpcAddr f l = case l of + LocalV3_ l3 -> LocalV3_ <$> rpcAddr f l3 + LocalV4_ l4 -> LocalV4_ <$> rpcAddr f l4 + +instance HasTokens Local (CQL.Set Text) where + tokens f l = case l of + LocalV3_ l3 -> LocalV3_ <$> tokens f l3 + LocalV4_ l4 -> LocalV4_ <$> tokens f l4 + +mURMUR3_PARTITIONER :: Text +mURMUR3_PARTITIONER = "org.apache.cassandra.dht.Murmur3Partitioner" + +peers :: QueryString R () (IP, Text, UUID, Maybe IP, Text, Text, IP, UUID, CQL.Set Text) +peers = "SELECT peer, data_center, host_id, preferred_ip, rack, \ + \release_version, rpc_address, schema_version, tokens FROM \ + \system.peers" + +peer :: QueryString R (Identity IP) (IP, Text, UUID, Maybe IP, Text, Text, IP, UUID, CQL.Set Text) +peer = "SELECT peer, data_center, host_id, preferred_ip, rack, \ + \release_version, rpc_address, schema_version, tokens FROM \ + \system.peers WHERE peer = ?" + +peersV2 :: QueryString R () (IP, Word32, Text, UUID, IP, Word32, Maybe IP, Maybe Word32, Text, Text, UUID, CQL.Set Text) +peersV2 = "SELECT peer, peer_port, data_center, host_id, native_address, \ + \native_port, preferred_ip, preferred_port, rack, release_version, \ + \schema_version, tokens FROM system.peers_v2" + +peerV2 :: QueryString R (Identity IP) (IP, Word32, Text, UUID, IP, Word32, Maybe IP, Maybe Word32, Text, Text, UUID, CQL.Set Text) +peerV2 = "SELECT peer, peer_port, data_center, host_id, native_address, \ + \native_port, preferred_ip, preferred_port, rack, release_version, \ + \schema_version, tokens FROM system.peers_v2 WHERE peer = ?" + +localV4 :: QueryString R () (Text, IP, Word32, Text, Text, Text, Int32, UUID, IP, Word32, Text, Text, Text, Text, IP, Word32, UUID, CQL.Set Text) +localV4 = "SELECT bootstrapped, broadcast_address, broadcast_port, cluster_name, \ + \cql_version, data_center, gossip_generation, host_id, listen_address, \ + \listen_port, native_protocol_version, partitioner, rack, \ + \release_version, rpc_address, rpc_port, schema_version, tokens \ + \FROM system.local WHERE key = 'local'" + +localV3 :: QueryString R () (Text, IP, Text, Text, Text, Int32, UUID, IP, Text, Text, Text, Text, IP, UUID, CQL.Set Text) +localV3 = "SELECT bootstrapped, broadcast_address, cluster_name, \ + \cql_version, data_center, gossip_generation, host_id, listen_address, \ + \native_protocol_version, partitioner, rack, \ + \release_version, rpc_address, schema_version, tokens \ + \FROM system.local WHERE key = 'local'" + +keyspaces :: QueryString R () (CQL.Keyspace, Bool, CQL.Map Text Text) +keyspaces = "SELECT keyspace_name, durable_writes, replication FROM \ + \system_schema.keyspaces" + +keyspace :: QueryString R (Identity CQL.Keyspace) (CQL.Keyspace, Bool, CQL.Map Text Text) +keyspace = "SELECT keyspace_name, durable_writes, replication FROM \ + \system_schema.keyspaces WHERE keyspace_name = ?" diff --git a/src/lib/Database/CQL/IO/Cluster/Host.hs b/src/lib/Database/CQL/IO/Cluster/Host.hs --- a/src/lib/Database/CQL/IO/Cluster/Host.hs +++ b/src/lib/Database/CQL/IO/Cluster/Host.hs @@ -7,20 +7,33 @@ module Database.CQL.IO.Cluster.Host where -import Control.Lens (Lens') +import Control.Lens (Lens', (^.)) import Database.CQL.Protocol (Response (..)) -import Database.CQL.IO.Cluster.Discovery +import Database.CQL.IO.Cluster.Discovery hiding (broadcastAddr, rack, hostId) +import qualified Database.CQL.IO.Cluster.Discovery as Disco (broadcastAddr, rack, hostId) import Data.IP import Data.Text (Text, unpack) import Network.Socket (SockAddr (..), PortNumber) +import Data.UUID (UUID, nil) -- | A Cassandra host known to the client. data Host = Host { _hostAddr :: !InetAddr + , _broadcastAddr :: !InetAddr + , _hostId :: !UUID , _dataCentre :: !Text , _rack :: !Text } +dummyHost :: InetAddr -> Host +dummyHost a = Host + { _hostAddr = a + , _broadcastAddr = a + , _hostId = nil + , _dataCentre = "" + , _rack = "" + } + instance Eq Host where a == b = _hostAddr a == _hostAddr b @@ -28,8 +41,42 @@ compare a b = compare (_hostAddr a) (_hostAddr b) peer2Host :: PortNumber -> Peer -> Host -peer2Host i p = Host (ip2inet i (peerRPC p)) (peerDC p) (peerRack p) +peer2Host i p = Host + { _hostAddr = ip2inet i (peerRPC p) + , _broadcastAddr = ip2inet 7000 (peerAddr p) + , _hostId = peerHostId p + , _dataCentre = peerDC p + , _rack = peerRack p + } +peer2HostV2 :: PeerV2 -> Host +peer2HostV2 peer2 = + let nativePortNum = fromIntegral . peerV2NativePort $ peer2 + broadcastPortNum = fromIntegral . peerV2Port $ peer2 + in Host + { _hostAddr = ip2inet nativePortNum (peerV2NativeAddr peer2) + , _broadcastAddr = ip2inet broadcastPortNum (peerV2Addr peer2) + , _hostId = peerV2HostId peer2 + , _dataCentre = peerV2DC peer2 + , _rack = peerV2Rack peer2 + } + +local2Host :: PortNumber -> Local -> Host +local2Host settingsPort curLocal = + let rpcPortNum = case curLocal of + LocalV3_ _ -> settingsPort + LocalV4_ l4 -> fromIntegral $ l4 ^. rpcPort + broadcastPortNum = case curLocal of + LocalV3_ _ -> 7000 + LocalV4_ l4 -> fromIntegral $ l4 ^. broadcastPort + in Host + { _hostAddr = ip2inet rpcPortNum (curLocal^.rpcAddr) + , _broadcastAddr = ip2inet broadcastPortNum (curLocal ^. Disco.broadcastAddr) + , _hostId = curLocal ^. Disco.hostId + , _dataCentre = curLocal^.dC + , _rack = curLocal ^. Disco.rack + } + updateHost :: Host -> Maybe (Text, Text) -> Host updateHost h (Just (dc, rk)) = h { _dataCentre = dc, _rack = rk } updateHost h Nothing = h @@ -44,23 +91,34 @@ -- cluster changes. data HostEvent = HostNew !Host -- ^ a new host has been added to the cluster - | HostGone !InetAddr -- ^ a host has been removed from the cluster - | HostUp !InetAddr -- ^ a host has been started - | HostDown !InetAddr -- ^ a host has been stopped + | HostGone !Host -- ^ a host has been removed from the cluster + | HostUp !Host -- ^ a host has been started + | HostDown !Host -- ^ a host has been stopped + | AddrDown !InetAddr -- ^ an address was reported down, we could not find the corresponding host. -- | The IP address and port number of a host. hostAddr :: Lens' Host InetAddr -hostAddr f ~(Host a c r) = fmap (\x -> Host x c r) (f a) +hostAddr f h = fmap (\x -> h {_hostAddr = x}) (f (_hostAddr h)) {-# INLINE hostAddr #-} +-- | The IP address and port number that will be used for events about a host. +broadcastAddr :: Lens' Host InetAddr +broadcastAddr f h = fmap (\x -> h {_broadcastAddr = x}) (f (_broadcastAddr h)) +{-# INLINE broadcastAddr #-} + +-- | The UUID hostId of a host. +hostId :: Lens' Host UUID +hostId f h = fmap (\x -> h {_hostId = x}) (f (_hostId h)) +{-# INLINE hostId #-} + -- | The data centre name (may be an empty string). dataCentre :: Lens' Host Text -dataCentre f ~(Host a c r) = fmap (\x -> Host a x r) (f c) +dataCentre f h = fmap (\x -> h {_dataCentre = x}) (f (_dataCentre h)) {-# INLINE dataCentre #-} -- | The rack name (may be an empty string). rack :: Lens' Host Text -rack f ~(Host a c r) = fmap (\x -> Host a c x) (f r) +rack f h = fmap (\x -> h {_rack = x}) (f (_rack h)) {-# INLINE rack #-} instance Show Host where @@ -69,6 +127,8 @@ . showString (unpack (_rack h)) . showString ":" . shows (_hostAddr h) + . showString ":" + . shows (_broadcastAddr h) $ "" ----------------------------------------------------------------------------- @@ -100,3 +160,12 @@ ip2inet p (IPv4 a) = InetAddr $ SockAddrInet p (toHostAddress a) ip2inet p (IPv6 a) = InetAddr $ SockAddrInet6 p 0 (toHostAddress6 a) 0 +inet2ip :: InetAddr -> IP +inet2ip (InetAddr (SockAddrInet _ a)) = IPv4 (fromHostAddress a) +inet2ip (InetAddr (SockAddrInet6 _ _ a6 _)) = IPv6 (fromHostAddress6 a6) +inet2ip x = error $ "inet2ip: expecting IPv4 or IPv6, got: " ++ show x + +inet2ipPort :: InetAddr -> (IP, PortNumber) +inet2ipPort (InetAddr (SockAddrInet p a)) = (IPv4 (fromHostAddress a), p) +inet2ipPort (InetAddr (SockAddrInet6 p _ a6 _)) = (IPv6 (fromHostAddress6 a6), p) +inet2ipPort x = error $ "inet2ipPort: expecting IPv4 or IPv6, got: " ++ show x diff --git a/src/lib/Database/CQL/IO/Cluster/Policies.hs b/src/lib/Database/CQL/IO/Cluster/Policies.hs --- a/src/lib/Database/CQL/IO/Cluster/Policies.hs +++ b/src/lib/Database/CQL/IO/Cluster/Policies.hs @@ -3,6 +3,7 @@ -- file, You can obtain one at http://mozilla.org/MPL/2.0/. {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE RankNTypes #-} module Database.CQL.IO.Cluster.Policies ( Policy (..) @@ -15,8 +16,11 @@ import Control.Lens ((^.), view, over, makeLenses) import Control.Monad import Data.Map.Strict (Map) +import Data.Maybe (isNothing) import Data.Word import Database.CQL.IO.Cluster.Host +import Database.CQL.IO.PrepQuery (ServerPrepQuery) +import Database.CQL.Protocol (Tuple) import System.Random.MWC import Prelude @@ -33,7 +37,7 @@ , onEvent :: HostEvent -> IO () -- ^ Event handler. Policies will be informed about cluster changes -- through this function. - , select :: IO (Maybe Host) + , select :: forall k a b. Tuple a => Maybe (ServerPrepQuery k a b, a) -> IO (Maybe Host) -- ^ Host selection. The driver will ask for a host to use in a query -- through this function. A policy which has no available nodes may -- return Nothing. @@ -42,6 +46,11 @@ , acceptable :: Host -> IO Bool -- ^ During startup and node discovery, the driver will ask the -- policy if a dicovered host should be ignored. + , priority :: Host -> IO Int + -- ^ During startup and node discovery, the driver will ask the + -- policy if a dicovered host should be ignored, and if not, for a + -- priority value for a control connection to that host. Lower values are + -- preferred before later values. Non-positive values are ignored. , hostCount :: IO Word -- ^ During query processing, the driver will ask the policy for -- a rough estimate of alive hosts. The number is used to repeatedly @@ -65,11 +74,10 @@ roundRobin = do h <- newTVarIO emptyHosts c <- newTVarIO 0 - return $ Policy (defSetup h) (defOnEvent h) (pickHost h c) - (defCurrent h) defAcceptable (defHostCount h) - (defDisplay h) + return (defPolicy h) + { select = pickHost h c } where - pickHost h c = atomically $ do + pickHost h c _ = atomically $ do m <- view alive <$> readTVar h if Map.null m then return Nothing @@ -83,11 +91,10 @@ random = do h <- newTVarIO emptyHosts g <- createSystemRandom - return $ Policy (defSetup h) (defOnEvent h) (pickHost h g) - (defCurrent h) defAcceptable (defHostCount h) - (defDisplay h) + return (defPolicy h) + { select = pickHost h g } where - pickHost h g = do + pickHost h g _ = do m <- view alive <$> readTVarIO h if Map.null m then return Nothing @@ -98,6 +105,18 @@ ----------------------------------------------------------------------------- -- Defaults +defPolicy :: TVar Hosts -> Policy +defPolicy h = Policy + { setup = defSetup h + , onEvent = defOnEvent h + , select = const $ pure Nothing + , current = defCurrent h + , acceptable = defAcceptable + , priority = defPriority + , hostCount = defHostCount h + , display = defDisplay h + } + emptyHosts :: Hosts emptyHosts = Hosts Map.empty Map.empty @@ -105,8 +124,11 @@ defDisplay h = show <$> readTVarIO h defAcceptable :: Host -> IO Bool -defAcceptable = const $ return True +defAcceptable = fmap (/= 0) . defPriority +defPriority :: Host -> IO Int +defPriority = const $ pure 1 + defSetup :: HostMap -> [Host] -> [Host] -> IO () defSetup r a b = do let ha = Map.fromList $ zip (map (view hostAddr) a) a @@ -123,27 +145,38 @@ defOnEvent :: HostMap -> HostEvent -> IO () defOnEvent r (HostNew h) = atomically $ do m <- readTVar r - when (Nothing == get (h^.hostAddr) m) $ + when (isNothing $ get (h^.hostAddr) m) $ writeTVar r (over alive (Map.insert (h^.hostAddr) h) m) -defOnEvent r (HostGone a) = atomically $ do +defOnEvent r (HostGone h) = atomically $ do m <- readTVar r + let a = (h^.hostAddr) if Map.member a (m^.alive) then writeTVar r (over alive (Map.delete a) m) else writeTVar r (over other (Map.delete a) m) -defOnEvent r (HostUp a) = atomically $ do +defOnEvent r (HostUp h) = atomically $ do m <- readTVar r + let a = (h^.hostAddr) case get a m of Nothing -> return () - Just h -> writeTVar r + Just _ -> writeTVar r $ over alive (Map.insert a h) . over other (Map.delete a) $ m -defOnEvent r (HostDown a) = atomically $ do +defOnEvent r (HostDown h) = atomically $ do m <- readTVar r + let a = (h^.hostAddr) case get a m of Nothing -> return () - Just h -> writeTVar r + Just _ -> writeTVar r + $ over other (Map.insert a h) + . over alive (Map.delete a) + $ m +defOnEvent r (AddrDown a) = atomically $ do + m <- readTVar r + case get a m of + Nothing -> return () + Just h -> writeTVar r $ over other (Map.insert a h) . over alive (Map.delete a) $ m diff --git a/src/lib/Database/CQL/IO/Connection.hs b/src/lib/Database/CQL/IO/Connection.hs --- a/src/lib/Database/CQL/IO/Connection.hs +++ b/src/lib/Database/CQL/IO/Connection.hs @@ -2,6 +2,7 @@ -- License, v. 2.0. If a copy of the MPL was not distributed with this -- file, You can obtain one at http://mozilla.org/MPL/2.0/. +{-# LANGUAGE CPP #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TemplateHaskell #-} @@ -12,6 +13,7 @@ , ConnId , ident , host + , protocol -- * Lifecycle , connect @@ -47,7 +49,9 @@ import Control.Monad.IO.Class import Data.ByteString.Builder import Data.Foldable (for_) +#if !MIN_VERSION_base(4,11,0) import Data.Semigroup ((<>)) +#endif import Data.Text.Lazy (fromStrict) import Data.Unique import Data.Vector (Vector, (!)) @@ -116,9 +120,23 @@ sta <- newTVarIO True sig <- signal rdr <- async (readLoop v g t tck h s syn sig sta lck) - Connection t h m v s sta syn lck rdr tck g sig . ConnId <$> newUnique + connId <- ConnId <$> newUnique + pure Connection + { _settings = t + , _host = h + , _tmanager = m + , _protocol = v + , _sock = s + , _status = sta + , _streams = syn + , _wLock = lck + , _reader = rdr + , _tickets = tck + , _logger = g + , _eventSig = sig + , _ident = connId + } initialise c - return c where sockOpen = Socket.open (t^.connectTimeout) (h^.hostAddr) (t^.tlsContext) @@ -127,6 +145,7 @@ startup c for_ (t^.defKeyspace) $ useKeyspace c + pure c `onException` close c @@ -202,7 +221,7 @@ checkAuth :: Connection -> IO () checkAuth c = unless (null (c^.settings.authenticators)) $ - logWarn' (c^.logger) (c^.host) $ + logWarn' (c^.logger) (c^.host) "Authentication configured but none required by the server." authenticate :: Connection -> Authenticate -> IO () diff --git a/src/lib/Database/CQL/IO/Connection/Socket.hs b/src/lib/Database/CQL/IO/Connection/Socket.hs --- a/src/lib/Database/CQL/IO/Connection/Socket.hs +++ b/src/lib/Database/CQL/IO/Connection/Socket.hs @@ -4,7 +4,6 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} -{-# LANGUAGE TemplateHaskell #-} -- | A thin wrapper of the Network.Socket API. module Database.CQL.IO.Connection.Socket @@ -32,8 +31,15 @@ import Database.CQL.IO.Cluster.Host import Database.CQL.IO.Exception (ConnectionError (..)) import Database.CQL.IO.Timeouts (Milliseconds (..)) -import Network.Socket (HostName, PortNumber, SockAddr (..), ShutdownCmd (..)) -import Network.Socket (Family (..), AddrInfo (..), AddrInfoFlag (..)) +import Network.Socket ( + AddrInfo (..), + AddrInfoFlag (..), + Family (..), + HostName, + PortNumber, + ShutdownCmd (..), + SockAddr (..) + ) import Network.Socket.ByteString.Lazy (sendAll) import OpenSSL.Session (SSL, SSLContext) import System.Timeout diff --git a/src/lib/Database/CQL/IO/Exception.hs b/src/lib/Database/CQL/IO/Exception.hs --- a/src/lib/Database/CQL/IO/Exception.hs +++ b/src/lib/Database/CQL/IO/Exception.hs @@ -126,6 +126,9 @@ -- | An error occurred during parsing of a response. This indicates -- a problem with the wire protocol and should be reported. ParseError :: String -> ProtocolError + -- | An error with the system.local table. This indicates a problem with the + -- assumptions about discovery during startup + SysLocalError :: Host -> String -> ProtocolError deriving instance Typeable ProtocolError instance Exception ProtocolError @@ -143,6 +146,9 @@ showString "unexpected query ID: " . shows i UnexpectedResponse h r -> showString "unexpected response: " . shows h . showString ": " . shows (f r) + SysLocalError h m -> + showString "error reading system.local from " . shows h . + showString " detail: " . showString m $ "" where f :: Response k a b -> Response k a NoShow diff --git a/src/lib/Database/CQL/IO/Jobs.hs b/src/lib/Database/CQL/IO/Jobs.hs --- a/src/lib/Database/CQL/IO/Jobs.hs +++ b/src/lib/Database/CQL/IO/Jobs.hs @@ -10,6 +10,7 @@ , runJob_ , tryRunJob , tryRunJob_ + , cancelJobRenumbered , cancelJobs , listJobs , listJobKeys @@ -47,6 +48,13 @@ toException = asyncExceptionToException fromException = asyncExceptionFromException +-- | The asynchronous exception used to cancel a job if it is obsolete due to a +-- host with the same hostId appearing at a different address. +data JobRenumbered = JobRenumbered deriving (Eq, Show, Typeable) +instance Exception JobRenumbered where + toException = asyncExceptionToException + fromException = asyncExceptionFromException + newJobs :: MonadIO m => m (Jobs k) newJobs = liftIO $ Jobs <$> newIORef Map.empty @@ -87,6 +95,15 @@ cancelJobs (Jobs d) = liftIO $ do jobs <- Map.elems <$> atomicModifyIORef' d (\m -> (Map.empty, m)) mapM_ (cancel . jobAsync) jobs + +cancelJobRenumbered :: (MonadIO m, Ord k) => Jobs k -> k -> m () +cancelJobRenumbered (Jobs ref) k = liftIO $ do + maybeCurJob <- atomicModifyIORef' ref $ \jobs -> + case Map.updateLookupWithKey (\_ _ -> Nothing) k jobs of + (Nothing, _) -> (jobs, Nothing) + (Just curJob, jobs') -> (jobs', Just curJob) + forM_ maybeCurJob $ \curJob -> + cancelWith (jobAsync curJob) JobRenumbered -- | List all running jobs. listJobs :: MonadIO m => Jobs k -> m [(k, Async ())] diff --git a/src/lib/Database/CQL/IO/Pool.hs b/src/lib/Database/CQL/IO/Pool.hs --- a/src/lib/Database/CQL/IO/Pool.hs +++ b/src/lib/Database/CQL/IO/Pool.hs @@ -2,6 +2,7 @@ -- License, v. 2.0. If a copy of the MPL was not distributed with this -- file, You can obtain one at http://mozilla.org/MPL/2.0/. +{-# LANGUAGE CPP #-} {-# LANGUAGE MultiWayIf #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -30,12 +31,18 @@ import Control.Lens ((^.), makeLenses, view) import Control.Monad.IO.Class import Control.Monad +#if !MIN_VERSION_base(4,8,0) import Data.Foldable (forM_, mapM_, find) +#else +import Data.Foldable (find) +#endif import Data.Function (on) import Data.Hashable import Data.IORef import Data.Sequence (Seq, ViewL (..), (|>), (><)) +#if !MIN_VERSION_base(4,11,0) import Data.Semigroup ((<>)) +#endif import Data.Time.Clock (UTCTime, NominalDiffTime, getCurrentTime, diffUTCTime) import Data.Vector (Vector, (!)) import Database.CQL.IO.Connection (Connection) @@ -207,10 +214,10 @@ return stale forM_ x $ \v -> ignore $ do logDebug (p^.logger) $ "Reaping idle connection: " <> string8 (show (value v)) - p^.destroyFn $ (value v) + p^.destroyFn $ value v stripe :: Pool -> IO Stripe -stripe p = ((p^.stripes) !) <$> ((`mod` (p^.settings.poolStripes)) . hash) <$> myThreadId +stripe p = ((p^.stripes) !) . ((`mod` (p^.settings.poolStripes)) . hash) <$> myThreadId {-# INLINE stripe #-} incrTimeouts :: Resource -> Resource diff --git a/src/lib/Database/CQL/IO/PrepQuery.hs b/src/lib/Database/CQL/IO/PrepQuery.hs --- a/src/lib/Database/CQL/IO/PrepQuery.hs +++ b/src/lib/Database/CQL/IO/PrepQuery.hs @@ -1,16 +1,22 @@ -- This Source Code Form is subject to the terms of the Mozilla Public -- License, v. 2.0. If a copy of the MPL was not distributed with this -- file, You can obtain one at http://mozilla.org/MPL/2.0/. +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE NamedFieldPuns #-} module Database.CQL.IO.PrepQuery - ( PrepQuery + ( PrepQuery (..) + , ServerPrepQuery (..) + , toLocal , prepared , queryString , PreparedQueries , new - , lookupQueryId - , lookupQueryString + , lookupQueryKeys + , lookupQueryLocal + , lookupQueryServer , insert , delete , queryStrings @@ -32,6 +38,7 @@ import Prelude import qualified Data.Map.Strict as M +import qualified Database.CQL.Protocol as Cql (MetaData (..)) ----------------------------------------------------------------------------- -- Prepared Query @@ -85,10 +92,13 @@ instance IsString (PrepQuery k a b) where fromString = prepared . fromString -newtype PrepQueryId = PrepQueryId (Digest SHA1) deriving (Eq, Ord) +newtype PrepQueryId = PrepQueryId (Digest SHA1) deriving (Eq, Ord, Show) prepared :: QueryString k a b -> PrepQuery k a b -prepared q = PrepQuery q $ PrepQueryId (hashlazy . encodeUtf8 . unQueryString $ q) +prepared q = PrepQuery + { pqStr = q + , pqId = PrepQueryId (hashlazy . encodeUtf8 . unQueryString $ q) + } queryString :: PrepQuery k a b -> QueryString k a b queryString = pqStr @@ -96,49 +106,78 @@ ----------------------------------------------------------------------------- -- Map of prepared queries to their query ID and query string -newtype QST = QST { unQST :: Text } -newtype QID = QID { unQID :: ByteString } deriving (Eq, Ord) +data ServerPrepQuery k a b = ServerPrepQuery + { spqStr :: !(QueryString k a b) + , spqId :: !(QueryId k a b) + , spqIdLocal :: !PrepQueryId + , spqMetadata :: !Cql.MetaData + } +toLocal :: ServerPrepQuery k a b -> PrepQuery k a b +toLocal ServerPrepQuery {..} = PrepQuery + { pqStr = spqStr + , pqId = spqIdLocal + } + +data SPQPhantom = forall k a b. SPQPhantom (ServerPrepQuery k a b) + +rePhantom :: SPQPhantom -> ServerPrepQuery k a b +rePhantom (SPQPhantom ServerPrepQuery{..}) = ServerPrepQuery + { spqStr = QueryString $ unQueryString spqStr + , spqId = QueryId $ unQueryId spqId + , spqIdLocal = spqIdLocal + , spqMetadata + } + data PreparedQueries = PreparedQueries - { queryMap :: !(TVar (Map PrepQueryId (QID, QST))) - , qid2Str :: !(TVar (Map QID QST)) + { queriesByLocalId :: !(TVar (Map PrepQueryId SPQPhantom)) + , queriesByServerId :: !(TVar (Map ByteString SPQPhantom)) } new :: IO PreparedQueries new = PreparedQueries <$> newTVarIO M.empty <*> newTVarIO M.empty -lookupQueryId :: PrepQuery k a b -> PreparedQueries -> STM (Maybe (QueryId k a b)) -lookupQueryId q m = do - qm <- readTVar (queryMap m) - return $ QueryId . unQID . fst <$> M.lookup (pqId q) qm +lookupQueryLocal :: PrepQuery k a b -> PreparedQueries -> STM (Maybe (ServerPrepQuery k a b)) +lookupQueryLocal q m = do + qm <- readTVar (queriesByLocalId m) + return . fmap rePhantom $ M.lookup (pqId q) qm -lookupQueryString :: QueryId k a b -> PreparedQueries -> STM (Maybe (QueryString k a b)) -lookupQueryString q m = do - qm <- readTVar (qid2Str m) - return $ QueryString . unQST <$> M.lookup (QID $ unQueryId q) qm +lookupQueryKeys :: PreparedQueries -> STM [PrepQueryId] +lookupQueryKeys m = do + qm <- readTVar (queriesByLocalId m) + return $ M.keys qm -insert :: PrepQuery k a b -> QueryId k a b -> PreparedQueries -> STM () -insert q i m = do - qq <- M.lookup (pqId q) <$> readTVar (queryMap m) - for_ qq (verify . snd) - modifyTVar' (queryMap m) $ - M.insert (pqId q) (QID $ unQueryId i, QST $ unQueryString (pqStr q)) - modifyTVar' (qid2Str m) $ - M.insert (QID $ unQueryId i) (QST $ unQueryString (pqStr q)) +lookupQueryServer :: QueryId k a b -> PreparedQueries -> STM (Maybe (ServerPrepQuery k a b)) +lookupQueryServer q m = do + qm <- readTVar (queriesByServerId m) + return . fmap rePhantom $ M.lookup (unQueryId q) qm + +insert :: PrepQuery k a b -> ServerPrepQuery k a b -> PreparedQueries -> STM () +insert q spq m = do + qq <- M.lookup (pqId q) <$> readTVar (queriesByLocalId m) + for_ qq verify + modifyTVar' (queriesByLocalId m) $ + M.insert (pqId q) (SPQPhantom spq) + modifyTVar' (queriesByServerId m) $ + M.insert (unQueryId $ spqId spq) (SPQPhantom spq) where - verify qs = - unless (unQST qs == unQueryString (pqStr q)) $ do - let a = unQST qs + verify spqPhantom = do + let ServerPrepQuery {..} = rePhantom spqPhantom + unless (spqStr == pqStr q) $ do let b = unQueryString (pqStr q) - throwSTM (HashCollision a b) + throwSTM (HashCollision (unQueryString spqStr) b) delete :: PrepQuery k a b -> PreparedQueries -> STM () delete q m = do - qid <- M.lookup (pqId q) <$> readTVar (queryMap m) - modifyTVar' (queryMap m) $ M.delete (pqId q) - case qid of + mSpqP <- M.lookup (pqId q) <$> readTVar (queriesByLocalId m) + modifyTVar' (queriesByLocalId m) $ M.delete (pqId q) + case mSpqP of Nothing -> return () - Just i -> modifyTVar' (qid2Str m) $ M.delete (fst i) + Just (SPQPhantom ServerPrepQuery {spqId}) -> + modifyTVar' (queriesByServerId m) $ M.delete (unQueryId spqId) queryStrings :: PreparedQueries -> STM [Text] -queryStrings m = map (unQST . snd) . M.elems <$> readTVar (queryMap m) +queryStrings m = map spqStrP . M.elems <$> readTVar (queriesByLocalId m) + where + spqStrP :: SPQPhantom -> Text + spqStrP (SPQPhantom ServerPrepQuery {spqStr}) = unQueryString spqStr diff --git a/src/lib/Database/CQL/IO/Protocol.hs b/src/lib/Database/CQL/IO/Protocol.hs --- a/src/lib/Database/CQL/IO/Protocol.hs +++ b/src/lib/Database/CQL/IO/Protocol.hs @@ -2,6 +2,7 @@ -- License, v. 2.0. If a copy of the MPL was not distributed with this -- file, You can obtain one at http://mozilla.org/MPL/2.0/. +{-# LANGUAGE CPP #-} {-# LANGUAGE OverloadedStrings #-} module Database.CQL.IO.Protocol where @@ -9,7 +10,9 @@ import Control.Monad.Catch import Data.ByteString.Lazy (ByteString) import Data.Maybe (fromMaybe) -import Data.Monoid ((<>)) +#if !MIN_VERSION_base(4,11,0) +import Data.Semigroup ((<>)) +#endif import Database.CQL.Protocol import Database.CQL.IO.Exception diff --git a/src/lib/Database/CQL/IO/Replication.hs b/src/lib/Database/CQL/IO/Replication.hs new file mode 100644 --- /dev/null +++ b/src/lib/Database/CQL/IO/Replication.hs @@ -0,0 +1,664 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} + +module Database.CQL.IO.Replication ( + SimpleReplicaMap, + NetworkTopologyReplicaMap, + ReplicationStrategy (..), + ReplicaJobKey (..), + dcReplicaMap, + simpleReplicaMap, + parseReplicationStrategy, + buildMasterReplicaMaps, + binBy, + cycle, + take, +) where + +import Bluefin.Consume (Consume, await, consumeStream) +import Bluefin.Eff (Eff, Effects, runPureEff, (:&), type (:>)) +import Bluefin.State (State, get, modify, put, runState) +import Bluefin.StateSource (newState, withStateSource) +import Bluefin.Stream (Stream, inFoldable, yield, yieldToList) +import Control.Lens ((^.)) +import Control.Monad (forever, replicateM, when) +import qualified Data.Attoparsec.Text as AP (decimal, parseOnly) +import Data.Foldable (Foldable (..), forM_) +import Data.Int (Int64) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map ( + alterF, + argSet, + delete, + elems, + empty, + findMin, + fromArgSet, + fromList, + insert, + insertLookupWithKey, + lookup, + map, + mapKeysMonotonic, + mapWithKey, + singleton, + size, + take, + ) +import Data.Semigroup (Arg (..), Sum (..)) +import Data.Sequence (Seq (..), (|>)) +import qualified Data.Sequence as Seq (empty) +import qualified Data.Set as Set ( + Set, + delete, + findMin, + fromList, + mapMonotonic, + singleton, + ) +import Data.Text (Text) +import Data.Word (Word64) +import qualified Database.CQL.Protocol as Cql (Map (..)) + +import Database.CQL.IO.Cluster.Host (Host (..), rack, hostId) + +import Data.Maybe (fromMaybe) +import Prelude hiding (cycle, take) +import Data.UUID (UUID) + +{- +>>> import Bluefin.Consume +>>> import Bluefin.Eff +>>> import Bluefin.Stream + +>>> runPureEff $ yieldToList $ \yOut -> consumeStream (\c -> binBy 2 c yOut) (\yIn -> inFoldable [1..4] yIn) +([[1,2],[3,4]],()) + -} +binBy :: (ea :> es, eb :> es) => Int -> Consume a ea -> Stream [a] eb -> Eff es () +binBy count source sink = + forever (replicateM count (await source) >>= yield sink) + +{- +>>> import Bluefin.Consume +>>> import Bluefin.Eff +>>> import Bluefin.Stream + +>>> runPureEff $ yieldToList $ \yOut -> consumeStream (\c -> take 4 c yOut) (\yIn -> inFoldable [1..10] yIn) +([1,2,3,4],()) + -} +take :: (ea :> es, eb :> es) => Int -> Consume a ea -> Stream a eb -> Eff es () +take count source sink = loop count + where + loop c | c <= 0 = pure () + loop c = do + await source >>= yield sink + loop (c - 1) + +{- +>>> import Bluefin.Consume +>>> import Bluefin.Eff +>>> import Bluefin.Stream + +>>> runPureEff $ yieldToList $ \yOut -> consumeStream (\c -> take 6 c yOut) (\yIn -> cycle [1..3] yIn) +([1,2,3,1,2,3],()) +-} + +cycle :: (Foldable f, ea :> es) => f a -> Stream a ea -> Eff es () +cycle f y = do + forever (inFoldable f y) + +provideStream :: + forall a (es :: Effects) r. + (forall (e :: Effects). Stream a e -> Eff (e :& es) r) -> + (forall (e :: Effects). Consume a e -> Eff (e :& es) r) -> + Eff es r +provideStream s c = consumeStream c s + +type TokenMap = Map Int64 UUID +type DCTokenMap = Map Text TokenMap +type DCHostCount = Map Text (Int, Map Text Int) +type SimpleReplicaMap = Map Int64 [UUID] +type NetworkTopologyReplicaMap = Map Text (Map Int64 [UUID]) + +data ReplicationStrategy = SimpleStrategy !Int | NetworkTopologyStrategy !(Map Text Int) + deriving (Eq, Ord, Show) + +data ReplicaJobKey = RebuildReplicaJob + deriving (Eq, Ord, Show) + +eitherToMaybe :: Either l r -> Maybe r +eitherToMaybe = either (const Nothing) Just + +parseReplicationStrategy :: Cql.Map Text Text -> Maybe ReplicationStrategy +parseReplicationStrategy (Cql.Map pairs) = do + let textMap = Map.fromList pairs + textClass <- Map.lookup "class" textMap + case textClass of + "org.apache.cassandra.locator.SimpleStrategy" -> parseSimple textMap + "org.apache.cassandra.locator.NetworkTopologyStrategy" -> parseTopology textMap + _ -> Nothing + where + maybeParseInt = eitherToMaybe . AP.parseOnly AP.decimal + parseSimple textMap = do + rfText <- Map.lookup "replication_factor" textMap + rfInt <- maybeParseInt rfText + pure $ SimpleStrategy rfInt + parseTopology textMap = do + let textMap' = Map.delete "class" textMap + NetworkTopologyStrategy <$> mapM maybeParseInt textMap' + +buildMasterReplicaMaps :: + TokenMap -> + DCTokenMap -> + DCHostCount -> + Map UUID Host -> + (SimpleReplicaMap, NetworkTopologyReplicaMap) +-- We can build a single replica map per DC, and one more for simple. +-- Because we produce the replica lists in order, whatever the replication +-- factor is, say , then we can just take the first n entries in the largish +-- map. +buildMasterReplicaMaps tokenMap dcTokenMap dcHostCount uuidMap = (masterSimpleMap, masterDCMap) + where + totalNodes :: Int + totalNodes = getSum $ foldMap' (Sum . fst) dcHostCount + masterSimpleMap :: SimpleReplicaMap + masterSimpleMap = simpleReplicaMap totalNodes tokenMap + masterDCMap :: NetworkTopologyReplicaMap + masterDCMap = Map.mapWithKey buildMapForDC dcHostCount + buildMapForDC dc (hostCount, rackSet) = fromMaybe Map.empty $ do + localTokenMap <- Map.lookup dc dcTokenMap + case Map.size rackSet of + x | x <= 0 -> pure Map.empty + -- When all the hosts are in their own rack, or all in the same rack, we + -- can ignore it. + x | x == 1 || x == hostCount -> pure $ simpleReplicaMap hostCount localTokenMap + numRacks -> pure $ dcReplicaMap hostCount numRacks localTokenMap uuidMap + +{- Shared helper functions for simpleReplicaMap and dcReplicaMap -} + +data Done = Done | NotDone +data Presence = Missing | Present + +{- | Use Map.alterF to update a map and store some output in a state cell. + The map is taken from a state cell and the result is placed back there. The + computed value is returned as the result of the entire computation. +-} +alterMapReturning :: + (e1 :> es, Ord k) => + State (Map k a) e1 -> + b -> + (forall e2. State b e2 -> Maybe a -> Eff (e2 :& es) (Maybe a)) -> + k -> + Eff es b +alterMapReturning curMapS initB f k = do + curMap <- get curMapS + (curMap', newB) <- runState initB $ \bState -> do + Map.alterF (f bState) k curMap + put curMapS curMap' + pure newB + +{- | Add a host to a hostMap state cell. + Takes a token offset, a Host, and a state cell for the hostMap. The token + offset is added to the set of active tokens for the Host. + Returns Present if the host was already in the map, or Missing if it + wasn't. +-} +addToHostMap :: + (e :> es) => Word64 -> UUID -> State (Map UUID (Set.Set Word64)) e -> Eff es Presence +addToHostMap tokenOffset host hostMapS = do + hostMap <- get hostMapS + let (mExists, hostMap') = Map.insertLookupWithKey addToken host (Set.singleton tokenOffset) hostMap + addToken _key = (<>) + put hostMapS hostMap' + pure $ maybe Missing (const Present) mExists + +{- | Remove a host from a hostMap in a state cell. + Takes a token offset, the host to be updated, and the token offset to be + removed. + Returns the smallest token offset in the Set.Set of token offsets for the + given host, after removing the supplied token offset. If a new minimum does + not exist, returns Nothing. +-} +dropFromHostMap :: + (e :> es) => Word64 -> UUID -> State (Map UUID (Set.Set Word64)) e -> Eff es (Maybe Word64) +dropFromHostMap tokenOffset host hostMapS = do + let dropToken _ Nothing = pure Nothing + dropToken hostNewMinS (Just windowSet) = do + let windowSet' = Set.delete tokenOffset windowSet + if null windowSet' + then pure Nothing + else do + put hostNewMinS $ Just (Set.findMin windowSet') + pure (Just windowSet') + alterMapReturning hostMapS Nothing dropToken host + +simpleReplicaMap :: Int -> TokenMap -> Map Int64 [UUID] +simpleReplicaMap 1 tokenMap = Map.map pure tokenMap +simpleReplicaMap n _ | n <= 0 = Map.empty +simpleReplicaMap n tokenMap = + Map.fromArgSet $ Set.fromList argList + where + argList :: [Arg Int64 [UUID]] + argList = + fst $ runPureEff $ yieldToList $ \yOut -> + provideStream (cycle (Map.argSet tokenMap)) $ \c -> argPipe c yOut + argPipe :: + forall e1 e2 eOuter. + (e1 :> eOuter, e2 :> eOuter) => + Consume (Arg Int64 UUID) e1 -> + Stream (Arg Int64 [UUID]) e2 -> + Eff eOuter () + argPipe c y = do + a@(Arg firstToken firstHost) <- await c + withStateSource $ \source -> do + relativeTokenS <- newState source (fromIntegral firstToken) + curArgS <- newState source a + argQueueS <- newState source Seq.empty + hostMapS <- newState source $ Map.singleton firstHost (Set.singleton 0) + smallestMapS <- newState source $ Map.singleton 0 firstHost + tokenLoop firstToken relativeTokenS curArgS argQueueS hostMapS smallestMapS + where + tokenLoop :: + forall e3 e4 e5 e6 e7 eLoop. + ( e1 :> eLoop + , e2 :> eLoop + , e3 :> eLoop + , e4 :> eLoop + , e5 :> eLoop + , e6 :> eLoop + , e7 :> eLoop + ) => + Int64 -> + State Word64 e3 -> + State (Arg Int64 UUID) e4 -> + State (Seq (Arg Int64 UUID)) e5 -> + State (Map UUID (Set.Set Word64)) e6 -> + State (Map Word64 UUID) e7 -> + Eff eLoop () + tokenLoop firstToken relativeTokenS curArgS argQueueS hostMapS smallestMapS = readArg + where + readArg :: Eff eLoop () + readArg = do + a@(Arg token _host) <- await c + relToken <- get relativeTokenS + {- Deal with wraparound. See similar comment in dcReplicaMap -} + when (fromIntegral token == relToken) $ do + (Arg curToken _host) <- get curArgS + let deltaShift = fromIntegral curToken - relToken + shrinkTokenDeltas deltaShift + put relativeTokenS $ fromIntegral curToken + modify argQueueS (|> a) + addToMaps a >>= yieldIfDone + + yieldIfDone :: Done -> Eff eLoop () + yieldIfDone Done = doYield >> dropHead + yieldIfDone NotDone = readArg + + doYield :: Eff eLoop () + doYield = do + (Arg curToken _host) <- get curArgS + smallestMap <- get smallestMapS + yield y (Arg curToken (Map.elems smallestMap)) + + checkSize :: Eff eLoop Done + checkSize = do + smallestMap <- get smallestMapS + if Map.size smallestMap == n + then pure Done + else pure NotDone + + dropHead :: Eff eLoop () + dropHead = do + aQueue <- get argQueueS + case aQueue of + Empty -> + error + "With replication factor of 2 or greater, arg queue \ + \should never be empty when dropHead is called." + a@(Arg nextToken _host) :<| aQueue' -> do + when (nextToken /= firstToken) $ do + stillDone <- deleteFromMaps + put curArgS a + put argQueueS aQueue' + case stillDone of + Done -> doYield >> dropHead + NotDone -> readArg + + addToMaps :: Arg Int64 UUID -> Eff eLoop Done + addToMaps (Arg token host) = do + relToken <- get relativeTokenS + let tokenDelta = fromIntegral token - relToken + wasPresent <- addToHostMap tokenDelta host hostMapS + case wasPresent of + Present -> pure NotDone + Missing -> do + modify smallestMapS (Map.insert tokenDelta host) + checkSize + + deleteFromMaps :: Eff eLoop Done + deleteFromMaps = do + (Arg curToken curHost) <- get curArgS + relToken <- get relativeTokenS + let curOffset = fromIntegral curToken - relToken + hostNewMin <- dropFromHostMap curOffset curHost hostMapS + dropFromSmallestMap curOffset + forM_ hostNewMin $ \hostNextBest -> do + modify smallestMapS (Map.insert hostNextBest curHost) + -- If we still have enough hosts, we can just yield them in the correct order. + checkSize + + dropFromSmallestMap :: Word64 -> Eff eLoop () + dropFromSmallestMap tokenOffset = modify smallestMapS (Map.delete tokenOffset) + + shrinkTokenDeltas :: Word64 -> Eff eLoop () + shrinkTokenDeltas deltaShift = do + modify smallestMapS (Map.mapKeysMonotonic (subtract deltaShift)) + modify hostMapS (Map.map (Set.mapMonotonic (subtract deltaShift))) + modify relativeTokenS (+ deltaShift) + +{- | Build a rack-aware ReplicaMap given the TokenMap for a single data + center +-} +dcReplicaMap :: + -- | The number of replicas in this data center + Int -> + -- | The number of unique racks in this data center + Int -> + -- | The TokenMap for this DC + TokenMap -> + -- | A map from UUID's to full host data + Map UUID Host -> + -- | The map from tokens to replicas. The order is one machine from each rack, + -- followed by any duplicates. + Map Int64 [UUID] +dcReplicaMap 1 _ tokenMap _ = Map.map pure tokenMap +dcReplicaMap n _ _ _ | n <= 0 = Map.empty +dcReplicaMap n rackCount tokenMap uuidMap = + {- The overall structure of this function is as follows: + - 1. Use bluefin streaming to look at a narrow slice of tokens. + - The token being examined is not part of the queue. + - Keep a queue of the other tokens currently being examined. + - Example diagram: + - 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 + - ^ ^---------------^ + - 2 is the current token (paired with a Host) + - 3-7 are the argQueue + - 2. Feed the source end of the pipeline with the *looped* entries from + - tokenMap. It is important that they are looped, so that we can get the + - correct replicas for the tokens at the high end of the map. This is + - why it is called a "ring" in Cassandra. The keys loop around. + - 3. Track values by 64-bit unsigned offset instead of as Int64's. Whenever + - the offsets would wrap around, we adjust the offset base and all the + - outstanding offsets so that the smallest offset is 0. + - 4. Use a group of mutually recursive functions to modify a handful of + - state variables: + - + - * curArgS: The current (Arg Int64 UUID), 2 in the diagram. + - * argQueueS: The queue of (Arg Int64 UUID)s currently in view, other + - than curArgS + - * hostMapS: A map from UUID's to the Set of currently in view token + - offsets for that host. We take advantage of the fact that the Set can + - be used as a priority queue for that host. + - * rackMapS, smallestPrimaryMapS, and smallestSecondaryMapS all use a + - Map of the following type: Map Word64 UUID. Similar to the Set + - mentioned above, this is a priority queue. These maps should have + - each host in them at most once. In every case, the key in these maps + - should be the smallest token offset for that host in hostMapS. + - * rackMapS: A map from rack name to a priority queue of hosts for this + - rack. The host with the smallest key is considered the primary for + - this rack. Any additional hosts are considired secondaries for this + - rack. + - * smallestPrimaryMapS: A priority queue of the hosts that are primaries + - for their rack. This should only ever have (min primaryCount n) + - entries at most. + - * smallestSecondaryMapS: A priority queue of the hosts that are not + - primaries for their rack. Every host in hostMapS should have an + - entry in exactly one of smallestPrimaryMapS or smallestSecondaryMapS + - 5. Use a handful of update-and-return modifiers for the maps above. When + - we remove the current token from hostMapS, we return the next lowest + - token, if any. When we remove the current token from rackMapS, we + - return the next lowest (token offset, host) pair for that rack, if + - any. + - 6. There are a handful of simple logic rules to maintain the invariants + - above: + - + - * When a (token, host) pair is added to a host that already + - existed in hostMapS, we don't update anything else. All of the other + - maps are keyed on the lowest value we have seen from a host so far. + - * When a (token, host) pair is a new (to the window) host, but the + - rack already had other hosts in it, that (token, host) pair gets + - added to smallestSecondaryMapS. + - * When a (token, host) pair is a new host, and the rack is new as + - well, add the (token, host) pair to smallestPrimaryMapS. + - * When a (token, host) pair is removed, remove the pair from + - smallestPrimaryMapS. It is an invariant that the current (token, + - host) pair should be the primary for the host's rack. + - * When a (token, host) pair is removed from its rack, and the rack is + - now empty, do nothing else. + - * When a (token, host) pair is removed from its rack and the same host + - is still the primary for the rack, with a larger token, add the new + - (token', host) pair to smallestPrimaryMapS + - * When a (token, host) pair is removed from its rack and a new host is + - now the primary for the rack, remove that new host from + - smallestSecondaryMapS. Add the new host to smallestPrimaryMapS. If + - we have a next smallest token for the removed host, add it to + - smallestSecondaryMapS. + - 7. Proceed, like an inch worm, in two phases: When we don't have enough + - data to list the replicas for the current token, read more data, + - adding it to arqQueueS and the appropriate maps. + - When we do have enough data, yield the replicas for this token, then + - drop the current (token, host) pair and update the maps + - appropriately. If we still have enough data, then repeat. If not, go + - back to reading. We grow on the read end just as far as necessary, + - and then we shrink on the write end as far as possible. + - /-\ + - ------- -/ \- ------- + - 1234567 45 67 4567890 + - 8. Stop when we encounter the first token as the head element of + - argQueueS. Because the queue contains the elements *after* the + - current token, when we encounter the first token as the head of the + - list, it is because we have completed wrapping around. + -} + Map.fromArgSet $ Set.fromList argList + where + primaryCount = min n rackCount + secondaryCount = max 0 (n - rackCount) + argList :: [Arg Int64 [UUID]] + argList = + fst $ runPureEff $ yieldToList $ \yOut -> + provideStream (cycle (Map.argSet tokenMap)) $ \c -> argPipe c yOut + argPipe :: + forall e1 e2 eOuter. + (e1 :> eOuter, e2 :> eOuter) => + Consume (Arg Int64 UUID) e1 -> + Stream (Arg Int64 [UUID]) e2 -> + Eff eOuter () + argPipe c y = do + a@(Arg firstToken firstUUID) <- await c + -- TODO: Add actual exception and catch it when building these + forM_ (Map.lookup firstUUID uuidMap) $ \firstHost -> do + withStateSource $ \source -> do + relativeTokenS <- newState source (fromIntegral firstToken) + curArgS <- newState source a + argQueueS <- newState source Seq.empty + hostMapS <- newState source $ Map.singleton (firstHost ^. hostId) (Set.singleton 0) + rackMapS <- newState source $ Map.singleton (firstHost ^. rack) (Map.singleton 0 firstUUID) + smallestPrimaryMapS <- newState source $ Map.singleton 0 firstUUID + smallestSecondaryMapS <- newState source Map.empty + tokenLoop firstToken relativeTokenS curArgS argQueueS hostMapS rackMapS smallestPrimaryMapS smallestSecondaryMapS + where + tokenLoop :: + forall e3 e4 e5 e6 e7 e8 e9 eLoop. + ( e1 :> eLoop + , e2 :> eLoop + , e3 :> eLoop + , e4 :> eLoop + , e5 :> eLoop + , e6 :> eLoop + , e7 :> eLoop + , e8 :> eLoop + , e9 :> eLoop + ) => + Int64 -> + State Word64 e3 -> + State (Arg Int64 UUID) e4 -> + State (Seq (Arg Int64 UUID)) e5 -> + State (Map UUID (Set.Set Word64)) e6 -> + State (Map Text (Map Word64 UUID)) e7 -> + State (Map Word64 UUID) e8 -> + State (Map Word64 UUID) e9 -> + Eff eLoop () + tokenLoop firstToken relativeTokenS curArgS argQueueS hostMapS rackMapS smallestPrimaryMapS smallestSecondaryMapS = + readArg + where + readArg :: Eff eLoop () + readArg = do + a@(Arg token _host) <- await c + relToken <- get relativeTokenS + {- To deal with wraparound, all the tokens are shifted so that the 0 + token is some token that we have seen before. When we read that token + in at the back end, we re-adjust so that the token at the beginning + is now the 0 token. In non-pathological cases, this should occur + once. -} + when (fromIntegral token == relToken) $ do + (Arg curToken _host) <- get curArgS + -- All the offsets are (t - relToken), we want them to be + -- (t - curToken). Subtracting (curToken - relToken) is the necessary + -- correction. + let deltaShift = fromIntegral curToken - relToken + shrinkTokenDeltas deltaShift + -- See the above note. To add further tokens, they need to be relative + -- to curToken. + put relativeTokenS $ fromIntegral curToken + modify argQueueS (|> a) + addToMaps a >>= yieldIfDone + + yieldIfDone :: Done -> Eff eLoop () + yieldIfDone Done = doYield >> dropHead + yieldIfDone NotDone = readArg + + doYield :: Eff eLoop () + doYield = do + (Arg curToken _host) <- get curArgS + smallestPrimaryMap <- get smallestPrimaryMapS + smallestSecondaryMap <- get smallestSecondaryMapS + let secondaries = Map.elems (Map.take secondaryCount smallestSecondaryMap) + yield y (Arg curToken (Map.elems smallestPrimaryMap ++ secondaries)) + + dropHead :: Eff eLoop () + dropHead = + do + aQueue <- get argQueueS + case aQueue of + Empty -> + error + "With replication factor of 2 or greater, arg queue \ + \should never be empty when dropHead is called." + a@(Arg nextToken _host) :<| aQueue' -> do + when (nextToken /= firstToken) $ do + stillDone <- deleteFromMaps + put curArgS a + put argQueueS aQueue' + yieldIfDone stillDone + + addToMaps :: Arg Int64 UUID -> Eff eLoop Done + addToMaps (Arg token host) = do + hostMap <- get hostMapS + relToken <- get relativeTokenS + let tokenDelta = fromIntegral token - relToken + (mExists, hostMap') = Map.insertLookupWithKey addToken host (Set.singleton tokenDelta) hostMap + addToken _key = (<>) + put hostMapS hostMap' + case mExists of + Just _ -> pure NotDone + Nothing -> handleNewHost tokenDelta host + + handleNewHost :: Word64 -> UUID -> Eff eLoop Done + handleNewHost tokenDelta host = do + newRack <- addToRackMap tokenDelta host + case newRack of + Missing -> modify smallestPrimaryMapS (Map.insert tokenDelta host) + Present -> modify smallestSecondaryMapS (Map.insert tokenDelta host) + checkSizes + + addToRackMap :: Word64 -> UUID -> Eff eLoop Presence + addToRackMap tokenDelta hostUUID = do + let addToken _prevRackS Nothing = + pure . Just $ Map.singleton tokenDelta hostUUID + addToken prevRackS (Just rackSmallest) = do + put prevRackS Present + pure . Just $ Map.insert tokenDelta hostUUID rackSmallest + mHost = Map.lookup hostUUID uuidMap + case mHost of + -- TODO: Add actual exception and catch it when building these + Nothing -> error $ "Couldn't find host: " ++ show hostUUID + Just host -> + alterMapReturning rackMapS Missing addToken (host ^. rack) + + checkSizes :: Eff eLoop Done + checkSizes = do + smallestPrimaryMap <- get smallestPrimaryMapS + smallestSecondaryMap <- get smallestSecondaryMapS + if Map.size smallestPrimaryMap >= primaryCount && Map.size smallestSecondaryMap >= secondaryCount + then pure Done + else pure NotDone + + deleteFromMaps :: Eff eLoop Done + deleteFromMaps = do + (Arg curToken curHost) <- get curArgS + relToken <- get relativeTokenS + let curOffset = fromIntegral curToken - relToken + hostNewMin <- dropFromHostMap curOffset curHost hostMapS + rackMapNewMin <- dropFromRackMap curOffset curHost hostNewMin + -- Delete the token from the smallestPrimaryMap + -- It has to be a primary now that it's at the front. + dropFromPrimaryMap curOffset + forM_ rackMapNewMin $ \(rackMinToken, rackMinHost) -> do + -- Whenever the rack has a new minimum, that (token, host) pair is + -- always a primary. + modify smallestPrimaryMapS (Map.insert rackMinToken rackMinHost) + -- If it was the same host that was the primary previously, we don't + -- have to adjust anything else. + when (rackMinHost /= curHost) $ do + -- Otherwise, remove the new host from the secondaries + modify smallestSecondaryMapS (Map.delete rackMinToken) + forM_ hostNewMin $ \hostNextBest -> do + -- And put the original host in the secondaries, if possible. + modify smallestSecondaryMapS (Map.insert hostNextBest curHost) + -- See if we still have enough hosts + checkSizes + + dropFromRackMap :: Word64 -> UUID -> Maybe Word64 -> Eff eLoop (Maybe (Word64, UUID)) + dropFromRackMap tokenDelta hostUUID hostNewMin = do + let dropToken _ Nothing = pure Nothing + dropToken rackNewMinS (Just rackSmallest) = do + let rackSmallest' = Map.delete tokenDelta rackSmallest + rackSmallest'' = case hostNewMin of + Nothing -> rackSmallest' + Just tokenDelta' -> Map.insert tokenDelta' hostUUID rackSmallest' + if null rackSmallest'' + then pure Nothing + else do + put rackNewMinS $ Just (Map.findMin rackSmallest'') + pure . Just $ rackSmallest'' + mHost = Map.lookup hostUUID uuidMap + case mHost of + -- TODO: Add actual exception and catch it when building these + Nothing -> error $ "Couldn't find host: " ++ show hostUUID + Just host -> + alterMapReturning rackMapS Nothing dropToken (host ^. rack) + + dropFromPrimaryMap :: Word64 -> Eff eLoop () + dropFromPrimaryMap tokenDelta = modify smallestPrimaryMapS (Map.delete tokenDelta) + + shrinkTokenDeltas :: Word64 -> Eff eLoop () + shrinkTokenDeltas deltaShift = do + let modifyKeys = Map.mapKeysMonotonic (subtract deltaShift) + modify smallestPrimaryMapS modifyKeys + modify smallestSecondaryMapS modifyKeys + modify rackMapS (Map.map modifyKeys) + modify hostMapS (Map.map (Set.mapMonotonic (subtract deltaShift))) diff --git a/src/lib/Database/CQL/IO/Settings.hs b/src/lib/Database/CQL/IO/Settings.hs --- a/src/lib/Database/CQL/IO/Settings.hs +++ b/src/lib/Database/CQL/IO/Settings.hs @@ -24,7 +24,6 @@ import Database.CQL.IO.Exception import Database.CQL.IO.Log import Database.CQL.IO.Pool as P -import Database.CQL.IO.Timeouts (Milliseconds (..)) import OpenSSL.Session (SSLContext, SomeSSLException) import qualified Data.HashMap.Strict as HashMap @@ -34,11 +33,13 @@ , _connSettings :: ConnectionSettings , _retrySettings :: RetrySettings , _logger :: Logger - , _protoVersion :: Version + , _minProtoVer :: Version + , _maxProtoVer :: Version , _portnumber :: PortNumber , _contacts :: NonEmpty String , _policyMaker :: IO Policy , _prepStrategy :: PrepareStrategy + , _maxControlConnections :: Int } -- | Strategy for the execution of 'PrepQuery's. @@ -80,8 +81,12 @@ -- -- * The load-balancing policy is 'random'. -- --- * The binary protocol version is 3. +-- * The minimum binary protocol version is 3. -- +-- * The preferred binary protocol version is 4. The initial connection is made +-- using V3. If the server supports V4, that connection is closed and a new +-- one opened using V4 +-- -- * The connection idle timeout is 60s. -- -- * The connection pool uses 4 stripes to mitigate thread contention. @@ -104,23 +109,30 @@ -- * Query preparation is done lazily. See 'PrepareStrategy'. defSettings :: Settings defSettings = Settings - P.defSettings - C.defSettings - defRetrySettings - nullLogger - V3 - 9042 - ("localhost" :| []) - random - LazyPrepare + { _poolSettings = P.defSettings + , _connSettings = C.defSettings + , _retrySettings = defRetrySettings + , _logger = nullLogger + , _minProtoVer = V3 + , _maxProtoVer = V4 + , _portnumber = 9042 + , _contacts = "localhost" :| [] + , _policyMaker = random + , _prepStrategy = LazyPrepare + , _maxControlConnections = 5 + } ----------------------------------------------------------------------------- -- Settings --- | Set the binary protocol version to use. -setProtocolVersion :: Version -> Settings -> Settings -setProtocolVersion v = set protoVersion v +-- | Set the minimum binary protocol version to use. +setMinProtocolVersion :: Version -> Settings -> Settings +setMinProtocolVersion = set minProtoVer +-- | Set the maximum binary protocol version to use. +setMaxProtocolVersion :: Version -> Settings -> Settings +setMaxProtocolVersion = set maxProtoVer + -- | Set the initial contact points (hosts) from which node discovery will -- start. setContacts :: String -> [String] -> Settings -> Settings @@ -145,6 +157,10 @@ -- | Set the 'Logger' to use for processing log messages emitted by the client. setLogger :: Logger -> Settings -> Settings setLogger v = set logger v + +-- | Set the max number of control connections +setMaxControlConnections :: Int -> Settings -> Settings +setMaxControlConnections = set maxControlConnections ----------------------------------------------------------------------------- -- Pool Settings diff --git a/src/test/Main.hs b/src/test/Main.hs --- a/src/test/Main.hs +++ b/src/test/Main.hs @@ -3,12 +3,14 @@ import Test.Tasty import Test.Database.CQL.IO import Test.Database.CQL.IO.Jobs +import Test.Database.CQL.IO.Replication main :: IO () main = do tree <- sequence [ Test.Database.CQL.IO.tests , pure Test.Database.CQL.IO.Jobs.tests + , pure Test.Database.CQL.IO.Replication.tests ] defaultMain $ testGroup "cql-io" tree diff --git a/src/test/Test/Database/CQL/IO.hs b/src/test/Test/Database/CQL/IO.hs --- a/src/test/Test/Database/CQL/IO.hs +++ b/src/test/Test/Database/CQL/IO.hs @@ -35,21 +35,22 @@ initSchema h fmap (testGroup "Database.CQL.IO") $ forM versions (\v -> do - c <- Client.init (settings h v) + c <- Client.init (settings h v v) return $ testGroup (show v) (cqlTests c)) versions :: [Version] versions = [V3, V4] -settings :: TestHost -> Version -> Settings -settings h v = setContacts h [] - . setProtocolVersion v +settings :: TestHost -> Version -> Version -> Settings +settings h vMin vMax = setContacts h [] + . setMinProtocolVersion vMin + . setMaxProtocolVersion vMax . setLogger (stdoutLogger LogInfo) $ defSettings initSchema :: TestHost -> IO () initSchema h = do - c <- Client.init (settings h V4) + c <- Client.init (settings h V3 V4) runClient c $ do dropKeyspace createKeyspace diff --git a/src/test/Test/Database/CQL/IO/Replication.hs b/src/test/Test/Database/CQL/IO/Replication.hs new file mode 100644 --- /dev/null +++ b/src/test/Test/Database/CQL/IO/Replication.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ImportQualifiedPost #-} +{-# LANGUAGE NoFieldSelectors #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Test.Database.CQL.IO.Replication (tests) where + +import Control.Monad (replicateM) +import Data.IP (IP (..), toIPv4, toIPv6) +import Data.Int (Int64) +import Data.List (sort) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map (fromList, unionWith, singleton, unionsWith, lookup, lookupGE, findMin, foldMapWithKey, size) +import Data.Maybe (mapMaybe) +import Data.Set (Set) +import Data.Set qualified as Set (fromList, lookupGE, findMin, size, unions) +import Data.Text (Text) +import Data.Traversable (forM) +import Data.UUID (UUID) +import Data.Word (Word64) +import Database.CQL.IO.Client (buildTokenMap) +import Database.CQL.IO.Cluster.Host (Host (..), ip2inet) +import Database.CQL.IO.Replication (buildMasterReplicaMaps) +import Formatting (sformat, (%), int) +import Test.QuickCheck (Arbitrary (..), chooseInt, chooseAny, elements, Property, (===), Large (..), Every (..), (.&&.), suchThat) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.QuickCheck (testProperty) + +data HostReplicationProblem = HostReplicationProblem + { simpleTokenMap :: Map Int64 UUID + , dcTokenMap :: Map Text (Map Int64 UUID) + , dcHostCount :: Map Text (Int, Map Text Int) + , fullHostMap :: Map UUID Host + , ownedTokenMap :: Map UUID (Set Int64) + } deriving (Eq, Ord, Show) + +instance Arbitrary IP where + arbitrary = do + version <- chooseInt (0, 1) + case version of + 0 -> do + suffix <- replicateM 3 (chooseInt (0, 255)) + pure . IPv4 . toIPv4 $ 10 : suffix + _ -> do + suffix <- replicateM 6 (chooseInt (0, 0xffff)) + pure . IPv6 . toIPv6 $ 0x2001 : 0x0db8 : suffix + +instance Arbitrary HostReplicationProblem where + arbitrary = do + dcCount <- chooseInt (1, 5) + rackCount <- chooseInt (1, 5) + perRackCount <- chooseInt (1, 10) + tokenCount <- elements [8, 16, 32, 64] + fullHosts <- forM [1..dcCount] (\dcNum -> do + let dcName = sformat ("DC" % int) dcNum + dcHosts <- forM [1..rackCount] $ \rackNum -> do + let rackName = sformat ("RAC" % int) rackNum + replicateM perRackCount (do + hostUUID <- chooseAny + hostIP <- arbitrary + tokens <- replicateM @_ @Int64 tokenCount chooseAny + let hostAddr = ip2inet 9042 hostIP + broadcastAddr = ip2inet 7000 hostIP + host = Host + { _hostAddr = hostAddr + , _broadcastAddr = broadcastAddr + , _hostId = hostUUID + , _dataCentre = dcName + , _rack = rackName + } + pure (host, Set.fromList tokens)) + pure $ concat dcHosts) + `suchThat` \possibleHostList -> + let expectedSize = dcCount * rackCount * perRackCount * tokenCount + allTokens = Set.unions $ map (Set.unions . map snd) possibleHostList + in Set.size allTokens == expectedSize + let fullHostsFlat = concat fullHosts + ownedTokenMap = Map.fromList $ [ (hostId, tokens) + | (Host {_hostId = hostId}, tokens) <- fullHostsFlat + ] + lookupMap = Map.fromList $ [ (hostId, host) + | (host@(Host {_hostId = hostId}), _tokens) <- fullHostsFlat + ] + (simpleMap, dcMap) = buildTokenMap ownedTokenMap lookupMap + mergeCounts :: (Int, Map Text Int) -> (Int, Map Text Int) -> (Int, Map Text Int) + mergeCounts (!count1, !racks1) (!count2, !racks2) = + let !count3 = count1 + count2 + !racks3 = Map.unionWith (+) racks1 racks2 + in (count3, racks3) + countMap = Map.unionsWith mergeCounts $ [ + Map.singleton dc (1, Map.singleton rack 1) + | (Host {_dataCentre = dc, _rack = rack}, _tokens) <- fullHostsFlat + ] + pure $ HostReplicationProblem + { simpleTokenMap = simpleMap + , dcTokenMap = dcMap + , dcHostCount = countMap + , fullHostMap = lookupMap + , ownedTokenMap = ownedTokenMap + } + +tokenOffsetDistance :: Int64 -> Set Int64 -> Word64 +tokenOffsetDistance queryVal tokenSet = + let qvW64 = fromIntegral queryVal + tokenW64 = case Set.lookupGE queryVal tokenSet of + Nothing -> fromIntegral $ Set.findMin tokenSet + Just owningToken -> fromIntegral owningToken + in tokenW64 - qvW64 + +distanceForHost :: Int64 -> Map UUID (Set Int64) -> UUID -> Word64 +distanceForHost queryVal ownedTokenMap hostId = + case Map.lookup hostId ownedTokenMap of + Nothing -> error ("distance requested for unknown host: " ++ show hostId) + Just tokenSet -> tokenOffsetDistance queryVal tokenSet + +correctlySorted :: Large Int64 -> HostReplicationProblem -> Property +correctlySorted (Large queryVal) hrp = + let (simpleReplicaMap, nonSimpleReplicaMap) = buildMasterReplicaMaps + hrp.simpleTokenMap + hrp.dcTokenMap + hrp.dcHostCount + hrp.fullHostMap + replicas = case Map.lookupGE queryVal simpleReplicaMap of + Just sortedReplicas -> snd sortedReplicas + Nothing -> snd $ Map.findMin simpleReplicaMap + replicaDistances = map (distanceForHost queryVal hrp.ownedTokenMap) replicas + in isSorted replicaDistances + .&&. allUnique replicas + .&&. Map.foldMapWithKey (\k v -> Every (dcSortedCorrectly queryVal hrp k v)) nonSimpleReplicaMap + +isSorted :: (Ord a, Eq a, Show a) => [a] -> Property +isSorted vals = sort vals === vals + +allUnique :: (Ord a, Eq a, Show a) => [a] -> Property +allUnique vals = length vals === Set.size (Set.fromList vals) + +hostRacks :: Map UUID Host -> [UUID] -> [Text] +hostRacks hostMap = + map (\x -> x._rack) . mapMaybe lookupHost + where lookupHost hostId = Map.lookup hostId hostMap + +dcSortedCorrectly :: Int64 -> HostReplicationProblem -> Text -> Map Int64 [UUID] -> Property +dcSortedCorrectly queryVal hrp dc localReplicaMap = + let replicas = case Map.lookupGE queryVal localReplicaMap of + Just sortedReplicas -> snd sortedReplicas + Nothing -> snd $ Map.findMin localReplicaMap + rackCount = case Map.lookup dc hrp.dcHostCount of + Nothing -> error ("could not find dc in dcHostCount map: " ++ show dc) + Just (_total, countByRack) -> Map.size countByRack + (primaryReplicas, secondaryReplicas) = splitAt rackCount replicas + rackSingleton :: UUID -> Map Text [UUID] + rackSingleton hostId = case Map.lookup hostId hrp.fullHostMap of + Nothing -> error ("could not find host in fullHostMap by UUID: " ++ show hostId) + Just fullHost -> Map.singleton fullHost._rack [hostId] + replicasByRack :: Map Text [UUID] + replicasByRack = Map.unionsWith (flip (++)) $ map rackSingleton (reverse replicas) + sortedDistances :: [UUID] -> Property + sortedDistances = isSorted . map (distanceForHost queryVal hrp.ownedTokenMap) + in sortedDistances primaryReplicas + .&&. sortedDistances secondaryReplicas + .&&. allUnique replicas + .&&. allUnique (hostRacks hrp.fullHostMap primaryReplicas) + .&&. foldMap (Every . sortedDistances) replicasByRack + +tests :: TestTree +tests = testGroup "Replication" + [ testGroup "Ordering" + [ testProperty "Simple and DC based replicas are sorted correctly." correctlySorted + ] + ]