nuthatch/src/Nuthatch/Protocol.hs

-- | Short Postgres transactions for mutation authority; store I/O stays outside them.
module Nuthatch.Protocol
  ( AdmitOutcome (..)
  , AdmissionIdentity (..)
  , MutationFingerprint (..)
  , AdmissionRequest (..)
  , AdmissionLimits (..)
  , DispatchTiming (..)
  , PriorObject (..)
  , DesiredObject (..)
  , ExistingDisposition (..)
  , MutationRecord (..)
  , MutationPhase (..)
  , TerminalResult (..)
  , mutationState
  , externalState
  , encodeMutationPhase
  , decodeMutationPhase
  , mrState
  , mrExternal
  , mrUserId
  , mrKind
  , mrOperationId
  , mrObjectKey
  , mrDeclaredSize
  , mrContentSha256
  , mrPriorSize
  , mrPriorSha256
  , desiredObject
  , ObjectObservation (..)
  , admitMutation
  , markApplying
  , markExternalPending
  , markExternalSettled
  , markSettledFromIdle
  , finalizeSuccess
  , finalizeFailure
  , claimActiveBatch
  , lookupMutation
  , mutationEpochExists
  , lookupObjectMeta
  , sameFingerprint
  , existingDisposition
  , growthBytes
  , quotaAllows
  , terminalExpiryEligible
  , terminalExpired
  , recoveryEligible
  , desiredMatches
  , priorMatches
  , finalizeFromObservation
  ) where

import           Control.Exception (throwIO)
import           Data.Maybe (listToMaybe)
import           Data.Text (Text)
import qualified Data.Text as T
import           Database.PostgreSQL.Simple (Connection, Only (..), execute, query)
import qualified Database.PostgreSQL.Simple.FromField as PG
import qualified Database.PostgreSQL.Simple.FromRow as PGRow
import qualified Database.PostgreSQL.Simple.ToField as PGToField
import qualified Database.PostgreSQL.Simple.ToRow as PGToRow
import           Nuthatch.Types

data MutationRecord = MutationRecord
  { mrAdmissionIdentity :: AdmissionIdentity
  , mrPoolId            :: PoolId
  , mrRelativePath      :: RelativePath
  , mrFingerprint       :: MutationFingerprint
  , mrPriorObject       :: PriorObject
  , mrGrowth            :: Integer
  , mrPhase             :: MutationPhase
  }
  deriving stock (Eq, Show)

data MutationPhase
  = PhasePrepared
  | PhaseApplying ExternalState
  | PhaseTerminal TerminalResult
  deriving stock (Eq, Show)

data TerminalResult
  = TerminalSucceeded
  | TerminalFailed
  | TerminalRejected
  deriving stock (Eq, Show)

mutationState :: MutationPhase -> MutationState
mutationState = \case
  PhasePrepared   -> Prepared
  PhaseApplying _ -> Applying
  PhaseTerminal TerminalSucceeded -> Succeeded
  PhaseTerminal TerminalFailed    -> Failed
  PhaseTerminal TerminalRejected  -> Rejected

externalState :: MutationPhase -> ExternalState
externalState = \case
  PhaseApplying external -> external
  _                      -> ExtIdle

encodeMutationPhase :: MutationPhase -> (MutationState, ExternalState)
encodeMutationPhase phase = (mutationState phase, externalState phase)

decodeMutationPhase :: MutationState -> ExternalState -> Maybe MutationPhase
decodeMutationPhase state external =
  case (state, external) of
    (Prepared, ExtIdle)  -> Just PhasePrepared
    (Applying, value)    -> Just (PhaseApplying value)
    (Succeeded, ExtIdle) ->
      Just (PhaseTerminal TerminalSucceeded)
    (Failed, ExtIdle) ->
      Just (PhaseTerminal TerminalFailed)
    (Rejected, ExtIdle) ->
      Just (PhaseTerminal TerminalRejected)
    _                    -> Nothing

mrState :: MutationRecord -> MutationState
mrState = mutationState . mrPhase

mrExternal :: MutationRecord -> ExternalState
mrExternal = externalState . mrPhase

-- | Unavailable must not drive terminal finalize.
data ObjectObservation
  = ObjectPresent Integer (Maybe ContentSha256)
  | ObjectAbsent
  | ObjectUnavailable Text
  deriving stock (Eq, Show)

data AdmitOutcome
  = AdmitPrepared MutationRecord
  | AdmitReplay MutationResult MutationRecord
  | AdmitRejected Text MutationRecord
  | AdmitConflict
  deriving stock (Show)

data AdmissionIdentity = AdmissionIdentity
  { admissionUserId      :: UserId
  , admissionKind        :: MutationKind
  , admissionOperationId :: OperationId
  }
  deriving stock (Eq, Show)

data MutationFingerprint = MutationFingerprint
  { fingerprintKind          :: MutationKind
  , fingerprintObjectKey     :: ObjectKey
  , fingerprintDeclaredSize  :: Integer
  , fingerprintContentSha256 :: Maybe ContentSha256
  }
  deriving stock (Eq, Show)

data AdmissionRequest = AdmissionRequest
  { admissionIdentity     :: AdmissionIdentity
  , admissionPoolId       :: PoolId
  , admissionRelativePath :: RelativePath
  , admissionFingerprint  :: MutationFingerprint
  }
  deriving stock (Eq, Show)

data AdmissionLimits = AdmissionLimits
  { admissionMaxObjectBytes       :: Integer
  , admissionIdempotencyWindowSec :: Integer
  }
  deriving stock (Eq, Show)

data DispatchTiming = DispatchTiming
  { dispatchSettlementBoundSec :: Integer
  , dispatchWorkerGraceSec     :: Integer
  }
  deriving stock (Eq, Show)

data PriorObject = PriorObject
  { priorObjectSize   :: Integer
  , priorObjectSha256 :: Maybe ContentSha256
  }
  deriving stock (Eq, Show)

data DesiredObject
  = DesiredPresent Integer (Maybe ContentSha256)
  | DesiredAbsent
  deriving stock (Eq, Show)

mrUserId :: MutationRecord -> UserId
mrUserId = admissionUserId . mrAdmissionIdentity

mrKind :: MutationRecord -> MutationKind
mrKind = admissionKind . mrAdmissionIdentity

mrOperationId :: MutationRecord -> OperationId
mrOperationId = admissionOperationId . mrAdmissionIdentity

mrObjectKey :: MutationRecord -> ObjectKey
mrObjectKey = fingerprintObjectKey . mrFingerprint

mrDeclaredSize :: MutationRecord -> Integer
mrDeclaredSize = fingerprintDeclaredSize . mrFingerprint

mrContentSha256 :: MutationRecord -> Maybe ContentSha256
mrContentSha256 = fingerprintContentSha256 . mrFingerprint

mrPriorSize :: MutationRecord -> Integer
mrPriorSize = priorObjectSize . mrPriorObject

mrPriorSha256 :: MutationRecord -> Maybe ContentSha256
mrPriorSha256 = priorObjectSha256 . mrPriorObject

desiredObject :: MutationRecord -> DesiredObject
desiredObject record =
  case mrKind record of
    Put -> DesiredPresent
      (mrDeclaredSize record)
      (mrContentSha256 record)
    Delete -> DesiredAbsent

data ExistingDisposition
  = NewMutationEpoch
  | ExistingReplay MutationResult MutationRecord
  | ExistingConflict
  deriving stock (Show)

data MutationRow = MutationRow
  { rowUserId        :: Text
  , rowKind          :: DbMutationKind
  , rowOperationId   :: Text
  , rowPoolId        :: Text
  , rowObjectKey     :: Text
  , rowRelativePath  :: Text
  , rowDeclaredSize  :: Integer
  , rowContentSha256 :: Maybe Text
  , rowPriorSize     :: Integer
  , rowPriorSha256   :: Maybe Text
  , rowGrowth        :: Integer
  , rowState         :: DbMutationState
  , rowExternal      :: DbExternalState
  }

data SettlementRow =
  SettlementRow DbMutationState DbExternalState Integer Integer

newtype DbMutationKind = DbMutationKind MutationKind

newtype DbMutationState = DbMutationState MutationState

newtype DbExternalState = DbExternalState ExternalState

newtype MutationInsertRow = MutationInsertRow MutationRecord

instance PG.FromField DbMutationKind where
  fromField field value = do
    label <- PG.fromField field value
    maybe
      (PG.returnError PG.ConversionFailed field "invalid mutation kind")
      (pure . DbMutationKind)
      (parseKind label)

instance PG.FromField DbMutationState where
  fromField field value = do
    label <- PG.fromField field value
    maybe
      (PG.returnError PG.ConversionFailed field "invalid mutation state")
      (pure . DbMutationState)
      (parseState label)

instance PG.FromField DbExternalState where
  fromField field value = do
    label <- PG.fromField field value
    maybe
      (PG.returnError PG.ConversionFailed field "invalid external state")
      (pure . DbExternalState)
      (parseExternal label)

instance PGRow.FromRow MutationRow where
  fromRow =
    MutationRow
      <$> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field

instance PGRow.FromRow SettlementRow where
  fromRow =
    SettlementRow
      <$> PGRow.field
      <*> PGRow.field
      <*> PGRow.field
      <*> PGRow.field

instance PGToRow.ToRow MutationInsertRow where
  toRow (MutationInsertRow record) =
    let (state, external) = encodeMutationPhase (mrPhase record)
    in
      [ PGToField.toField (unUserId (mrUserId record))
      , PGToField.toField (kindLabel (mrKind record))
      , PGToField.toField (unOperationId (mrOperationId record))
      , PGToField.toField (unPoolId (mrPoolId record))
      , PGToField.toField (unObjectKey (mrObjectKey record))
      , PGToField.toField (unRelativePath (mrRelativePath record))
      , PGToField.toField (mrDeclaredSize record)
      , PGToField.toField
          (unContentSha256 <$> mrContentSha256 record)
      , PGToField.toField (mrPriorSize record)
      , PGToField.toField
          (unContentSha256 <$> mrPriorSha256 record)
      , PGToField.toField (mrGrowth record)
      , PGToField.toField (stateLabel state)
      , PGToField.toField (externalLabel external)
      ]

rowToRecord :: MutationRow -> Either Text MutationRecord
rowToRecord MutationRow{..} = do
  let DbMutationKind kind = rowKind
      DbMutationState state = rowState
      DbExternalState ext = rowExternal
  phase <- maybe (Left "invalid mutation phase") Right (decodeMutationPhase state ext)
  pure MutationRecord
    { mrAdmissionIdentity =
        AdmissionIdentity (UserId rowUserId) kind (OperationId rowOperationId)
    , mrPoolId = PoolId rowPoolId
    , mrRelativePath = RelativePath rowRelativePath
    , mrFingerprint =
        MutationFingerprint
          kind
          (ObjectKey rowObjectKey)
          rowDeclaredSize
          (ContentSha256 <$> rowContentSha256)
    , mrPriorObject =
        PriorObject rowPriorSize (ContentSha256 <$> rowPriorSha256)
    , mrGrowth = rowGrowth
    , mrPhase = phase
    }

decodeRows :: [MutationRow] -> IO [MutationRecord]
decodeRows rows =
  either (throwIO . userError . T.unpack) pure (traverse rowToRecord rows)

lookupMutation :: Connection -> UserId -> MutationKind -> OperationId -> IO (Maybe MutationRecord)
lookupMutation conn uid kind op = do
  rows <- query conn
    "SELECT user_id, kind, operation_id, pool_id, object_key, relative_path, \
    \declared_size, content_sha256, prior_size, prior_sha256, growth_bytes, \
    \state, external_state FROM mutations \
    \WHERE user_id = ? AND kind = ? AND operation_id = ?"
    (unUserId uid, kindLabel kind, unOperationId op)
  listToMaybe <$> decodeRows rows

mutationEpochExists
  :: Connection
  -> UserId
  -> MutationKind
  -> OperationId
  -> Integer
  -> IO Bool
mutationEpochExists conn uid kind op idempotencyWindow = do
  let identity =
        unUserId uid <> "\NUL" <> kindLabel kind <> "\NUL" <> unOperationId op
  _ <- query conn
    "SELECT pg_advisory_xact_lock(hashtextextended(?, 0))"
    (Only identity) :: IO [Only ()]
  rows <- query conn
    "SELECT 1 FROM mutations \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? \
    \AND (state IN ('prepared', 'applying') \
    \  OR external_state <> 'idle' \
    \  OR updated_at >= now() - (? * interval '1 second'))"
    (unUserId uid, kindLabel kind, unOperationId op, idempotencyWindow)
      :: IO [Only Int]
  pure (not (null rows))

lookupObjectMeta :: Connection -> ObjectKey -> IO (Maybe (Integer, ContentSha256))
lookupObjectMeta conn (ObjectKey key) = do
  rows <- query conn
    "SELECT size_bytes, content_sha256 FROM objects_meta WHERE object_key = ?"
    (Only key)
  pure $ case rows of
    [(sz, sha)] -> Just (sz, ContentSha256 sha)
    _           -> Nothing

sameFingerprint :: MutationRecord -> MutationFingerprint -> Bool
sameFingerprint rec fingerprint =
  mrFingerprint rec == fingerprint

existingDisposition
  :: Maybe MutationRecord
  -> MutationFingerprint
  -> ExistingDisposition
existingDisposition existing fingerprint =
  case existing of
    Nothing -> NewMutationEpoch
    Just rec
      | sameFingerprint rec fingerprint ->
          ExistingReplay (stateToResult (mrState rec)) rec
      | otherwise -> ExistingConflict

growthBytes :: Integer -> Integer -> Integer
growthBytes priorSize desiredSize = max 0 (desiredSize - priorSize)

quotaAllows :: Integer -> Integer -> Integer -> Integer -> Bool
quotaAllows limitBytes usedBytes reservedBytes growth =
  usedBytes + reservedBytes + growth <= limitBytes

terminalExpiryEligible :: MutationState -> ExternalState -> Bool
terminalExpiryEligible state external =
  state `elem` [Succeeded, Failed, Rejected] && external == ExtIdle

terminalExpired :: MutationState -> ExternalState -> Integer -> Integer -> Bool
terminalExpired state external elapsedSeconds windowSeconds =
  terminalExpiryEligible state external && elapsedSeconds > windowSeconds

recoveryEligible :: MutationState -> Integer -> Integer -> Bool
recoveryEligible state elapsedSeconds graceSeconds =
  state `elem` [Prepared, Applying] && elapsedSeconds > graceSeconds

admitMutation
  :: Connection
  -> AdmissionRequest
  -> AdmissionLimits
  -> IO AdmitOutcome
admitMutation conn AdmissionRequest{..} limits =
  if fingerprintKind admissionFingerprint
      /= admissionKind admissionIdentity
    then pure AdmitConflict
    else admitConsistentMutation
      conn
      admissionIdentity
      admissionPoolId
      admissionRelativePath
      admissionFingerprint
      limits

admitConsistentMutation
  :: Connection
  -> AdmissionIdentity
  -> PoolId
  -> RelativePath
  -> MutationFingerprint
  -> AdmissionLimits
  -> IO AdmitOutcome
admitConsistentMutation
  conn
  (AdmissionIdentity uid kind op)
  pool
  rel
  fingerprint@(MutationFingerprint _ key size msha)
  AdmissionLimits{..} = do
  let maxObject = admissionMaxObjectBytes
      idempotencyWindow = admissionIdempotencyWindowSec
  let identity =
        unUserId uid <> "\NUL" <> kindLabel kind <> "\NUL" <> unOperationId op
  _ <- query conn
    "SELECT pg_advisory_xact_lock(hashtextextended(?, 0))"
    (Only identity) :: IO [Only ()]
  _ <- execute conn
    "DELETE FROM mutations \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? \
    \AND state IN ('succeeded', 'failed', 'rejected') \
    \AND external_state = 'idle' \
    \AND updated_at < now() - (? * interval '1 second')"
    (unUserId uid, kindLabel kind, unOperationId op, idempotencyWindow)
  existing <- lookupMutation conn uid kind op
  case existingDisposition existing fingerprint of
    ExistingReplay result rec -> pure (AdmitReplay result rec)
    ExistingConflict -> pure AdmitConflict
    NewMutationEpoch -> do
      priorMeta <- lookupObjectMeta conn key
      let priorSize = maybe 0 fst priorMeta
          priorSha = snd <$> priorMeta
          desiredSize = case kind of
            Put    -> size
            Delete -> 0
          growth = growthBytes priorSize desiredSize
          baseRec phase =
            MutationRecord
              { mrAdmissionIdentity = AdmissionIdentity uid kind op
              , mrPoolId = pool
              , mrRelativePath = rel
              , mrFingerprint = fingerprint
              , mrPriorObject = PriorObject priorSize priorSha
              , mrGrowth = 0
              , mrPhase = phase
              }

      poolRows <- query conn
        "SELECT 1 FROM pools WHERE pool_id = ? FOR UPDATE"
        (Only (unPoolId pool)) :: IO [Only Int]

      if null poolRows
        then pure AdmitConflict
        else if kind == Put && (size <= 0 || size > maxObject)
        then do
          rec <- insertTerminalRejected conn uid kind op pool key rel size msha priorSize priorSha
          pure (AdmitRejected "invalid object size" rec)
        else do
          locked <- query conn
            "SELECT 1 FROM key_locks WHERE object_key = ? FOR UPDATE"
            (Only (unObjectKey key)) :: IO [Only Int]
          if not (null locked)
            then do
              rec <- insertTerminalRejected conn uid kind op pool key rel size msha priorSize priorSha
              pure (AdmitRejected "key locked by another mutation" rec)
            else do
              okQuota <- query conn
                "SELECT 1 FROM pools \
                \WHERE pool_id = ? AND used_bytes + reserved_bytes + ? <= limit_bytes"
                (unPoolId pool, growth) :: IO [Only Int]
              if null okQuota
                then do
                  rec <- insertTerminalRejected conn uid kind op pool key rel size msha priorSize priorSha
                  pure (AdmitRejected "quota exceeded" rec)
                else do
                  let preparedRecord = MutationRecord
                        { mrAdmissionIdentity = AdmissionIdentity uid kind op
                        , mrPoolId = pool
                        , mrRelativePath = rel
                        , mrFingerprint = fingerprint
                        , mrPriorObject = PriorObject priorSize priorSha
                        , mrGrowth = growth
                        , mrPhase = PhasePrepared
                        }
                  n <- execute conn
                    "INSERT INTO mutations \
                    \(user_id, kind, operation_id, pool_id, object_key, relative_path, \
                    \ declared_size, content_sha256, prior_size, prior_sha256, growth_bytes, \
                    \ state, external_state) \
                    \VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)"
                    (MutationInsertRow preparedRecord)
                  if n /= 1
                    then pure AdmitConflict
                    else do
                      lockN <- execute conn
                        "INSERT INTO key_locks (object_key, user_id, kind, operation_id) \
                        \VALUES (?,?,?,?) ON CONFLICT DO NOTHING"
                        (unObjectKey key, unUserId uid, kindLabel kind, unOperationId op)
                      if lockN /= 1
                        then do
                          _ <- execute conn
                            "UPDATE mutations SET state = 'rejected', growth_bytes = 0, \
                            \updated_at = now() \
                            \WHERE user_id = ? AND kind = ? AND operation_id = ?"
                            (unUserId uid, kindLabel kind, unOperationId op)
                          pure $ AdmitRejected "key locked by another mutation"
                            (baseRec (PhaseTerminal TerminalRejected))
                        else do
                          _ <- execute conn
                            "UPDATE pools SET reserved_bytes = reserved_bytes + ? WHERE pool_id = ?"
                            (growth, unPoolId pool)
                          pure (AdmitPrepared preparedRecord)

insertTerminalRejected
  :: Connection
  -> UserId
  -> MutationKind
  -> OperationId
  -> PoolId
  -> ObjectKey
  -> RelativePath
  -> Integer
  -> Maybe ContentSha256
  -> Integer
  -> Maybe ContentSha256
  -> IO MutationRecord
insertTerminalRejected conn uid kind op pool key rel size msha priorSize priorSha = do
  let record = MutationRecord
        { mrAdmissionIdentity = AdmissionIdentity uid kind op
        , mrPoolId = pool
        , mrRelativePath = rel
        , mrFingerprint =
            MutationFingerprint kind key size msha
        , mrPriorObject = PriorObject priorSize priorSha
        , mrGrowth = 0
        , mrPhase = PhaseTerminal TerminalRejected
        }
  _ <- execute conn
    "INSERT INTO mutations \
    \(user_id, kind, operation_id, pool_id, object_key, relative_path, \
    \ declared_size, content_sha256, prior_size, prior_sha256, growth_bytes, \
    \ state, external_state) \
    \VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?) \
    \ON CONFLICT DO NOTHING"
    (MutationInsertRow record)
  pure record

stateToResult :: MutationState -> MutationResult
stateToResult Prepared  = MutationPending
stateToResult Applying  = MutationPending
stateToResult Succeeded = MutationSucceeded
stateToResult Failed    = MutationFailed
stateToResult Rejected  = MutationRejected

markApplying :: Connection -> MutationRecord -> IO Bool
markApplying conn rec = do
  n <- execute conn
    "UPDATE mutations SET state = 'applying', updated_at = now() \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? AND state = 'prepared'"
    (unUserId (mrUserId rec), kindLabel (mrKind rec), unOperationId (mrOperationId rec))
  pure (n == 1)

markExternalPending
  :: Connection
  -> MutationRecord
  -> DispatchTiming
  -> IO Bool
markExternalPending conn rec DispatchTiming{..} = do
  n <- execute conn
    "UPDATE mutations SET external_state = 'pending', \
    \dispatch_started_at = now(), \
    \request_deadline_at = now() + (? * interval '1 second'), \
    \worker_lease_until = now() + (? * interval '1 second'), \
    \updated_at = now() \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? \
    \AND state = 'applying' AND external_state = 'idle'"
    ( dispatchSettlementBoundSec
    , dispatchWorkerGraceSec
    , unUserId (mrUserId rec)
    , kindLabel (mrKind rec)
    , unOperationId (mrOperationId rec)
    )
  pure (n == 1)

markExternalSettled :: Connection -> MutationRecord -> IO Bool
markExternalSettled conn rec = do
  n <- execute conn
    "UPDATE mutations SET external_state = 'settled', settled_at = now(), \
    \updated_at = now() \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? \
    \AND state = 'applying' AND external_state = 'pending'"
    (unUserId (mrUserId rec), kindLabel (mrKind rec), unOperationId (mrOperationId rec))
  pure (n == 1)

markSettledFromIdle :: Connection -> MutationRecord -> IO Bool
markSettledFromIdle conn rec = do
  n <- execute conn
    "UPDATE mutations SET external_state = 'settled', settled_at = now(), \
    \updated_at = now() \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? \
    \AND state = 'applying' AND external_state = 'idle'"
    (unUserId (mrUserId rec), kindLabel (mrKind rec), unOperationId (mrOperationId rec))
  pure (n == 1)

finalizeSuccess :: Connection -> MutationRecord -> Integer -> Maybe ContentSha256 -> IO Bool
finalizeSuccess conn rec finalSize finalSha = do
  rows <- query conn
    "SELECT state, external_state, growth_bytes, prior_size FROM mutations \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? FOR UPDATE"
    (unUserId (mrUserId rec), kindLabel (mrKind rec), unOperationId (mrOperationId rec))
      :: IO [SettlementRow]
  case rows of
    [SettlementRow
      (DbMutationState state)
      (DbExternalState external)
      growth
      prior]
      | decodeMutationPhase state external
          == Just (PhaseApplying ExtSettled) -> do
          let pool = mrPoolId rec
          _ <- execute conn
            "UPDATE pools SET used_bytes = used_bytes - ? + ?, \
            \reserved_bytes = reserved_bytes - ? WHERE pool_id = ?"
            (prior, finalSize, growth, unPoolId pool)
          case mrKind rec of
            Put -> do
              let sha = maybe "" unContentSha256 finalSha
              _ <- execute conn
                "INSERT INTO objects_meta (object_key, pool_id, size_bytes, content_sha256, updated_at) \
                \VALUES (?,?,?,?, now()) \
                \ON CONFLICT (object_key) DO UPDATE SET \
                \  size_bytes = EXCLUDED.size_bytes, \
                \  content_sha256 = EXCLUDED.content_sha256, \
                \  updated_at = now()"
                (unObjectKey (mrObjectKey rec), unPoolId pool, finalSize, sha)
              pure ()
            Delete -> do
              _ <- execute conn
                "DELETE FROM objects_meta WHERE object_key = ?"
                (Only (unObjectKey (mrObjectKey rec)))
              pure ()
          _ <- execute conn
            "UPDATE mutations SET state = 'succeeded', external_state = 'idle', updated_at = now() \
            \WHERE user_id = ? AND kind = ? AND operation_id = ?"
            (unUserId (mrUserId rec), kindLabel (mrKind rec), unOperationId (mrOperationId rec))
          _ <- execute conn
            "DELETE FROM key_locks WHERE object_key = ? AND user_id = ? AND kind = ? AND operation_id = ?"
            ( unObjectKey (mrObjectKey rec)
            , unUserId (mrUserId rec)
            , kindLabel (mrKind rec)
            , unOperationId (mrOperationId rec)
            )
          pure True
    _ -> pure False

finalizeFailure :: Connection -> MutationRecord -> IO Bool
finalizeFailure conn rec = do
  rows <- query conn
    "SELECT state, external_state, growth_bytes, prior_size FROM mutations \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? FOR UPDATE"
    (unUserId (mrUserId rec), kindLabel (mrKind rec), unOperationId (mrOperationId rec))
      :: IO [SettlementRow]
  case rows of
    [SettlementRow
      (DbMutationState state)
      (DbExternalState external)
      growth
      _]
      | decodeMutationPhase state external
          == Just (PhaseApplying ExtSettled) -> do
          _ <- execute conn
            "UPDATE pools SET reserved_bytes = reserved_bytes - ? WHERE pool_id = ?"
            (growth, unPoolId (mrPoolId rec))
          _ <- execute conn
            "UPDATE mutations SET state = 'failed', external_state = 'idle', updated_at = now() \
            \WHERE user_id = ? AND kind = ? AND operation_id = ?"
            (unUserId (mrUserId rec), kindLabel (mrKind rec), unOperationId (mrOperationId rec))
          _ <- execute conn
            "DELETE FROM key_locks WHERE object_key = ? AND user_id = ? AND kind = ? AND operation_id = ?"
            ( unObjectKey (mrObjectKey rec)
            , unUserId (mrUserId rec)
            , kindLabel (mrKind rec)
            , unOperationId (mrOperationId rec)
            )
          pure True
    _ -> pure False

-- | updated_at advances on state transitions only, not during a long put body.
claimActiveBatch :: Connection -> Int -> Integer -> IO [MutationRecord]
claimActiveBatch conn limit graceSec = do
  rows <- query conn
    "SELECT user_id, kind, operation_id, pool_id, object_key, relative_path, \
    \declared_size, content_sha256, prior_size, prior_sha256, growth_bytes, \
    \state, external_state FROM mutations \
    \WHERE state IN ('prepared', 'applying') \
    \AND ( \
    \  (external_state = 'idle' \
    \    AND updated_at < now() - (? * interval '1 second')) \
    \  OR (external_state = 'settled' \
    \    AND updated_at < now() - (? * interval '1 second')) \
    \  OR (external_state = 'pending' \
    \    AND worker_lease_until < now() \
    \    AND request_deadline_at < now()) \
    \) \
    \ORDER BY updated_at ASC \
    \FOR UPDATE SKIP LOCKED \
    \LIMIT ?"
    (graceSec, graceSec, limit)
  decodeRows rows

desiredMatches :: MutationRecord -> ObjectObservation -> Bool
desiredMatches rec observation =
  case (desiredObject rec, observation) of
    (_, ObjectUnavailable _) -> False
    (DesiredAbsent, ObjectAbsent) -> True
    (DesiredPresent size expectedSha, ObjectPresent observedSize observedSha) ->
      size == observedSize
        && case (expectedSha, observedSha) of
             (Just want, Just got) -> want == got
             (Just _, Nothing)     -> False
             (Nothing, _)          -> True
    _ -> False

priorMatches :: MutationRecord -> ObjectObservation -> Bool
priorMatches rec obs =
  let PriorObject priorSize priorSha = mrPriorObject rec
      priorEmpty = priorSize == 0 && priorSha == Nothing
  in case obs of
       ObjectUnavailable _ -> False
       ObjectAbsent -> priorEmpty
       ObjectPresent sz msha ->
         sz == priorSize
           && ( case (priorSha, msha) of
                  (Just want, Just got) -> want == got
                  (Nothing, Nothing)    -> True
                  (Just _, Nothing)     -> False
                  (Nothing, Just _)     -> priorSize == sz
              )

finalizeFromObservation
  :: Connection
  -> MutationRecord
  -> ObjectObservation
  -> IO (Maybe MutationResult)
finalizeFromObservation conn rec observed =
  case observed of
    ObjectUnavailable _ -> pure Nothing
    _ | desiredMatches rec observed -> do
      let (finalSize, finalSha) =
            case mrKind rec of
              Put ->
                ( mrDeclaredSize rec
                , mrContentSha256 rec
                )
              Delete -> (0, Nothing)
      ok <- finalizeSuccess conn rec finalSize finalSha
      pure $ if ok then Just MutationSucceeded else Nothing
    _ | priorMatches rec observed -> do
      ok <- finalizeFailure conn rec
      pure $ if ok then Just MutationFailed else Nothing
    _ -> pure Nothing