WPB-25915 add timeout and duration metric for conversation migration

changelog.d/0-release-notes/WPB-25915

@@ -0,0 +1 @@
+The background-worker migration timeout configuration was renamed from `migrateConversationsOptions` to `migrationOptions`. The old key is no longer read, so update any custom Helm overrides to the new name to avoid unexpected defaults.

changelog.d/5-internal/WPB-25915

@@ -0,0 +1 @@
+Add timeout and duration metrics to Cassandra-to-PostgreSQL migration options

charts/wire-server/templates/background-worker/configmap.yaml

@@ -84,8 +84,8 @@ data:
     migrateConversationCodes: {{ .migrateConversationCodes }}
     migrateTeamFeatures: {{ .migrateTeamFeatures }}
     migrateDomainRegistration: {{ .migrateDomainRegistration }}
-    migrateConversationsOptions:
-{{toYaml .migrateConversationsOptions | indent 6 }}
+    migrationOptions:
+{{toYaml .migrationOptions | indent 6 }}
 
     backendNotificationPusher:
 {{toYaml .backendNotificationPusher | indent 6 }}

charts/wire-server/values.yaml

@@ -952,9 +952,10 @@ background-worker:
     # `settings.postgresMigration.conversation` with `migration-to-postgresql`
     # before setting this to `true`.
     migrateConversations: false
-    migrateConversationsOptions:
+    migrationOptions:
       pageSize: 10000
       parallelism: 2
+      timeout: 5s
     # This will start the migration of conversation codes.
     # It's important to set `settings.postgresMigration.conversationCodes` to `migration-to-postgresql`
     # before starting the migration.

docs/src/developer/reference/config-options.md

@@ -1922,6 +1922,16 @@ time. For conversations, this is necessary for channel search and management of
 channels from the team-management UI. It is highly recommended to take a backup
 of the affected Cassandra data before triggering a migration.
 
+The `background-worker.config.migrationOptions.timeout` defaults to 5s. It
+sets an upper bound for a single migration attempt after it has acquired the
+migration lock, so a stuck item does not keep the migration run blocked
+indefinitely.
+
+Although `migrationOptions` is a shared background-worker setting, the
+migrations run sequentially. That means the same configured timeout can be
+used as a practical per-migration limit, because only one migration type is
+executing at a time in the worker.
+
 Migrations are independent and can be run separately, in batches, or all at
 once. This is expected, because migrations will be released over time. The
 pattern below applies per `postgresMigration` setting. A single setting may
@@ -2096,6 +2106,13 @@ migrateConversationCodes: false
 migrateTeamFeatures: false
 migrateDomainRegistration: false
 
+# migration settings
+migrationOptions:
+  pageSize: 10000
+  parallelism: 2
+  # Required migration timeout for a single migration attempt.
+  timeout: 5s
+
 # Background jobs consumer
 backgroundJobs:
   concurrency: 8     # in-flight jobs per process
@@ -2106,6 +2123,15 @@ backgroundJobs:
 federationDomain: example.org
 ```
 
+The `migrationOptions.timeout` setting limits how long a single migration
+attempt may run after it has acquired the migration lock. If the timeout is
+exceeded, that migration attempt is aborted and treated as failed.
+
+Choose a value that is comfortably above the normal time for one conversation
+migration, but still low enough to catch a genuinely stuck migration in a
+reasonable time. A good starting point for most deployments is `5s`,
+then adjust based on observed migration durations.
+
 Secrets
 
 - Set `background-worker.secrets.pgPassword` to pass the PostgreSQL password. The chart mounts it to `/etc/wire/background-worker/secrets/pgPassword` and sets `postgresqlPassword` accordingly.

libs/types-common/default.nix

@@ -36,6 +36,7 @@
 , optparse-applicative
 , pem
 , polysemy
+, polysemy-time
 , protobuf
 , QuickCheck
 , quickcheck-instances
@@ -95,6 +96,7 @@ mkDerivation {
     optparse-applicative
     pem
     polysemy
+    polysemy-time
     protobuf
     QuickCheck
     quickcheck-instances

libs/types-common/src/Data/Misc.hs

@@ -104,6 +104,8 @@ import Data.Time
 import GHC.TypeLits (Nat)
 import GHC.TypeNats (KnownNat)
 import Imports
+import Polysemy.Time
+import Polysemy.Time.Data.TimeUnit
 import Servant (FromHttpApiData (..))
 import Test.QuickCheck (Arbitrary (arbitrary), chooseInteger)
 import Test.QuickCheck qualified as QC
@@ -269,6 +271,16 @@ instance Cql Milliseconds where
 newtype Duration = Duration {duration :: DiffTime}
   deriving (Eq, Show)
 
+instance TimeUnit Duration where
+  nanos = NanoSeconds 1_000_000_000
+
+  toNanos (Duration dt) =
+    NanoSeconds . fromInteger $
+      diffTimeToPicoseconds dt `div` 1_000
+
+  fromNanos (NanoSeconds ns) =
+    Duration $ picosecondsToDiffTime (fromIntegral ns * 1_000)
+
 diffTimeParser :: Parser DiffTime
 diffTimeParser = do
   n :: DiffTime <- Atto.rational

libs/types-common/src/Util/Timeout.hs

@@ -26,12 +26,13 @@ import Data.Aeson.Types
 import Data.Scientific
 import Data.Time.Clock
 import Imports
+import Polysemy.Time
 import Test.QuickCheck (Arbitrary (arbitrary), choose)
 
 newtype Timeout = Timeout
   { timeoutDiff :: NominalDiffTime
   }
-  deriving newtype (Eq, Enum, Ord, Num, Real, Fractional, RealFrac, Show)
+  deriving newtype (Eq, Enum, Ord, Num, Real, Fractional, RealFrac, Show, TimeUnit)
 
 instance Arbitrary Timeout where
   arbitrary = Timeout . fromIntegral <$> choose (60 :: Int, 10 * 24 * 3600)

libs/types-common/types-common.cabal

@@ -125,6 +125,7 @@ library
     , optparse-applicative   >=0.10
     , pem
     , polysemy
+    , polysemy-time
     , protobuf               >=0.2
     , QuickCheck             >=2.9
     , quickcheck-instances   >=0.3.16

libs/wire-subsystems/src/Wire/CodeStore/Migration.hs

@@ -22,16 +22,25 @@ import Data.ByteString.Conversion
 import Data.Code (Key, Value)
 import Data.Conduit
 import Data.Conduit.List qualified as C
+import Data.IORef qualified as IORef
 import Data.Id (ConvId)
 import Data.Misc (HttpsUrl)
+import Data.Text qualified as T
+import Data.Time
 import Hasql.Pool qualified as Hasql
 import Imports
 import Polysemy
+import Polysemy.Async
+import Polysemy.Conc (interpretRace)
+import Polysemy.Conc qualified as Conc
+import Polysemy.Conc.Effect.Race hiding (Timeout)
 import Polysemy.Input
+import Polysemy.Resource (Resource, bracket, resourceToIOFinal)
 import Polysemy.State
 import Polysemy.TinyLog
 import Prometheus qualified
 import System.Logger qualified as Log
+import UnliftIO qualified
 import Wire.API.Password
 import Wire.CodeStore
 import Wire.CodeStore.Cassandra.Queries qualified as Cql
@@ -47,6 +56,9 @@ type EffectStack =
     Input ClientState,
     Input Hasql.Pool,
     Input (Either HttpsUrl (Map Text HttpsUrl)),
+    Resource,
+    Async,
+    Race,
     TinyLog,
     Embed IO,
     Final IO
@@ -60,15 +72,16 @@ migrateCodesLoop ::
   Prometheus.Counter ->
   Prometheus.Counter ->
   Prometheus.Counter ->
+  Prometheus.Vector Text Prometheus.Histogram ->
   IO ()
-migrateCodesLoop migOpts cassClient pgPool logger migCounter migFinished migFailed =
+migrateCodesLoop migOpts cassClient pgPool logger migCounter migFinished migFailed migDuration =
   migrationLoop
     logger
     "conversation codes"
     migFinished
     migFailed
     (interpreter cassClient pgPool logger "conversation codes")
-    (migrateAllCodes migOpts migCounter)
+    (migrateAllCodes migOpts migCounter migDuration)
 
 interpreter :: ClientState -> Hasql.Pool -> Log.Logger -> ByteString -> Sem EffectStack a -> IO (Int, a)
 interpreter cassClient pgPool logger name =
@@ -77,6 +90,9 @@ interpreter cassClient pgPool logger name =
     . loggerToTinyLog logger
     . mapLogger (Log.field "migration" (Log.val name) .)
     . raiseUnder
+    . interpretRace
+    . asyncToIOFinal
+    . resourceToIOFinal
     . runInputConst (Right mempty)
     . runInputConst pgPool
     . runInputConst cassClient
@@ -88,26 +104,60 @@ migrateAllCodes ::
     Member (Embed IO) r,
     Member (Input ClientState) r,
     Member TinyLog r,
-    Member (State Int) r
+    Member (State Int) r,
+    Member Resource r,
+    Member Race r
   ) =>
   MigrationOptions ->
   Prometheus.Counter ->
+  Prometheus.Vector Text Prometheus.Histogram ->
   ConduitM () Void (Sem r) ()
-migrateAllCodes migOpts migCounter = do
+migrateAllCodes migOpts migCounter migDuration = do
   lift $ info $ Log.msg (Log.val "migrateAllCodes")
   withCount (paginateSem Cql.selectAllCodes (paramsP LocalQuorum () migOpts.pageSize) x5)
     .| logRetrievedPage migOpts.pageSize id
-    .| C.mapM_ (traverse_ (\row@(key, _, _, _, _) -> handleErrors (toByteString' key) (migrateCodeRow migCounter row)))
+    .| C.mapM_ (traverse_ (\row@(key, _, _, _, _) -> handleErrors (toByteString' key) (migrateCodeRow migOpts migCounter migDuration row)))
 
 migrateCodeRow ::
   ( Member (Input (Either HttpsUrl (Map Text HttpsUrl))) r,
-    PGConstraints r
+    PGConstraints r,
+    Member TinyLog r,
+    Member Resource r,
+    Member Race r
   ) =>
+  MigrationOptions ->
   Prometheus.Counter ->
+  Prometheus.Vector Text Prometheus.Histogram ->
   (Key, Value, Int32, ConvId, Maybe Password) ->
   Sem r ()
-migrateCodeRow migCounter (k, v, ttl, cnv, mPw) =
+migrateCodeRow migOpts migCounter migDuration (k, v, ttl, cnv, mPw) =
   when (ttl > 0) $ do
     let (code, _) = toCode k (v, ttl, cnv, mPw)
-    Postgres.interpretCodeStoreToPostgres $ createCode code mPw
-    liftIO $ Prometheus.incCounter migCounter
+        keyText = T.pack (show k)
+    outcomeRef <- liftIO $ IORef.newIORef @Text "error"
+    bracket
+      (liftIO getCurrentTime)
+      (observeDuration migDuration outcomeRef)
+      ( const $ do
+          timeoutResult <- Conc.timeout (migOpts.timeout <$ handleTimeout) migOpts.timeout $ Postgres.interpretCodeStoreToPostgres $ createCode code mPw
+          case timeoutResult of
+            Left timedOutAfter -> do
+              markOutcome outcomeRef "timeout"
+              liftIO . UnliftIO.throwIO $ MigrationTimedOut keyText timedOutAfter
+            Right () -> do
+              markOutcome outcomeRef "success"
+              liftIO $ Prometheus.incCounter migCounter
+      )
+  where
+    handleTimeout =
+      err $
+        Log.msg (Log.val "conversation code migration timed out")
+          . Log.field "key" (show k)
+          . Log.field "timeout" (show migOpts.timeout)
+
+    markOutcome ref outcome = liftIO $ IORef.writeIORef ref outcome
+
+    observeDuration metric outcomeRef start = do
+      outcome <- liftIO $ IORef.readIORef outcomeRef
+      end <- liftIO getCurrentTime
+      liftIO $ Prometheus.withLabel metric outcome (`Prometheus.observe` realToFrac (diffUTCTime end start))