Make object recovery settlement-safe

test/RecoveryPhase.hs

@@ -0,0 +1,239 @@
+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
+        "settlement failed"
+        "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 "wp09" "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"