441 lines
17 KiB
Haskell
441 lines
17 KiB
Haskell
-- | 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
|
|
)
|