-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2026 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Test.Migration.Util where

import Control.Applicative
import Control.Concurrent (threadDelay)
import Control.Monad.Reader
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import GHC.Stack
import SetupHelpers hiding (deleteUser)
import Testlib.Prelude
import Text.Regex.TDFA ((=~))
import UnliftIO

waitForMigration :: (HasCallStack) => String -> String -> App ()
waitForMigration :: HasCallStack => String -> String -> App ()
waitForMigration String
domain String
metricName =
  App () -> (() -> App ()) -> Maybe () -> App ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe App ()
failWithContext () -> App ()
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe () -> App ()) -> App (Maybe ()) -> App ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> App () -> App (Maybe ())
forall (m :: * -> *) a.
MonadUnliftIO m =>
Int -> m a -> m (Maybe a)
timeout Int
30_000_000 App ()
go
  where
    failWithContext :: App ()
failWithContext = do
      String -> Service -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> Service -> App Response
getMetrics String
domain Service
BackgroundWorker App Response -> (Response -> App ()) -> App ()
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp -> do
        Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
        String -> App ()
forall a. HasCallStack => String -> App a
assertFailure String
"Timed out waiting for postgresql migration"
    go :: App ()
go = do
      metrics <-
        String -> Service -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> Service -> App Response
getMetrics String
domain Service
BackgroundWorker App Response -> (Response -> App Text) -> App Text
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp -> do
          Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
          Text -> App Text
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> App Text) -> Text -> App Text
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
Text.decodeUtf8 Response
resp.body
      let (_, _, _, finishedMatches) :: (Text, Text, Text, [Text]) = (metrics =~ Text.pack (metricName <> "\\ ([0-9]+\\.[0-9]+)$"))
      when (finishedMatches /= [Text.pack "1.0"]) $ do
        liftIO $ threadDelay 100_000
        go

assertMigrationSuccessful :: (HasCallStack) => String -> String -> App ()
assertMigrationSuccessful :: HasCallStack => String -> String -> App ()
assertMigrationSuccessful String
domain String
failedMetricName = do
  String -> Service -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> Service -> App Response
getMetrics String
domain Service
BackgroundWorker App Response -> (Response -> App ()) -> App ()
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp -> do
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
    let metrics :: Text
metrics = ByteString -> Text
Text.decodeUtf8 Response
resp.body
        (Text
_, Text
_, Text
_, [Text]
failedMatches) :: (Text, Text, Text, [Text]) = (Text
metrics Text -> Text -> (Text, Text, Text, [Text])
forall source source1 target.
(RegexMaker Regex CompOption ExecOption source,
 RegexContext Regex source1 target) =>
source1 -> source -> target
=~ String -> Text
Text.pack (String
failedMetricName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"\\ ([0-9]+\\.[0-9]+)$"))
    [Text]
failedMatches [Text] -> [Text] -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` [String -> Text
Text.pack String
"0.0"]

-- | Retry a request on transient errors, bounded by a cumulative cap so
-- genuine failures still surface.
--
-- Retryable statuses:
-- * 500 — transient server errors under load: during the user-migration phase
--   the background worker holds per-user Postgres advisory locks and brig's
--   dual Cassandra/Postgres reads time out acquiring their shared locks; spar
--   rethrows such brig errors on SCIM reads. Also covers spar-side Cassandra
--   timeouts under burst load.
-- * 422 — observed transiently during migration test runs.
-- * 521 / 525 / 533 — federation cold-start errors when a freshly
--   (re)started dynamic backend is hit with a parallel fan-out (connection
--   refused, SSL, unreachable / unexpected federation response).
--
-- Safe to wrap unconditionally only for idempotent reads — a retried mutation
-- can duplicate effects. Test.Migration.Conversation deliberately also wraps
-- federation-facing mutations (conversation create, membership changes)
-- during migration cold-start, where the transient 521/525/533 window would
-- otherwise flake.
retryTransient :: App Response -> App Response
retryTransient :: App Response -> App Response
retryTransient App Response
action = Int -> Int -> App Response
go (Int
0 :: Int) (Int
100_000 :: Int)
  where
    go :: Int -> Int -> App Response
go Int
spent Int
delay = do
      resp <- App Response
action
      if resp.status `elem` [500, 422, 521, 525, 533] && spent < maxCumulative
        then do
          liftIO $ threadDelay delay
          go (spent + delay) (min 2_000_000 (delay * 2))
        else pure resp
    maxCumulative :: Int
maxCumulative = Int
30_000_000