nuthatch/test/RecoveryPhase.hs

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           Nuthatch.App (dispatchAndFinalize, newEnv)
import           Nuthatch.Config
import           Nuthatch.Db
import qualified Nuthatch.Db as Db
import           Nuthatch.Protocol
import           Nuthatch.Redis
import           Nuthatch.S3
import           Nuthatch.Types
import           Servant.Types.SourceT
import           System.Environment (getEnv)
import           System.IO (BufferMode (LineBuffering), hSetBuffering, stdout)

main :: IO ()
main = do
  hSetBuffering stdout LineBuffering
  cfg <- loadConfig
  phase <- getEnv "NUTHATCH_RECOVERY_PHASE"
  withDb cfg $ \database -> do
    migrate database "sql" cfg
    withRedis cfg $ \redis -> do
      store <- newS3 cfg
      runPhase cfg database redis store phase

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
      recoveryUser
      Put
      (phaseOperation name)
      sharedPool
      (phaseKey name)
      (RelativePath (name <> ".bin"))
      size
      (Just identity)
      (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
      recoveryUser
      Delete
      (phaseOperation name)
      sharedPool
      key
      (RelativePath (name <> ".bin"))
      0
      Nothing
      (cfgMaxObjectBytes cfg)
      (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
      (cfgObjectRequestTimeoutSec cfg * 2)
      (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"