nuthatch/src/Nuthatch/Protocol.hs

-- | Short Postgres transactions for mutation authority; store I/O stays outside them.
module Nuthatch.Protocol
  ( AdmitOutcome (..)
  , MutationRecord (..)
  , ObjectObservation (..)
  , admitMutation
  , markApplying
  , markExternalPending
  , markExternalSettled
  , markSettledFromIdle
  , finalizeSuccess
  , finalizeFailure
  , claimActiveBatch
  , lookupMutation
  , lookupObjectMeta
  , desiredMatches
  , priorMatches
  ) where

import           Data.Maybe (listToMaybe, mapMaybe)
import           Data.Text (Text)
import           Database.PostgreSQL.Simple (Connection, Only (..), execute, query)
import           Nuthatch.Types

data MutationRecord = MutationRecord
  { mrUserId        :: UserId
  , mrKind          :: MutationKind
  , mrOperationId   :: OperationId
  , mrPoolId        :: PoolId
  , mrObjectKey     :: ObjectKey
  , mrRelativePath  :: RelativePath
  , mrDeclaredSize  :: Integer
  , mrContentSha256 :: Maybe ContentSha256
  , mrPriorSize     :: Integer
  , mrPriorSha256   :: Maybe ContentSha256
  , mrGrowth        :: Integer
  , mrState         :: MutationState
  , mrExternal      :: ExternalState
  }
  deriving stock (Show)

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

type ResultRow =
  ( Text, Text, Text, Text, Text, Text, Integer, Maybe Text
  , Integer, Maybe Text, Integer, Text, Text
  )

rowToRecord :: ResultRow -> Maybe MutationRecord
rowToRecord (user, kindT, op, pool, key, rel, size, csha, prior, psha, growth, stateT, extT) = do
  kind <- parseKind kindT
  state <- parseState stateT
  ext <- parseExternal extT
  pure MutationRecord
    { mrUserId = UserId user
    , mrKind = kind
    , mrOperationId = OperationId op
    , mrPoolId = PoolId pool
    , mrObjectKey = ObjectKey key
    , mrRelativePath = RelativePath rel
    , mrDeclaredSize = size
    , mrContentSha256 = ContentSha256 <$> csha
    , mrPriorSize = prior
    , mrPriorSha256 = ContentSha256 <$> psha
    , mrGrowth = growth
    , mrState = state
    , mrExternal = ext
    }

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)
  pure (listToMaybe rows >>= rowToRecord)

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 -> MutationKind -> ObjectKey -> Integer -> Maybe ContentSha256 -> Bool
sameFingerprint rec kind key size msha =
  mrKind rec == kind
    && mrObjectKey rec == key
    && mrDeclaredSize rec == size
    && mrContentSha256 rec == msha

admitMutation
  :: Connection
  -> UserId
  -> MutationKind
  -> OperationId
  -> PoolId
  -> ObjectKey
  -> RelativePath
  -> Integer
  -> Maybe ContentSha256
  -> Integer
  -> IO AdmitOutcome
admitMutation conn uid kind op pool key rel size msha maxObject = do
  existing <- lookupMutation conn uid kind op
  case existing of
    Just rec
      | sameFingerprint rec kind key size msha ->
          pure $ AdmitReplay (stateToResult (mrState rec)) rec
      | otherwise -> pure AdmitConflict
    Nothing -> do
      priorMeta <- lookupObjectMeta conn key
      let priorSize = maybe 0 fst priorMeta
          priorSha = snd <$> priorMeta
          desiredSize = case kind of
            Put    -> size
            Delete -> 0
          growth = max 0 (desiredSize - priorSize)
          baseRec state =
            MutationRecord
              { mrUserId = uid
              , mrKind = kind
              , mrOperationId = op
              , mrPoolId = pool
              , mrObjectKey = key
              , mrRelativePath = rel
              , mrDeclaredSize = size
              , mrContentSha256 = msha
              , mrPriorSize = priorSize
              , mrPriorSha256 = priorSha
              , mrGrowth = 0
              , mrState = state
              , mrExternal = ExtIdle
              }

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

      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
                  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 (?,?,?,?,?,?,?,?,?,?,?,'prepared','idle')"
                    ( unUserId uid
                    , kindLabel kind
                    , unOperationId op
                    , unPoolId pool
                    , unObjectKey key
                    , unRelativePath rel
                    , size
                    , unContentSha256 <$> msha
                    , priorSize
                    , unContentSha256 <$> priorSha
                    , growth
                    )
                  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 Rejected)
                        else do
                          _ <- execute conn
                            "UPDATE pools SET reserved_bytes = reserved_bytes + ? WHERE pool_id = ?"
                            (growth, unPoolId pool)
                          pure $ AdmitPrepared MutationRecord
                            { mrUserId = uid
                            , mrKind = kind
                            , mrOperationId = op
                            , mrPoolId = pool
                            , mrObjectKey = key
                            , mrRelativePath = rel
                            , mrDeclaredSize = size
                            , mrContentSha256 = msha
                            , mrPriorSize = priorSize
                            , mrPriorSha256 = priorSha
                            , mrGrowth = growth
                            , mrState = Prepared
                            , mrExternal = ExtIdle
                            }

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
  _ <- 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 (?,?,?,?,?,?,?,?,?,?,0,'rejected','idle') \
    \ON CONFLICT DO NOTHING"
    ( unUserId uid
    , kindLabel kind
    , unOperationId op
    , unPoolId pool
    , unObjectKey key
    , unRelativePath rel
    , size
    , unContentSha256 <$> msha
    , priorSize
    , unContentSha256 <$> priorSha
    )
  pure MutationRecord
    { mrUserId = uid
    , mrKind = kind
    , mrOperationId = op
    , mrPoolId = pool
    , mrObjectKey = key
    , mrRelativePath = rel
    , mrDeclaredSize = size
    , mrContentSha256 = msha
    , mrPriorSize = priorSize
    , mrPriorSha256 = priorSha
    , mrGrowth = 0
    , mrState = Rejected
    , mrExternal = ExtIdle
    }

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 -> IO Bool
markExternalPending conn rec = do
  n <- execute conn
    "UPDATE mutations SET external_state = 'pending', 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)

markExternalSettled :: Connection -> MutationRecord -> IO Bool
markExternalSettled conn rec = do
  n <- execute conn
    "UPDATE mutations SET external_state = 'settled', 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', 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 [(Text, Text, Integer, Integer)]
  case rows of
    [(st, ext, growth, prior)]
      | st == "applying" && ext == "settled" -> 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 FROM mutations \
    \WHERE user_id = ? AND kind = ? AND operation_id = ? FOR UPDATE"
    (unUserId (mrUserId rec), kindLabel (mrKind rec), unOperationId (mrOperationId rec))
      :: IO [(Text, Text, Integer)]
  case rows of
    [(st, ext, growth)]
      | st == "applying" && ext == "settled" -> 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 updated_at < now() - (? * interval '1 second') \
    \ORDER BY updated_at ASC \
    \FOR UPDATE SKIP LOCKED \
    \LIMIT ?"
    (graceSec, limit)
  pure (mapMaybe rowToRecord rows)

desiredMatches :: MutationRecord -> ObjectObservation -> Bool
desiredMatches rec = \case
  ObjectUnavailable _ -> False
  ObjectAbsent ->
    case mrKind rec of
      Delete -> True
      Put    -> False
  ObjectPresent sz msha ->
    case mrKind rec of
      Delete -> sz == 0
      Put ->
        sz == mrDeclaredSize rec
          && case (mrContentSha256 rec, msha) of
               (Just want, Just got) -> want == got
               -- Some stores omit custom metadata on HEAD.
               (Just _, Nothing)     -> True
               (Nothing, _)          -> True

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