cql-io-2.0.0: src/lib/Database/CQL/IO/Client.hs
-- 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 BangPatterns #-}
{-# 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