packages feed

cql-io 1.1.1 → 2.0.0

raw patch · 22 files changed

+3241/−951 lines, 22 filesdep +QuickCheckdep +attoparsecdep +bluefindep −cryptonitedep ~HsOpenSSLdep ~asyncdep ~auto-updatenew-uploader

Dependencies added: QuickCheck, attoparsec, bluefin, crypton, formatting, persist, random-shuffle, semver, tasty-quickcheck

Dependencies removed: cryptonite

Dependency ranges changed: HsOpenSSL, async, auto-update, base, bytestring, containers, cql, data-default-class, exceptions, hashable, iproute, lens, mtl, mwc-random, network, retry, semigroups, stm, text, time, transformers, unliftio-core, unordered-containers, uuid, vector

Files

AUTHORS view
@@ -1,6 +1,11 @@ Maintainers ----------- +- Kyle Butt <kyle@iteratee.net>++Prior Maintainers+-----------------+ - Toralf Wittner <tw@dtex.org> - Roman S. Borschel <roman@pkaboo.org>
CHANGELOG view
@@ -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`.
README.md view
@@ -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 =======
cql-io.cabal view
@@ -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 <tw@dtex.org>,-                      Roman S. Borschel <roman@pkaboo.org>-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 <kyle@iteratee.net>+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
src/api/Database/CQL/IO.hs view
@@ -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
src/lib/Database/CQL/IO/Batch.hs view
@@ -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 ()
src/lib/Database/CQL/IO/Client.hs view
@@ -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
src/lib/Database/CQL/IO/Cluster/Discovery.hs view
@@ -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 = ?"
src/lib/Database/CQL/IO/Cluster/Host.hs view
@@ -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
src/lib/Database/CQL/IO/Cluster/Policies.hs view
@@ -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
src/lib/Database/CQL/IO/Connection.hs view
@@ -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 ()
src/lib/Database/CQL/IO/Connection/Socket.hs view
@@ -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
src/lib/Database/CQL/IO/Exception.hs view
@@ -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
src/lib/Database/CQL/IO/Jobs.hs view
@@ -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 ())]
src/lib/Database/CQL/IO/Pool.hs view
@@ -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
src/lib/Database/CQL/IO/PrepQuery.hs view
@@ -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
src/lib/Database/CQL/IO/Protocol.hs view
@@ -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
+ src/lib/Database/CQL/IO/Replication.hs view
@@ -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 <n>, 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)))
src/lib/Database/CQL/IO/Settings.hs view
@@ -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
src/test/Main.hs view
@@ -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
src/test/Test/Database/CQL/IO.hs view
@@ -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
+ src/test/Test/Database/CQL/IO/Replication.hs view
@@ -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+    ]+  ]