341 lines
11 KiB
Haskell
341 lines
11 KiB
Haskell
module Main (main) where
|
|
|
|
import Control.Concurrent (threadDelay)
|
|
import Control.Monad (unless)
|
|
import qualified Data.ByteString as BS
|
|
import qualified Data.ByteString.Lazy as BL
|
|
import Data.Text (Text)
|
|
import qualified Data.Text as T
|
|
import Database.PostgreSQL.Simple
|
|
( FromRow
|
|
, Only (..)
|
|
, Query
|
|
, ToRow
|
|
)
|
|
import Nuthatch.App.Env (newEnv)
|
|
import Nuthatch.Config
|
|
import Nuthatch.Db
|
|
import qualified Nuthatch.Db as Db
|
|
import Nuthatch.Mutation.Service (dispatchAndFinalize)
|
|
import Nuthatch.Protocol
|
|
import Nuthatch.Redis
|
|
import Nuthatch.S3
|
|
import Nuthatch.Types
|
|
import Servant.Types.SourceT
|
|
import System.Environment (getArgs, getEnv)
|
|
import System.IO (BufferMode (LineBuffering), hSetBuffering, stdout)
|
|
import Text.Read (readMaybe)
|
|
|
|
main :: IO ()
|
|
main = do
|
|
hSetBuffering stdout LineBuffering
|
|
cfg <- loadConfig
|
|
withDb cfg $ \database -> do
|
|
migrate database "sql" cfg
|
|
withRedis cfg $ \redis -> do
|
|
store <- newS3 cfg
|
|
getArgs >>= \case
|
|
[] -> getEnv "NUTHATCH_RECOVERY_PHASE" >>=
|
|
runPhase cfg database redis store
|
|
["wait-phase", operation, expected] ->
|
|
waitForPhase database operation expected
|
|
["terminal", operation, expected] ->
|
|
verifyTerminal database operation expected
|
|
["active", operation, expectedReservation] ->
|
|
maybe
|
|
(fail "invalid expected reservation")
|
|
(verifyActive database operation)
|
|
(readMaybe expectedReservation)
|
|
["object", expectation, key] ->
|
|
verifyObject store expectation (ObjectKey (T.pack key))
|
|
["cache-invalidated"] ->
|
|
cacheGet redis "nuthatch:list:recovery-user:shared::" >>=
|
|
requireEqual "cache invalidation" Nothing
|
|
["database-available"] ->
|
|
waitForQuery database "SELECT 1" ((), Only (1 :: Int))
|
|
["log-exhausted"] -> do
|
|
logBytes <- BS.readFile "/run/nuthatch-recovery/app.log"
|
|
unless ("\"result\":\"exhausted\"" `BS.isInfixOf` logBytes) $
|
|
fail "recovery log omitted exhausted retry result"
|
|
_ -> fail "unknown recovery probe mode"
|
|
|
|
waitForPhase :: Db -> String -> String -> IO ()
|
|
waitForPhase database operation expected =
|
|
waitForQuery database
|
|
"SELECT state || ':' || external_state FROM mutations WHERE operation_id = ?"
|
|
(Only (T.pack operation), Only (T.pack expected))
|
|
|
|
verifyTerminal :: Db -> String -> String -> IO ()
|
|
verifyTerminal database operation expected = do
|
|
waitForQuery database
|
|
"SELECT state FROM mutations WHERE operation_id = ?"
|
|
(Only (T.pack operation), Only (T.pack expected))
|
|
locks <- Db.query database
|
|
"SELECT count(*) FROM key_locks WHERE operation_id = ?"
|
|
(Only (T.pack operation))
|
|
:: IO [Only Int]
|
|
requireEqual "terminal key lock" [Only 0] locks
|
|
reservations <- Db.query database
|
|
"SELECT reserved_bytes FROM pools WHERE pool_id = 'shared'"
|
|
()
|
|
:: IO [Only Integer]
|
|
requireEqual "terminal reservation" [Only 0] reservations
|
|
|
|
verifyActive :: Db -> String -> Integer -> IO ()
|
|
verifyActive database operation expectedReservation = do
|
|
locks <- Db.query database
|
|
"SELECT count(*) FROM key_locks WHERE operation_id = ?"
|
|
(Only (T.pack operation))
|
|
:: IO [Only Int]
|
|
requireEqual "active key lock" [Only 1] locks
|
|
reservations <- Db.query database
|
|
"SELECT reserved_bytes FROM pools WHERE pool_id = 'shared'"
|
|
()
|
|
:: IO [Only Integer]
|
|
requireEqual "active reservation"
|
|
[Only expectedReservation]
|
|
reservations
|
|
|
|
verifyObject :: S3 -> String -> ObjectKey -> IO ()
|
|
verifyObject store expectation key = do
|
|
observed <- headObject store key
|
|
case (expectation, observed) of
|
|
("present", Right (Just _)) -> pure ()
|
|
("absent", Right Nothing) -> pure ()
|
|
_ -> fail
|
|
("object expectation failed: " <> expectation <> " " <> show observed)
|
|
|
|
waitForQuery
|
|
:: (Eq row, Show row, ToRow parameters, FromRow row)
|
|
=> Db
|
|
-> Query
|
|
-> (parameters, row)
|
|
-> IO ()
|
|
waitForQuery database queryText (parameters, expected) =
|
|
go (120 :: Int)
|
|
where
|
|
go remaining = do
|
|
rows <- Db.query database queryText parameters
|
|
if rows == [expected]
|
|
then pure ()
|
|
else
|
|
if remaining <= 1
|
|
then fail
|
|
("timed out waiting for database value: " <> show rows)
|
|
else threadDelay 250000 >> go (remaining - 1)
|
|
|
|
requireEqual :: (Eq value, Show value) => String -> value -> value -> IO ()
|
|
requireEqual label expected actual =
|
|
unless (actual == expected) $
|
|
fail (label <> ": expected " <> show expected <> ", got " <> show actual)
|
|
|
|
runPhase :: Config -> Db -> Redis -> S3 -> String -> IO ()
|
|
runPhase cfg database redis store phase =
|
|
case phase of
|
|
"before-dispatch" ->
|
|
preparePut database redis "before-dispatch" 13 phaseIdentity >>= checkpoint
|
|
"during-put" -> do
|
|
rec <- preparePut database redis "during-put" slowSize slowIdentity
|
|
beginDispatch cfg database rec
|
|
result <- putObjectStream
|
|
store
|
|
(mrObjectKey rec)
|
|
slowSize
|
|
(Just slowIdentity)
|
|
65536
|
|
(const (pure ()))
|
|
(slowZeroSource slowSize)
|
|
fail ("during-put completed before process loss: " <> show result)
|
|
"live-timeout" -> do
|
|
rec <- preparePut database redis "live-timeout" slowSize slowIdentity
|
|
result <- dispatchAndFinalize
|
|
(newEnv cfg database redis store)
|
|
recoverySession
|
|
Shared
|
|
rec
|
|
"request timed out"
|
|
(putObjectStream
|
|
store
|
|
(mrObjectKey rec)
|
|
slowSize
|
|
(Just slowIdentity)
|
|
65536
|
|
(const (pure ()))
|
|
(slowZeroSource slowSize)
|
|
>>= pure . either (Left . T.pack . show) Right)
|
|
case result of
|
|
Left "request timed out" -> checkpoint rec
|
|
other -> fail ("live timeout returned unexpected result: " <> show other)
|
|
"after-dispatch" -> do
|
|
rec <- preparePut database redis "after-dispatch" 13 phaseIdentity
|
|
beginDispatch cfg database rec
|
|
expectRight =<< putObject
|
|
store
|
|
(mrObjectKey rec)
|
|
13
|
|
(Just phaseIdentity)
|
|
"recovered-put"
|
|
checkpoint rec
|
|
"after-settlement" -> do
|
|
rec <- preparePut database redis "after-settlement" 13 phaseIdentity
|
|
beginDispatch cfg database rec
|
|
expectRight =<< putObject
|
|
store
|
|
(mrObjectKey rec)
|
|
13
|
|
(Just phaseIdentity)
|
|
"recovered-put"
|
|
settled <- withTransaction database $ \connection ->
|
|
markExternalSettled connection rec
|
|
unless settled (fail "could not record settlement")
|
|
checkpoint rec
|
|
"delete-settled" -> do
|
|
rec <- prepareDelete cfg database redis store
|
|
beginDispatch cfg database rec
|
|
expectRight =<< deleteObject store (mrObjectKey rec)
|
|
settled <- withTransaction database $ \connection ->
|
|
markExternalSettled connection rec
|
|
unless settled (fail "could not record delete settlement")
|
|
checkpoint rec
|
|
_ -> fail ("unknown recovery phase: " <> phase)
|
|
|
|
preparePut
|
|
:: Db
|
|
-> Redis
|
|
-> Text
|
|
-> Integer
|
|
-> ContentSha256
|
|
-> IO MutationRecord
|
|
preparePut database redis name size identity = do
|
|
cacheRecoveryListing redis
|
|
outcome <- withTransaction database $ \connection ->
|
|
admitMutation
|
|
connection
|
|
AdmissionRequest
|
|
{ admissionIdentity =
|
|
AdmissionIdentity recoveryUser Put (phaseOperation name)
|
|
, admissionPoolId = sharedPool
|
|
, admissionRelativePath = RelativePath (name <> ".bin")
|
|
, admissionFingerprint =
|
|
MutationFingerprint Put (phaseKey name) size (Just identity)
|
|
}
|
|
(AdmissionLimits (1024 * 1024 * 1024) 3600)
|
|
prepared outcome
|
|
|
|
prepareDelete :: Config -> Db -> Redis -> S3 -> IO MutationRecord
|
|
prepareDelete cfg database redis store = do
|
|
let name = "delete-settled"
|
|
key = phaseKey name
|
|
body = "delete-me"
|
|
size = toInteger (BL.length body)
|
|
cacheRecoveryListing redis
|
|
expectRight =<< putObject store key size (Just deleteIdentity) body
|
|
_ <- Db.execute database
|
|
"UPDATE pools SET used_bytes = ? WHERE pool_id = ?"
|
|
(size, unPoolId sharedPool)
|
|
_ <- Db.execute database
|
|
"INSERT INTO objects_meta \
|
|
\(object_key, pool_id, size_bytes, content_sha256) VALUES (?,?,?,?)"
|
|
( unObjectKey key
|
|
, unPoolId sharedPool
|
|
, size
|
|
, unContentSha256 deleteIdentity
|
|
)
|
|
outcome <- withTransaction database $ \connection ->
|
|
admitMutation
|
|
connection
|
|
AdmissionRequest
|
|
{ admissionIdentity =
|
|
AdmissionIdentity recoveryUser Delete (phaseOperation name)
|
|
, admissionPoolId = sharedPool
|
|
, admissionRelativePath = RelativePath (name <> ".bin")
|
|
, admissionFingerprint =
|
|
MutationFingerprint Delete key 0 Nothing
|
|
}
|
|
AdmissionLimits
|
|
{ admissionMaxObjectBytes = cfgMaxObjectBytes cfg
|
|
, admissionIdempotencyWindowSec = cfgIdempotencyWindowSec cfg
|
|
}
|
|
prepared outcome
|
|
|
|
beginDispatch :: Config -> Db -> MutationRecord -> IO ()
|
|
beginDispatch cfg database rec = do
|
|
started <- withTransaction database $ \connection -> do
|
|
applying <- markApplying connection rec
|
|
pending <- markExternalPending
|
|
connection
|
|
rec
|
|
DispatchTiming
|
|
{ dispatchSettlementBoundSec = cfgObjectRequestTimeoutSec cfg * 2
|
|
, dispatchWorkerGraceSec = cfgRecoveryGraceSec cfg
|
|
}
|
|
pure (applying && pending)
|
|
unless started (fail "could not begin dispatch")
|
|
|
|
cacheRecoveryListing :: Redis -> IO ()
|
|
cacheRecoveryListing redis =
|
|
cacheSet redis "nuthatch:list:recovery-user:shared::" 600 "stale"
|
|
|
|
prepared :: AdmitOutcome -> IO MutationRecord
|
|
prepared outcome =
|
|
case outcome of
|
|
AdmitPrepared rec -> pure rec
|
|
other -> fail ("phase admission failed: " <> show other)
|
|
|
|
checkpoint :: MutationRecord -> IO a
|
|
checkpoint rec = do
|
|
putStrLn
|
|
("recovery checkpoint "
|
|
<> T.unpack (unOperationId (mrOperationId rec)))
|
|
spin
|
|
|
|
spin :: IO a
|
|
spin = threadDelay 1000000 >> spin
|
|
|
|
slowZeroSource :: Integer -> SourceT IO BS.ByteString
|
|
slowZeroSource total =
|
|
SourceT $ \consume -> consume (go total)
|
|
where
|
|
go remaining
|
|
| remaining <= 0 = Stop
|
|
| otherwise = Effect $ do
|
|
threadDelay 200000
|
|
let size = fromInteger (min remaining 65536)
|
|
pure (Yield (BS.replicate size 0) (go (remaining - toInteger size)))
|
|
|
|
expectRight :: Either S3Error a -> IO a
|
|
expectRight result =
|
|
case result of
|
|
Right value -> pure value
|
|
Left err -> fail ("object operation failed: " <> show err)
|
|
|
|
phaseOperation :: Text -> OperationId
|
|
phaseOperation name = OperationId ("recovery-" <> name)
|
|
|
|
phaseKey :: Text -> ObjectKey
|
|
phaseKey name = ObjectKey ("shared/recovery-" <> name <> ".bin")
|
|
|
|
recoveryUser :: UserId
|
|
recoveryUser = UserId "recovery-user"
|
|
|
|
recoverySession :: SessionUser
|
|
recoverySession =
|
|
SessionUser "recovery" "recovery" recoveryUser (Just "Recovery")
|
|
|
|
slowSize :: Integer
|
|
slowSize = 20 * 1024 * 1024
|
|
|
|
phaseIdentity :: ContentSha256
|
|
phaseIdentity =
|
|
ContentSha256
|
|
"1cfe64e6951d3efd9cfe80b53c65cb55d0e87e2a3c72474e6afd8a89fb0e2f69"
|
|
|
|
deleteIdentity :: ContentSha256
|
|
deleteIdentity =
|
|
ContentSha256
|
|
"af77c0879d9cc69681600bb17fa8fedccfe7ae605b65974fee792581ed63a84c"
|
|
|
|
slowIdentity :: ContentSha256
|
|
slowIdentity =
|
|
ContentSha256
|
|
"cd52d81e25f372e6fa4db2c0dfceb59862c1969cab17096da352b34950c973cc"
|