{-# OPTIONS_GHC -Wno-ambiguous-fields #-}

module Testlib.ResourcePool
  ( ResourcePool,
    BackendResource (..),
    DynamicBackendConfig (..),
    resourceServiceMap,
    backendResources,
    createBackendResourcePool,
    acquireResources,
  )
where

import Control.Concurrent
import Control.Monad.Catch
import Control.Monad.Codensity
import Control.Monad.IO.Class
import Data.Foldable (for_)
import Data.Functor
import Data.IORef
import qualified Data.Set as Set
import Data.String
import Data.Tuple
import Database.CQL.IO
import GHC.Stack (HasCallStack)
import Network.AMQP.Extended
import Network.RabbitMqAdmin
import System.IO
import Testlib.Ports
import Testlib.Types
import Prelude

resourceServiceMap :: BackendResource -> ServiceMap
resourceServiceMap :: BackendResource -> ServiceMap
resourceServiceMap BackendResource
resource =
  let g :: Service -> HostPort
g Service
srv = String -> Word16 -> HostPort
HostPort String
"127.0.0.1" (BackendResource -> forall a. Num a => Service -> a
berInternalServicePorts BackendResource
resource Service
srv)
   in ServiceMap
        { brig :: HostPort
brig = Service -> HostPort
g Service
Brig,
          backgroundWorker :: HostPort
backgroundWorker = Service -> HostPort
g Service
BackgroundWorker,
          cannon :: HostPort
cannon = Service -> HostPort
g Service
Cannon,
          cargohold :: HostPort
cargohold = Service -> HostPort
g Service
Cargohold,
          federatorInternal :: HostPort
federatorInternal = Service -> HostPort
g Service
FederatorInternal,
          federatorExternal :: HostPort
federatorExternal = String -> Word16 -> HostPort
HostPort String
"127.0.0.1" BackendResource
resource.berFederatorExternal,
          galley :: HostPort
galley = Service -> HostPort
g Service
Galley,
          gundeck :: HostPort
gundeck = Service -> HostPort
g Service
Gundeck,
          nginz :: HostPort
nginz = Service -> HostPort
g Service
Nginz,
          proxy :: HostPort
proxy = Service -> HostPort
g Service
WireProxy,
          spar :: HostPort
spar = Service -> HostPort
g Service
Spar,
          stern :: HostPort
stern = Service -> HostPort
g Service
Stern,
          wireServerEnterprise :: HostPort
wireServerEnterprise = Service -> HostPort
g Service
WireServerEnterprise,
          rabbitMqVHost :: Text
rabbitMqVHost = String -> Text
forall a. IsString a => String -> a
fromString BackendResource
resource.berVHost
        }

acquireResources :: forall m a. (Ord a, MonadIO m, MonadMask m, HasCallStack) => Int -> ResourcePool a -> Codensity m [a]
acquireResources :: forall (m :: * -> *) a.
(Ord a, MonadIO m, MonadMask m, HasCallStack) =>
Int -> ResourcePool a -> Codensity m [a]
acquireResources Int
n ResourcePool a
pool = (forall b. ([a] -> m b) -> m b) -> Codensity m [a]
forall k (m :: k -> *) a.
(forall (b :: k). (a -> m b) -> m b) -> Codensity m a
Codensity ((forall b. ([a] -> m b) -> m b) -> Codensity m [a])
-> (forall b. ([a] -> m b) -> m b) -> Codensity m [a]
forall a b. (a -> b) -> a -> b
$ \[a] -> m b
f -> m (Set a) -> (Set a -> m ()) -> (Set a -> m b) -> m b
forall (m :: * -> *) a c b.
(HasCallStack, MonadMask m) =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket m (Set a)
acquire Set a -> m ()
release ((Set a -> m b) -> m b) -> (Set a -> m b) -> m b
forall a b. (a -> b) -> a -> b
$ \Set a
s -> do
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ (a -> IO ()) -> Set a -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ResourcePool a
pool.onAcquire Set a
s
  [a] -> m b
f ([a] -> m b) -> [a] -> m b
forall a b. (a -> b) -> a -> b
$ Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
s
  where
    release :: Set.Set a -> m ()
    release :: Set a -> m ()
release Set a
s =
      IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        IORef (Set a) -> (Set a -> (Set a, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef ResourcePool a
pool.resources ((Set a -> (Set a, ())) -> IO ())
-> (Set a -> (Set a, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ (,()) (Set a -> (Set a, ())) -> (Set a -> Set a) -> Set a -> (Set a, ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set a
s
        QSemN -> Int -> IO ()
signalQSemN ResourcePool a
pool.sem (Set a -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Set a
s)

    acquire :: m (Set.Set a)
    acquire :: m (Set a)
acquire = IO (Set a) -> m (Set a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Set a) -> m (Set a)) -> IO (Set a) -> m (Set a)
forall a b. (a -> b) -> a -> b
$ do
      QSemN -> Int -> IO ()
waitQSemN ResourcePool a
pool.sem Int
n
      IORef (Set a) -> (Set a -> (Set a, Set a)) -> IO (Set a)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef ResourcePool a
pool.resources ((Set a -> (Set a, Set a)) -> IO (Set a))
-> (Set a -> (Set a, Set a)) -> IO (Set a)
forall a b. (a -> b) -> a -> b
$ (Set a, Set a) -> (Set a, Set a)
forall a b. (a, b) -> (b, a)
swap ((Set a, Set a) -> (Set a, Set a))
-> (Set a -> (Set a, Set a)) -> Set a -> (Set a, Set a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Set a -> (Set a, Set a)
forall a. Int -> Set a -> (Set a, Set a)
Set.splitAt Int
n

createBackendResourcePool :: [BackendResource] -> RabbitMqAdminOpts -> ClientState -> IO (ResourcePool BackendResource)
createBackendResourcePool :: [BackendResource]
-> RabbitMqAdminOpts
-> ClientState
-> IO (ResourcePool BackendResource)
createBackendResourcePool [BackendResource]
resources RabbitMqAdminOpts
rabbitmq ClientState
cassClient =
  let cleanupBackend :: BackendResource -> IO ()
      cleanupBackend :: BackendResource -> IO ()
cleanupBackend BackendResource
resource = do
        RabbitMqAdminOpts -> BackendResource -> IO ()
deleteAllRabbitMQQueues RabbitMqAdminOpts
rabbitmq BackendResource
resource
        ClientState -> Client () -> IO ()
forall (m :: * -> *) a. MonadIO m => ClientState -> Client a -> m a
runClient ClientState
cassClient (Client () -> IO ()) -> Client () -> IO ()
forall a b. (a -> b) -> a -> b
$ BackendResource -> Client ()
deleteAllDynamicBackendConfigs BackendResource
resource
   in QSemN
-> IORef (Set BackendResource)
-> (BackendResource -> IO ())
-> ResourcePool BackendResource
forall a. QSemN -> IORef (Set a) -> (a -> IO ()) -> ResourcePool a
ResourcePool
        (QSemN
 -> IORef (Set BackendResource)
 -> (BackendResource -> IO ())
 -> ResourcePool BackendResource)
-> IO QSemN
-> IO
     (IORef (Set BackendResource)
      -> (BackendResource -> IO ()) -> ResourcePool BackendResource)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO QSemN
newQSemN ([BackendResource] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [BackendResource]
resources)
        IO
  (IORef (Set BackendResource)
   -> (BackendResource -> IO ()) -> ResourcePool BackendResource)
-> IO (IORef (Set BackendResource))
-> IO ((BackendResource -> IO ()) -> ResourcePool BackendResource)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set BackendResource -> IO (IORef (Set BackendResource))
forall a. a -> IO (IORef a)
newIORef ([BackendResource] -> Set BackendResource
forall a. Ord a => [a] -> Set a
Set.fromList [BackendResource]
resources)
        IO ((BackendResource -> IO ()) -> ResourcePool BackendResource)
-> IO (BackendResource -> IO ())
-> IO (ResourcePool BackendResource)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (BackendResource -> IO ()) -> IO (BackendResource -> IO ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BackendResource -> IO ()
cleanupBackend

deleteAllRabbitMQQueues :: RabbitMqAdminOpts -> BackendResource -> IO ()
deleteAllRabbitMQQueues :: RabbitMqAdminOpts -> BackendResource -> IO ()
deleteAllRabbitMQQueues RabbitMqAdminOpts
opts BackendResource
resource = do
  AdminAPI (AsClientT IO)
client <- RabbitMqAdminOpts -> IO (AdminAPI (AsClientT IO))
mkRabbitMqAdminClientEnv RabbitMqAdminOpts
opts {vHost = fromString resource.berVHost}
  Page Queue
queuesPage <- AdminAPI (AsClientT IO)
-> AsClientT IO
   :- ("api"
       :> ("queues"
           :> (Capture "vhost" Text
               :> (QueryParam' '[Required, Strict] "name" Text
                   :> (QueryParam' '[Required, Strict] "use_regex" Bool
                       :> (QueryParam' '[Required, Strict] "page_size" Int
                           :> (QueryParam' '[Required, Strict] "page" Int
                               :> Get '[JSON] (Page Queue))))))))
forall {k} (route :: k).
AdminAPI route
-> route
   :- ("api"
       :> ("queues"
           :> (Capture "vhost" Text
               :> (QueryParam' '[Required, Strict] "name" Text
                   :> (QueryParam' '[Required, Strict] "use_regex" Bool
                       :> (QueryParam' '[Required, Strict] "page_size" Int
                           :> (QueryParam' '[Required, Strict] "page" Int
                               :> Get '[JSON] (Page Queue))))))))
listQueuesByVHost AdminAPI (AsClientT IO)
client (String -> Text
forall a. IsString a => String -> a
fromString BackendResource
resource.berVHost) (String -> Text
forall a. IsString a => String -> a
fromString String
"") Bool
False Int
100 Int
1
  [Queue] -> (Queue -> IO NoContent) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ Page Queue
queuesPage.items ((Queue -> IO NoContent) -> IO ())
-> (Queue -> IO NoContent) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Queue
queue ->
    AdminAPI (AsClientT IO)
-> AsClientT IO
   :- ("api"
       :> ("queues"
           :> (Capture "vhost" Text
               :> (Capture "queue" Text :> DeleteNoContent))))
forall {k} (route :: k).
AdminAPI route
-> route
   :- ("api"
       :> ("queues"
           :> (Capture "vhost" Text
               :> (Capture "queue" Text :> DeleteNoContent))))
deleteQueue AdminAPI (AsClientT IO)
client (String -> Text
forall a. IsString a => String -> a
fromString BackendResource
resource.berVHost) Queue
queue.name

deleteAllDynamicBackendConfigs :: BackendResource -> Client ()
deleteAllDynamicBackendConfigs :: BackendResource -> Client ()
deleteAllDynamicBackendConfigs BackendResource
resource = PrepQuery W () () -> QueryParams () -> Client ()
forall (m :: * -> *) a (q :: * -> * -> * -> *).
(MonadClient m, Tuple a, RunQ q) =>
q W a () -> QueryParams a -> m ()
write PrepQuery W () ()
cql (Consistency -> () -> QueryParams ()
forall a. Consistency -> a -> QueryParams a
defQueryParams Consistency
LocalQuorum ())
  where
    cql :: PrepQuery W () ()
    cql :: PrepQuery W () ()
cql = String -> PrepQuery W () ()
forall a. IsString a => String -> a
fromString (String -> PrepQuery W () ()) -> String -> PrepQuery W () ()
forall a b. (a -> b) -> a -> b
$ String
"TRUNCATE " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> BackendResource
resource.berBrigKeyspace String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
".federation_remotes"

backendResources :: [DynamicBackendConfig] -> [BackendResource]
backendResources :: [DynamicBackendConfig] -> [BackendResource]
backendResources [DynamicBackendConfig]
dynConfs =
  ([DynamicBackendConfig] -> [Int] -> [(DynamicBackendConfig, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [DynamicBackendConfig]
dynConfs [Int
1 ..])
    [(DynamicBackendConfig, Int)]
-> ((DynamicBackendConfig, Int) -> BackendResource)
-> [BackendResource]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> ( \(DynamicBackendConfig
dynConf, Int
i) ->
            let name :: BackendName
name = Int -> BackendName
DynamicBackend Int
i
             in BackendResource
                  { berName :: BackendName
berName = BackendName
name,
                    berBrigKeyspace :: String
berBrigKeyspace = String
"brig_test_dyn_" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i,
                    berGalleyKeyspace :: String
berGalleyKeyspace = String
"galley_test_dyn_" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i,
                    berSparKeyspace :: String
berSparKeyspace = String
"spar_test_dyn_" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i,
                    berGundeckKeyspace :: String
berGundeckKeyspace = String
"gundeck_test_dyn_" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i,
                    berElasticsearchIndex :: String
berElasticsearchIndex = String
"directory_dyn_" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_test",
                    berPostgresqlDBName :: String
berPostgresqlDBName = String
"dyn-" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i,
                    berFederatorInternal :: Word16
berFederatorInternal = PortNamespace -> Int -> Word16
forall a. Num a => PortNamespace -> Int -> a
portForDyn (Service -> PortNamespace
ServiceInternal Service
FederatorInternal) Int
i,
                    berFederatorExternal :: Word16
berFederatorExternal = DynamicBackendConfig
dynConf.federatorExternalPort,
                    berDomain :: String
berDomain = DynamicBackendConfig
dynConf.domain,
                    berAwsUserJournalQueue :: String
berAwsUserJournalQueue = String
"integration-user-events" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. (Show a, Num a) => a -> String
suffix Int
i String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
".fifo",
                    berAwsPrekeyTable :: String
berAwsPrekeyTable = String
"integration-brig-prekeys" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. (Show a, Num a) => a -> String
suffix Int
i,
                    berAwsS3Bucket :: String
berAwsS3Bucket = String
"dummy-bucket" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. (Show a, Num a) => a -> String
suffix Int
i,
                    berAwsQueueName :: String
berAwsQueueName = String
"integration-gundeck-events" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. (Show a, Num a) => a -> String
suffix Int
i,
                    berBrigInternalEvents :: String
berBrigInternalEvents = String
"integration-brig-events-internal" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. (Show a, Num a) => a -> String
suffix Int
i,
                    berEmailSMSSesQueue :: String
berEmailSMSSesQueue = String
"integration-brig-events" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. (Show a, Num a) => a -> String
suffix Int
i,
                    berEmailSMSEmailSender :: String
berEmailSMSEmailSender = String
"backend-integration" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. (Show a, Num a) => a -> String
suffix Int
i String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"@wire.com",
                    berGalleyJournal :: String
berGalleyJournal = String
"integration-team-events" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. (Show a, Num a) => a -> String
suffix Int
i String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
".fifo",
                    berVHost :: String
berVHost = DynamicBackendConfig
dynConf.domain,
                    berNginzSslPort :: Word16
berNginzSslPort = PortNamespace -> Int -> Word16
forall a. Num a => PortNamespace -> Int -> a
portForDyn PortNamespace
NginzSSL Int
i,
                    berNginzHttp2Port :: Word16
berNginzHttp2Port = PortNamespace -> Int -> Word16
forall a. Num a => PortNamespace -> Int -> a
portForDyn PortNamespace
NginzHttp2 Int
i,
                    berInternalServicePorts :: forall a. Num a => Service -> a
berInternalServicePorts = BackendName -> Service -> a
forall a. Num a => BackendName -> Service -> a
internalServicePorts BackendName
name,
                    berEnableService :: Service -> Bool
berEnableService = Bool -> Service -> Bool
forall a b. a -> b -> a
const Bool
True,
                    berMlsPrivateKeyPaths :: Value
berMlsPrivateKeyPaths = DynamicBackendConfig
dynConf.mlsPrivateKeyPaths
                  }
        )
  where
    suffix :: (Show a, Num a) => a -> String
    suffix :: forall a. (Show a, Num a) => a -> String
suffix a
i = a -> String
forall a. Show a => a -> String
show (a -> String) -> a -> String
forall a b. (a -> b) -> a -> b
$ a
i a -> a -> a
forall a. Num a => a -> a -> a
+ a
2