ntf server: use multiple APNS clients for each provider

Author: epoberezkin

src/Simplex/Messaging/Notifications/Server.hs

@@ -641,13 +641,15 @@ showServer' :: SMPServer -> Text
 showServer' = decodeLatin1 . strEncode . host
 
 pushNotification :: NtfPushServer -> Maybe T.Text -> OwnServer -> NtfTknRec -> PushNotification -> M ()
-pushNotification s srvHost_ isOwn tkn@NtfTknRec {ntfTknId, token = token@(DeviceToken pp _)} ntf =
+pushNotification s srvHost_ isOwn tkn@NtfTknRec {token = token@(DeviceToken pp _)} ntf =
   ifM
     (pushProviderAllowed token)
-    (getOrCreatePushWorker s (srvHost_, pp, hash (unEntityId ntfTknId) `mod` pushWorkersPerServer) isOwn >>= atomically . (`writeTBQueue` (tkn, ntf)))
+    (getOrCreatePushWorker s (srvHost_, pp, pushTokenShardId 8 tkn) isOwn >>= atomically . (`writeTBQueue` (tkn, ntf)))
     (logWarn "skipping disabled APNS test push provider")
-  where
-    pushWorkersPerServer = 8
+
+pushTokenShardId :: Int -> NtfTknRec -> Int
+pushTokenShardId pushWorkersPerServer NtfTknRec {ntfTknId} =
+  hash (unEntityId ntfTknId) `mod` pushWorkersPerServer
 
 pushProviderAllowed :: DeviceToken -> M Bool
 pushProviderAllowed (DeviceToken PPApnsTest _) = asks (allowTestPushProvider . config)
@@ -707,7 +709,7 @@ runPushWorker s srvHost_ isOwn q = forever $ do
       | otherwise = liftIO $ logError "bad notification token status"
     deliverNotification :: NtfPostgresStore -> PushProvider -> NtfTknRec -> PushNotification -> IO (Either PushProviderError ())
     deliverNotification st pp tkn@NtfTknRec {ntfTknId} ntf' = do
-      (deliver, clientVar) <- getPushClient s pp
+      (deliver, clientVar) <- getPushClient s pp tokenShardId
       runExceptT (deliver tkn ntf') >>= \case
         Right _ -> pure $ Right ()
         Left e -> case e of
@@ -720,11 +722,12 @@ runPushWorker s srvHost_ isOwn q = forever $ do
             err e
           PPPermanentError -> err e
       where
+        tokenShardId = pushTokenShardId 8 tkn
         retryDeliver :: PushClientVar -> Text -> IO (Either PushProviderError ())
         retryDeliver oldVar reason = do
           logWarn $ "retrying push (" <> tshow pp <> ", " <> tshow ntfTknId <> "): " <> reason
-          atomically $ removeSessVar oldVar pp (pushClients s)
-          (deliver, _) <- getPushClient s pp
+          atomically $ removeSessVar oldVar (pp, tokenShardId) (pushClients s)
+          (deliver, _) <- getPushClient s pp tokenShardId
           runExceptT (deliver tkn ntf') >>= \case
             Right _ -> pure $ Right ()
             Left e -> case e of

src/Simplex/Messaging/Notifications/Server/Env.hs

@@ -30,11 +30,9 @@ module Simplex.Messaging.Notifications.Server.Env
 import Control.Concurrent (ThreadId)
 import qualified Control.Exception as E
 import Control.Logger.Simple
-import Control.Monad
 import Control.Monad.Except
 import Control.Monad.Trans.Except
 import Crypto.Random
-import Data.Functor (($>))
 import Data.Int (Int64)
 import Simplex.Messaging.Agent.RetryInterval
 import Data.List.NonEmpty (NonEmpty)
@@ -66,7 +64,6 @@ import Simplex.Messaging.Transport.Credentials (genCredentials, tlsCredentials)
 import Simplex.Messaging.Transport.Server (AddHTTP, ServerCredentials, TransportServerConfig, loadFingerprint, loadServerCredential)
 import Simplex.Messaging.Util (liftEitherWith, tshow)
 import Simplex.Messaging.Util ()
-import System.Exit (exitFailure)
 import System.Mem.Weak (Weak)
 import UnliftIO.STM
 
@@ -177,7 +174,7 @@ data NtfPushServer = NtfPushServer
   { pushWorkers :: TMap (Maybe T.Text, PushProvider, Int) PushWorkerVar, -- Int is the worker shard
     pushWorkerSeq :: TVar Int,
     pushQSize :: Natural,
-    pushClients :: TMap PushProvider PushClientVar,
+    pushClients :: TMap (PushProvider, Int) PushClientVar,
     pushClientSeq :: TVar Int,
     apnsConfig :: APNSPushClientConfig
   }
@@ -203,27 +200,27 @@ newNtfPushServer pushQSize apnsConfig = do
 -- | Single-flight access to the per-provider push client with bounded retry.
 -- The returned PushClientVar is the handle retryDeliver passes to removeSessVar to evict
 -- this specific instance before re-fetching.
-getPushClient :: NtfPushServer -> PushProvider -> IO (PushProviderClient, PushClientVar)
-getPushClient s@NtfPushServer {apnsConfig = APNSPushClientConfig {reconnectInterval}} pp =
+getPushClient :: NtfPushServer -> PushProvider -> Int -> IO (PushProviderClient, PushClientVar)
+getPushClient s@NtfPushServer {apnsConfig = APNSPushClientConfig {reconnectInterval}} pp tokenShardId =
   withRetryIntervalCount reconnectInterval $ \n _delay loop -> do
     ts <- getCurrentTime
-    E.try (atomically (getSessVar (pushClientSeq s) pp (pushClients s) ts) >>= either (newPushClient s pp) waitForPushClient) >>= \case
+    E.try (atomically (getSessVar (pushClientSeq s) (pp, tokenShardId) (pushClients s) ts) >>= either (newPushClient s pp tokenShardId) waitForPushClient) >>= \case
       Right result -> pure result
       Left e
         | n < 2 -> do
             logError $ "getPushClient error (" <> tshow pp <> "): " <> tshow (e :: E.SomeException)
             loop
         | otherwise -> E.throwIO e
 
-newPushClient :: NtfPushServer -> PushProvider -> PushClientVar -> IO (PushProviderClient, PushClientVar)
-newPushClient NtfPushServer {pushClients, apnsConfig} pp v = do
+newPushClient :: NtfPushServer -> PushProvider -> Int -> PushClientVar -> IO (PushProviderClient, PushClientVar)
+newPushClient NtfPushServer {pushClients, apnsConfig} pp tokenShardId v = do
   r <- E.try $ case apnsProviderHost pp of
     Nothing -> pure $ \_ _ -> pure ()
     Just host -> apnsPushProviderClient <$> createAPNSPushClient host apnsConfig
   atomically $ do
     putTMVar (sessionVar v) r
     case r of
-      Left _ -> removeSessVar v pp pushClients
+      Left _ -> removeSessVar v (pp, tokenShardId) pushClients
       Right _ -> pure ()
   either E.throwIO (\c -> pure (c, v)) r