packages feed

cql-io-2.0.0: src/lib/Database/CQL/IO/Replication.hs

{-# 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)))