natskell 1.0.0.1 → 1.1.0.0
raw patch · 43 files changed
+4257/−627 lines, 43 filesPVP ok
version bump matches the API change (PVP)
API changes (from Hackage documentation)
+ API: ConnectionClosed :: ClientExitReason -> ConnectionState
+ API: ConnectionClosing :: ClientExitReason -> ConnectionState
+ API: ConnectionConnected :: ConnectionState
+ API: ConnectionConnecting :: ConnectionState
+ API: ConnectionReconnecting :: ConnectionState
+ API: connectionState :: Client -> IO ConnectionState
+ API: data ConnectionState
+ API: withFlushTimeout :: NominalDiffTime -> FlushOption
+ API: withPingTimeout :: NominalDiffTime -> PingOption
+ Client: ConnectionEventClosed :: ClientExitReason -> ConnectionEvent
+ Client: ConnectionEventDisconnected :: ConnectionEvent
+ Client: ConnectionEventReconnected :: ConnectionEvent
+ Client: ServerErrorAuthentication :: ServerErrorKind
+ Client: ServerErrorConnection :: ServerErrorKind
+ Client: ServerErrorPermission :: ServerErrorKind
+ Client: ServerErrorProtocol :: ServerErrorKind
+ Client: ServerErrorResourceLimit :: ServerErrorKind
+ Client: ServerErrorUnknown :: ServerErrorKind
+ Client: data ConnectionEvent
+ Client: data ServerErrorKind
+ Client: serverErrorKind :: ServerError -> ServerErrorKind
+ Client: withConnectionEventHandler :: (ConnectionEvent -> IO ()) -> ConfigOption
+ Client: withServerErrorHandler :: (ServerError -> IO ()) -> ConfigOption
+ JetStream.API.Consumer: withConsumerBackoff :: [NominalDiffTime] -> ConsumerConfigOption
+ JetStream.API.Consumer: withConsumerMaxRequestBatch :: Int -> ConsumerConfigOption
+ JetStream.API.Message: data NakDelay
+ JetStream.API.Message: nakDelay :: NominalDiffTime -> Maybe NakDelay
+ JetStream.API.Stream: withMaxMessageSize :: Int32 -> StreamConfigOption
Files
- client/API.hs +8/−0
- client/Client.hs +5/−0
- client/Client/Implementation.hs +560/−200
- internal/Client/API.hs +17/−1
- internal/Lib/Exception.hs +18/−0
- internal/Lib/Logger.hs +6/−2
- internal/Lib/WorkerPool.hs +15/−6
- internal/Plumbing/Network/Connection.hs +2/−0
- internal/Plumbing/Network/Connection/Core.hs +22/−13
- internal/Plumbing/Network/Connection/Tcp.hs +16/−20
- internal/Plumbing/Network/Connection/Tls.hs +13/−14
- internal/Plumbing/Network/Connection/Types.hs +1/−0
- internal/Plumbing/Network/ConnectionAPI.hs +2/−0
- internal/Plumbing/Pipeline/Streaming/Parser.hs +20/−4
- internal/Plumbing/Queue/API.hs +49/−3
- internal/Plumbing/Queue/TransactionalQueue.hs +89/−20
- internal/Policy/Engine.hs +433/−149
- internal/Policy/Handshake/Nats.hs +71/−36
- internal/Policy/Router/Nats.hs +40/−12
- internal/Policy/State/Store.hs +526/−53
- internal/Policy/State/Types.hs +67/−16
- internal/Policy/Subscription/Store.hs +61/−9
- internal/Policy/Subscription/Types.hs +21/−1
- jetstream/JetStream/Consumer.hs +8/−3
- jetstream/JetStream/Consumer/API.hs +4/−0
- jetstream/JetStream/Consumer/Types.hs +90/−1
- jetstream/JetStream/Message.hs +50/−4
- jetstream/JetStream/Message/API.hs +10/−0
- jetstream/JetStream/Message/Types.hs +31/−0
- jetstream/JetStream/Stream.hs +20/−4
- jetstream/JetStream/Stream/API.hs +2/−0
- jetstream/JetStream/Stream/Types.hs +29/−3
- jetstream/JetStream/Types.hs +7/−1
- natskell.cabal +5/−1
- test/Integration/ClientSpec.hs +646/−10
- test/PublicAPI/Main.hs +56/−7
- test/Unit/ConnectionSpec.hs +806/−26
- test/Unit/JetStream/ConsumerSpec.hs +97/−2
- test/Unit/JetStream/MessageSpec.hs +155/−1
- test/Unit/JetStream/ProtocolSpec.hs +1/−0
- test/Unit/JetStream/StreamSpec.hs +76/−0
- test/Unit/StreamingSpec.hs +9/−1
- test/Unit/SubscriptionStoreSpec.hs +93/−4
client/API.hs view
@@ -10,6 +10,7 @@ , Subscription , subscriptionSid , NatsError (..)+ , ConnectionState (..) , Subject , Payload , Headers@@ -29,6 +30,7 @@ , newInbox , ping , flush+ , connectionState , reset , close , subject@@ -42,11 +44,14 @@ , withHeaders , withRequestTimeout , withRequestHeaders+ , withPingTimeout+ , withFlushTimeout ) where import Client.API ( Client , CloseOption+ , ConnectionState (..) , FlushOption , Headers , Message@@ -62,6 +67,7 @@ , Subscription , UnsubscribeOption , close+ , connectionState , flush , headers , newInbox@@ -77,7 +83,9 @@ , subscribeOnce , subscriptionSid , unsubscribe+ , withFlushTimeout , withHeaders+ , withPingTimeout , withQueueGroup , withReplyTo , withRequestHeaders
client/Client.hs view
@@ -32,6 +32,8 @@ , withMessageLimit , withPendingDeliveryLimits , withErrorHandler+ , withServerErrorHandler+ , withConnectionEventHandler , withBufferLimit , withExitAction , LogLevel (..)@@ -52,8 +54,11 @@ , TLSCertData , TLSConfig (..) , ClientExitReason (..)+ , ConnectionEvent (..) , ServerError+ , ServerErrorKind (..) , serverErrorReason+ , serverErrorKind , ConnectError (..) , ConnectAttemptError (..) , ConnectFailure (..)
client/Client/Implementation.hs view
@@ -1,3 +1,4 @@+{-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} -- | High-level client implementation for NATS.@@ -34,6 +35,8 @@ , withMessageLimit , withPendingDeliveryLimits , withErrorHandler+ , withServerErrorHandler+ , withConnectionEventHandler , withBufferLimit , withExitAction , LogLevel (..)@@ -54,8 +57,11 @@ , TLSCertData , TLSConfig (..) , ClientExitReason (..)+ , ConnectionEvent (..) , ServerError+ , ServerErrorKind (..) , serverErrorReason+ , serverErrorKind , ConnectError (..) , ConnectAttemptError (..) , ConnectFailure (..)@@ -92,22 +98,32 @@ , Subscription (..) , UnsubscribeConfig (..) )-import Control.Concurrent (forkIO)+import Control.Concurrent+ ( forkIOWithUnmask+ , killThread+ , myThreadId+ ) import Control.Concurrent.STM import Control.Exception ( SomeException- , displayException+ , finally , mask+ , mask_ , onException ) import Control.Monad (forM_, void, when) import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as BC import Data.Char (isAsciiUpper)-import Data.Maybe (fromMaybe)+import Data.Maybe (fromMaybe, isJust) import Data.Time.Clock (NominalDiffTime) import Data.Version (showVersion)-import Engine (closeClient, resetClient, runEngine)+import Engine+ ( closeClient+ , interruptConnection+ , resetClient+ , runEngine+ ) import Lib.CallOption (CallOption, applyCallOptions) import Lib.Logger ( LogEntry (..)@@ -119,52 +135,81 @@ , renderLogEntry ) import Network.Connection (connectionApi)-import Network.ConnectionAPI (newConn)+import Network.ConnectionAPI (ConnectionAPI, newConn)+import qualified Network.ConnectionAPI as Connection import Parser.Attoparsec (parserApiWithMessageLimit) import qualified Paths_natskell as Package import Pipeline.Broadcasting (broadcastingApi) import Pipeline.Streaming (streamingApi) import Publish (defaultPublishConfig) import Publish.Config (PublishConfig (..))-import Queue.API (QueueItem (QueueItem))+import Queue.API+ ( QueueItem (QueueBatch, QueueConnectionScoped, QueueItem, QueuePayloadBound)+ , TryEnqueueResult (..)+ ) import Queue.TransactionalQueue (newQueue) import State.Store ( ClientState+ , PingResult (..)+ , PublishEnqueueResult (..) , config , enqueue+ , enqueueOnGenerationTracked+ , enqueuePublishOnConnected+ , enqueuePublishOnConnectedGenerationTracked+ , isManagedThread+ , markClosed , newClientState , nextInbox , nextSid- , pushPingAction- , readServerInfo+ , readConnectionGeneration , readStatus+ , registerManagedThread+ , registerPingWaiterAndEnqueue , runClient+ , setClosing , setConnectName+ , startCallbackWorker+ , stopCallbackWorker+ , tryEnqueue+ , tryEnqueueOnGeneration+ , unregisterManagedThread+ , waitForCallbackWorker , waitForClosed+ , waitForConnected+ , waitForConnectionGenerationLoss , waitForInitialConnection , waitForNotRunning+ , withSubscriptionGate ) import State.Types ( ClientConfig (..) , ClientExitReason (..)- , ClientStatus (..) , ConnectAttemptError (..) , ConnectError (..) , ConnectFailure (..)+ , ConnectionEvent (..)+ , ConnectionState (..) , ServerError+ , ServerErrorKind (..) , TLSCertData , TLSConfig (..) , TLSPrivateKey , TLSPublicKey , defaultTLSConfig+ , serverErrorKind , serverErrorReason ) import Subscription.Store ( SubscriptionStore+ , awaitCallbackDrain , awaitNoTrackedExpiries+ , closeStore+ , enqueueControl , hasTrackedExpiries , newSubscriptionStore , register+ , registerWithDispatchHooks , startExpiryWorker , startWorkers , unregister@@ -172,11 +217,13 @@ import Subscription.Types ( PendingLimits (..) , SubscribeConfig (..)+ , SubscriptionKind (..) , SubscriptionMeta (SubscriptionMeta) , defaultPendingLimits+ , isOneShotSubscription )+import System.Timeout (timeout) import qualified Types.Connect as Connect-import qualified Types.Info as Info import qualified Types.Msg as Msg import Types.Ping (Ping (..)) import qualified Types.Pub as Pub@@ -185,20 +232,26 @@ import Validators.Validators (validate) data ClientOptions = ClientOptions- { optionConnectConfig :: Connect.Connect- , optionAuth :: Auth- , optionTlsConfig :: Maybe TLSConfig- , optionLoggerConfig :: LoggerConfig- , optionConnectionAttempts :: Int+ { optionConnectConfig :: Connect.Connect+ , optionAuth :: Auth+ , optionTlsConfig :: Maybe TLSConfig+ , optionLoggerConfig :: LoggerConfig+ , optionConnectionAttempts :: Int , optionConnectTimeoutMicros :: Int- , optionCallbackConcurrency :: Int- , optionMessageLimit :: Int- , optionPendingLimits :: PendingLimits- , optionErrorHandler :: NatsError -> IO ()- , optionExitAction :: ClientExitReason -> IO ()- , optionConnectOptions :: [(String, Int)]+ , optionCallbackConcurrency :: Int+ , optionMessageLimit :: Int+ , optionPendingLimits :: PendingLimits+ , optionErrorHandler :: NatsError -> IO ()+ , optionServerErrorHandler :: ServerError -> IO ()+ , optionConnectionEventHandler :: ConnectionEvent -> IO ()+ , optionExitAction :: ClientExitReason -> IO ()+ , optionConnectOptions :: [(String, Int)] } +data SubscriptionQueueResult = SubscriptionQueued | SubscriptionReconnect++data UnsubscribeQueueResult = UnsubscribeQueued | UnsubscribeReconnect | UnsubscribeReset+ -- | An opaque NATS server endpoint. -- -- Keeping this representation private allows endpoint schemes and transports@@ -234,7 +287,7 @@ -- | Compatibility connection entry point using raw @(host, port)@ tuples. newClient :: [(String, Int)] -> [ConfigOption] -> IO (Either ConnectError Client)-newClient servers configOptions = do+newClient servers configOptions = mask $ \restoreInitialWait -> do loggerConfig' <- defaultLogger ctx <- newLogContext let defaultOptions = applyCallOptions configOptions ClientOptions@@ -248,6 +301,8 @@ , optionMessageLimit = 1024 * 1024 , optionPendingLimits = defaultPendingLimits , optionErrorHandler = const (pure ())+ , optionServerErrorHandler = const (pure ())+ , optionConnectionEventHandler = const (pure ()) , optionExitAction = const (pure ()) , optionConnectOptions = servers }@@ -260,6 +315,8 @@ , connectConfig = optionConnectConfig defaultOptions , loggerConfig = optionLoggerConfig defaultOptions , tlsConfig = optionTlsConfig defaultOptions+ , serverErrorHandler = optionServerErrorHandler defaultOptions+ , connectionEventHandler = optionConnectionEventHandler defaultOptions , exitAction = optionExitAction defaultOptions , connectOptions = optionConnectOptions defaultOptions }@@ -276,37 +333,76 @@ setConnectName clientState (Connect.name (optionConnectConfig defaultOptions)) logStaticConfiguration clientState defaultOptions - startWorkers+ callbackThreads <- startWorkers (callbackConcurrency clientConfig) store (do waitForClosed clientState+ awaitCallbackDrain store awaitNoTrackedExpiries store) (handleCallbackError clientState) - startExpiryWorker store $+ expiryThread <- startExpiryWorker store $ shouldStopExpiryWorker clientState store - void . forkIO $- runEngine- connectionApi- streamingApi- broadcastingApi- (parserApiWithMessageLimit (messageLimit clientConfig))- clientState- store- configuredAuth+ mapM_ (registerManagedThread clientState) (expiryThread : callbackThreads)+ callbackWorker <- startCallbackWorker clientState + engineDone <- newEmptyTMVarIO+ engineThread <- forkIOWithUnmask $ \unmask ->+ do+ threadId <- myThreadId+ registerManagedThread clientState threadId+ unmask+ ( runEngine+ connectionApi+ streamingApi+ broadcastingApi+ (parserApiWithMessageLimit (messageLimit clientConfig))+ clientState+ store+ configuredAuth+ )+ `finally` do+ unregisterManagedThread clientState threadId+ atomically (void (tryPutTMVar engineDone ()))++ let auxiliaryThreads = expiryThread : callbackThreads+ stopAuxiliaryThreads excluded =+ mapM_ killThread (filter (/= excluded) auxiliaryThreads)+ finishClose excluded = do+ atomically (readTMVar engineDone)+ Connection.close connectionApi conn+ atomically (awaitCallbackDrain store)+ stopAuxiliaryThreads excluded+ atomically (waitForCallbackWorker callbackWorker)+ scheduleFinish excluded = do+ _ <- forkIOWithUnmask $ \unmask -> unmask (finishClose excluded)+ pure ()+ closeAndFinish = mask $ \restore -> do+ current <- myThreadId+ closeClient connectionApi clientState store+ `onException` scheduleFinish current+ managed <- isManagedThread clientState current+ if managed+ then scheduleFinish current+ else restore (finishClose current) `onException` scheduleFinish current+ let client = Client { publish = \subject payload publishOptions -> do let cfg = applyCallOptions publishOptions defaultPublishConfig- publishClient clientState subject payload cfg+ publishClient+ clientState+ (enqueueControl store . optionErrorHandler defaultOptions)+ subject+ payload+ cfg , subscribe = \subject subscribeOptions callback -> do let cfg = applyCallOptions subscribeOptions defaultSubscribeConfig- subscribeClient clientState store False subject cfg callback+ subscribeClient clientState store StandardSubscription subject cfg callback , subscribeOnce = \subject subscribeOptions callback -> do let cfg = applyCallOptions subscribeOptions defaultSubscribeConfig- subscribeClient clientState store True subject cfg callback+ subscribeClient clientState store OneShotSubscription subject cfg callback , request = \subject payload requestOptions -> do let cfg = applyCallOptions requestOptions defaultRequestConfig requestClient clientState store subject payload cfg@@ -315,21 +411,50 @@ UnsubscribeConfig -> unsubscribeClient clientState store subscription , newInbox = nextInbox clientState , ping = \options ->- case applyCallOptions options PingConfig of- PingConfig -> flushClient clientState+ case applyCallOptions options defaultPingConfig of+ PingConfig -> flushClient connectionApi clientState defaultRoundTripTimeout+ PingConfigTimeout timeoutSeconds ->+ flushClient connectionApi clientState timeoutSeconds , flush = \options ->- case applyCallOptions options FlushConfig of- FlushConfig -> flushClient clientState+ case applyCallOptions options defaultFlushConfig of+ FlushConfig -> flushClient connectionApi clientState defaultRoundTripTimeout+ FlushConfigTimeout timeoutSeconds ->+ flushClient connectionApi clientState timeoutSeconds+ , connectionState = readStatus clientState , reset = \options -> case applyCallOptions options ResetConfig of ResetConfig -> resetClient connectionApi clientState store , close = \options -> case applyCallOptions options CloseConfig of- CloseConfig -> closeClient connectionApi clientState store+ CloseConfig -> closeAndFinish } - initialResult <- atomically (waitForInitialConnection clientState)- pure (client <$ initialResult)+ let abortInitialWait = do+ setClosing clientState ExitClosedByUser+ closeStore store+ void (interruptConnection connectionApi clientState)+ killThread engineThread+ atomically (readTMVar engineDone)+ Connection.close connectionApi conn+ void (markClosed clientState ExitClosedByUser)+ stopCallbackWorker callbackWorker+ current <- myThreadId+ stopAuxiliaryThreads current+ initialResult <-+ restoreInitialWait+ (atomically $+ waitForInitialConnection clientState+ `orElse` (Left (ConnectAttemptsExhausted []) <$ readTMVar engineDone))+ `onException` abortInitialWait+ case initialResult of+ Left err -> do+ atomically (readTMVar engineDone)+ Connection.close connectionApi conn+ stopCallbackWorker callbackWorker+ current <- myThreadId+ stopAuxiliaryThreads current+ pure (Left err)+ Right () -> pure (Right client) type ConfigOption = CallOption ClientOptions @@ -433,8 +558,9 @@ withConnectionAttempts attempts config = config { optionConnectionAttempts = max 1 attempts } --- | Set the maximum time for INFO, TLS, CONNECT, and PONG negotiation on each--- server attempt. Values below one microsecond are clamped to one.+-- | Set the total time budget for TCP acquisition, INFO, TLS, CONNECT, PONG,+-- and reconnect resubscription/readiness on each server attempt. Values below+-- one microsecond are clamped to one. withConnectTimeoutMicros :: Int -> ConfigOption withConnectTimeoutMicros timeoutMicros config = config { optionConnectTimeoutMicros = max 1 timeoutMicros }@@ -468,6 +594,19 @@ withErrorHandler handler config = config { optionErrorHandler = handler } +-- | Receive protocol @-ERR@ values on a bounded, dedicated serial worker.+-- Protocol processing, including PONG handling, does not wait for the handler.+withServerErrorHandler :: (ServerError -> IO ()) -> ConfigOption+withServerErrorHandler handler config =+ config { optionServerErrorHandler = handler }++-- | Receive completed disconnect, reconnect, and close transitions on the+-- same serial callback worker as server errors. Lifecycle delivery is reserved+-- independently from the bounded server-error backlog.+withConnectionEventHandler :: (ConnectionEvent -> IO ()) -> ConfigOption+withConnectionEventHandler handler config =+ config { optionConnectionEventHandler = handler }+ -- | Compatibility alias for 'withMessageLimit'. withBufferLimit :: Int -> ConfigOption withBufferLimit = withMessageLimit@@ -504,6 +643,15 @@ defaultRequestConfig :: RequestConfig defaultRequestConfig = RequestConfig 2 Nothing +defaultPingConfig :: PingConfig+defaultPingConfig = PingConfig++defaultFlushConfig :: FlushConfig+defaultFlushConfig = FlushConfig++defaultRoundTripTimeout :: NominalDiffTime+defaultRoundTripTimeout = 2+ logStaticConfiguration :: ClientState -> ClientOptions -> IO () logStaticConfiguration client options = runClient client $ do@@ -520,9 +668,9 @@ logMessage Warn "tls certificate verification is disabled" handleCallbackError :: ClientState -> SomeException -> IO ()-handleCallbackError client err =+handleCallbackError client _ = runClient client $- logMessage Error ("callback failed: " ++ displayException err)+ logMessage Error "callback failed" handleSlowConsumer :: ClientState -> (NatsError -> IO ()) -> IO () handleSlowConsumer client handler = do@@ -536,8 +684,8 @@ tracked <- hasTrackedExpiries store pure $ case status of- Closed _ -> not tracked- _ -> False+ ConnectionClosed _ -> not tracked+ _ -> False toMessage :: Msg.Msg -> Message toMessage msg =@@ -551,11 +699,12 @@ publishClient :: ClientState+ -> (NatsError -> IO ()) -> Msg.Subject -> Msg.Payload -> PublishConfig -> IO (Either NatsError ())-publishClient client subject messagePayload cfg = do+publishClient client errorHandler subject messagePayload cfg = do runClient client $ logMessage Debug ("publishing to subject: " ++ show subject) let publishMessage =@@ -566,70 +715,81 @@ , Pub.replyTo = publishReplyTo cfg , Pub.headers = publishHeaders cfg }- validation <- canPublish client publishMessage+ actualSize = Pub.messageSize publishMessage+ queuedPublish =+ QueuePayloadBound+ actualSize+ (errorHandler . NatsPayloadTooLarge actualSize)+ (QueueItem publishMessage)+ validation <- validatePublish client publishMessage case validation of Left err -> pure (Left err) Right () -> do- enqueue client (QueueItem publishMessage)- pure (Right ())+ enqueueResult <- enqueuePublishOnConnected client actualSize queuedPublish+ publishEnqueueResult client actualSize enqueueResult -canPublish :: ClientState -> Pub.Pub -> IO (Either NatsError ())-canPublish client publishMessage =+validatePublish :: ClientState -> Pub.Pub -> IO (Either NatsError ())+validatePublish client publishMessage = case validate publishMessage of Left reason -> do runClient client $ logMessage Error ("rejecting invalid publish: " ++ BC.unpack reason) pure (Left (NatsValidationError reason)) Right () -> do- statusResult <- runningResult client- case statusResult of- Left err -> pure (Left err)- Right () -> do- serverInfo <- readServerInfo client- let actual = Pub.messageSize publishMessage- clientMaximum = messageLimit (config client)- maximumSize =- maybe- clientMaximum- (min clientMaximum . Info.max_payload)- serverInfo- if actual > maximumSize- then do- runClient client $- logMessage Error- ( "rejecting publish: message size "- ++ show actual- ++ " exceeds effective limit "- ++ show maximumSize- )- pure (Left (NatsPayloadTooLarge actual maximumSize))- else- pure (Right ())+ let actualSize = Pub.messageSize publishMessage+ clientMaximum = messageLimit (config client)+ if actualSize > clientMaximum+ then rejectPayloadTooLarge client actualSize clientMaximum+ else pure (Right ()) +publishEnqueueResult+ :: ClientState+ -> Int+ -> PublishEnqueueResult+ -> IO (Either NatsError ())+publishEnqueueResult client actualSize enqueueResult =+ case enqueueResult of+ PublishEnqueued _ -> pure (Right ())+ PublishTooLarge maximumSize ->+ rejectPayloadTooLarge client actualSize maximumSize+ PublishConnectionClosed reason ->+ pure (Left (NatsConnectionClosed reason))++rejectPayloadTooLarge :: ClientState -> Int -> Int -> IO (Either NatsError a)+rejectPayloadTooLarge client actualSize maximumSize = do+ runClient client $+ logMessage Error+ ( "rejecting publish: message size "+ ++ show actualSize+ ++ " exceeds effective limit "+ ++ show maximumSize+ )+ pure (Left (NatsPayloadTooLarge actualSize maximumSize))+ subscribeClient :: ClientState -> SubscriptionStore- -> Bool+ -> SubscriptionKind -> Msg.Subject -> SubscribeConfig -> (Message -> IO ()) -> IO (Either NatsError Subscription)-subscribeClient client store isReply subject cfg callback =- subscribeRawClient client store isReply subject cfg (maybe (pure ()) (callback . toMessage))+subscribeClient client store subscriptionKind subject cfg callback =+ subscribeRawClient client store subscriptionKind subject cfg (maybe (pure ()) (callback . toMessage)) subscribeRawClient :: ClientState -> SubscriptionStore- -> Bool+ -> SubscriptionKind -> Msg.Subject -> SubscribeConfig -> (Maybe Msg.Msg -> IO ()) -> IO (Either NatsError Subscription)-subscribeRawClient client store isReply subject cfg callback = do+subscribeRawClient client store subscriptionKind subject cfg callback = do subscribeRawClientWithOverflow client store- isReply+ subscriptionKind subject cfg callback@@ -638,47 +798,103 @@ subscribeRawClientWithOverflow :: ClientState -> SubscriptionStore- -> Bool+ -> SubscriptionKind -> Msg.Subject -> SubscribeConfig -> (Maybe Msg.Msg -> IO ()) -> IO () -> IO (Either NatsError Subscription)-subscribeRawClientWithOverflow client store isReply subject cfg callback onDropped = do- runClient client $- logMessage Debug ("subscribing to subject: " ++ show subject)- statusResult <- runningResult client- let queueGroup = subscribeQueueGroup cfg- sid <- nextSid client- let- subscriptionMessage =- Sub.Sub- { Sub.subject = subject- , Sub.queueGroup = queueGroup- , Sub.sid = sid- }- case validate subscriptionMessage of- Left reason -> do- runClient client $- logMessage Error ("rejecting invalid subscription: " ++ BC.unpack reason)- pure (Left (NatsValidationError reason))- Right () -> case statusResult of- Left err -> pure (Left err)+subscribeRawClientWithOverflow client store subscriptionKind subject cfg callback onDropped =+ mask $ \restore -> do+ runClient client $+ logMessage Debug ("subscribing to subject: " ++ show subject)+ let queueGroup = subscribeQueueGroup cfg+ sid <- nextSid client+ registeredGeneration <- newEmptyTMVarIO+ committedGeneration <- newEmptyTMVarIO+ let subscriptionMessage =+ Sub.Sub+ { Sub.subject = subject+ , Sub.queueGroup = queueGroup+ , Sub.sid = sid+ }+ meta = SubscriptionMeta subject queueGroup subscriptionKind+ subscription = Subscription sid+ cleanup =+ cleanupResumableSubscription+ client+ store+ registeredGeneration+ committedGeneration+ subscription+ deliver maybeMessage = do+ case maybeMessage of+ Nothing -> cleanup+ Just _ -> pure ()+ callback maybeMessage+ commands =+ QueueItem subscriptionMessage :+ [ QueueItem+ Unsub.Unsub+ { Unsub.sid = sid+ , Unsub.maxMsg = Just 1+ }+ | isOneShotSubscription meta+ ]+ queuedCommands = QueueConnectionScoped (QueueBatch commands)+ case validate subscriptionMessage of+ Left reason -> do+ runClient client $+ logMessage Error ("rejecting invalid subscription: " ++ BC.unpack reason)+ pure (Left (NatsValidationError reason)) Right () -> do- let meta =- SubscriptionMeta subject queueGroup isReply- register store sid meta cfg callback onDropped- enqueue client $- QueueItem- subscriptionMessage- when isReply $- enqueue client- (QueueItem- Unsub.Unsub- { Unsub.sid = sid- , Unsub.maxMsg = Just 1- })- pure (Right (Subscription sid))+ let waitForServerRegistration = do+ statusResult <- runningResult client+ case statusResult of+ Left err -> cleanup >> pure (Left err)+ Right () -> pure (Right subscription)+ acquireGate = do+ statusResult <- runningResult client+ case statusResult of+ Left err -> pure (Left err)+ Right () -> do+ gated <- withSubscriptionGate client $ do+ status <- readStatus client+ case status of+ ConnectionConnected -> do+ generation <- readConnectionGeneration client+ atomically (void (tryPutTMVar registeredGeneration generation))+ register store sid meta cfg deliver onDropped+ enqueueResult <-+ enqueueOnGenerationTracked+ client+ generation+ committedGeneration+ queuedCommands+ case enqueueResult of+ Right _ ->+ pure (Just (Right SubscriptionQueued))+ Left _ -> do+ nextStatus <- readStatus client+ case nextStatus of+ ConnectionClosing reason ->+ pure (Just (Left (NatsConnectionClosed reason)))+ ConnectionClosed reason ->+ pure (Just (Left (NatsConnectionClosed reason)))+ _ ->+ pure (Just (Right SubscriptionReconnect))+ _ -> pure Nothing+ maybe acquireGate pure gated+ subscribeOperation = do+ queueResult <- acquireGate+ case queueResult of+ Left err -> cleanup >> pure (Left err)+ Right SubscriptionQueued -> do+ runClient client $+ logMessage Debug "subscription commands queued"+ pure (Right subscription)+ Right SubscriptionReconnect -> waitForServerRegistration+ restore subscribeOperation `onException` cleanup requestClient :: ClientState@@ -687,10 +903,25 @@ -> Msg.Payload -> RequestConfig -> IO (Either NatsError Message)-requestClient client store requestSubject requestPayload cfg =+requestClient client store requestSubject requestPayload cfg = do+ timedResult <-+ timeout+ (durationMicros (requestTimeout cfg))+ (requestBeforeDeadline client store requestSubject requestPayload cfg)+ pure (fromMaybe (Left NatsRequestTimedOut) timedResult)++requestBeforeDeadline+ :: ClientState+ -> SubscriptionStore+ -> Msg.Subject+ -> Msg.Payload+ -> RequestConfig+ -> IO (Either NatsError Message)+requestBeforeDeadline client store requestSubject requestPayload cfg = mask $ \restore -> do response <- newEmptyTMVarIO- deadline <- registerDelay (durationMicros (requestTimeout cfg))+ accepted <- newEmptyTMVarIO+ committedGeneration <- newEmptyTMVarIO inbox <- nextInbox client let subscriptionConfig = SubscribeConfig Nothing Nothing deliver Nothing = atomically (void (tryPutTMVar response (Left NatsRequestTimedOut)))@@ -701,53 +932,103 @@ then Left NatsNoResponders else Right message rejectSlowConsumer =- atomically (void (tryPutTMVar response (Left NatsSlowConsumer)))- subscriptionResult <-- subscribeRawClientWithOverflow- client- store- True- inbox- subscriptionConfig- deliver- rejectSlowConsumer- case subscriptionResult of- Left err -> pure (Left err)- Right subscription -> do- let publishConfig = PublishConfig (requestHeaders cfg) (Just inbox)- cleanup = void (unsubscribeClient client store subscription)- publishResult <- publishClient client requestSubject requestPayload publishConfig- case publishResult of- Left err -> cleanup >> pure (Left err)+ void (tryPutTMVar response (Left NatsSlowConsumer))+ sid <- nextSid client+ let subscription = Subscription sid+ subscriptionMessage =+ Sub.Sub+ { Sub.subject = inbox+ , Sub.queueGroup = Nothing+ , Sub.sid = sid+ }+ unsubscribeMessage =+ Unsub.Unsub+ { Unsub.sid = sid+ , Unsub.maxMsg = Just 1+ }+ publishMessage =+ Pub.Pub+ { Pub.subject = requestSubject+ , Pub.payload =+ if BS.null requestPayload then Nothing else Just requestPayload+ , Pub.replyTo = Just inbox+ , Pub.headers = requestHeaders cfg+ }+ commands = QueueConnectionScoped . QueueBatch $+ [ QueueItem subscriptionMessage+ , QueueItem unsubscribeMessage+ , QueueItem publishMessage+ ]+ meta = SubscriptionMeta inbox Nothing RequestReplySubscription+ case validate subscriptionMessage of+ Left reason -> pure (Left (NatsValidationError reason))+ Right () -> do+ publishValidation <- validatePublish client publishMessage+ case publishValidation of+ Left err -> pure (Left err) Right () -> do- result <- restore (awaitRequest client deadline response) `onException` cleanup- cleanup- pure result+ registerWithDispatchHooks+ store+ sid+ meta+ subscriptionConfig+ (void (tryPutTMVar accepted ()))+ rejectSlowConsumer+ deliver+ (pure ())+ let cleanupCommitted = do+ committed <- atomically (tryReadTMVar committedGeneration)+ case committed of+ Nothing -> unregister store sid+ Just generation ->+ cleanupRequestSubscription client store generation subscription+ let publishSize = Pub.messageSize publishMessage+ queued <- restore+ ( enqueuePublishOnConnectedGenerationTracked+ client+ committedGeneration+ publishSize+ commands+ )+ `onException` cleanupCommitted+ case queued of+ PublishConnectionClosed reason -> do+ unregister store sid+ pure (Left (NatsConnectionClosed reason))+ PublishTooLarge maximumSize -> do+ unregister store sid+ rejectPayloadTooLarge client publishSize maximumSize+ PublishEnqueued generation -> do+ let cleanup =+ cleanupRequestSubscription client store generation subscription+ result <- restore (awaitRequest client generation accepted response)+ `onException` cleanup+ cleanup+ pure result awaitRequest :: ClientState- -> TVar Bool+ -> Int+ -> TMVar () -> TMVar (Either NatsError Message) -> IO (Either NatsError Message)-awaitRequest client deadline response = do- outcome <- atomically $- (Just <$> readTMVar response)- `orElse` (do- expired <- readTVar deadline- check expired- pure (Just (Left NatsRequestTimedOut)))- `orElse` (Nothing <$ waitForNotRunning client)- case outcome of- Just result -> pure result- Nothing -> do- status <- readStatus client- pure (Left (closedError status))+awaitRequest client generation accepted response =+ atomically $+ readTMVar response+ `orElse`+ do+ acceptedReply <- tryReadTMVar accepted+ case acceptedReply of+ Just () -> retry+ Nothing ->+ Left . NatsConnectionClosed+ <$> waitForConnectionGenerationLoss client generation durationMicros :: NominalDiffTime -> Int durationMicros duration = fromInteger (min (toInteger (maxBound :: Int)) micros) where- micros = max 0 (floor (realToFrac duration * (1000000 :: Double)))+ micros = max 0 (floor (toRational duration * 1000000)) isNoResponders :: Message -> Bool isNoResponders message =@@ -767,58 +1048,137 @@ -> SubscriptionStore -> Subscription -> IO (Either NatsError ())-unsubscribeClient client store (Subscription sid) = do- runClient client $- logMessage Debug ("unsubscribing SID: " ++ show sid)- unregister store sid- statusResult <- runningResult client- case statusResult of- Left err -> pure (Left err)- Right () -> do- enqueue client $- QueueItem+unsubscribeClient client store (Subscription sid) =+ mask $ \restore -> do+ runClient client $+ logMessage Debug ("unsubscribing SID: " ++ show sid)+ let command = QueueConnectionScoped . QueueItem $ Unsub.Unsub { Unsub.sid = sid , Unsub.maxMsg = Nothing }- pure (Right ())+ unregisterAndTry = do+ unregister store sid+ tryEnqueue client command >>= \case+ TryEnqueued -> pure (Right UnsubscribeQueued)+ TryQueueClosed -> pure (Right UnsubscribeReconnect)+ TryQueueFull -> pure (Right UnsubscribeReset)+ queueResult <- withSubscriptionGate client $ do+ status <- readStatus client+ case status of+ ConnectionConnected -> do+ unregister store sid+ enqueueResult <-+ restore (enqueue client command)+ `onException` interruptConnection connectionApi client+ case enqueueResult of+ Right () -> pure (Right UnsubscribeQueued)+ Left _ -> do+ nextStatus <- readStatus client+ case nextStatus of+ ConnectionClosing reason ->+ pure (Left (NatsConnectionClosed reason))+ ConnectionClosed reason ->+ pure (Left (NatsConnectionClosed reason))+ _ -> pure (Right UnsubscribeReconnect)+ ConnectionConnecting -> unregisterAndTry+ ConnectionReconnecting -> unregisterAndTry+ ConnectionClosing reason -> do+ unregister store sid+ pure (Left (NatsConnectionClosed reason))+ ConnectionClosed reason -> do+ unregister store sid+ pure (Left (NatsConnectionClosed reason))+ case queueResult of+ Left err -> pure (Left err)+ Right UnsubscribeQueued -> pure (Right ())+ Right UnsubscribeReconnect -> pure (Right ())+ Right UnsubscribeReset -> do+ interruptConnection connectionApi client+ pure (Right ()) -pingClient :: ClientState -> IO () -> IO ()-pingClient client action = do- runClient client $- logMessage Debug "sending ping to server"- pushPingAction client action- enqueue client (QueueItem Ping)+cleanupResumableSubscription+ :: ClientState+ -> SubscriptionStore+ -> TMVar Int+ -> TMVar Int+ -> Subscription+ -> IO ()+cleanupResumableSubscription client store registeredGeneration committedGeneration (Subscription sid) =+ mask_ $ do+ resetRequired <- withSubscriptionGate client $ do+ unregister store sid+ status <- readStatus client+ case status of+ ConnectionConnected -> do+ generation <- readConnectionGeneration client+ registered <- atomically (tryReadTMVar registeredGeneration)+ committed <- atomically (tryReadTMVar committedGeneration)+ let serverMayKnow =+ isJust committed || maybe False (< generation) registered+ if serverMayKnow+ then do+ result <- tryEnqueueOnGeneration client generation unsubscribeCommand+ pure (result == TryQueueFull)+ else pure False+ _ -> pure False+ when resetRequired (void (interruptConnection connectionApi client))+ where+ unsubscribeCommand = QueueConnectionScoped . QueueItem $+ Unsub.Unsub { Unsub.sid = sid, Unsub.maxMsg = Nothing } -flushClient :: ClientState -> IO (Either NatsError ())-flushClient client = do- statusResult <- runningResult client- case statusResult of- Left err -> pure (Left err)- Right () -> do- ponged <- newEmptyTMVarIO- pingClient client (atomically (void (tryPutTMVar ponged ())))- outcome <- atomically $- (True <$ readTMVar ponged)- `orElse` (False <$ waitForNotRunning client)- if outcome- then pure (Right ())- else do- status <- readStatus client- pure (Left (closedError status))+cleanupRequestSubscription :: ClientState -> SubscriptionStore -> Int -> Subscription -> IO ()+cleanupRequestSubscription client store generation (Subscription sid) = mask_ $ do+ unregister store sid+ result <- tryEnqueueOnGeneration client generation . QueueConnectionScoped . QueueItem $+ Unsub.Unsub { Unsub.sid = sid, Unsub.maxMsg = Nothing }+ case result of+ TryEnqueued -> pure ()+ TryQueueClosed -> pure ()+ TryQueueFull -> void (interruptConnection connectionApi client) +flushClient :: ConnectionAPI -> ClientState -> NominalDiffTime -> IO (Either NatsError ())+flushClient connectionApi' client timeoutSeconds =+ mask $ \restore -> do+ let reset = interruptConnection connectionApi' client+ timedResult <-+ restore (timeout (durationMicros timeoutSeconds) (awaitPong client))+ `onException` reset+ case timedResult of+ Nothing -> reset >> pure (Left NatsRequestTimedOut)+ Just result -> pure result+ where+ awaitPong client' = do+ waiter <- newEmptyTMVarIO+ registered <- registerPingWaiterAndEnqueue+ client'+ waiter+ (QueueConnectionScoped (QueueItem Ping))+ case registered of+ Left reason -> pure (Left (NatsConnectionClosed reason))+ Right _ -> do+ result <- atomically $+ (Just <$> readTMVar waiter)+ `orElse` (Nothing <$ waitForNotRunning client')+ case result of+ Just PingReceived -> pure (Right ())+ Just PingConnectionLost ->+ Left . closedError <$> readStatus client'+ Nothing -> do+ status <- readStatus client'+ pure (Left (closedError status))+ runningResult :: ClientState -> IO (Either NatsError ()) runningResult client = do- status <- readStatus client+ result <- atomically (waitForConnected client) pure $- case status of- Running -> Right ()- Closing reason -> Left (NatsConnectionClosed reason)- Closed reason -> Left (NatsConnectionClosed reason)+ case result of+ Left reason -> Left (NatsConnectionClosed reason)+ Right () -> Right () -closedError :: ClientStatus -> NatsError+closedError :: ConnectionState -> NatsError closedError status = case status of- Closing reason -> NatsConnectionClosed reason- Closed reason -> NatsConnectionClosed reason- Running -> NatsConnectionClosed ExitResetRequested+ ConnectionClosing reason -> NatsConnectionClosed reason+ ConnectionClosed reason -> NatsConnectionClosed reason+ _ -> NatsConnectionClosed ExitResetRequested
internal/Client/API.hs view
@@ -10,6 +10,7 @@ , Subscription (..) , subscriptionSid , NatsError (..)+ , ConnectionState (..) , Subject , Payload , Headers@@ -33,13 +34,15 @@ , withHeaders , withRequestTimeout , withRequestHeaders+ , withPingTimeout+ , withFlushTimeout ) where import qualified Data.ByteString as BS import Data.Time.Clock (NominalDiffTime) import Lib.CallOption (CallOption) import Publish.Config (PublishConfig (..))-import State.Types (ClientExitReason)+import State.Types (ClientExitReason, ConnectionState (..)) import Subscription.Types (SubscribeConfig (..)) import Types.Msg (Headers, Payload, SID, Subject) @@ -65,6 +68,8 @@ -- ^ Round trip a PING through the server. , flush :: [FlushOption] -> IO (Either NatsError ()) -- ^ Wait until all previously buffered writes reach the server.+ , connectionState :: IO ConnectionState+ -- ^ Read the client's current connection lifecycle state. , reset :: [ResetOption] -> IO () -- ^ Reset the client connection state. , close :: [CloseOption] -> IO ()@@ -119,10 +124,13 @@ data UnsubscribeConfig = UnsubscribeConfig type UnsubscribeOption = CallOption UnsubscribeConfig +-- Internal configuration constructors are not covered by public API stability. data PingConfig = PingConfig+ | PingConfigTimeout NominalDiffTime type PingOption = CallOption PingConfig data FlushConfig = FlushConfig+ | FlushConfigTimeout NominalDiffTime type FlushOption = CallOption FlushConfig data ResetConfig = ResetConfig@@ -155,3 +163,11 @@ -- | Set headers on an outgoing request. withRequestHeaders :: Headers -> RequestOption withRequestHeaders messageHeaders cfg = cfg { requestHeaders = Just messageHeaders }++-- | Set the maximum time to wait for a ping response.+withPingTimeout :: NominalDiffTime -> PingOption+withPingTimeout timeout _ = PingConfigTimeout (max 0 timeout)++-- | Set the maximum time to wait for a flush response.+withFlushTimeout :: NominalDiffTime -> FlushOption+withFlushTimeout timeout _ = FlushConfigTimeout (max 0 timeout)
+ internal/Lib/Exception.hs view
@@ -0,0 +1,18 @@+{-# LANGUAGE TypeApplications #-}++module Lib.Exception+ ( trySync+ ) where++import Control.Exception++-- | Convert synchronous failures to values without swallowing cancellation.+trySync :: IO a -> IO (Either SomeException a)+trySync action = do+ result <- try @SomeException action+ case result of+ Left err ->+ case fromException err :: Maybe SomeAsyncException of+ Just _ -> throwIO err+ Nothing -> pure (Left err)+ Right value -> pure (Right value)
internal/Lib/Logger.hs view
@@ -14,6 +14,7 @@ import Control.Monad import Control.Monad.Reader import Data.Maybe (catMaybes, fromMaybe)+import Lib.Exception (trySync) import Lib.Logger.Types import Lib.LoggerAPI (LoggerAPI (LoggerAPI)) @@ -31,13 +32,13 @@ instance MonadLogger AppM where logEntry entry = AppM . ReaderT $ \(LoggerEnv (LoggerConfig minLvl out lock) _) ->- when (leLevel entry >= minLvl) (withLogLock lock (out entry))+ when (leLevel entry >= minLvl) (safeLog lock (out entry)) getLogContext = AppM . ReaderT $ \(LoggerEnv _ ctxVar) -> readTVarIO ctxVar logMessage lvl msg = AppM . ReaderT $ \(LoggerEnv (LoggerConfig minLvl out lock) ctxVar) -> when (lvl >= minLvl) $ do ctx <- readTVarIO ctxVar- withLogLock lock $+ safeLog lock $ out LogEntry { leLevel = lvl , leMessage = msg@@ -45,6 +46,9 @@ , leConnectName = lcConnectName ctx , leServer = lcServer ctx }++safeLog :: TMVar () -> IO () -> IO ()+safeLog lock action = void (trySync (withLogLock lock action)) renderLogEntry :: LogEntry -> String renderLogEntry entry =
internal/Lib/WorkerPool.hs view
@@ -2,7 +2,7 @@ ( startWorkerPool ) where -import Control.Concurrent (forkIO)+import Control.Concurrent (ThreadId, forkIOWithUnmask) import Control.Concurrent.STM ( STM , TQueue@@ -12,13 +12,19 @@ , orElse , readTQueue )-import Control.Exception (SomeException, try)-import Control.Monad (replicateM_, void)+import Control.Exception+ ( SomeAsyncException+ , SomeException+ , fromException+ , throwIO+ , try+ )+import Control.Monad (replicateM) -startWorkerPool :: Int -> TQueue (IO ()) -> STM () -> (SomeException -> IO ()) -> IO ()+startWorkerPool :: Int -> TQueue (IO ()) -> STM () -> (SomeException -> IO ()) -> IO [ThreadId] startWorkerPool concurrency queue stopSignal onError = do let workerCount = max 1 concurrency- replicateM_ workerCount (void (forkIO worker))+ replicateM workerCount (forkIOWithUnmask (\unmask -> unmask worker)) where worker = do let loop = do@@ -35,7 +41,10 @@ Just job -> do result <- (try job :: IO (Either SomeException ())) case result of- Left err -> onError err+ Left err ->+ case fromException err :: Maybe SomeAsyncException of+ Just _ -> throwIO err+ Nothing -> onError err Right () -> pure () loop loop
internal/Plumbing/Network/Connection.hs view
@@ -13,6 +13,7 @@ { newConn = Core.newConn , reader = ReaderAPI { readData = Core.readData+ , bufferRead = Core.bufferRead , closeReader = Core.closeReader , openReader = Core.openReader }@@ -24,6 +25,7 @@ } , open = Core.openConn , close = Core.closeConn+ , abort = Core.abortConn , connectTcp = Tcp.connectTcp , configure = Tls.configureTransport }
internal/Plumbing/Network/Connection/Core.hs view
@@ -1,5 +1,3 @@-{-# LANGUAGE TypeApplications #-}- module Network.Connection.Core ( ReadError , WriteError@@ -16,20 +14,22 @@ , openWriter , openConn , closeConn+ , abortConn , pointTransport , currentTransport , bufferRead , enableReadWorker ) where -import Control.Concurrent (forkIO)+import Control.Concurrent (forkIOWithUnmask) import Control.Concurrent.STM-import Control.Exception+import Control.Exception (SomeException, finally, onException) import Control.Monad import Data.ByteString (ByteString) import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LBS import Data.Maybe (isJust, isNothing)+import Lib.Exception (trySync) import Network.Connection.Types newConn :: IO Conn@@ -56,7 +56,9 @@ else do writeTVar (readWorkerRunning conn) True return True- when shouldStart . void $ forkIO (readWorkerLoop conn)+ when shouldStart $+ void (forkIOWithUnmask (\unmask -> unmask (readWorkerLoop conn)))+ `onException` atomically (writeTVar (readWorkerRunning conn) False) readWorkerLoop :: Conn -> IO () readWorkerLoop conn = do@@ -76,7 +78,7 @@ case current of Nothing -> return () Just currentTransport' -> do- result <- try @SomeException (transportRead currentTransport' readChunkSize)+ result <- trySync (transportRead currentTransport' readChunkSize) shouldContinue <- enqueueReadResult conn result case result of Left _ -> return ()@@ -137,8 +139,8 @@ return $ Right chunk else do resultVar <- newEmptyTMVarIO- _ <- forkIO $ do- result <- try @SomeException (transportRead currentTransport' n)+ _ <- forkIOWithUnmask $ \unmask -> unmask $ do+ result <- trySync (transportRead currentTransport' n) atomically . void $ tryPutTMVar resultVar result result <- atomically $ (Left "Read operation is blocked" <$ readTMVar (readBlock conn))@@ -169,7 +171,7 @@ case current of Nothing -> return $ Left "Transport not initialized" Just currentTransport' -> do- result <- try @SomeException $ do+ result <- trySync $ do transportWrite currentTransport' bytes transportFlush currentTransport' case result of@@ -186,7 +188,7 @@ case current of Nothing -> return $ Left "Transport not initialized" Just currentTransport' -> do- result <- try @SomeException $ do+ result <- trySync $ do transportWriteLazy currentTransport' bytes transportFlush currentTransport' case result of@@ -211,11 +213,18 @@ closeConn conn = do closeReader conn closeWriter conn- current <- currentTransport conn+ current <- atomically (tryTakeTMVar (transport conn)) case current of- Nothing -> return ()- Just currentTransport' -> void $ try @SomeException (transportClose currentTransport')+ Nothing -> return ()+ Just currentTransport' -> void $ trySync (transportClose currentTransport') waitForReadWorkerStopped conn++abortConn :: Conn -> IO ()+abortConn conn = do+ current <- atomically (tryTakeTMVar (transport conn))+ case current of+ Nothing -> pure ()+ Just currentTransport' -> void $ trySync (transportAbort currentTransport') openConn :: Conn -> IO () openConn conn = do
internal/Plumbing/Network/Connection/Tcp.hs view
@@ -1,11 +1,10 @@-{-# LANGUAGE TypeApplications #-}- module Network.Connection.Tcp ( connectTcp ) where -import Control.Exception+import Control.Exception (displayException, mask, onException) import qualified Data.ByteString.Lazy as LBS+import Lib.Exception (trySync) import Network.Connection.Core (pointTransport) import qualified Network.Connection.Tls as Tls import Network.Connection.Types (Conn, Transport (..))@@ -21,25 +20,22 @@ , transportWriteLazy = NSB.sendMany sock . LBS.toChunks , transportFlush = pure () , transportClose = NS.close sock+ , transportAbort = NS.close sock , transportUpgrade = Just (Tls.upgradeTcp sock) } -openTcpTransport :: String -> Int -> IO (Either String Transport)-openTcpTransport host port = do- result <- try @SomeException $ do- (sock, _) <- TCP.connectSock host (show port)- NS.setSocketOption sock NS.NoDelay 1- NS.setSocketOption sock NS.Cork 0- pure (tcpTransport sock)- case result of- Left err -> return (Left (show err))- Right transport -> return (Right transport)- connectTcp :: Conn -> String -> Int -> IO (Either String ()) connectTcp conn host port = do- result <- openTcpTransport host port- case result of- Left err -> return (Left err)- Right transport -> do- pointTransport conn transport- return (Right ())+ result <- trySync (mask $ \restore -> do+ (sock, _) <- TCP.connectSock host (show port)+ let transport = tcpTransport sock+ restore+ (do+ NS.setSocketOption sock NS.NoDelay 1+ NS.setSocketOption sock NS.Cork 0)+ `onException` NS.close sock+ pointTransport conn transport `onException` transportClose transport)+ pure $+ case result of+ Left err -> Left (displayException err)+ Right () -> Right ()
internal/Plumbing/Network/Connection/Tls.hs view
@@ -1,5 +1,3 @@-{-# LANGUAGE TypeApplications #-}- module Network.Connection.Tls ( configureTransport , upgradeTcp@@ -12,6 +10,7 @@ import Data.Maybe (fromMaybe) import Data.X509.CertificateStore (makeCertificateStore) import Data.X509.Memory (readSignedObjectFromMemory)+import Lib.Exception (trySync) import Network.Connection.Core ( bufferRead , currentTransport@@ -37,29 +36,31 @@ , transportWriteLazy = TLS.sendData ctx , transportFlush = TLS.contextFlush ctx , transportClose = do- void $ try @SomeException (TLS.bye ctx)- void $ try @SomeException (TLS.contextClose ctx)+ void $ trySync (TLS.bye ctx)+ void $ trySync (TLS.contextClose ctx)+ , transportAbort = void (trySync (TLS.contextClose ctx)) , transportUpgrade = Nothing } upgradeTcp :: NS.Socket -> TLS.ClientParams -> IO (Either String Transport)-upgradeTcp sock params = do+upgradeTcp sock params = mask $ \restore -> do let backend = TLS.Backend { TLS.backendSend = NSB.sendAll sock , TLS.backendRecv = NSB.recv sock , TLS.backendFlush = pure () , TLS.backendClose = NS.close sock }- result <- try @SomeException $ do- ctx <- TLS.contextNew backend params- TLS.handshake ctx+ result <- trySync $ do+ ctx <- TLS.contextNew backend params `onException` NS.close sock+ restore (TLS.handshake ctx)+ `onException` void (trySync (TLS.contextClose ctx)) pure (tlsTransport ctx) case result of Left err -> return $ Left (show err) Right transport -> return $ Right transport upgradeToTLS :: Conn -> TLS.ClientParams -> IO (Either String ())-upgradeToTLS conn params = do+upgradeToTLS conn params = mask_ $ do current <- currentTransport conn case current of Nothing -> return $ Left "Transport not initialized"@@ -72,6 +73,7 @@ Left err -> return $ Left err Right newTransport -> do pointTransport conn newTransport+ `onException` transportClose newTransport return $ Right () upgradeToTLSWithConfig :: Conn -> String -> TLSConfig -> IO (Either String ())@@ -102,7 +104,7 @@ buildTlsParams :: String -> TLSConfig -> IO (Either String TLS.ClientParams) buildTlsParams host tlsConfig = do- result <- try @SomeException $ do+ result <- trySync $ do systemStore <- getSystemCertificateStore let verificationHost = fromMaybe host (tlsServerName tlsConfig) base = TLS.defaultParamsClient verificationHost (BC.pack verificationHost)@@ -143,8 +145,5 @@ } } case result of- Left err ->- case fromException err :: Maybe SomeAsyncException of- Just _ -> throwIO err- Nothing -> pure (Left (displayException err))+ Left err -> pure (Left (displayException err)) Right params -> pure (Right params)
internal/Plumbing/Network/Connection/Types.hs view
@@ -21,6 +21,7 @@ , transportWriteLazy :: LBS.ByteString -> IO () , transportFlush :: IO () , transportClose :: IO ()+ , transportAbort :: IO () , transportUpgrade :: Maybe (TLS.ClientParams -> IO (Either String Transport)) }
internal/Plumbing/Network/ConnectionAPI.hs view
@@ -19,6 +19,7 @@ data ReaderAPI reader = ReaderAPI { readData :: reader -> Int -> IO (Either ReadError ByteString)+ , bufferRead :: reader -> ByteString -> IO () , closeReader :: reader -> IO () , openReader :: reader -> IO () }@@ -36,6 +37,7 @@ , writer :: WriterAPI Conn , open :: Conn -> IO () , close :: Conn -> IO ()+ , abort :: Conn -> IO () , connectTcp :: Conn -> String -> Int -> IO (Either String ()) , configure :: Conn -> TransportOption -> IO (Either String ()) }
internal/Plumbing/Pipeline/Streaming/Parser.hs view
@@ -2,6 +2,7 @@ import Conduit import Data.ByteString import qualified Data.ByteString as BS+import qualified Data.List as List import Lib.Logger.Types (LogLevel (..), MonadLogger (..)) import Parser.API ( ParseStep (DropPrefix, Emit, NeedMore, Reject)@@ -40,13 +41,28 @@ lift . logMessage Debug $ "message spans frame, waiting for more data" loop bs DropPrefix n reason -> do- lift . logMessage Error $ ("dropping invalid prefix: " ++ reason)- lift . logMessage Debug $ ("invalid prefix: " ++ show bs)- lift . logMessage Error $ ("dropping " ++ show n ++ " bytes")+ lift . logMessage Error $+ "dropping invalid protocol bytes: category="+ ++ reasonCategory reason+ ++ " dropped_bytes="+ ++ show n+ ++ " buffered_bytes="+ ++ show (length bs) handleChunk (drop n bs) Reject reason ->- lift . logMessage Error $ ("parser rejected inbound data: " ++ reason)+ lift . logMessage Error $+ "parser rejected inbound data: category="+ ++ reasonCategory reason+ ++ " buffered_bytes="+ ++ show (length bs) Emit message rest -> do lift . logMessage Debug $ "parsed message" yield message handleChunk rest++ reasonCategory reason+ | "unknown protocol prefix" `List.isPrefixOf` reason = "unknown-prefix"+ | "invalid INFO" `List.isPrefixOf` reason = "invalid-info"+ | "invalid MSG" `List.isPrefixOf` reason = "invalid-msg"+ | "invalid HMSG" `List.isPrefixOf` reason = "invalid-hmsg"+ | otherwise = "malformed-frame"
internal/Plumbing/Queue/API.hs view
@@ -3,18 +3,64 @@ module Queue.API ( Queue (..) , QueueItem (..)+ , TryEnqueueResult (..)+ , isConnectionScoped+ , isWithinPayloadLimit+ , payloadLimitRejections ) where +import Control.Concurrent.STM (STM) import Transformers.Transformers (Transformer (..)) data QueueItem where QueueItem :: Transformer m => m -> QueueItem+ QueueBatch :: [QueueItem] -> QueueItem+ QueueConnectionScoped :: QueueItem -> QueueItem+ QueuePayloadBound :: Int -> (Int -> IO ()) -> QueueItem -> QueueItem instance Transformer QueueItem where- transform (QueueItem item) = transform item+ transform (QueueItem item) = transform item+ transform (QueueBatch items) = mconcat (map transform items)+ transform (QueueConnectionScoped item) = transform item+ transform (QueuePayloadBound _ _ item) = transform item +isConnectionScoped :: QueueItem -> Bool+isConnectionScoped (QueueConnectionScoped _) = True+isConnectionScoped (QueuePayloadBound _ _ item) = isConnectionScoped item+isConnectionScoped _ = False++isWithinPayloadLimit :: Int -> QueueItem -> Bool+isWithinPayloadLimit maximumSize (QueuePayloadBound actualSize _ item) =+ actualSize <= maximumSize && isWithinPayloadLimit maximumSize item+isWithinPayloadLimit maximumSize (QueueConnectionScoped item) =+ isWithinPayloadLimit maximumSize item+isWithinPayloadLimit maximumSize (QueueBatch items) =+ all (isWithinPayloadLimit maximumSize) items+isWithinPayloadLimit _ (QueueItem _) = True++payloadLimitRejections :: Int -> QueueItem -> [IO ()]+payloadLimitRejections maximumSize (QueuePayloadBound actualSize reject item)+ | actualSize > maximumSize = [reject maximumSize]+ | otherwise = payloadLimitRejections maximumSize item+payloadLimitRejections maximumSize (QueueConnectionScoped item) =+ payloadLimitRejections maximumSize item+payloadLimitRejections maximumSize (QueueBatch items) =+ concatMap (payloadLimitRejections maximumSize) items+payloadLimitRejections _ (QueueItem _) = []++data TryEnqueueResult = TryEnqueued | TryQueueClosed | TryQueueFull+ deriving (Eq, Show)+ data Queue = Queue { enqueue :: QueueItem -> IO (Either String ())+ , enqueueSTM :: QueueItem -> STM (Either String ())+ , tryEnqueue :: QueueItem -> IO TryEnqueueResult+ , tryEnqueueSTM :: QueueItem -> STM TryEnqueueResult , dequeue :: IO (Either String QueueItem)- , close :: IO ()- , open :: IO ()+ , close :: IO ()+ , open :: IO ()+ , openSTM :: STM ()+ , discardConnectionScoped :: IO ()+ , closeAndDiscardConnectionScoped :: STM ()+ , closeAndDiscardAll :: STM ()+ , discardOversizedPayloads :: Int -> STM [IO ()] }
internal/Plumbing/Queue/TransactionalQueue.hs view
@@ -1,10 +1,20 @@ module Queue.TransactionalQueue ( newQueue+ , newQueueWithCapacity ) where import Control.Concurrent.STM import qualified Control.Monad-import Queue.API (Queue (Queue), QueueItem)+import Data.List (partition)+import Numeric.Natural (Natural)+import Queue.API+ ( Queue (Queue)+ , QueueItem+ , TryEnqueueResult (..)+ , isConnectionScoped+ , isWithinPayloadLimit+ , payloadLimitRejections+ ) data TransactionalQueue = TransactionalQueue { transactionalQueueItems :: TBQueue QueueItem@@ -12,8 +22,11 @@ } newQueue :: IO Queue-newQueue = do- queueItems <- newTBQueueIO 1000+newQueue = newQueueWithCapacity 1000++newQueueWithCapacity :: Natural -> IO Queue+newQueueWithCapacity capacity = do+ queueItems <- newTBQueueIO (max 1 capacity) queueClosed <- newEmptyTMVarIO let transactionalQueue = TransactionalQueue@@ -23,28 +36,51 @@ pure $ Queue (enqueue transactionalQueue)+ (enqueueSTM transactionalQueue)+ (tryEnqueue transactionalQueue)+ (tryEnqueueSTM transactionalQueue) (dequeue transactionalQueue) (close transactionalQueue) (open transactionalQueue)+ (openSTM transactionalQueue)+ (discardConnectionScoped transactionalQueue)+ (closeAndDiscardConnectionScoped transactionalQueue)+ (closeAndDiscardAll transactionalQueue)+ (discardOversizedPayloads transactionalQueue) enqueue :: TransactionalQueue -> QueueItem -> IO (Either String ())-enqueue transactionalQueue item = do- isOpen <- atomically $ isEmptyTMVar (transactionalQueueClosed transactionalQueue)- if isOpen- then atomically $ do- writeTBQueue (transactionalQueueItems transactionalQueue) item- pure (Right ())- else pure (Left "Queue is closed")+enqueue transactionalQueue = atomically . enqueueSTM transactionalQueue +enqueueSTM :: TransactionalQueue -> QueueItem -> STM (Either String ())+enqueueSTM transactionalQueue item =+ (do+ isOpen <- isEmptyTMVar (transactionalQueueClosed transactionalQueue)+ check isOpen+ writeTBQueue (transactionalQueueItems transactionalQueue) item+ pure (Right ()))+ `orElse`+ (Left "Queue is closed" <$ readTMVar (transactionalQueueClosed transactionalQueue))++tryEnqueue :: TransactionalQueue -> QueueItem -> IO TryEnqueueResult+tryEnqueue transactionalQueue = atomically . tryEnqueueSTM transactionalQueue++tryEnqueueSTM :: TransactionalQueue -> QueueItem -> STM TryEnqueueResult+tryEnqueueSTM transactionalQueue item = do+ isOpen <- isEmptyTMVar (transactionalQueueClosed transactionalQueue)+ if not isOpen+ then pure TryQueueClosed+ else do+ isFull <- isFullTBQueue (transactionalQueueItems transactionalQueue)+ if isFull+ then pure TryQueueFull+ else writeTBQueue (transactionalQueueItems transactionalQueue) item >> pure TryEnqueued+ dequeue :: TransactionalQueue -> IO (Either String QueueItem)-dequeue transactionalQueue =- atomically $- (Right <$> readTBQueue (transactionalQueueItems transactionalQueue))- `orElse`- (do- isOpen <- isEmptyTMVar (transactionalQueueClosed transactionalQueue)- check (not isOpen)- pure (Left "Queue is closed"))+dequeue transactionalQueue = atomically $ do+ isOpen <- isEmptyTMVar (transactionalQueueClosed transactionalQueue)+ if isOpen+ then Right <$> readTBQueue (transactionalQueueItems transactionalQueue)+ else pure (Left "Queue is closed") close :: TransactionalQueue -> IO () close transactionalQueue = atomically $ do@@ -52,5 +88,38 @@ putTMVar (transactionalQueueClosed transactionalQueue) () open :: TransactionalQueue -> IO ()-open transactionalQueue =- Control.Monad.void (atomically (tryTakeTMVar (transactionalQueueClosed transactionalQueue)))+open = atomically . openSTM++openSTM :: TransactionalQueue -> STM ()+openSTM transactionalQueue =+ Control.Monad.void (tryTakeTMVar (transactionalQueueClosed transactionalQueue))++discardConnectionScoped :: TransactionalQueue -> IO ()+discardConnectionScoped transactionalQueue = atomically $ do+ discardConnectionScopedSTM transactionalQueue++closeAndDiscardConnectionScoped :: TransactionalQueue -> STM ()+closeAndDiscardConnectionScoped transactionalQueue = do+ _ <- tryTakeTMVar (transactionalQueueClosed transactionalQueue)+ putTMVar (transactionalQueueClosed transactionalQueue) ()+ discardConnectionScopedSTM transactionalQueue++closeAndDiscardAll :: TransactionalQueue -> STM ()+closeAndDiscardAll transactionalQueue = do+ _ <- tryTakeTMVar (transactionalQueueClosed transactionalQueue)+ putTMVar (transactionalQueueClosed transactionalQueue) ()+ Control.Monad.void (flushTBQueue (transactionalQueueItems transactionalQueue))++discardConnectionScopedSTM :: TransactionalQueue -> STM ()+discardConnectionScopedSTM transactionalQueue = do+ queued <- flushTBQueue (transactionalQueueItems transactionalQueue)+ mapM_+ (writeTBQueue (transactionalQueueItems transactionalQueue))+ (filter (not . isConnectionScoped) queued)++discardOversizedPayloads :: TransactionalQueue -> Int -> STM [IO ()]+discardOversizedPayloads transactionalQueue maximumSize = do+ queued <- flushTBQueue (transactionalQueueItems transactionalQueue)+ let (retained, discarded) = partition (isWithinPayloadLimit maximumSize) queued+ mapM_ (writeTBQueue (transactionalQueueItems transactionalQueue)) retained+ pure (concatMap (payloadLimitRejections maximumSize) discarded)
internal/Policy/Engine.hs view
@@ -1,78 +1,118 @@ module Engine ( closeClient+ , interruptConnection , resetClient , runEngine ) where import Auth.Types-import Control.Concurrent (forkIO)+import Control.Concurrent+ ( MVar+ , ThreadId+ , forkIOWithUnmask+ , killThread+ , myThreadId+ , newMVar+ , withMVar+ ) import Control.Concurrent.STM-import Control.Monad (forM_, unless, void, when)+import Control.Exception+ ( SomeException+ , finally+ , mask+ , mask_+ , onException+ )+import Control.Monad (unless, void, when) import Data.Foldable (for_) import Data.Maybe (listToMaybe)+import GHC.Clock (getMonotonicTimeNSec) import Handshake.Nats ( HandshakeError (..) , performHandshake )+import Lib.Exception (trySync) import Lib.Logger (LogLevel (..), MonadLogger (..)) import Network.ConnectionAPI ( Conn , ConnectionAPI+ , WriterAPI (..)+ , abort , close , closeReader- , closeWriter , connectTcp , open , reader , writer )-import Parser.API (ParsedMessage, ParserAPI)+import Parser.API+ ( ParsedMessage (ParsedPing)+ , ParserAPI+ ) import Pipeline.Broadcasting.API (BroadcastingAPI (BroadcastingAPI)) import Pipeline.Streaming.API (StreamingAPI (StreamingAPI))-import Queue.API (QueueItem (QueueItem))+import Queue.API (QueueItem (QueueBatch, QueueItem)) import Router.Nats (RouteDirective (..), routeMessage) import State.Store ( ClientState+ , activateConnectionAttempt+ , beginConnectionAttempt , closeQueue , config+ , connectedOnce , connection- , enqueue , failInitialConnection , incrementAttemptIndex+ , isManagedThread , markClosed- , markConnected- , markConnectionReady- , openQueue+ , markConnectionReadyWithRejectedPublishes , queue , readAttemptIndex- , readConnectionGeneration , readStatus+ , registerManagedThread , runClient , setClosing , setEndpoint+ , setReconnecting+ , unregisterManagedThread , waitForClosed , waitForConnectionGenerationAfter+ , withSubscriptionGate ) import State.Types ( ClientConfig (..) , ClientExitReason (..)- , ClientStatus (..) , ConnectAttemptError (..) , ConnectError (..) , ConnectFailure (..)+ , ConnectionState (..)+ , serverErrorKind ) import Subscription.Store ( SubscriptionStore , active- , awaitCallbackDrain- , awaitNoTrackedExpiries+ , closeStore )-import Subscription.Types (SubscriptionMeta (SubscriptionMeta))+import Subscription.Types+ ( SubscriptionMeta (SubscriptionMeta)+ , isOneShotSubscription+ , isResumableSubscription+ )+import System.Timeout (timeout)+import Transformers.Transformers (Transformer (transform))+import Types.Pong (Pong (Pong)) import qualified Types.Sub as Sub+import qualified Types.Unsub as Unsub data ConnectionRunResult = ConnectionDisconnected | ConnectionExit ClientExitReason +data ReaderRuntime = ReaderRuntime+ { runtimeReaderThread :: ThreadId+ , runtimeReaderDone :: TMVar ()+ , runtimeExit :: TMVar ClientExitReason+ }+ runEngine :: ConnectionAPI -> StreamingAPI@@ -83,16 +123,39 @@ -> Auth -> IO () runEngine connectionApi streamingApi broadcastingApi parserApi state store auth =- case connectOptions (config state) of- [] -> do- failInitialConnection state ConnectNoServers- finalize (ExitRetriesExhausted (Just "No servers provided"))- _ ->- loop (connectionAttempts (config state)) []+ mask $ \restore -> do+ result <- trySync (restore engine `finally` cleanupEngineResources)+ case result of+ Left err -> handleEngineFailure err+ Right () -> pure () where StreamingAPI runStreaming = streamingApi BroadcastingAPI runBroadcasting = broadcastingApi + engine =+ case connectOptions (config state) of+ [] -> do+ failInitialConnection state ConnectNoServers+ finalize (ExitRetriesExhausted (Just "No servers provided"))+ _ ->+ loop (connectionAttempts (config state)) []++ handleEngineFailure :: SomeException -> IO ()+ handleEngineFailure _ = mask_ $ do+ let reason = ExitRetriesExhausted (Just "client engine failed")+ failInitialConnection state (ConnectAttemptsExhausted [])+ setClosing state reason+ finalize reason++ cleanupEngineResources = mask_ $ do+ let ignoreSync action = void (trySync action)+ conn = connection state+ ignoreSync (closeQueue state)+ ignoreSync (closeReader (reader connectionApi) conn)+ ignoreSync (closeWriter (writer connectionApi) conn)+ ignoreSync (abort connectionApi conn)+ ignoreSync (close connectionApi conn)+ loop remaining attemptErrors | remaining <= 0 = do runClient state $@@ -105,63 +168,150 @@ | otherwise = do status <- readStatus state case status of- Running -> do- attemptErr <- runAttempt- nextStatus <- readStatus state- case nextStatus of- Running -> do- when (remaining > 1) $ do- runClient state $- logMessage Info "retrying client connection"- incrementAttemptIndex state- loop (remaining - 1) (attemptErr : attemptErrors)- Closing reason ->- finalize reason- Closed _ ->- pure ()- Closing reason ->+ ConnectionConnecting -> continueLoop remaining attemptErrors+ ConnectionConnected -> continueLoop remaining attemptErrors+ ConnectionReconnecting -> continueLoop remaining attemptErrors+ ConnectionClosing reason -> finalize reason- Closed _ ->+ ConnectionClosed _ -> pure ()+ where+ continueLoop remaining' attemptErrors' = do+ attemptErr <- runAttempt+ nextStatus <- readStatus state+ case nextStatus of+ ConnectionConnecting -> retryLoop attemptErr+ ConnectionConnected -> retryLoop attemptErr+ ConnectionReconnecting -> retryLoop attemptErr+ ConnectionClosing reason ->+ finalize reason+ ConnectionClosed _ ->+ pure ()+ where+ retryLoop attemptErr = do+ when (remaining' > 1) $ do+ runClient state $+ logMessage Info "retrying client connection"+ incrementAttemptIndex state+ loop (remaining' - 1) (attemptErr : attemptErrors') runAttempt = do- openQueue state- open connectionApi (connection state) let cfg = config state endpoints = connectOptions cfg attemptIndex <- readAttemptIndex state let endpoint@(host, _port) = endpoints !! (attemptIndex `mod` length endpoints)- setEndpoint state endpoint- transportResult <- acquireTransport endpoint- case transportResult of- Left err -> do+ attemptTimeout = max 1 (connectTimeoutMicros cfg)+ attempt <- beginConnectionAttempt state+ case attempt of+ Nothing ->+ pure (ConnectAttemptError endpoint (ConnectTransportFailure "connection attempt cancelled"))+ Just attemptEpoch -> do+ attemptStarted <- toInteger <$> getMonotonicTimeNSec+ setEndpoint state endpoint+ acquired <- timeout attemptTimeout (acquireTransport endpoint)+ case acquired of+ Nothing -> do+ runClient state $+ logMessage Error "transport acquisition timed out"+ terminateTransport (connection state)+ pure (ConnectAttemptError endpoint (ConnectTransportFailure "transport acquisition timed out"))+ Just (Left err) -> do+ terminateTransport (connection state)+ pure (ConnectAttemptError endpoint (ConnectTransportFailure err))+ Just (Right conn) -> do+ open connectionApi conn+ activated <- activateConnectionAttempt state attemptEpoch+ if not activated+ then do+ terminateTransport conn+ pure (ConnectAttemptError endpoint (ConnectTransportFailure "connection attempt cancelled"))+ else do+ remaining <- remainingAttemptMicros attemptStarted attemptTimeout+ initialized <-+ if remaining <= 0+ then pure Nothing+ else timeout remaining $+ performHandshake connectionApi parserApi state auth conn host+ case initialized of+ Nothing -> do+ runClient state $+ logMessage Error "connection handshake timed out"+ terminateTransport conn+ pure (ConnectAttemptError endpoint ConnectHandshakeTimeout)+ Just (Left handshakeError) -> do+ let failure = handshakeFailure handshakeError+ runClient state $+ logMessage Error ("connection initialization failed: " ++ failureCategory failure)+ terminateTransport conn+ pure (ConnectAttemptError endpoint failure)+ Just (Right ()) ->+ finishConnection+ endpoint+ conn+ attemptEpoch+ attemptStarted+ attemptTimeout++ finishConnection endpoint conn attemptEpoch attemptStarted attemptTimeout = mask $ \restore -> do+ writeLock <- newMVar ()+ readerRuntime <- startConnectionReader conn writeLock+ remaining <- remainingAttemptMicros attemptStarted attemptTimeout+ readinessResult <-+ if remaining <= 0+ then pure Nothing+ else restore (timeout remaining (completeReadiness writeLock conn attemptEpoch))+ `onException` terminateReader conn readerRuntime+ case readinessResult of+ Nothing -> do runClient state $- logMessage Error ("connection attempt failed: " ++ show err)- close connectionApi (connection state)+ logMessage Error "connection readiness timed out"+ terminateReader conn readerRuntime+ terminateTransport conn+ pure+ ( ConnectAttemptError+ endpoint+ (ConnectTransportFailure "connection readiness timed out")+ )+ Just (Left err) -> do+ runClient state $+ logMessage Error+ ("resubscription failed: " ++ failureCategory (ConnectTransportFailure err))+ void (setReconnecting state)+ terminateReader conn readerRuntime+ terminateTransport conn pure (ConnectAttemptError endpoint (ConnectTransportFailure err))- Right conn -> do- handshakeResult <- performHandshake connectionApi parserApi state auth conn host- case handshakeResult of- Left err -> do+ Just (Right False) -> do+ closeQueue state+ terminateReader conn readerRuntime+ terminateTransport conn+ pure (ConnectAttemptError endpoint (ConnectTransportFailure "connection attempt cancelled"))+ Just (Right True) -> do+ connectionResult <- restore (runConnection conn writeLock readerRuntime)+ `onException` terminateReader conn readerRuntime+ case connectionResult of+ ConnectionDisconnected -> void (setReconnecting state)+ ConnectionExit reason -> setClosing state reason+ close connectionApi conn+ case connectionResult of+ ConnectionDisconnected -> do runClient state $- logMessage Error ("connection initialization failed: " ++ show err)- close connectionApi conn- pure (ConnectAttemptError endpoint (handshakeFailure err))- Right () -> do- resubscribeIfNeeded- markConnectionReady state- connectionResult <- runConnection conn- close connectionApi conn- case connectionResult of- ConnectionDisconnected -> do- runClient state $- logMessage Info "connection disconnected"- pure (ConnectAttemptError endpoint (ConnectTransportFailure "connection disconnected"))- ConnectionExit reason -> do- setClosing state reason- pure (ConnectAttemptError endpoint (ConnectProtocolFailure (show reason)))+ logMessage Info "connection disconnected"+ pure (ConnectAttemptError endpoint (ConnectTransportFailure "connection disconnected"))+ ConnectionExit reason ->+ pure (ConnectAttemptError endpoint (ConnectProtocolFailure (exitCategory reason))) + terminateTransport conn = do+ abort connectionApi conn+ close connectionApi conn++ remainingAttemptMicros started timeoutMicros = do+ now <- toInteger <$> getMonotonicTimeNSec+ let deadline = started + toInteger timeoutMicros * 1000+ remainingNanos = max 0 (deadline - now)+ roundedMicros = (remainingNanos + 999) `div` 1000+ pure (fromInteger (min (toInteger (maxBound :: Int)) roundedMicros))+ acquireTransport (host, port) = do let conn = connection state result <- connectTcp connectionApi conn host port@@ -173,91 +323,229 @@ handshakeFailure (HandshakeTLSError err) = ConnectTLSFailure err handshakeFailure (HandshakeProtocolError err) = ConnectProtocolFailure err handshakeFailure (HandshakeAuthError err) = ConnectAuthenticationFailure (show err)- handshakeFailure HandshakeTimeout = ConnectHandshakeTimeout - resubscribeIfNeeded = do- alreadyConnected <- markConnected state- when alreadyConnected $ do- activeSubscriptions <- active store- let resumable = filter isResumable activeSubscriptions- unless (null resumable) $- runClient state- (logMessage Info ("resubscribing " ++ show (length resumable) ++ " subscriptions"))- forM_ resumable $ \(sid, SubscriptionMeta subject queueGroup _) ->- enqueue state $- QueueItem- Sub.Sub- { Sub.subject = subject- , Sub.queueGroup = queueGroup- , Sub.sid = sid+ failureCategory (ConnectTransportFailure _) = "transport failure"+ failureCategory (ConnectTLSFailure _) = "tls failure"+ failureCategory (ConnectProtocolFailure _) = "protocol failure"+ failureCategory (ConnectAuthenticationFailure _) = "authentication failure"+ failureCategory ConnectHandshakeTimeout = "handshake timeout"++ exitCategory ExitClosedByUser = "closed by user"+ exitCategory (ExitRetriesExhausted _) = "retries exhausted"+ exitCategory (ExitServerError serverError) = "server error: " ++ show (serverErrorKind serverError)+ exitCategory (ExitInboundMessageTooLarge _ _) = "inbound message too large"+ exitCategory ExitResetRequested = "reset requested"++ completeReadiness writeLock conn attemptEpoch =+ withSubscriptionGate state $+ (do+ result <- resubscribeIfNeeded writeLock conn+ case result of+ Left err -> do+ abort connectionApi conn+ pure (Left err)+ Right () -> do+ (accepted, rejectedPublishes) <-+ markConnectionReadyWithRejectedPublishes state attemptEpoch+ sequence_ rejectedPublishes+ unless accepted (abort connectionApi conn)+ pure (Right accepted))+ `onException` abort connectionApi conn++ resubscribeIfNeeded writeLock conn = do+ alreadyConnected <- connectedOnce state+ if alreadyConnected+ then do+ activeSubscriptions <- active store+ let resumable = filter isResumable activeSubscriptions+ result <- writeSubscriptions writeLock conn resumable+ let resumableCount = length resumable+ case result of+ Right () ->+ unless (resumableCount == 0) $+ runClient state+ (logMessage Info ("resubscribed " ++ show resumableCount ++ " subscriptions"))+ Left _ -> pure ()+ pure result+ else pure (Right ())++ writeSubscriptions _ _ [] = pure (Right ())+ writeSubscriptions writeLock conn subscriptions = do+ let (chunk, remaining) = splitAt resubscribeChunkSize subscriptions+ result <- writeSerialized+ writeLock+ conn+ (transform (QueueBatch (map subscriptionCommand chunk)))+ case result of+ Left err -> pure (Left err)+ Right () -> writeSubscriptions writeLock conn remaining++ writeSerialized writeLock conn bytes =+ withMVar writeLock $ \_ ->+ writeDataLazy (writer connectionApi) conn bytes++ resubscribeChunkSize = 64++ subscriptionCommand (sid, meta@(SubscriptionMeta subject queueGroup _)) =+ QueueBatch $+ QueueItem+ Sub.Sub+ { Sub.subject = subject+ , Sub.queueGroup = queueGroup+ , Sub.sid = sid+ } :+ [ QueueItem+ Unsub.Unsub+ { Unsub.sid = sid+ , Unsub.maxMsg = Just 1 }+ | isOneShotSubscription meta+ ] finalize reason = do+ closeStore store result <- markClosed state reason for_ result (exitAction (config state)) - runConnection :: Conn -> IO ConnectionRunResult- runConnection conn = do+ startManaged :: IO () -> IO () -> IO ThreadId+ startManaged body cleanup =+ forkIOWithUnmask $ \unmask -> do+ threadId <- myThreadId+ registerManagedThread state threadId+ unmask body+ `finally` (cleanup `finally` unregisterManagedThread state threadId)++ startConnectionReader :: Conn -> MVar () -> IO ReaderRuntime+ startConnectionReader conn writeLock = do exitVar <- newEmptyTMVarIO readerDone <- newEmptyTMVarIO+ let stopReader = closeReader (reader connectionApi) conn+ stopWriter = closeWriter (writer connectionApi) conn+ stopQueue = closeQueue state+ signalReaderDone = atomically (void (tryPutTMVar readerDone ()))+ signalExit reason = atomically (void (tryPutTMVar exitVar reason))+ cleanup =+ ( do+ void (setReconnecting state)+ stopQueue+ stopWriter+ abort connectionApi conn+ ) `finally` signalReaderDone+ route message = do+ directive <- routeForConnection conn writeLock message+ case directive of+ RouteContinue -> pure ()+ RouteReconnect -> do+ signalExit ExitResetRequested+ stopQueue+ stopReader+ stopWriter+ RouteExit reason -> do+ setClosing state reason+ signalExit reason+ stopQueue+ stopReader+ stopWriter+ readerThread <- startManaged+ (runClient state $ do+ logMessage Debug "starting streaming thread"+ runStreaming+ (streamBufferLimit (messageLimit (config state)))+ (reader connectionApi)+ conn+ parserApi+ route+ logMessage Debug "streaming thread exited")+ cleanup+ pure+ ReaderRuntime+ { runtimeReaderThread = readerThread+ , runtimeReaderDone = readerDone+ , runtimeExit = exitVar+ }++ routeForConnection :: Conn -> MVar () -> ParsedMessage -> IO RouteDirective+ routeForConnection conn writeLock message@(ParsedPing _) = do+ status <- readStatus state+ case status of+ ConnectionConnected -> routeMessage state store message+ ConnectionClosing _ -> pure RouteReconnect+ ConnectionClosed _ -> pure RouteReconnect+ _ -> do+ runClient state (logMessage Debug "routing pre-ready PING")+ result <- writeSerialized writeLock conn (transform Pong)+ pure $ case result of+ Left _ -> RouteReconnect+ Right () -> RouteContinue+ routeForConnection _ _ message = routeMessage state store message++ terminateReader :: Conn -> ReaderRuntime -> IO ()+ terminateReader conn runtime = do+ closeReader (reader connectionApi) conn+ closeWriter (writer connectionApi) conn+ abort connectionApi conn+ killThread (runtimeReaderThread runtime)+ atomically (readTMVar (runtimeReaderDone runtime))++ runConnection :: Conn -> MVar () -> ReaderRuntime -> IO ConnectionRunResult+ runConnection conn writeLock readerRuntime = mask $ \restore -> do writerDone <- newEmptyTMVarIO let stopReader = closeReader (reader connectionApi) conn stopWriter = closeWriter (writer connectionApi) conn stopQueue = closeQueue state- signalReaderDone = atomically (void (tryPutTMVar readerDone ())) signalWriterDone = atomically (void (tryPutTMVar writerDone ()))- signalExit reason = atomically (void (tryPutTMVar exitVar reason)) - void . forkIO $ do- runClient state $ do+ writerThread <- startManaged+ (runClient state $ do logMessage Debug "starting broadcasting thread" runBroadcasting (queue state)- (writer connectionApi)+ (serializedWriter writeLock) conn- logMessage Debug "broadcasting thread exited"- stopQueue- stopReader- signalWriterDone+ logMessage Debug "broadcasting thread exited")+ ((stopQueue >> stopReader) `finally` signalWriterDone) - void . forkIO $ do- runClient state $ do- logMessage Debug "starting streaming thread"- runStreaming- (streamBufferLimit (messageLimit (config state)))- (reader connectionApi)- conn- parserApi- (\message -> do- directive <- routeMessage state store message- case directive of- RouteContinue ->- pure ()- RouteExit reason -> do- setClosing state reason- signalExit reason- stopQueue- stopReader- stopWriter)- logMessage Debug "streaming thread exited"- stopQueue- stopWriter- signalReaderDone+ let terminateChildren = do+ stopQueue+ stopReader+ stopWriter+ abort connectionApi conn+ killThread writerThread+ killThread (runtimeReaderThread readerRuntime)+ awaitChildren (runtimeReaderDone readerRuntime) writerDone - outcome <- atomically $- (Left <$> readTMVar exitVar)- `orElse` (Right () <$ readTMVar readerDone)- `orElse` (Right () <$ readTMVar writerDone)- atomically $ do- readTMVar readerDone- readTMVar writerDone+ outcome <- restore+ (atomically $+ (Left <$> readTMVar (runtimeExit readerRuntime))+ `orElse` (Right () <$ readTMVar (runtimeReaderDone readerRuntime))+ `orElse` (Right () <$ readTMVar writerDone))+ `onException` terminateChildren+ terminateChildren case outcome of- Left reason -> pure (ConnectionExit reason)- Right () -> pure ConnectionDisconnected+ Left ExitResetRequested -> pure ConnectionDisconnected+ Left reason -> pure (ConnectionExit reason)+ Right () -> pure ConnectionDisconnected + where+ awaitChildren readerDone writerDone = atomically $ do+ readTMVar readerDone+ readTMVar writerDone++ serializedWriter :: MVar () -> WriterAPI Conn+ serializedWriter writeLock =+ let baseWriter = writer connectionApi+ in WriterAPI+ { writeData = \conn bytes ->+ withMVar writeLock $ \_ -> writeData baseWriter conn bytes+ , writeDataLazy = \conn bytes ->+ withMVar writeLock $ \_ -> writeDataLazy baseWriter conn bytes+ , closeWriter = closeWriter baseWriter+ , openWriter = openWriter baseWriter+ }+ isResumable :: (a, SubscriptionMeta) -> Bool- isResumable (_, SubscriptionMeta _ _ isReply) = not isReply+ isResumable = isResumableSubscription . snd streamBufferLimit :: Int -> Int streamBufferLimit maximumMessageSize@@ -267,35 +555,31 @@ controlLineAllowance = 64 * 1024 closeClient :: ConnectionAPI -> ClientState -> SubscriptionStore -> IO ()-closeClient connectionApi state store =- shutdownClient connectionApi state store ExitClosedByUser "closing client connection"--resetClient :: ConnectionAPI -> ClientState -> SubscriptionStore -> IO ()-resetClient connectionApi state _store = do- generation <- readConnectionGeneration state- runClient state $- logMessage Info "resetting client connection"+closeClient connectionApi state store = mask_ $ do+ setClosing state ExitClosedByUser+ closeStore store closeQueue state closeReader (reader connectionApi) (connection state) closeWriter (writer connectionApi) (connection state)- atomically $- waitForConnectionGenerationAfter state generation+ abort connectionApi (connection state)+ void . trySync . runClient state $+ logMessage Info "closing client connection"++resetClient :: ConnectionAPI -> ClientState -> SubscriptionStore -> IO ()+resetClient connectionApi state _store = mask $ \restore -> do+ invalidatedGeneration <- interruptConnection connectionApi state+ void . trySync . runClient state $+ logMessage Info "resetting client connection"+ current <- myThreadId+ managed <- isManagedThread state current+ unless managed . restore . atomically $+ waitForConnectionGenerationAfter state invalidatedGeneration `orElse` waitForClosed state -shutdownClient- :: ConnectionAPI- -> ClientState- -> SubscriptionStore- -> ClientExitReason- -> String- -> IO ()-shutdownClient connectionApi state store reason message = do- setClosing state reason- runClient state $- logMessage Info message- closeQueue state+interruptConnection :: ConnectionAPI -> ClientState -> IO Int+interruptConnection connectionApi state = mask_ $ do+ invalidatedGeneration <- setReconnecting state closeReader (reader connectionApi) (connection state) closeWriter (writer connectionApi) (connection state)- atomically $ waitForClosed state- atomically $ awaitNoTrackedExpiries store- atomically $ awaitCallbackDrain store+ abort connectionApi (connection state)+ pure invalidatedGeneration
internal/Policy/Handshake/Nats.hs view
@@ -6,11 +6,13 @@ import Auth.Resolver (applyAuthPatch, buildAuthPatch) import Auth.Types import qualified Data.ByteString as BS+import Data.List (isPrefixOf) import Data.Maybe (fromMaybe, isJust) import Network.ConnectionAPI ( Conn , ConnectionAPI , TransportOption (..)+ , bufferRead , configure , readData , reader@@ -26,11 +28,14 @@ import State.Store ( ClientState , config+ , notifyServerError , setServerInfo , updateLogContextFromInfo )-import State.Types (ClientConfig (..))-import System.Timeout (timeout)+import State.Types+ ( ClientConfig (..)+ , serverErrorFromProtocol+ ) import Transformers.Transformers (Transformer (transform)) import qualified Types.Connect as Connect import qualified Types.Err as Err@@ -43,13 +48,13 @@ | HandshakeTLSError String | HandshakeProtocolError String | HandshakeAuthError AuthError- | HandshakeTimeout deriving (Eq, Show) +handshakeControlFrameLimit :: Int+handshakeControlFrameLimit = 64 * 1024+ performHandshake :: ConnectionAPI -> ParserAPI ParsedMessage -> ClientState -> Auth -> Conn -> String -> IO (Either HandshakeError ())-performHandshake connectionApi parserApi state auth conn host = do- result <- timeout (max 1 (connectTimeoutMicros (config state))) handshake- pure (fromMaybe (Left HandshakeTimeout) result)+performHandshake connectionApi parserApi state auth conn host = handshake where handshake = do infoResult <- readInitialInfo@@ -107,9 +112,9 @@ awaitPong mempty readInitialInfo :: IO (Either HandshakeError (I.Info, BS.ByteString))- readInitialInfo = go mempty+ readInitialInfo = readMore mempty where- go acc = do+ readMore acc = do result <- readData (reader connectionApi) conn 4096 case result of Left err ->@@ -117,25 +122,33 @@ Right chunk -> if BS.null chunk then pure (Left (HandshakeTransportError "read returned empty chunk before INFO"))- else do- let bytes = acc <> chunk- case parse parserApi bytes of- NeedMore ->- go bytes- DropPrefix n _ ->- go (BS.drop n bytes)- Reject reason ->- pure (Left (HandshakeProtocolError reason))- Emit (ParsedInfo info) rest ->- pure (Right (info, rest))- Emit (ParsedErr err) _- | isAuthenticationError err ->- pure (Left (HandshakeAuthError (AuthError (show err))))- | otherwise ->- pure (Left (HandshakeProtocolError ("server error before INFO: " ++ show err)))- Emit _ rest ->- go rest+ else consumeInfoFrames (acc <> chunk) + consumeInfoFrames bytes =+ case parse parserApi bytes of+ NeedMore ->+ retainIncompleteFrame "INFO" bytes readMore+ DropPrefix n _+ | n <= 0 ->+ pure (Left (HandshakeProtocolError "malformed protocol frame"))+ | otherwise ->+ continueInfoWith (BS.drop n bytes)+ Reject reason ->+ pure (Left (HandshakeProtocolError (reasonCategory reason)))+ Emit (ParsedInfo info) rest ->+ pure (Right (info, rest))+ Emit (ParsedErr err) _ -> do+ notifyServerError state (serverErrorFromProtocol err)+ if isAuthenticationError err+ then pure (Left (HandshakeAuthError (AuthError (show err))))+ else pure (Left (HandshakeProtocolError "server error before INFO"))+ Emit _ rest ->+ continueInfoWith rest++ continueInfoWith rest+ | BS.null rest = readMore mempty+ | otherwise = consumeInfoFrames rest+ awaitPong :: BS.ByteString -> IO (Either HandshakeError ()) awaitPong acc = do result <- readData (reader connectionApi) conn 4096@@ -152,12 +165,13 @@ consumePongFrames bytes = case parse parserApi bytes of NeedMore ->- awaitPong bytes+ retainIncompleteFrame "PONG" bytes awaitPong DropPrefix _ reason ->- pure (Left (HandshakeProtocolError reason))+ pure (Left (HandshakeProtocolError (reasonCategory reason))) Reject reason ->- pure (Left (HandshakeProtocolError reason))- Emit (ParsedPong _) _ ->+ pure (Left (HandshakeProtocolError (reasonCategory reason)))+ Emit (ParsedPong _) rest -> do+ bufferRead (reader connectionApi) conn rest pure (Right ()) Emit (ParsedOk _) rest -> continueWith rest@@ -170,21 +184,42 @@ setServerInfo state info updateLogContextFromInfo state info continueWith rest- Emit (ParsedErr err) _- | isAuthenticationError err ->- pure (Left (HandshakeAuthError (AuthError (show err))))- | otherwise ->- pure (Left (HandshakeProtocolError (show err)))+ Emit (ParsedErr err) _ -> do+ notifyServerError state (serverErrorFromProtocol err)+ if isAuthenticationError err+ then pure (Left (HandshakeAuthError (AuthError (show err))))+ else pure (Left (HandshakeProtocolError "server error before PONG")) Emit message _ ->- pure (Left (HandshakeProtocolError ("unexpected frame before PONG: " ++ show message)))+ pure (Left (HandshakeProtocolError ("unexpected " ++ frameKind message ++ " frame before PONG"))) continueWith rest | BS.null rest = awaitPong mempty | otherwise = consumePongFrames rest + retainIncompleteFrame frame bytes continue+ | BS.length bytes > handshakeControlFrameLimit =+ pure . Left . HandshakeProtocolError $+ frame ++ " control frame exceeds 65536 byte limit"+ | otherwise = continue bytes+ isAuthenticationError (Err.ErrAuthViolation _) = True isAuthenticationError (Err.ErrAuthTimeout _) = True isAuthenticationError (Err.ErrAuthExpired _) = True isAuthenticationError (Err.ErrAuthRevoked _) = True isAuthenticationError (Err.ErrAccountAuthExpired _) = True isAuthenticationError _ = False++ reasonCategory reason+ | "unknown protocol prefix" `isPrefixOf` reason = "unknown protocol prefix"+ | "invalid INFO" `isPrefixOf` reason = "invalid INFO frame"+ | "invalid MSG" `isPrefixOf` reason = "invalid MSG frame"+ | "invalid HMSG" `isPrefixOf` reason = "invalid HMSG frame"+ | otherwise = "malformed protocol frame"++ frameKind (ParsedMsg _) = "MSG"+ frameKind (ParsedInfo _) = "INFO"+ frameKind (ParsedPing _) = "PING"+ frameKind (ParsedPong _) = "PONG"+ frameKind (ParsedOk _) = "+OK"+ frameKind (ParsedErr _) = "-ERR"+ frameKind (ParsedMessageTooLarge _ _) = "oversized message"
internal/Policy/Router/Nats.hs view
@@ -3,23 +3,30 @@ , routeMessage ) where +import Control.Monad (when) import Control.Monad.IO.Class (liftIO)+import qualified Data.ByteString as BS import Lib.Logger (LogLevel (..), MonadLogger (..)) import Parser.API ( ParsedMessage (ParsedErr, ParsedInfo, ParsedMessageTooLarge, ParsedMsg, ParsedOk, ParsedPing, ParsedPong) )-import Queue.API (QueueItem (QueueItem))+import Queue.API+ ( QueueItem (QueueConnectionScoped, QueueItem)+ ) import State.Store ( ClientState+ , ServerErrorEnqueueResult (..) , enqueue+ , notifyServerError+ , resolveNextPing , runClient- , runNextPingAction , setServerInfo , updateLogContextFromInfo ) import State.Types ( ClientExitReason (ExitInboundMessageTooLarge, ExitServerError) , serverErrorFromProtocol+ , serverErrorKind ) import Subscription.Store ( DispatchResult (DispatchDropped, DispatchMissing, DispatchQueued)@@ -31,6 +38,7 @@ import Types.Pong (Pong (..)) data RouteDirective = RouteContinue+ | RouteReconnect | RouteExit ClientExitReason deriving (Eq, Show) @@ -39,7 +47,16 @@ runClient state $ case parsed of ParsedMsg msg -> do- logMessage Debug ("routing MSG: " ++ show msg)+ logMessage Debug+ ( "routing MSG: subject="+ ++ show (Msg.subject msg)+ ++ " sid="+ ++ show (Msg.sid msg)+ ++ " payload_bytes="+ ++ show (maybe 0 BS.length (Msg.payload msg))+ ++ " header_count="+ ++ show (maybe 0 length (Msg.headers msg))+ ) result <- liftIO $ dispatchMessage store msg case result of DispatchQueued ->@@ -63,27 +80,38 @@ ) pure (RouteExit (ExitInboundMessageTooLarge actual maximumSize)) ParsedInfo info -> do- logMessage Debug ("routing INFO: " ++ show info)+ logMessage Debug "routing INFO" liftIO $ setServerInfo state info liftIO $ updateLogContextFromInfo state info pure RouteContinue ParsedPing _ -> do logMessage Debug "routing PING"- liftIO $ enqueue state (QueueItem Pong)- pure RouteContinue+ result <- liftIO $ enqueue state (QueueConnectionScoped (QueueItem Pong))+ pure $+ case result of+ Left _ -> RouteReconnect+ Right () -> RouteContinue ParsedPong _ -> do logMessage Debug "routing PONG"- liftIO $ runNextPingAction state+ liftIO $ resolveNextPing state pure RouteContinue ParsedOk okMsg -> do logMessage Debug ("routing OK: " ++ show okMsg) pure RouteContinue ParsedErr err -> do- logMessage Debug ("routing ERR: " ++ show err)+ let serverError = serverErrorFromProtocol err+ kind = serverErrorKind serverError+ logMessage Debug ("routing ERR: kind=" ++ show kind)+ enqueueResult <- liftIO $ notifyServerError state serverError+ let accepted = enqueueResult == ServerErrorQueued+ when (enqueueResult == ServerErrorDroppedReport) $+ logMessage Warn "server error handler queue full; dropping event" if Err.isFatal err then do- logMessage Error ("fatal server error: " ++ show err)- pure (RouteExit (ExitServerError (serverErrorFromProtocol err)))- else do- logMessage Warn ("server error: " ++ show err)+ logMessage Error ("fatal server error: kind=" ++ show kind)+ pure (RouteExit (ExitServerError serverError))+ else if accepted+ then do+ logMessage Warn ("server error: kind=" ++ show kind) pure RouteContinue+ else pure RouteContinue
internal/Policy/State/Store.hs view
@@ -9,10 +9,15 @@ , queue , connection , enqueue+ , tryEnqueue , openQueue , closeQueue- , pushPingAction- , runNextPingAction+ , discardConnectionScoped+ , PingResult (..)+ , registerPingWaiter+ , registerPingWaiterAndEnqueue+ , resolveNextPing+ , clearPendingPings , nextSid , nextInbox , readServerInfo@@ -20,27 +25,57 @@ , updateLogContextFromInfo , setConnectName , setEndpoint+ , ServerErrorEnqueueResult (..)+ , notifyServerError+ , CallbackWorker+ , startCallbackWorker+ , callbackWorkerThreadId+ , waitForCallbackWorker+ , stopCallbackWorker , readStatus+ , waitForConnected+ , beginConnectionAttempt+ , activateConnectionAttempt+ , setReconnecting , setClosing , markClosed , waitForClosed , waitForNotRunning , waitForServerInfo- , markConnected+ , connectedOnce , markConnectionReady+ , markConnectionReadyWithRejectedPublishes , waitForInitialConnection , failInitialConnection , readConnectionGeneration , waitForConnectionGenerationAfter+ , waitForConnectionGenerationLoss+ , enqueueOnGenerationTracked+ , PublishEnqueueResult (..)+ , enqueuePublishOnConnected+ , enqueuePublishOnConnectedGenerationTracked+ , tryEnqueueOnGeneration , readAttemptIndex , incrementAttemptIndex+ , registerManagedThread+ , unregisterManagedThread+ , isManagedThread+ , withSubscriptionGate ) where +import Control.Concurrent+ ( ThreadId+ , forkIOWithUnmask+ , killThread+ , myThreadId+ ) import Control.Concurrent.STM-import Control.Monad (unless, void)+import Control.Exception (bracket_, finally)+import Control.Monad (unless, void, when) import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as BC-import Data.Maybe (fromMaybe)+import Data.List (delete)+import Lib.Exception (trySync) import Lib.Logger ( LogLevel (..) , MonadLogger (..)@@ -50,16 +85,22 @@ import Lib.LoggerAPI (runWithLogger, updateLogContext) import Network.ConnectionAPI (Conn) import Nuid (Nuid, newNuidIO, nextNuid)-import Queue.API (Queue (Queue), QueueItem, close, open)+import qualified Queue.API as Queue+import Queue.API (Queue, QueueItem, TryEnqueueResult) import Sid (SIDCounter, initialSIDCounter, nextSID) import State.Types+import qualified Types.Info as Info import Types.Info (Info, client_id) import Types.Msg (SID, Subject) data ClientState = ClientState { clientConfig :: ClientConfig , clientQueue :: Queue- , clientPings :: TQueue (IO ())+ , clientPings :: TQueue (Int, TMVar PingResult)+ , clientCallbacks :: TQueue ClientCallback+ , clientPendingServerCallbacks :: TVar Int+ , clientPendingServerCallbackBytes :: TVar Int+ , clientServerCallbackOverflowed :: TVar Bool , clientConnectedOnce :: TVar Bool , clientConnectionGeneration :: TVar Int , clientInitialConnection :: TMVar (Either ConnectError ())@@ -67,14 +108,41 @@ , clientInboxNuid :: TVar Nuid , clientServerInfo :: TVar (Maybe Info) , clientAttemptIndex :: TVar Int- , clientStatus :: TVar ClientStatus+ , clientAttemptEpoch :: TVar Int+ , clientStatus :: TVar ConnectionState , clientConnection :: Conn , clientLogContext :: TVar LogContext+ , clientManagedThreads :: TVar [ThreadId]+ , clientSubscriptionGate :: TMVar () } +data PingResult = PingReceived | PingConnectionLost+ deriving (Eq, Show)++data PublishEnqueueResult = PublishEnqueued Int+ | PublishTooLarge Int+ | PublishConnectionClosed ClientExitReason+ deriving (Eq, Show)++data ClientCallback = CallbackServerError ServerError+ | CallbackConnectionEvent ConnectionEvent++data ServerErrorEnqueueResult = ServerErrorQueued | ServerErrorDropped | ServerErrorDroppedReport+ deriving (Eq, Show)++data CallbackWorker = CallbackWorker+ { callbackWorkerThreadId :: ThreadId+ , callbackWorkerDone :: TMVar ()+ }+ newClientState :: ClientConfig -> Queue -> Conn -> TVar LogContext -> IO ClientState newClientState cfg queue conn logContext = do+ Queue.close queue pings <- newTQueueIO+ callbacks <- newTQueueIO+ pendingServerCallbacks <- newTVarIO 0+ pendingServerCallbackBytes <- newTVarIO 0+ serverCallbackOverflowed <- newTVarIO False connectedOnce <- newTVarIO False connectionGeneration <- newTVarIO 0 initialConnection <- newEmptyTMVarIO@@ -82,11 +150,18 @@ inboxNuid <- newTVarIO =<< newNuidIO serverInfo <- newTVarIO Nothing attemptIndex <- newTVarIO 0- status <- newTVarIO Running+ attemptEpoch <- newTVarIO 0+ status <- newTVarIO ConnectionConnecting+ managedThreads <- newTVarIO []+ subscriptionGate <- newTMVarIO () pure ClientState { clientConfig = cfg , clientQueue = queue , clientPings = pings+ , clientCallbacks = callbacks+ , clientPendingServerCallbacks = pendingServerCallbacks+ , clientPendingServerCallbackBytes = pendingServerCallbackBytes+ , clientServerCallbackOverflowed = serverCallbackOverflowed , clientConnectedOnce = connectedOnce , clientConnectionGeneration = connectionGeneration , clientInitialConnection = initialConnection@@ -94,9 +169,12 @@ , clientInboxNuid = inboxNuid , clientServerInfo = serverInfo , clientAttemptIndex = attemptIndex+ , clientAttemptEpoch = attemptEpoch , clientStatus = status , clientConnection = conn , clientLogContext = logContext+ , clientManagedThreads = managedThreads+ , clientSubscriptionGate = subscriptionGate } runClient :: ClientState -> AppM a -> IO a@@ -112,32 +190,182 @@ connection :: ClientState -> Conn connection = clientConnection -enqueue :: ClientState -> QueueItem -> IO ()-enqueue client item = do- let Queue queueEnqueue _ _ _ = clientQueue client- result <- queueEnqueue item- case result of- Left err ->- runClient client $- logMessage Error ("enqueueing item failed: " ++ err)- Right () ->- pure ()+enqueue :: ClientState -> QueueItem -> IO (Either String ())+enqueue client = Queue.enqueue (clientQueue client) +tryEnqueue :: ClientState -> QueueItem -> IO TryEnqueueResult+tryEnqueue client = Queue.tryEnqueue (clientQueue client)+ openQueue :: ClientState -> IO ()-openQueue = open . clientQueue+openQueue = Queue.open . clientQueue closeQueue :: ClientState -> IO ()-closeQueue = close . clientQueue+closeQueue = Queue.close . clientQueue -pushPingAction :: ClientState -> IO () -> IO ()-pushPingAction client action =- atomically $ writeTQueue (clientPings client) action+discardConnectionScoped :: ClientState -> IO ()+discardConnectionScoped = Queue.discardConnectionScoped . clientQueue -runNextPingAction :: ClientState -> IO ()-runNextPingAction client = do- action <- atomically $ tryReadTQueue (clientPings client)- fromMaybe (pure ()) action+registerPingWaiter :: ClientState -> TMVar PingResult -> IO (Either ClientExitReason ())+registerPingWaiter client waiter = atomically $ do+ status <- readTVar (clientStatus client)+ case status of+ ConnectionConnecting -> retry+ ConnectionReconnecting -> retry+ ConnectionConnected -> do+ generation <- readTVar (clientConnectionGeneration client)+ writeTQueue (clientPings client) (generation, waiter)+ pure (Right ())+ ConnectionClosing reason -> pure (Left reason)+ ConnectionClosed reason -> pure (Left reason) +registerPingWaiterAndEnqueue+ :: ClientState+ -> TMVar PingResult+ -> QueueItem+ -> IO (Either ClientExitReason Int)+registerPingWaiterAndEnqueue client waiter item = do+ connected <- atomically (waitForConnectedGeneration client)+ case connected of+ Left reason -> pure (Left reason)+ Right generation -> atomically $ do+ result <- enqueueOnGenerationSTM client generation item+ case result of+ Left reason -> pure (Left reason)+ Right () -> do+ writeTQueue (clientPings client) (generation, waiter)+ pure (Right generation)++enqueueOnGenerationTracked+ :: ClientState+ -> Int+ -> TMVar Int+ -> QueueItem+ -> IO (Either ClientExitReason Int)+enqueueOnGenerationTracked client expectedGeneration committedGeneration item =+ atomically $ do+ result <- enqueueOnGenerationSTM client expectedGeneration item+ case result of+ Left reason -> pure (Left reason)+ Right () -> do+ putTMVar committedGeneration expectedGeneration+ pure (Right expectedGeneration)++enqueuePublishOnConnected+ :: ClientState+ -> Int+ -> QueueItem+ -> IO PublishEnqueueResult+enqueuePublishOnConnected client actualSize item =+ atomically (enqueuePublishOnConnectedSTM client Nothing actualSize item)++enqueuePublishOnConnectedGenerationTracked+ :: ClientState+ -> TMVar Int+ -> Int+ -> QueueItem+ -> IO PublishEnqueueResult+enqueuePublishOnConnectedGenerationTracked client committedGeneration actualSize item =+ atomically $+ enqueuePublishOnConnectedSTM client (Just committedGeneration) actualSize item++enqueuePublishOnConnectedSTM+ :: ClientState+ -> Maybe (TMVar Int)+ -> Int+ -> QueueItem+ -> STM PublishEnqueueResult+enqueuePublishOnConnectedSTM client committedGeneration actualSize item = do+ status <- readTVar (clientStatus client)+ case status of+ ConnectionConnecting -> retry+ ConnectionReconnecting -> retry+ ConnectionConnected -> do+ generation <- readTVar (clientConnectionGeneration client)+ serverInfo <- readTVar (clientServerInfo client)+ let maximumSize = effectivePayloadLimit client serverInfo+ if actualSize > maximumSize+ then pure (PublishTooLarge maximumSize)+ else do+ queued <- Queue.enqueueSTM (clientQueue client) item+ case queued of+ Left _ -> pure (PublishConnectionClosed ExitResetRequested)+ Right () -> do+ mapM_ (`putTMVar` generation) committedGeneration+ pure (PublishEnqueued generation)+ ConnectionClosing reason -> pure (PublishConnectionClosed reason)+ ConnectionClosed reason -> pure (PublishConnectionClosed reason)++effectivePayloadLimit :: ClientState -> Maybe Info -> Int+effectivePayloadLimit client =+ maybe+ (messageLimit (clientConfig client))+ (min (messageLimit (clientConfig client)) . Info.max_payload)++waitForConnectedGeneration :: ClientState -> STM (Either ClientExitReason Int)+waitForConnectedGeneration client = do+ status <- readTVar (clientStatus client)+ case status of+ ConnectionConnecting -> retry+ ConnectionReconnecting -> retry+ ConnectionConnected ->+ Right <$> readTVar (clientConnectionGeneration client)+ ConnectionClosing reason -> pure (Left reason)+ ConnectionClosed reason -> pure (Left reason)++enqueueOnGenerationSTM+ :: ClientState+ -> Int+ -> QueueItem+ -> STM (Either ClientExitReason ())+enqueueOnGenerationSTM client expectedGeneration item = do+ status <- readTVar (clientStatus client)+ generation <- readTVar (clientConnectionGeneration client)+ case status of+ ConnectionConnecting -> pure (Left ExitResetRequested)+ ConnectionReconnecting -> pure (Left ExitResetRequested)+ ConnectionConnected -> do+ if generation /= expectedGeneration+ then pure (Left ExitResetRequested)+ else Queue.enqueueSTM (clientQueue client) item >>= \case+ Right () -> pure (Right ())+ Left _ -> pure (Left ExitResetRequested)+ ConnectionClosing reason -> pure (Left reason)+ ConnectionClosed reason -> pure (Left reason)++tryEnqueueOnGeneration+ :: ClientState+ -> Int+ -> QueueItem+ -> IO TryEnqueueResult+tryEnqueueOnGeneration client expectedGeneration item = atomically $ do+ status <- readTVar (clientStatus client)+ generation <- readTVar (clientConnectionGeneration client)+ if status == ConnectionConnected && generation == expectedGeneration+ then Queue.tryEnqueueSTM (clientQueue client) item+ else pure Queue.TryQueueClosed++resolveNextPing :: ClientState -> IO ()+resolveNextPing client =+ atomically $ do+ generation <- readTVar (clientConnectionGeneration client)+ resolveNextPingSTM client generation++resolveNextPingSTM :: ClientState -> Int -> STM ()+resolveNextPingSTM client expectedGeneration = do+ currentGeneration <- readTVar (clientConnectionGeneration client)+ when (expectedGeneration == currentGeneration) resolveCurrent+ where+ resolveCurrent = do+ pending <- tryReadTQueue (clientPings client)+ case pending of+ Nothing -> pure ()+ Just (generation, waiter) -> do+ when (generation == expectedGeneration) $+ void (tryPutTMVar waiter PingReceived)++clearPendingPings :: ClientState -> IO ()+clearPendingPings = atomically . clearPendingPingsSTM+ nextSid :: ClientState -> IO SID nextSid client = atomically $ do@@ -176,43 +404,215 @@ updateLogContext loggerApi (clientLogContext client) $ \ctx -> ctx { lcServer = Just (host ++ ":" ++ show port) } -readStatus :: ClientState -> IO ClientStatus+notifyServerError :: ClientState -> ServerError -> IO ServerErrorEnqueueResult+notifyServerError client = atomically . enqueueServerErrorSTM client++enqueueServerErrorSTM :: ClientState -> ServerError -> STM ServerErrorEnqueueResult+enqueueServerErrorSTM client serverError = do+ pending <- readTVar (clientPendingServerCallbacks client)+ pendingBytes <- readTVar (clientPendingServerCallbackBytes client)+ let reasonBytes = BS.length (serverErrorReason serverError)+ hasCapacity =+ pending < serverCallbackLimit+ && reasonBytes <= serverCallbackByteLimit - pendingBytes+ if hasCapacity+ then do+ modifyTVar' (clientPendingServerCallbacks client) (+ 1)+ modifyTVar' (clientPendingServerCallbackBytes client) (+ reasonBytes)+ writeTQueue (clientCallbacks client) (CallbackServerError serverError)+ pure ServerErrorQueued+ else do+ alreadyReported <- readTVar (clientServerCallbackOverflowed client)+ if alreadyReported+ then pure ServerErrorDropped+ else do+ writeTVar (clientServerCallbackOverflowed client) True+ pure ServerErrorDroppedReport++enqueueCallbackSTM :: ClientState -> ClientCallback -> STM Bool+enqueueCallbackSTM client event =+ case event of+ CallbackServerError serverError ->+ (== ServerErrorQueued) <$> enqueueServerErrorSTM client serverError+ CallbackConnectionEvent _ -> do+ writeTQueue (clientCallbacks client) event+ pure True++serverCallbackLimit :: Int+serverCallbackLimit = 256++serverCallbackByteLimit :: Int+serverCallbackByteLimit = 1024 * 1024++startCallbackWorker :: ClientState -> IO CallbackWorker+startCallbackWorker client = do+ done <- newEmptyTMVarIO+ thread <- forkIOWithUnmask $ \unmask -> do+ threadId <- myThreadId+ let signalDone = atomically (void (tryPutTMVar done ()))+ registerManagedThread client threadId+ (unmask loop `finally` unregisterManagedThread client threadId)+ `finally` signalDone+ pure (CallbackWorker thread done)+ where+ loop = do+ next <- atomically $+ (Just <$> readTQueue (clientCallbacks client))+ `orElse` do+ status <- readTVar (clientStatus client)+ case status of+ ConnectionClosed _ -> do+ empty <- isEmptyTQueue (clientCallbacks client)+ check empty+ pure Nothing+ _ -> retry+ case next of+ Nothing -> pure ()+ Just (CallbackServerError serverError) -> do+ atomically $ do+ modifyTVar' (clientPendingServerCallbacks client) (subtract 1)+ modifyTVar'+ (clientPendingServerCallbackBytes client)+ (subtract (BS.length (serverErrorReason serverError)))+ pending <- readTVar (clientPendingServerCallbacks client)+ pendingBytes <- readTVar (clientPendingServerCallbackBytes client)+ when+ ( pending <= serverCallbackLimit `div` 2+ && pendingBytes <= serverCallbackByteLimit `div` 2+ )+ (writeTVar (clientServerCallbackOverflowed client) False)+ result <- trySync (serverErrorHandler (clientConfig client) serverError)+ case result of+ Right () -> pure ()+ Left _ ->+ runClient client $+ logMessage Error "server error handler failed"+ loop+ Just (CallbackConnectionEvent event) -> do+ result <- trySync (connectionEventHandler (clientConfig client) event)+ case result of+ Right () -> pure ()+ Left _ ->+ runClient client $+ logMessage Error "connection event handler failed"+ loop++waitForCallbackWorker :: CallbackWorker -> STM ()+waitForCallbackWorker = readTMVar . callbackWorkerDone++stopCallbackWorker :: CallbackWorker -> IO ()+stopCallbackWorker worker = do+ killThread (callbackWorkerThreadId worker)+ atomically (waitForCallbackWorker worker)++readStatus :: ClientState -> IO ConnectionState readStatus = readTVarIO . clientStatus +waitForConnected :: ClientState -> STM (Either ClientExitReason ())+waitForConnected client = do+ status <- readTVar (clientStatus client)+ case status of+ ConnectionConnecting -> retry+ ConnectionConnected -> pure (Right ())+ ConnectionReconnecting -> retry+ ConnectionClosing reason -> pure (Left reason)+ ConnectionClosed reason -> pure (Left reason)++setReconnecting :: ClientState -> IO Int+setReconnecting client =+ atomically $ do+ generation <- readTVar (clientConnectionGeneration client)+ previous <- readTVar (clientStatus client)+ case previous of+ ConnectionClosing _ -> pure ()+ ConnectionClosed _ -> pure ()+ ConnectionConnected -> do+ writeTVar (clientStatus client) ConnectionReconnecting+ void $ enqueueCallbackSTM client+ (CallbackConnectionEvent ConnectionEventDisconnected)+ _ -> writeTVar (clientStatus client) ConnectionReconnecting+ modifyTVar' (clientAttemptEpoch client) (+ 1)+ clearPendingPingsSTM client+ Queue.closeAndDiscardConnectionScoped (clientQueue client)+ pure generation+ setClosing :: ClientState -> ClientExitReason -> IO () setClosing client reason =- atomically . modifyTVar' (clientStatus client) $ \case- Closed result -> Closed result- Closing result -> Closing result- Running -> Closing reason+ atomically $ do+ modifyTVar' (clientStatus client) $ \case+ ConnectionClosed result -> ConnectionClosed result+ ConnectionClosing result -> ConnectionClosing result+ ConnectionConnecting -> ConnectionClosing reason+ ConnectionConnected -> ConnectionClosing reason+ ConnectionReconnecting -> ConnectionClosing reason+ modifyTVar' (clientAttemptEpoch client) (+ 1)+ clearPendingPingsSTM client+ Queue.closeAndDiscardAll (clientQueue client) +beginConnectionAttempt :: ClientState -> IO (Maybe Int)+beginConnectionAttempt client = atomically $ do+ status <- readTVar (clientStatus client)+ case status of+ ConnectionConnecting -> begin+ ConnectionReconnecting -> begin+ ConnectionConnected -> pure Nothing+ ConnectionClosing _ -> pure Nothing+ ConnectionClosed _ -> pure Nothing+ where+ begin = do+ modifyTVar' (clientAttemptEpoch client) (+ 1)+ epoch <- readTVar (clientAttemptEpoch client)+ pure (Just epoch)++activateConnectionAttempt :: ClientState -> Int -> IO Bool+activateConnectionAttempt client attemptEpoch = atomically $ do+ status <- readTVar (clientStatus client)+ currentEpoch <- readTVar (clientAttemptEpoch client)+ if currentEpoch /= attemptEpoch+ then pure False+ else case status of+ ConnectionConnecting -> activate+ ConnectionReconnecting -> activate+ ConnectionConnected -> pure False+ ConnectionClosing _ -> pure False+ ConnectionClosed _ -> pure False+ where+ activate = pure True+ markClosed :: ClientState -> ClientExitReason -> IO (Maybe ClientExitReason) markClosed client fallbackReason = atomically $ do+ clearPendingPingsSTM client+ Queue.closeAndDiscardAll (clientQueue client) status <- readTVar (clientStatus client) case status of- Closed _ ->+ ConnectionClosed _ -> pure Nothing- Closing reason -> do- writeTVar (clientStatus client) (Closed reason)+ ConnectionClosing reason -> do+ writeTVar (clientStatus client) (ConnectionClosed reason)+ void $ enqueueCallbackSTM client+ (CallbackConnectionEvent (ConnectionEventClosed reason)) pure (Just reason)- Running -> do- writeTVar (clientStatus client) (Closed fallbackReason)+ _ -> do+ writeTVar (clientStatus client) (ConnectionClosed fallbackReason)+ void $ enqueueCallbackSTM client+ (CallbackConnectionEvent (ConnectionEventClosed fallbackReason)) pure (Just fallbackReason) waitForClosed :: ClientState -> STM () waitForClosed client = do status <- readTVar (clientStatus client) case status of- Closed _ -> pure ()- _ -> retry+ ConnectionClosed _ -> pure ()+ _ -> retry waitForNotRunning :: ClientState -> STM () waitForNotRunning client = do status <- readTVar (clientStatus client) case status of- Running -> retry- _ -> pure ()+ ConnectionClosing _ -> pure ()+ ConnectionClosed _ -> pure ()+ _ -> retry waitForServerInfo :: ClientState -> STM () waitForServerInfo client = do@@ -221,18 +621,48 @@ Just _ -> pure () Nothing -> retry -markConnected :: ClientState -> IO Bool-markConnected client =- atomically $ do- alreadyConnected <- readTVar (clientConnectedOnce client)- writeTVar (clientConnectedOnce client) True- pure alreadyConnected+connectedOnce :: ClientState -> IO Bool+connectedOnce = readTVarIO . clientConnectedOnce -markConnectionReady :: ClientState -> IO ()-markConnectionReady client =+markConnectionReady :: ClientState -> Int -> IO Bool+markConnectionReady client attemptEpoch = do+ (accepted, rejectedPublishes) <-+ markConnectionReadyWithRejectedPublishes client attemptEpoch+ sequence_ rejectedPublishes+ pure accepted++markConnectionReadyWithRejectedPublishes+ :: ClientState+ -> Int+ -> IO (Bool, [IO ()])+markConnectionReadyWithRejectedPublishes client attemptEpoch = atomically $ do- modifyTVar' (clientConnectionGeneration client) (+ 1)- void (tryPutTMVar (clientInitialConnection client) (Right ()))+ status <- readTVar (clientStatus client)+ currentEpoch <- readTVar (clientAttemptEpoch client)+ if currentEpoch /= attemptEpoch+ then pure (False, [])+ else case status of+ ConnectionConnecting -> markReady+ ConnectionReconnecting -> markReady+ ConnectionConnected -> pure (False, [])+ ConnectionClosing _ -> pure (False, [])+ ConnectionClosed _ -> pure (False, [])+ where+ markReady = do+ wasConnected <- readTVar (clientConnectedOnce client)+ serverInfo <- readTVar (clientServerInfo client)+ rejectedPublishes <-+ Queue.discardOversizedPayloads+ (clientQueue client)+ (effectivePayloadLimit client serverInfo)+ modifyTVar' (clientConnectionGeneration client) (+ 1)+ writeTVar (clientConnectedOnce client) True+ writeTVar (clientStatus client) ConnectionConnected+ Queue.openSTM (clientQueue client)+ void (tryPutTMVar (clientInitialConnection client) (Right ()))+ when wasConnected . void $ enqueueCallbackSTM client+ (CallbackConnectionEvent ConnectionEventReconnected)+ pure (True, rejectedPublishes) waitForInitialConnection :: ClientState -> STM (Either ConnectError ()) waitForInitialConnection = readTMVar . clientInitialConnection@@ -247,12 +677,55 @@ waitForConnectionGenerationAfter :: ClientState -> Int -> STM () waitForConnectionGenerationAfter client generation = do+ status <- readTVar (clientStatus client) current <- readTVar (clientConnectionGeneration client)- unless (current > generation) retry+ unless (status == ConnectionConnected && current > generation) retry +waitForConnectionGenerationLoss :: ClientState -> Int -> STM ClientExitReason+waitForConnectionGenerationLoss client expectedGeneration = do+ status <- readTVar (clientStatus client)+ generation <- readTVar (clientConnectionGeneration client)+ case status of+ ConnectionConnected+ | generation == expectedGeneration -> retry+ | otherwise -> pure ExitResetRequested+ ConnectionClosing reason -> pure reason+ ConnectionClosed reason -> pure reason+ ConnectionConnecting -> pure ExitResetRequested+ ConnectionReconnecting -> pure ExitResetRequested+ readAttemptIndex :: ClientState -> IO Int readAttemptIndex = readTVarIO . clientAttemptIndex incrementAttemptIndex :: ClientState -> IO () incrementAttemptIndex client = atomically (modifyTVar' (clientAttemptIndex client) (+ 1))++registerManagedThread :: ClientState -> ThreadId -> IO ()+registerManagedThread client threadId =+ atomically $ modifyTVar' (clientManagedThreads client) (threadId :)++unregisterManagedThread :: ClientState -> ThreadId -> IO ()+unregisterManagedThread client threadId =+ atomically $ modifyTVar' (clientManagedThreads client) (delete threadId)++isManagedThread :: ClientState -> ThreadId -> IO Bool+isManagedThread client threadId =+ elem threadId <$> readTVarIO (clientManagedThreads client)++withSubscriptionGate :: ClientState -> IO a -> IO a+withSubscriptionGate client =+ bracket_+ (atomically (takeTMVar (clientSubscriptionGate client)))+ (atomically (putTMVar (clientSubscriptionGate client) ()))++clearPendingPingsSTM :: ClientState -> STM ()+clearPendingPingsSTM client = loop+ where+ loop = do+ waiter <- tryReadTQueue (clientPings client)+ case waiter of+ Nothing -> pure ()+ Just (_, result) -> do+ void (tryPutTMVar result PingConnectionLost)+ loop
internal/Policy/State/Types.hs view
@@ -11,28 +11,34 @@ , ConnectAttemptError (..) , ConnectFailure (..) , ServerError+ , ServerErrorKind (..) , serverErrorReason+ , serverErrorKind , serverErrorFromProtocol , ClientExitReason (..)- , ClientStatus (..)+ , ConnectionEvent (..)+ , ConnectionState (..) ) where import Data.ByteString (ByteString)+import qualified Data.ByteString as BS import Lib.Logger.Types (LoggerConfig) import Types.Connect (Connect) import qualified Types.Err as Err import Types.TLS data ClientConfig = ClientConfig- { connectionAttempts :: Int- , connectTimeoutMicros :: Int- , callbackConcurrency :: Int- , messageLimit :: Int- , connectConfig :: Connect- , loggerConfig :: LoggerConfig- , tlsConfig :: Maybe TLSConfig- , exitAction :: ClientExitReason -> IO ()- , connectOptions :: [(String, Int)]+ { connectionAttempts :: Int+ , connectTimeoutMicros :: Int+ , callbackConcurrency :: Int+ , messageLimit :: Int+ , connectConfig :: Connect+ , loggerConfig :: LoggerConfig+ , tlsConfig :: Maybe TLSConfig+ , serverErrorHandler :: ServerError -> IO ()+ , connectionEventHandler :: ConnectionEvent -> IO ()+ , exitAction :: ClientExitReason -> IO ()+ , connectOptions :: [(String, Int)] } -- | Why an initial connection could not be established.@@ -59,15 +65,51 @@ -- -- Its wire representation is intentionally opaque so that new server error -- categories do not force changes to the public type.-newtype ServerError = ServerError ByteString+data ServerError = ServerError ServerErrorKind ByteString+ deriving (Eq)++instance Show ServerError where+ showsPrec precedence (ServerError _ reason) =+ showParen (precedence > applicationPrecedence) $+ showString "ServerError " . showsPrec (applicationPrecedence + 1) reason+ where+ applicationPrecedence = 10++-- | Stable categories for server-reported protocol errors.+data ServerErrorKind = ServerErrorAuthentication | ServerErrorPermission | ServerErrorProtocol | ServerErrorResourceLimit | ServerErrorConnection | ServerErrorUnknown deriving (Eq, Show) serverErrorReason :: ServerError -> ByteString-serverErrorReason (ServerError reason) = reason+serverErrorReason (ServerError _ reason) = reason +serverErrorKind :: ServerError -> ServerErrorKind+serverErrorKind (ServerError kind _) = kind+ serverErrorFromProtocol :: Err.Err -> ServerError-serverErrorFromProtocol = ServerError . Err.errReason+serverErrorFromProtocol err =+ ServerError (classifyServerError err) (BS.copy (Err.errReason err)) +classifyServerError :: Err.Err -> ServerErrorKind+classifyServerError err =+ case err of+ Err.ErrAuthViolation _ -> ServerErrorAuthentication+ Err.ErrAuthTimeout _ -> ServerErrorAuthentication+ Err.ErrAuthExpired _ -> ServerErrorAuthentication+ Err.ErrAuthRevoked _ -> ServerErrorAuthentication+ Err.ErrAccountAuthExpired _ -> ServerErrorAuthentication+ Err.ErrPermViolation _ -> ServerErrorPermission+ Err.ErrUnknownOp _ -> ServerErrorProtocol+ Err.ErrInvalidProtocol _ -> ServerErrorProtocol+ Err.ErrInvalidSubject _ -> ServerErrorProtocol+ Err.ErrMaxControlLineEx _ -> ServerErrorResourceLimit+ Err.ErrMaxConnsEx _ -> ServerErrorResourceLimit+ Err.ErrSlowConsumer _ -> ServerErrorResourceLimit+ Err.ErrMaxPayload _ -> ServerErrorResourceLimit+ Err.ErrRoutePortConn _ -> ServerErrorConnection+ Err.ErrTlsRequired _ -> ServerErrorConnection+ Err.ErrStaleConn _ -> ServerErrorConnection+ Err.ErrErr _ -> ServerErrorUnknown+ data ClientExitReason = ExitClosedByUser | ExitRetriesExhausted (Maybe String) | ExitServerError ServerError@@ -75,7 +117,16 @@ | ExitResetRequested deriving (Eq, Show) -data ClientStatus = Running- | Closing ClientExitReason- | Closed ClientExitReason+-- | A completed connection lifecycle transition.+data ConnectionEvent = ConnectionEventDisconnected+ | ConnectionEventReconnected+ | ConnectionEventClosed ClientExitReason+ deriving (Eq, Show)++-- | Current lifecycle state of a client connection.+data ConnectionState = ConnectionConnecting+ | ConnectionConnected+ | ConnectionReconnecting+ | ConnectionClosing ClientExitReason+ | ConnectionClosed ClientExitReason deriving (Eq, Show)
internal/Policy/Subscription/Store.hs view
@@ -3,20 +3,28 @@ , DispatchResult (..) , newSubscriptionStore , register+ , registerWithAcceptance+ , registerWithDispatchHooks , unregister , dispatchMessage , active+ , closeStore , startWorkers , awaitCallbackDrain+ , enqueueControl , startExpiryWorker , awaitNoTrackedExpiries , hasTrackedExpiries ) where -import Control.Concurrent (forkIO, threadDelay)+import Control.Concurrent+ ( ThreadId+ , forkIOWithUnmask+ , threadDelay+ ) import Control.Concurrent.STM import Control.Exception (SomeException, finally)-import Control.Monad (unless, void, when)+import Control.Monad (unless, when) import qualified Data.Heap as Heap import qualified Data.Map as Map import Data.Time.Clock@@ -35,6 +43,8 @@ data SubscriptionCallback = SubscriptionCallback { deliverMessage :: Maybe M.Msg -> IO () , dropMessage :: IO ()+ , acceptMessage :: STM ()+ , rejectMessage :: STM () } data DispatchResult = DispatchMissing@@ -70,7 +80,40 @@ SubscriptionState Map.empty Heap.empty Map.empty Map.empty register :: SubscriptionStore -> SID -> SubscriptionMeta -> SubscribeConfig -> (Maybe M.Msg -> IO ()) -> IO () -> IO ()-register store sid meta cfg callback onDropped = do+register store sid meta cfg =+ registerWithAcceptance store sid meta cfg (pure ())++registerWithAcceptance+ :: SubscriptionStore+ -> SID+ -> SubscriptionMeta+ -> SubscribeConfig+ -> STM ()+ -> (Maybe M.Msg -> IO ())+ -> IO ()+ -> IO ()+registerWithAcceptance store sid meta cfg onAccepted callback onDropped = do+ registerWithDispatchHooks+ store+ sid+ meta+ cfg+ onAccepted+ (pure ())+ callback+ onDropped++registerWithDispatchHooks+ :: SubscriptionStore+ -> SID+ -> SubscriptionMeta+ -> SubscribeConfig+ -> STM ()+ -> STM ()+ -> (Maybe M.Msg -> IO ())+ -> IO ()+ -> IO ()+registerWithDispatchHooks store sid meta cfg onAccepted onRejected callback onDropped = do expiryAt <- case expiry cfg of Nothing -> pure Nothing Just ttl -> Just . addUTCTime ttl <$> getCurrentTime@@ -80,12 +123,12 @@ { subscriptionCallbacks = Map.insert sid- (SubscriptionCallback callback onDropped)+ (SubscriptionCallback callback onDropped onAccepted onRejected) (subscriptionCallbacks state) , subscriptionMeta = Map.insert sid meta (subscriptionMeta state) } in case expiryAt of- Just expiresAt | isReply meta ->+ Just expiresAt | tracksSubscriptionExpiry meta -> trackSubscriptionExpiry sid expiresAt stateWithCallback _ -> stateWithCallback@@ -104,8 +147,9 @@ Nothing -> pure (DispatchMissing, pure ()) Just callback -> do+ acceptMessage callback let shouldRemove =- maybe False isReply (Map.lookup sid (subscriptionMeta state))+ maybe False isOneShotSubscription (Map.lookup sid (subscriptionMeta state)) nextState = if shouldRemove then removeSubscriptionLocal sid state@@ -118,6 +162,7 @@ enqueueDeliverySTM store messageBytes (deliverMessage callback (Just msg)) pure (DispatchQueued, pure ()) else do+ rejectMessage callback alreadySlow <- readTVar (storeSlowConsumer store) unless alreadySlow $ do writeTVar (storeSlowConsumer store) True@@ -130,7 +175,11 @@ active store = Map.toList . subscriptionMeta <$> readTVarIO (storeState store) -startWorkers :: Int -> SubscriptionStore -> STM () -> (SomeException -> IO ()) -> IO ()+closeStore :: SubscriptionStore -> IO ()+closeStore store =+ atomically (writeTVar (storeState store) emptySubscriptionState)++startWorkers :: Int -> SubscriptionStore -> STM () -> (SomeException -> IO ()) -> IO [ThreadId] startWorkers concurrency store = startWorkerPool (max 1 concurrency) (storeCallbackQueue store) @@ -139,9 +188,12 @@ pending <- readTVar (storeCallbackPending store) check (pending == 0) -startExpiryWorker :: SubscriptionStore -> IO Bool -> IO ()+enqueueControl :: SubscriptionStore -> IO () -> IO ()+enqueueControl store = atomically . enqueueControlCallbackSTM store++startExpiryWorker :: SubscriptionStore -> IO Bool -> IO ThreadId startExpiryWorker store shouldStop =- void . forkIO $ loop+ forkIOWithUnmask (\unmask -> unmask loop) where loop = do stop <- shouldStop
internal/Policy/Subscription/Types.hs view
@@ -1,6 +1,10 @@ module Subscription.Types ( SubscribeConfig (..)+ , SubscriptionKind (..) , SubscriptionMeta (..)+ , isOneShotSubscription+ , isResumableSubscription+ , tracksSubscriptionExpiry , PendingLimits (..) , defaultPendingLimits ) where@@ -13,12 +17,28 @@ , subscribeQueueGroup :: Maybe Subject } +data SubscriptionKind = StandardSubscription | OneShotSubscription | RequestReplySubscription+ deriving (Eq, Show)+ data SubscriptionMeta = SubscriptionMeta { subject :: Subject , queueGroup :: Maybe Subject- , isReply :: Bool+ , kind :: SubscriptionKind } deriving (Eq, Show)++isOneShotSubscription :: SubscriptionMeta -> Bool+isOneShotSubscription meta =+ case kind meta of+ StandardSubscription -> False+ OneShotSubscription -> True+ RequestReplySubscription -> True++isResumableSubscription :: SubscriptionMeta -> Bool+isResumableSubscription meta = kind meta /= RequestReplySubscription++tracksSubscriptionExpiry :: SubscriptionMeta -> Bool+tracksSubscriptionExpiry meta = kind meta == OneShotSubscription data PendingLimits = PendingLimits { pendingMessageLimit :: Int
jetstream/JetStream/Consumer.hs view
@@ -19,6 +19,7 @@ , consumerNamesRequest , consumerPauseRequest , consumerResetRequest+ , validateConsumerConfigRequest ) import JetStream.Error ( JetStreamError (JetStreamDecodeError)@@ -93,11 +94,15 @@ case consumerSubject context stream action target of Left err -> Left err- Right subject ->- Right+ Right subject -> do+ let config =+ applyConsumerTarget target . applyConsumerKind kind $+ consumerConfigRequest options+ validateConsumerConfigRequest config+ pure ( subject , actionValue- , applyConsumerTarget target . applyConsumerKind kind $ consumerConfigRequest options+ , config ) where actionValue =
jetstream/JetStream/Consumer/API.hs view
@@ -31,6 +31,8 @@ , consumerConfigMaxDeliver , consumerConfigMaxWaiting , consumerConfigMaxAckPending+ , consumerConfigBackoff+ , consumerConfigMaxRequestBatch , consumerConfigInactiveThreshold , consumerConfigIdleHeartbeat , consumerConfigHeadersOnly@@ -80,6 +82,7 @@ , ReplayPolicy (..) , withConsumerAckPolicy , withConsumerAckWait+ , withConsumerBackoff , withConsumerDescription , withConsumerDeliverGroup , withConsumerDeliverPolicy@@ -89,6 +92,7 @@ , withConsumerInactiveThreshold , withConsumerListOffset , withConsumerMaxAckPending+ , withConsumerMaxRequestBatch , withConsumerMaxDeliver , withConsumerMaxWaiting , withConsumerMemoryStorage
jetstream/JetStream/Consumer/Types.hs view
@@ -25,6 +25,7 @@ , applyConsumerKind , applyConsumerTarget , consumerConfigRequest+ , validateConsumerConfigRequest , consumerActionValue , consumerListRequest , consumerNamesRequest@@ -32,6 +33,7 @@ , consumerResetRequest , withConsumerAckPolicy , withConsumerAckWait+ , withConsumerBackoff , withConsumerDescription , withConsumerDeliverGroup , withConsumerDurableName@@ -42,6 +44,7 @@ , withConsumerInactiveThreshold , withConsumerListOffset , withConsumerMaxAckPending+ , withConsumerMaxRequestBatch , withConsumerMaxDeliver , withConsumerMaxWaiting , withConsumerMemoryStorage@@ -55,11 +58,12 @@ import Data.Aeson import Data.Aeson.Types (Pair, Parser) import qualified Data.ByteString as BS+import Data.Int (Int64) import Data.Maybe (catMaybes) import qualified Data.Text as T import Data.Time.Clock (NominalDiffTime, UTCTime) import Data.Word (Word64)-import JetStream.Error (JetStreamError)+import JetStream.Error (JetStreamError (JetStreamDecodeError)) import JetStream.Types ( AckPolicy (..) , CallOption@@ -72,6 +76,7 @@ , applyCallOptions , byteStringToJSON , diffTimeNanosToJSON+ , diffTimeToNanoseconds , parseByteString ) @@ -109,6 +114,8 @@ , consumerConfigRequestMaxDeliver :: Maybe Int , consumerConfigRequestMaxWaiting :: Maybe Int , consumerConfigRequestMaxAckPending :: Maybe Int+ , consumerConfigRequestBackoff :: Maybe [NominalDiffTime]+ , consumerConfigRequestMaxRequestBatch :: Maybe Int , consumerConfigRequestInactiveThreshold :: Maybe NominalDiffTime , consumerConfigRequestIdleHeartbeat :: Maybe NominalDiffTime , consumerConfigRequestHeadersOnly :: Maybe Bool@@ -156,6 +163,8 @@ , consumerConfigRequestMaxDeliver = Nothing , consumerConfigRequestMaxWaiting = Nothing , consumerConfigRequestMaxAckPending = Nothing+ , consumerConfigRequestBackoff = Nothing+ , consumerConfigRequestMaxRequestBatch = Nothing , consumerConfigRequestInactiveThreshold = Nothing , consumerConfigRequestIdleHeartbeat = Nothing , consumerConfigRequestHeadersOnly = Nothing@@ -163,6 +172,42 @@ , consumerConfigRequestMemoryStorage = Nothing } +validateConsumerConfigRequest :: ConsumerConfigRequest -> Either JetStreamError ()+validateConsumerConfigRequest config = do+ validateMaxRequestBatch+ validateBackoffRange+ validateBackoffLength+ where+ validateMaxRequestBatch =+ case consumerConfigRequestMaxRequestBatch config of+ Just maxRequestBatch+ | maxRequestBatch < 0 ->+ Left (JetStreamDecodeError+ "consumer max request batch must be zero or greater")+ _ ->+ Right ()+ validateBackoffRange =+ case consumerConfigRequestBackoff config of+ Nothing ->+ Right ()+ Just backoff ->+ mapM_ validateBackoffDuration backoff+ validateBackoffDuration duration =+ let nanoseconds = diffTimeToNanoseconds duration+ maxNanoseconds = toInteger (maxBound :: Int64)+ in if nanoseconds >= 0 && nanoseconds <= maxNanoseconds+ then Right ()+ else Left (JetStreamDecodeError+ "consumer backoff durations must floor to nanoseconds between 0 and 9223372036854775807")+ validateBackoffLength =+ case (consumerConfigRequestMaxDeliver config, consumerConfigRequestBackoff config) of+ (Just maxDeliver, Just backoff)+ | maxDeliver > 0 && not (null (drop maxDeliver backoff)) ->+ Left (JetStreamDecodeError+ "consumer backoff length cannot exceed positive max deliver")+ _ ->+ Right ()+ applyConsumerTarget :: ConsumerTarget -> ConsumerConfigRequest -> ConsumerConfigRequest applyConsumerTarget EphemeralConsumer config = config@@ -233,6 +278,15 @@ withConsumerAckWait ackWait config = config { consumerConfigRequestAckWait = Just ackWait } +-- | Configure acknowledgment-timeout redelivery delays (NATS Server 2.7.1+).+-- This overrides AckWait and does not control explicit NAK redelivery. Values+-- are floored to whole nanoseconds. The first delay becomes AckWait, the final+-- delay repeats after the schedule, and the schedule cannot be longer than a+-- positive MaxDeliver. An empty list leaves the server default.+withConsumerBackoff :: [NominalDiffTime] -> ConsumerConfigOption+withConsumerBackoff backoff config =+ config { consumerConfigRequestBackoff = nonEmpty backoff }+ withConsumerMaxDeliver :: Int -> ConsumerConfigOption withConsumerMaxDeliver maxDeliver config = config { consumerConfigRequestMaxDeliver = Just maxDeliver }@@ -245,6 +299,12 @@ withConsumerMaxAckPending maxAckPending config = config { consumerConfigRequestMaxAckPending = Just maxAckPending } +-- | Limit a pull consumer request batch (NATS Server 2.7.0+). Zero leaves the+-- limit unset. This option is only meaningful for pull consumers.+withConsumerMaxRequestBatch :: Int -> ConsumerConfigOption+withConsumerMaxRequestBatch maxRequestBatch config =+ config { consumerConfigRequestMaxRequestBatch = nonZero maxRequestBatch }+ withConsumerInactiveThreshold :: NominalDiffTime -> ConsumerConfigOption withConsumerInactiveThreshold inactiveThreshold config = config { consumerConfigRequestInactiveThreshold = Just inactiveThreshold }@@ -280,6 +340,8 @@ , consumerConfigMaxDeliver :: Maybe Int , consumerConfigMaxWaiting :: Maybe Int , consumerConfigMaxAckPending :: Maybe Int+ , consumerConfigBackoff :: Maybe [NominalDiffTime]+ , consumerConfigMaxRequestBatch :: Maybe Int , consumerConfigInactiveThreshold :: Maybe NominalDiffTime , consumerConfigIdleHeartbeat :: Maybe NominalDiffTime , consumerConfigHeadersOnly :: Maybe Bool@@ -401,6 +463,8 @@ , maybePair "max_deliver" (consumerConfigRequestMaxDeliver config) , maybePair "max_waiting" (consumerConfigRequestMaxWaiting config) , maybePair "max_ack_pending" (consumerConfigRequestMaxAckPending config)+ , maybeJsonPair "backoff" (durationListToJSON <$> consumerConfigRequestBackoff config)+ , maybePair "max_batch" (consumerConfigRequestMaxRequestBatch config) , maybeJsonPair "inactive_threshold" (diffTimeNanosToJSON <$> consumerConfigRequestInactiveThreshold config) , maybeJsonPair "idle_heartbeat" (diffTimeNanosToJSON <$> consumerConfigRequestIdleHeartbeat config) , maybePair "headers_only" (consumerConfigRequestHeadersOnly config)@@ -425,6 +489,8 @@ , maybePair "max_deliver" (consumerConfigMaxDeliver config) , maybePair "max_waiting" (consumerConfigMaxWaiting config) , maybePair "max_ack_pending" (consumerConfigMaxAckPending config)+ , maybeJsonPair "backoff" (durationListToJSON <$> (consumerConfigBackoff config >>= nonEmpty))+ , maybePair "max_batch" (consumerConfigMaxRequestBatch config >>= nonZero) , maybeJsonPair "inactive_threshold" (diffTimeNanosToJSON <$> consumerConfigInactiveThreshold config) , maybeJsonPair "idle_heartbeat" (diffTimeNanosToJSON <$> consumerConfigIdleHeartbeat config) , maybePair "headers_only" (consumerConfigHeadersOnly config)@@ -449,6 +515,8 @@ <*> obj .:? "max_deliver" <*> obj .:? "max_waiting" <*> obj .:? "max_ack_pending"+ <*> parseOptionalDurationListField obj "backoff"+ <*> parseOptionalNonZeroIntField obj "max_batch" <*> parseOptionalDurationField obj "inactive_threshold" <*> parseOptionalDurationField obj "idle_heartbeat" <*> obj .:? "headers_only"@@ -636,6 +704,27 @@ parseOptionalDurationField :: Object -> Key -> Parser (Maybe NominalDiffTime) parseOptionalDurationField obj key = fmap nanosToDiffTime <$> obj .:? key++parseOptionalDurationListField :: Object -> Key -> Parser (Maybe [NominalDiffTime])+parseOptionalDurationListField obj key = do+ durations <- obj .:? key+ pure (durations >>= nonEmpty . fmap nanosToDiffTime)++parseOptionalNonZeroIntField :: Object -> Key -> Parser (Maybe Int)+parseOptionalNonZeroIntField obj key =+ fmap (>>= nonZero) (obj .:? key)++durationListToJSON :: [NominalDiffTime] -> Value+durationListToJSON =+ toJSON . fmap diffTimeToNanoseconds++nonEmpty :: [value] -> Maybe [value]+nonEmpty [] = Nothing+nonEmpty values = Just values++nonZero :: (Eq value, Num value) => value -> Maybe value+nonZero 0 = Nothing+nonZero value = Just value nanosToDiffTime :: Integer -> NominalDiffTime nanosToDiffTime nanoseconds =
jetstream/JetStream/Message.hs view
@@ -49,12 +49,15 @@ ( JetStreamContext (..) , requestTimeoutMicros )+import JetStream.Protocol.Headers (statusError)+import JetStream.Protocol.Request (requestMsg) import qualified JetStream.Protocol.Subject as Subject import JetStream.Types ( AckPolicy (AckNone) , ConsumerName , DeliverPolicy (DeliverByStartSequence) , JetStreamRequestOption+ , Payload , StreamName , Subject )@@ -70,10 +73,14 @@ consumePushMessages context deliverSubject (pushConsumeConfig options) handler , createOrderedConsumer = \stream options requestOptions -> createOrderedConsumerHandle context consumerAPI stream (orderedConsumerConfig options) requestOptions- , ack = \message _ -> ackMessage (contextClient context) message- , nak = \message _ -> nakMessage (contextClient context) message- , inProgress = \message _ -> inProgressMessage (contextClient context) message- , term = \message _ -> termMessage (contextClient context) message+ , ack = sendDisposition context ConfirmWhenRequested (ackPayload Ack)+ , ackSync = sendDisposition context ConfirmAlways (ackPayload Ack)+ , nak = sendDisposition context ConfirmWhenRequested (ackPayload Nak)+ , nakWithDelay = \message delay ->+ sendDisposition context ConfirmWhenRequested (nakDelayPayload delay) message+ , inProgress = sendDisposition context ConfirmWhenRequested (ackPayload InProgress)+ , term = sendDisposition context ConfirmWhenRequested (ackPayload Term)+ , termSync = sendDisposition context ConfirmAlways (ackPayload Term) } consumePushMessages@@ -373,6 +380,45 @@ pure (Left JetStreamNoReply) Just reply -> mapNatsResult <$> Nats.publish natsClient reply (ackPayload verb) []++data ConfirmationMode = ConfirmWhenRequested | ConfirmAlways++sendDisposition+ :: JetStreamContext+ -> ConfirmationMode+ -> Payload+ -> Message+ -> [JetStreamRequestOption]+ -> IO (Either JetStreamError ())+sendDisposition context confirmationMode body message requestOptions =+ case messageReplyTo message of+ Nothing ->+ pure (Left JetStreamNoReply)+ Just reply+ | confirmationRequired confirmationMode requestOptions ->+ confirmDisposition context reply body requestOptions+ | otherwise ->+ mapNatsResult <$> Nats.publish (contextClient context) reply body []++confirmDisposition+ :: JetStreamContext+ -> Subject+ -> Payload+ -> [JetStreamRequestOption]+ -> IO (Either JetStreamError ())+confirmDisposition context reply body requestOptions = do+ response <- requestMsg context reply body [] requestOptions+ pure $ do+ message <- response+ case statusError message of+ Just err ->+ Left err+ Nothing ->+ Right ()++confirmationRequired :: ConfirmationMode -> [JetStreamRequestOption] -> Bool+confirmationRequired ConfirmAlways _ = True+confirmationRequired ConfirmWhenRequested requestOptions = not (null requestOptions) mapNatsResult :: Either Nats.NatsError () -> Either JetStreamError () mapNatsResult =
jetstream/JetStream/Message/API.hs view
@@ -4,9 +4,14 @@ , consumePush , createOrderedConsumer , ack+ , ackSync , nak+ , nakWithDelay+ , NakDelay+ , nakDelay , inProgress , term+ , termSync , FetchOption , FetchWait (..) , Headers@@ -56,6 +61,7 @@ , Message (..) , MessageAPI , MessageMetadata+ , NakDelay , OrderedConsumer , OrderedConsumerOption , PullResponse@@ -63,6 +69,7 @@ , PushConsumeOption , PushSubscription , ack+ , ackSync , consumePush , createOrderedConsumer , fetch@@ -83,12 +90,15 @@ , messageStatus , messageSubject , nak+ , nakDelay+ , nakWithDelay , orderedConsumerInfo , pullResponseMessages , pullResponseStatus , stopOrderedConsumer , stopPushSubscription , term+ , termSync , withFetchBatch , withFetchWait , withOrderedConsumerDeliverPolicy
jetstream/JetStream/Message/Types.hs view
@@ -3,6 +3,7 @@ module JetStream.Message.Types ( MessageAPI (..) , AckVerb (..)+ , NakDelay (..) , FetchOption , FetchWait (..) , Headers@@ -26,6 +27,8 @@ , isStatusMessage , messageMetadata , nakPayload+ , nakDelay+ , nakDelayPayload , orderedConsumerConfig , pullRequestPayload , pushConsumeConfig@@ -46,6 +49,7 @@ import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as BC import Data.Char (isDigit, toLower)+import Data.Int (Int64) import Data.Maybe (isJust) import Data.Time.Clock (NominalDiffTime, UTCTime) import qualified Data.Time.Clock.POSIX as Time@@ -72,9 +76,12 @@ , consumePush :: Subject -> [PushConsumeOption] -> [JetStreamRequestOption] -> (Message -> IO ()) -> IO (Either JetStreamError PushSubscription) , createOrderedConsumer :: StreamName -> [OrderedConsumerOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError OrderedConsumer) , ack :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ())+ , ackSync :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) , nak :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ())+ , nakWithDelay :: Message -> NakDelay -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) , inProgress :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) , term :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ())+ , termSync :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) } data PullRequest = PullRequest@@ -213,6 +220,21 @@ data AckVerb = Ack | Nak | InProgress | Term deriving (Eq, Show) +-- | A positive, whole-nanosecond JetStream redelivery delay.+newtype NakDelay = NakDelay Int64+ deriving (Eq, Show)++-- | Build a delayed-NAK duration. Durations below one nanosecond, negative+-- durations, and values outside the JetStream signed 64-bit wire range are+-- rejected. Sub-nanosecond precision is rounded down.+nakDelay :: NominalDiffTime -> Maybe NakDelay+nakDelay duration+ | nanoseconds < 1 = Nothing+ | nanoseconds > toInteger (maxBound :: Int64) = Nothing+ | otherwise = Just (NakDelay (fromInteger nanoseconds))+ where+ nanoseconds = floor (toRational duration * 1000000000)+ ackPayload :: AckVerb -> Payload ackPayload Ack = "+ACK" ackPayload Nak = "-NAK"@@ -221,6 +243,15 @@ nakPayload :: Payload nakPayload = ackPayload Nak++nakDelayPayload :: NakDelay -> Payload+nakDelayPayload (NakDelay nanoseconds) =+ BS.concat+ [ nakPayload+ , " {\"delay\":"+ , BC.pack (show nanoseconds)+ , "}"+ ] inProgressPayload :: Payload inProgressPayload = ackPayload InProgress
jetstream/JetStream/Stream.hs view
@@ -14,22 +14,25 @@ import JetStream.Stream.API import JetStream.Stream.Types ( StreamAPI (..)+ , StreamConfigRequest , purgeStreamRequest , streamConfigRequest , streamListRequest , streamMessageDeleteRequest , streamMessageGetRequest , streamNamesRequest+ , validateStreamConfigRequest )+import JetStream.Types (JetStreamRequestOption, Subject) streamAPI :: JetStreamContext -> StreamAPI streamAPI context = StreamAPI { create = \name subjects options requestOptions -> let config = streamConfigRequest name subjects options- in Request.requestJSON context+ in requestStreamConfig context (Subject.streamCreateSubject context name)- (Just (toJSON config))+ config requestOptions , createOrUpdate = \name subjects options requestOptions -> do updateResult <- update (streamAPI context) name subjects options requestOptions@@ -41,9 +44,9 @@ pure other , update = \name subjects options requestOptions -> let config = streamConfigRequest name subjects options- in Request.requestJSON context+ in requestStreamConfig context (Subject.streamUpdateSubject context name)- (Just (toJSON config))+ config requestOptions , info = \streamName requestOptions -> Request.requestJSON context@@ -77,6 +80,19 @@ Request.requestJSON context (Subject.streamNamesSubject context) . Just . toJSON . streamNamesRequest }++requestStreamConfig+ :: JetStreamContext+ -> Subject+ -> StreamConfigRequest+ -> [JetStreamRequestOption]+ -> IO (Either JetStreamError StreamInfo)+requestStreamConfig context subject config requestOptions =+ case validateStreamConfigRequest config of+ Left err ->+ pure (Left err)+ Right () ->+ Request.requestJSON context subject (Just (toJSON config)) requestOptions isStreamNotFound :: JetStreamError -> Bool isStreamNotFound (JetStreamApiFailure err) =
jetstream/JetStream/Stream/API.hs view
@@ -27,6 +27,7 @@ , streamConfigMaxMessages , streamConfigMaxBytes , streamConfigMaxAge+ , streamConfigMaxMessageSize , streamConfigReplicas , streamConfigDuplicateWindow , streamConfigAllowDirect@@ -37,6 +38,7 @@ , withMaxMessages , withMaxBytes , withMaxAge+ , withMaxMessageSize , withReplicas , withDuplicateWindow , withAllowDirect
jetstream/JetStream/Stream/Types.hs view
@@ -8,13 +8,16 @@ , DiscardPolicy (..) , StreamConfig (..) , StreamConfigOption+ , StreamConfigRequest , streamConfigRequest+ , validateStreamConfigRequest , withRetention , withStorage , withDiscard , withMaxMessages , withMaxBytes , withMaxAge+ , withMaxMessageSize , withReplicas , withDuplicateWindow , withAllowDirect@@ -51,10 +54,11 @@ import Data.Aeson.Types (Pair, Parser) import qualified Data.ByteString as BS import qualified Data.ByteString.Base64 as Base64+import Data.Int (Int32) import Data.Maybe (catMaybes) import Data.Time.Clock (NominalDiffTime, UTCTime) import Data.Word (Word64)-import JetStream.Error (JetStreamError)+import JetStream.Error (JetStreamError (JetStreamDecodeError)) import JetStream.Types ( CallOption , DiscardPolicy (..)@@ -66,6 +70,7 @@ , Subject , applyCallOptions , byteStringToJSON+ , diffTimeToNanoseconds , parseByteString ) @@ -98,6 +103,7 @@ , streamConfigRequestMaxMessages :: Maybe Integer , streamConfigRequestMaxBytes :: Maybe Integer , streamConfigRequestMaxAge :: Maybe NominalDiffTime+ , streamConfigRequestMaxMessageSize :: Maybe Int32 , streamConfigRequestReplicas :: Maybe Int , streamConfigRequestDuplicateWindow :: Maybe NominalDiffTime , streamConfigRequestAllowDirect :: Maybe Bool@@ -118,11 +124,21 @@ , streamConfigRequestMaxMessages = Nothing , streamConfigRequestMaxBytes = Nothing , streamConfigRequestMaxAge = Nothing+ , streamConfigRequestMaxMessageSize = Nothing , streamConfigRequestReplicas = Nothing , streamConfigRequestDuplicateWindow = Nothing , streamConfigRequestAllowDirect = Nothing } +validateStreamConfigRequest :: StreamConfigRequest -> Either JetStreamError ()+validateStreamConfigRequest config =+ case streamConfigRequestMaxMessageSize config of+ Just maxMessageSize+ | maxMessageSize < (-1) ->+ Left (JetStreamDecodeError "stream max message size must be -1 or greater")+ _ ->+ Right ()+ withRetention :: RetentionPolicy -> StreamConfigOption withRetention retention config = config { streamConfigRequestRetention = Just retention }@@ -147,6 +163,12 @@ withMaxAge maxAge config = config { streamConfigRequestMaxAge = Just maxAge } +-- | Limit the total encoded bytes of one stored message, including headers.+-- Use @-1@ for unlimited.+withMaxMessageSize :: Int32 -> StreamConfigOption+withMaxMessageSize maxMessageSize config =+ config { streamConfigRequestMaxMessageSize = Just maxMessageSize }+ withReplicas :: Int -> StreamConfigOption withReplicas replicas config = config { streamConfigRequestReplicas = Just replicas }@@ -168,6 +190,7 @@ , streamConfigMaxMessages :: Integer , streamConfigMaxBytes :: Integer , streamConfigMaxAge :: NominalDiffTime+ , streamConfigMaxMessageSize :: Int32 , streamConfigReplicas :: Int , streamConfigDuplicateWindow :: Maybe NominalDiffTime , streamConfigAllowDirect :: Bool@@ -363,6 +386,7 @@ , maybePair "max_msgs" (streamConfigRequestMaxMessages config) , maybePair "max_bytes" (streamConfigRequestMaxBytes config) , maybeDurationPair "max_age" (streamConfigRequestMaxAge config)+ , maybePair "max_msg_size" (streamConfigRequestMaxMessageSize config) , maybePair "num_replicas" (streamConfigRequestReplicas config) , maybeDurationPair "duplicate_window" (streamConfigRequestDuplicateWindow config) , maybePair "allow_direct" (streamConfigRequestAllowDirect config)@@ -379,6 +403,7 @@ , Just ("max_msgs" .= streamConfigMaxMessages config) , Just ("max_bytes" .= streamConfigMaxBytes config) , Just ("max_age" .= durationToNanoseconds (streamConfigMaxAge config))+ , Just ("max_msg_size" .= streamConfigMaxMessageSize config) , Just ("num_replicas" .= streamConfigReplicas config) , maybeDurationPair "duplicate_window" (streamConfigDuplicateWindow config) , Just ("allow_direct" .= streamConfigAllowDirect config)@@ -396,6 +421,7 @@ <*> value .: "max_msgs" <*> value .: "max_bytes" <*> parseDurationField value "max_age"+ <*> value .:? "max_msg_size" .!= (-1) <*> value .: "num_replicas" <*> parseOptionalDurationField value "duplicate_window" <*> value .: "allow_direct"@@ -638,8 +664,8 @@ <*> value .: "time" durationToNanoseconds :: NominalDiffTime -> Integer-durationToNanoseconds duration =- round (realToFrac duration * (1000000000 :: Double))+durationToNanoseconds =+ diffTimeToNanoseconds nanosecondsToDuration :: Integer -> NominalDiffTime nanosecondsToDuration nanoseconds =
jetstream/JetStream/Types.hs view
@@ -23,6 +23,7 @@ , byteStringToJSON , parseByteString , diffTimeNanosToJSON+ , diffTimeToNanoseconds , applyCallOptions ) where @@ -259,7 +260,12 @@ diffTimeNanosToJSON :: NominalDiffTime -> Value diffTimeNanosToJSON diffTime =- Number (fromInteger (floor (realToFrac diffTime * (1000000000 :: Double))))+ Number (fromInteger (diffTimeToNanoseconds diffTime))++-- | Convert seconds to whole nanoseconds using exact arithmetic and floor.+diffTimeToNanoseconds :: NominalDiffTime -> Integer+diffTimeToNanoseconds =+ floor . (* 1000000000) . toRational instance FromJSON AccountInfo where parseJSON value@(Object obj) =
natskell.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: natskell-version: 1.0.0.1+version: 1.1.0.0 synopsis: A NATS client library written in Haskell tested-with: GHC ==8.8,@@ -110,7 +110,9 @@ jetstream library natskell-internal+ -- Private implementation library, also reused by the test suites. import: shared+ visibility: private build-depends: conduit >= 1.3 && < 1.4, attoparsec >= 0.14 && < 0.15,@@ -154,6 +156,7 @@ Lib.Logger Lib.Logger.Types Lib.LoggerAPI+ Lib.Exception Lib.CallOption Lib.WorkerPool Parser.Attoparsec@@ -253,6 +256,7 @@ JetStream.MessageSpec JetStream.ProtocolSpec JetStream.PublishSpec+ JetStream.StreamSpec test-suite fuzz-test import: shared
test/Integration/ClientSpec.hs view
@@ -43,10 +43,21 @@ pingWithCallback client action = void . forkIO $ void (ping client []) >> action +awaitClosed :: Client -> IO ()+awaitClosed client = do+ status <- connectionState client+ case status of+ ConnectionClosed _ -> pure ()+ _ -> threadDelay 1000 >> awaitClosed client+ isValidationFailure :: Either NatsError a -> Bool isValidationFailure (Left (NatsValidationError _)) = True isValidationFailure _ = False +isJetStreamDecodeFailure :: Either JetStream.JetStreamError a -> Bool+isJetStreamDecodeFailure (Left (JetStream.JetStreamDecodeError _)) = True+isJetStreamDecodeFailure _ = False+ tooLongMSG = "MSG a b 5000\r\n" <> BS.replicate 5000 _x <> "\r\n" headerBlock :: [(BS.ByteString, BS.ByteString)] -> BS.ByteString@@ -95,6 +106,7 @@ , capturedSubject :: BS.ByteString , capturedReplyTo :: BS.ByteString , capturedPayload :: BS.ByteString+ , capturedWire :: BS.ByteString } deriving (Eq, Show) @@ -169,8 +181,11 @@ pure value completeHandshake :: Socket -> IO BS.ByteString-completeHandshake sock = do- sendAll sock defaultINFO+completeHandshake sock = completeHandshakeWithInfo sock defaultINFO++completeHandshakeWithInfo :: Socket -> BS.ByteString -> IO BS.ByteString+completeHandshakeWithInfo sock info = do+ sendAll sock info bytes <- expectClientBytes sock@@ -179,6 +194,14 @@ sendAll sock "PONG\r\n" pure bytes +infoWithMaxPayload :: Int -> BS.ByteString+infoWithMaxPayload maximumPayload =+ BS.concat+ [ "INFO {\"server_id\":\"limit-server\",\"version\":\"1.0.0\",\"go\":\"go1\",\"host\":\"127.0.0.1\",\"port\":4222,\"max_payload\":"+ , C.pack (show maximumPayload)+ , ",\"proto\":1}\r\n"+ ]+ expectQueuedSub :: Socket -> BS.ByteString -> BS.ByteString -> BS.ByteString -> IO () expectQueuedSub sock subject queueGroup sid = expectClientCommand sock (BS.concat ["SUB ", subject, " ", queueGroup, " ", sid, "\r\n"])@@ -253,6 +276,7 @@ , capturedSubject = subject' , capturedReplyTo = reply , capturedPayload = payloadBytes+ , capturedWire = bytes } where wireLines = C.split '\n' bytes@@ -328,11 +352,11 @@ protoStreamInfoResponse :: BS.ByteString protoStreamInfoResponse =- "{\"config\":{\"name\":\"PROTO_STREAM\",\"subjects\":[\"PROTO.SUBJECT\"],\"retention\":\"limits\",\"storage\":\"memory\",\"discard\":\"new\",\"max_msgs\":1,\"max_bytes\":-1,\"max_age\":0,\"num_replicas\":1,\"allow_direct\":false},\"created\":\"2024-01-01T00:00:00Z\",\"state\":{\"messages\":0,\"bytes\":0,\"first_seq\":0,\"first_ts\":\"2024-01-01T00:00:00Z\",\"last_seq\":0,\"last_ts\":\"2024-01-01T00:00:00Z\",\"consumer_count\":0}}"+ "{\"config\":{\"name\":\"PROTO_STREAM\",\"subjects\":[\"PROTO.SUBJECT\"],\"retention\":\"limits\",\"storage\":\"memory\",\"discard\":\"new\",\"max_msgs\":1,\"max_bytes\":-1,\"max_age\":0,\"max_msg_size\":64,\"num_replicas\":1,\"allow_direct\":false},\"created\":\"2024-01-01T00:00:00Z\",\"state\":{\"messages\":0,\"bytes\":0,\"first_seq\":0,\"first_ts\":\"2024-01-01T00:00:00Z\",\"last_seq\":0,\"last_ts\":\"2024-01-01T00:00:00Z\",\"consumer_count\":0}}" protoConsumerInfoResponse :: BS.ByteString protoConsumerInfoResponse =- "{\"stream_name\":\"PROTO_STREAM\",\"name\":\"PROTO_CONSUMER\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\",\"filter_subjects\":[\"PROTO.A\",\"PROTO.B\"]}}"+ "{\"stream_name\":\"PROTO_STREAM\",\"name\":\"PROTO_CONSUMER\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\",\"filter_subjects\":[\"PROTO.A\",\"PROTO.B\"],\"backoff\":[100000001,250000000],\"max_batch\":32}}" protoPushConsumerInfoResponse :: BS.ByteString protoPushConsumerInfoResponse =@@ -442,6 +466,133 @@ sendAll serv "PONG\r\n" expectMVar "flush did not resolve after PONG" resultVar `shouldReturn` Right ()+ it "bounds flush when the server does not send PONG" $ \(serv, client, _, _) -> do+ resultVar <- newEmptyMVar+ void . forkIO $+ flush client [withFlushTimeout 0.05] >>= putMVar resultVar+ expectClientCommand serv "PING\r\n"++ timeout 500000 (takeMVar resultVar)+ `shouldReturn` Just (Left NatsRequestTimedOut)+ it "reports the terminal close reason to a pending ping" $ \(serv, client, _, _) -> do+ resultVar <- newEmptyMVar+ void . forkIO $ ping client [withPingTimeout 5] >>= putMVar resultVar+ expectClientCommand serv "PING\r\n"++ close client []++ expectMVar "pending ping did not receive terminal close" resultVar+ `shouldReturn` Left (NatsConnectionClosed ExitClosedByUser)+ it "cleans up a one-shot subscription cancelled after queue commit" $ \_ -> do+ queued <- newEmptyMVar+ release <- newEmptyMVar+ let capture entry =+ when (leMessage entry == "subscription commands queued") $ do+ putMVar queued ()+ takeMVar release+ bracket+ (startClientWith [withLogAction capture])+ stopClientSafely $ \(serv, client, _, _) -> do+ finished <- newEmptyMVar+ subscriptionThread <- forkIO $+ void (subscribeOnce client "CANCEL.ONCE" [] (const (pure ())))+ `finally` putMVar finished ()+ expectMVar "subscribe did not commit its command batch" queued+ expectClientCommand serv "SUB CANCEL.ONCE 1\r\nUNSUB 1 1\r\n"++ killThread subscriptionThread++ expectMVar "cancelled subscribe thread did not stop" finished+ expectClientCommand serv "UNSUB 1\r\n"+ it "removes an expired one-shot subscription from the server" $ \(serv, client, _, _) -> do+ Right subscription <-+ subscribeOnce client "EXPIRING.ONCE" [withSubscriptionExpiry 0.01] (const (pure ()))+ let sid = subscriptionSid subscription+ expectClientCommand serv $+ BS.concat ["SUB EXPIRING.ONCE ", sid, "\r\nUNSUB ", sid, " 1\r\n"]++ expectClientCommand serv (BS.concat ["UNSUB ", sid, "\r\n"])+ it "does not let a server error handler block a following PONG" $ \_ -> do+ observed <- newEmptyMVar+ releaseHandler <- newEmptyMVar+ let handler serverError = do+ putMVar observed (serverErrorReason serverError)+ takeMVar releaseHandler+ bracket+ (startClientWith [withServerErrorHandler handler])+ stopClientSafely $ \(serv, client, _, _) -> do+ pingResult <- newEmptyMVar+ void . forkIO $ ping client [] >>= putMVar pingResult+ expectClientCommand serv "PING\r\n"+ sendAll+ serv+ "-ERR 'Permissions Violation For Publish To PRIVATE'\r\nPONG\r\n"++ expectMVar "server error callback did not run" observed+ `shouldReturn` "Permissions Violation For Publish To PRIVATE"+ expectMVar "PONG waited for the server error handler" pingResult+ `shouldReturn` Right ()+ putMVar releaseHandler ()+ it "allows a server error handler to close its own client" $ \_ -> do+ clientRef <- newEmptyMVar+ handlerReturned <- newEmptyMVar+ let handler _ = do+ client <- readMVar clientRef+ close client []+ putMVar handlerReturned ()+ bracket+ (startClientWith [withServerErrorHandler handler])+ stopClientSafely $ \(serv, client, _, _) -> do+ putMVar clientRef client++ sendAll serv "-ERR 'Permissions Violation For Publish To PRIVATE'\r\n"++ void (expectMVar "server error handler deadlocked in close" handlerReturned)+ timeout 1000000 (awaitClosed client) `shouldReturn` Just ()+ it "allows a message callback to close its own client" $ \_ -> do+ clientRef <- newEmptyMVar+ callbackReturned <- newEmptyMVar+ bracket startClient stopClientSafely $ \(serv, client, _, _) -> do+ putMVar clientRef client+ Right subscription <- subscribe client "CLOSE.FROM.CALLBACK" [] $ \_ -> do+ callbackClient <- readMVar clientRef+ close callbackClient []+ putMVar callbackReturned ()+ expectClientCommand serv+ ("SUB CLOSE.FROM.CALLBACK " <> subscriptionSid subscription <> "\r\n")++ sendAll serv (msgFrame "CLOSE.FROM.CALLBACK" (subscriptionSid subscription) "close")++ void (expectMVar "message callback deadlocked in close" callbackReturned)+ timeout 1000000 (awaitClosed client) `shouldReturn` Just ()+ it "does not put message or malformed bytes in debug logs" $ \_ -> do+ logMessages <- newIORef ([] :: [String])+ let capture entry =+ atomicModifyIORef' logMessages $ \messages ->+ (leMessage entry : messages, ())+ bracket+ (startClientWith [withLogAction capture])+ stopClientSafely $ \(serv, client, _, _) -> do+ delivered <- newEmptyMVar+ Right subscription <- subscribe client "SECRET.LOG" [] (putMVar delivered)+ let subscriptionId = subscriptionSid subscription+ expectClientCommand serv ("SUB SECRET.LOG " <> subscriptionId <> "\r\n")+ sendAll serv $+ hmsgFrame+ "SECRET.LOG"+ subscriptionId+ [("x-private-header", "secret-header-sentinel")]+ "secret-payload-sentinel"+ <> "secret-malformed-sentinel"+ void (expectMVar "secret message was not delivered" delivered)+ sendAll serv "PING\r\n"+ expectClientCommand serv "PONG\r\n"++ messages <- readIORef logMessages+ let combined = unlines messages+ combined `shouldNotContain` "secret-header-sentinel"+ combined `shouldNotContain` "secret-payload-sentinel"+ combined `shouldNotContain` "secret-malformed-sentinel" it "subscribes with a queue group" $ \(serv, client, _, _) -> do Right subscription <- subscribe client "JOBS" [withQueueGroup "WORKERS"] (const (pure ())) expectQueuedSub serv "JOBS" "WORKERS" (subscriptionSid subscription)@@ -727,6 +878,39 @@ expectClientCommand serv "PONG\r\n" putMVar releaseCallback () + around (withClientWith [withCallbackConcurrency 1]) $ do+ it "returns an accepted reply after disconnect while its callback is queued" $ \(serv, client, _, _) -> do+ callbackStarted <- newEmptyMVar+ releaseCallback <- newEmptyMVar+ Right subscription <- subscribe client "SLOW.REQUEST.BARRIER" [] $ \_ -> do+ putMVar callbackStarted ()+ takeMVar releaseCallback+ let sid = subscriptionSid subscription+ expectClientCommand serv (BS.concat ["SUB SLOW.REQUEST.BARRIER ", sid, "\r\n"])+ sendAll serv (msgFrame "SLOW.REQUEST.BARRIER" sid "busy")+ expectMVar "blocking callback did not start" callbackStarted++ requestResult <- newEmptyMVar+ void . forkIO $+ request client "SERVICE.ACCEPTED" "request" [withRequestTimeout 5]+ >>= putMVar requestResult+ captured <- capturePublish serv "SERVICE.ACCEPTED"+ sendAll serv $+ msgFrame (capturedInbox captured) (capturedSid captured) "reply"+ <> "PING\r\n"+ expectClientCommand serv "PONG\r\n"+ Network.Socket.close serv++ timeout 1000000 (awaitClosed client) `shouldReturn` Just ()+ tryTakeMVar requestResult `shouldReturn` Nothing+ putMVar releaseCallback ()+ reply <- expectMVar "accepted request reply did not complete" requestResult+ case reply of+ Left err ->+ expectationFailure ("accepted request failed after disconnect: " ++ show err)+ Right message ->+ payload message `shouldBe` "reply"+ it "cleans up no responders replies for expiring requests" $ do bracket startClient stopClientSafely $ \(serv, client, _, _) -> do replyBox <- newEmptyMVar@@ -770,20 +954,33 @@ Nothing -> expectationFailure "client did not connect" Just client -> do+ requestResult <- newEmptyMVar void . forkIO $- void (request client "SERVICE.RECONNECT" "ping" [withRequestTimeout 5])+ request client "SERVICE.RECONNECT" "ping" [withRequestTimeout 5]+ >>= putMVar requestResult captured <- capturePublish firstConn "SERVICE.RECONNECT"+ let replySetup = BS.concat+ [ "SUB "+ , capturedInbox captured+ , " "+ , capturedSid captured+ , "\r\nUNSUB "+ , capturedSid captured+ , " 1\r\n"+ ]+ capturedWire captured `shouldSatisfy` BS.isInfixOf replySetup Network.Socket.close firstConn+ expectMVar "request did not fail at its connection boundary" requestResult+ `shouldReturn` Left (NatsConnectionClosed ExitResetRequested) (secondConn, _) <- accept sock void (completeHandshake secondConn)- let replySubCommand = BS.concat ["SUB ", capturedInbox captured] pingBox <- newEmptyMVar pingWithCallback client (putMVar pingBox ()) expectNoClientBytesBefore secondConn- (BS.isInfixOf replySubCommand)+ (BS.isInfixOf (capturedInbox captured)) "PING\r\n"- "reply inbox was resubscribed"+ "reply inbox protocol crossed the reconnect boundary" sendAll secondConn "PONG\r\n" expectMVar "PING failed after reconnect" pingBox close client []@@ -801,12 +998,14 @@ ["PROTO.SUBJECT"] [ JetStream.withStorage JetStream.MemoryStorage , JetStream.withMaxMessages 1+ , JetStream.withMaxMessageSize 64 , JetStream.withDiscard JetStream.DiscardNew ] [] putMVar done result captured <- capturePublish serv "$JS.API.STREAM.CREATE.PROTO_STREAM" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"max_msgs\":1"+ capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"max_msg_size\":64" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"discard\":\"new\"" replyToCapturedPublish serv captured protoStreamInfoResponse result <- timeout 1000000 (takeMVar done)@@ -815,8 +1014,9 @@ expectationFailure "stream create did not complete" Just (Left err) -> expectationFailure ("stream create failed: " ++ show err)- Just (Right info) ->+ Just (Right info) -> do JetStream.streamConfigName (JetStream.streamInfoConfig info) `shouldBe` "PROTO_STREAM"+ JetStream.streamConfigMaxMessageSize (JetStream.streamInfoConfig info) `shouldBe` 64 it "sends only the selected consumer filter representation" $ \(serv, client, _, _) -> do let Right jetStream = newJetStream client []@@ -832,12 +1032,16 @@ , JetStream.withConsumerAckPolicy JetStream.AckExplicit , JetStream.withConsumerFilter (JetStream.ConsumerFilterSubjects ["PROTO.A", "PROTO.B"])+ , JetStream.withConsumerBackoff [0.100000001, 0.25]+ , JetStream.withConsumerMaxRequestBatch 32 ] [] putMVar done result captured <- capturePublish serv "$JS.API.CONSUMER.DURABLE.CREATE.PROTO_STREAM.PROTO_CONSUMER" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"filter_subjects\":[\"PROTO.A\",\"PROTO.B\"]" capturedPayload captured `shouldSatisfy` not . BS.isInfixOf "\"filter_subject\":"+ capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"backoff\":[100000001,250000000]"+ capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"max_batch\":32" replyToCapturedPublish serv captured protoConsumerInfoResponse result <- timeout 1000000 (takeMVar done) case result of@@ -845,9 +1049,65 @@ expectationFailure "consumer create did not complete" Just (Left err) -> expectationFailure ("consumer create failed: " ++ show err)- Just (Right info) ->+ Just (Right info) -> do JetStream.consumerInfoName info `shouldBe` "PROTO_CONSUMER"+ JetStream.consumerConfigBackoff (JetStream.consumerInfoConfig info)+ `shouldBe` Just [0.100000001, 0.25]+ JetStream.consumerConfigMaxRequestBatch (JetStream.consumerInfoConfig info)+ `shouldBe` Just 32 + it "rejects invalid resource configs before publishing" $ \(serv, client, _, _) -> do+ let Right jetStream = newJetStream client []+ invalidStream <- JetStream.create+ (JetStream.streams jetStream)+ "INVALID_STREAM"+ ["INVALID.STREAM"]+ [JetStream.withMaxMessageSize (-2)]+ []+ invalidStream `shouldSatisfy` isJetStreamDecodeFailure++ invalidBatch <- JetStream.putConsumer+ (JetStream.consumers jetStream)+ "PROTO_STREAM"+ JetStream.ConsumerCreate+ (JetStream.DurableConsumer "INVALID_BATCH")+ JetStream.PullConsumer+ [JetStream.withConsumerMaxRequestBatch (-1)]+ []+ invalidBatch `shouldSatisfy` isJetStreamDecodeFailure++ invalidBackoff <- JetStream.putConsumer+ (JetStream.consumers jetStream)+ "PROTO_STREAM"+ JetStream.ConsumerCreate+ (JetStream.DurableConsumer "INVALID_BACKOFF")+ JetStream.PullConsumer+ [JetStream.withConsumerBackoff [10000000000]]+ []+ invalidBackoff `shouldSatisfy` isJetStreamDecodeFailure++ invalidBackoffLength <- JetStream.putConsumer+ (JetStream.consumers jetStream)+ "PROTO_STREAM"+ JetStream.ConsumerCreate+ (JetStream.DurableConsumer "INVALID_BACKOFF_LENGTH")+ JetStream.PullConsumer+ [ JetStream.withConsumerMaxDeliver 1+ , JetStream.withConsumerBackoff [1, 2]+ ]+ []+ invalidBackoffLength `shouldSatisfy` isJetStreamDecodeFailure++ pingBox <- newEmptyMVar+ pingWithCallback client (putMVar pingBox ())+ expectNoClientBytesBefore+ serv+ (BS.isInfixOf "$JS.API.")+ "PING\r\n"+ "invalid JetStream config published a request"+ sendAll serv "PONG\r\n"+ expectMVar "PING failed after local JetStream validation" pingBox+ it "sends push consumer create requests with a delivery subject" $ \(serv, client, _, _) -> do let Right jetStream = newJetStream client [] done <- newEmptyMVar@@ -1211,6 +1471,284 @@ close client [] Network.Socket.close secondConn Network.Socket.close sock+ it "waits reconnecting operations until the connection is ready" $ do+ (p, sock) <- openFreePort+ listen sock 2+ clientVar <- newEmptyMVar+ void . forkIO $ do+ client <-+ newClientOrFail+ [("127.0.0.1", p)]+ [ withConnectionAttempts 2+ , withConnectName "state-client"+ ]+ putMVar clientVar client+ (firstConn, _) <- accept sock+ void (completeHandshake firstConn)+ client <- expectMVar "client did not connect" clientVar+ connectionState client `shouldReturn` ConnectionConnected++ Network.Socket.close firstConn+ (secondConn, _) <- accept sock+ stateBeforeHandshake <- timeout 1000000 $ do+ let awaitReconnecting = do+ state <- connectionState client+ if state == ConnectionReconnecting+ then pure state+ else threadDelay 1000 >> awaitReconnecting+ awaitReconnecting+ stateBeforeHandshake `shouldBe` Just ConnectionReconnecting+ publishResult <- newEmptyMVar+ void . forkIO $+ publish client "RECOVERED.SUBJECT" "ready" [] >>= putMVar publishResult+ timeout 50000 (takeMVar publishResult) `shouldReturn` Nothing++ void (completeHandshake secondConn)+ expectMVar "publish did not recover after reconnect" publishResult+ `shouldReturn` Right ()+ expectClientCommand secondConn "PUB RECOVERED.SUBJECT 5\r\nready\r\n"+ connectionState client `shouldReturn` ConnectionConnected+ close client []+ finalState <- connectionState client+ case finalState of+ ConnectionClosed ExitClosedByUser -> pure ()+ other -> expectationFailure ("unexpected final connection state: " ++ show other)+ Network.Socket.close secondConn+ Network.Socket.close sock+ it "reports disconnect, reconnect, and terminal close in order" $ do+ (p, sock) <- openFreePort+ listen sock 2+ observed <- newTQueueIO+ clientVar <- newEmptyMVar+ void . forkIO $ do+ client <-+ newClientOrFail+ [("127.0.0.1", p)]+ [ withConnectionAttempts 2+ , withConnectionEventHandler (atomically . writeTQueue observed)+ ]+ putMVar clientVar client+ (firstConn, _) <- accept sock+ void (completeHandshake firstConn)+ client <- expectMVar "client did not connect" clientVar+ atomically (tryReadTQueue observed) `shouldReturn` Nothing++ Network.Socket.close firstConn+ (secondConn, _) <- accept sock+ void (completeHandshake secondConn)+ timeout 1000000 (atomically (readTQueue observed))+ `shouldReturn` Just ConnectionEventDisconnected+ timeout 1000000 (atomically (readTQueue observed))+ `shouldReturn` Just ConnectionEventReconnected++ close client []+ timeout 1000000 (atomically (readTQueue observed))+ `shouldReturn` Just (ConnectionEventClosed ExitClosedByUser)+ Network.Socket.close secondConn+ Network.Socket.close sock+ it "allows a disconnect callback to close its own client" $ do+ (p, sock) <- openFreePort+ listen sock 2+ clientRef <- newEmptyMVar+ clientVar <- newEmptyMVar+ handlerReturned <- newEmptyMVar+ observed <- newTQueueIO+ let handler event = do+ atomically (writeTQueue observed event)+ when (event == ConnectionEventDisconnected) $ do+ client <- readMVar clientRef+ close client []+ putMVar handlerReturned ()+ void . forkIO $ do+ client <-+ newClientOrFail+ [("127.0.0.1", p)]+ [ withConnectionAttempts 2+ , withConnectionEventHandler handler+ ]+ putMVar clientVar client+ (firstConn, _) <- accept sock+ void (completeHandshake firstConn)+ client <- expectMVar "client did not connect" clientVar+ putMVar clientRef client++ Network.Socket.close firstConn+ void (expectMVar "disconnect callback deadlocked in close" handlerReturned)+ timeout 1000000 (awaitClosed client) `shouldReturn` Just ()+ timeout 1000000 (atomically (readTQueue observed))+ `shouldReturn` Just ConnectionEventDisconnected+ timeout 1000000 (atomically (readTQueue observed))+ `shouldReturn` Just (ConnectionEventClosed ExitClosedByUser)+ Network.Socket.close sock+ it "resets after ping cancellation without poisoning the next PONG" $ do+ (p, sock) <- openFreePort+ listen sock 2+ clientVar <- newEmptyMVar+ void . forkIO $ do+ client <-+ newClientOrFail+ [("127.0.0.1", p)]+ [ withConnectionAttempts 2+ , withConnectName "cancelled-ping-client"+ ]+ putMVar clientVar client+ (firstConn, _) <- accept sock+ void (completeHandshake firstConn)+ client <- expectMVar "client did not connect" clientVar+ abandonedResult <- newEmptyMVar+ pingThread <- forkIO $+ ping client [withPingTimeout 5] >>= putMVar abandonedResult+ expectClientCommand firstConn "PING\r\n"++ killThread pingThread+ (secondConn, _) <- accept sock+ void (completeHandshake secondConn)+ nextResult <- newEmptyMVar+ void . forkIO $+ ping client [withPingTimeout 1] >>= putMVar nextResult+ expectClientCommand secondConn "PING\r\n"+ timeout 50000 (takeMVar nextResult) `shouldReturn` Nothing+ sendAll secondConn "PONG\r\n"+ expectMVar "next ping did not receive its own PONG" nextResult+ `shouldReturn` Right ()+ tryTakeMVar abandonedResult `shouldReturn` Nothing++ close client []+ Network.Socket.close firstConn+ Network.Socket.close secondConn+ Network.Socket.close sock+ it "classifies an established TCP connection timeout as handshake timeout" $ do+ (p, sock) <- openFreePort+ listen sock 1+ clientResult <- newEmptyMVar+ void . forkIO $+ newClient+ [("127.0.0.1", p)]+ [ withConnectionAttempts 1+ , withConnectTimeoutMicros 50000+ ]+ >>= putMVar clientResult+ (serverConn, _) <- accept sock++ result <- expectMVar "client did not finish after handshake timeout" clientResult+ case result of+ Left (ConnectAttemptsExhausted [ConnectAttemptError _ ConnectHandshakeTimeout]) ->+ pure ()+ Left err ->+ expectationFailure ("unexpected connection timeout: " ++ show err)+ Right client -> do+ close client []+ expectationFailure "client connected without a handshake"+ Network.Socket.close serverConn+ Network.Socket.close sock+ it "closes the socket when initial connection waiting is cancelled" $ do+ (p, sock) <- openFreePort+ listen sock 1+ accepted <- newEmptyMVar+ void . forkIO $ do+ (serverConn, _) <- accept sock+ putMVar accepted serverConn++ clientFinished <- newEmptyMVar+ clientThread <- forkIO $+ void+ ( newClient+ [("127.0.0.1", p)]+ [ withConnectionAttempts 1+ , withConnectTimeoutMicros (5 * 1000000)+ ]+ )+ `finally` putMVar clientFinished ()+ serverConn <- expectMVar "client never opened its initial socket" accepted+ killThread clientThread+ void (expectMVar "cancelled client did not stop" clientFinished)+ timeout 1000000 (recv serverConn 1) `shouldReturn` Just BS.empty+ Network.Socket.close serverConn+ Network.Socket.close sock+ it "revalidates a waiting publish against the reconnect target" $ do+ (p, sock) <- openFreePort+ listen sock 2+ clientVar <- newEmptyMVar+ void . forkIO $+ newClientOrFail+ [("127.0.0.1", p)]+ [ withConnectionAttempts 2+ , withMessageLimit 4096+ ]+ >>= putMVar clientVar+ (firstConn, _) <- accept sock+ void (completeHandshakeWithInfo firstConn (infoWithMaxPayload 4096))+ client <- expectMVar "publish boundary client did not connect" clientVar++ Network.Socket.close firstConn+ (secondConn, _) <- accept sock+ publishResult <- newEmptyMVar+ void . forkIO $+ publish client "LIMIT.PUBLISH" (BS.replicate 2048 _x) []+ >>= putMVar publishResult+ timeout 50000 (takeMVar publishResult) `shouldReturn` Nothing++ void (completeHandshakeWithInfo secondConn (infoWithMaxPayload 1024))+ expectMVar "publish did not revalidate on reconnect" publishResult+ `shouldReturn` Left (NatsPayloadTooLarge 2048 1024)+ pingResult <- newEmptyMVar+ void . forkIO $ ping client [] >>= putMVar pingResult+ expectNoClientBytesBefore+ secondConn+ (BS.isInfixOf "LIMIT.PUBLISH")+ "PING\r\n"+ "oversized publish reached the reconnect target"+ sendAll secondConn "PONG\r\n"+ expectMVar "ping failed after rejected reconnect publish" pingResult+ `shouldReturn` Right ()++ close client []+ Network.Socket.close secondConn+ Network.Socket.close sock+ it "revalidates a waiting request against the reconnect target" $ do+ (p, sock) <- openFreePort+ listen sock 2+ clientVar <- newEmptyMVar+ void . forkIO $+ newClientOrFail+ [("127.0.0.1", p)]+ [ withConnectionAttempts 2+ , withMessageLimit 4096+ ]+ >>= putMVar clientVar+ (firstConn, _) <- accept sock+ void (completeHandshakeWithInfo firstConn (infoWithMaxPayload 4096))+ client <- expectMVar "request boundary client did not connect" clientVar++ Network.Socket.close firstConn+ (secondConn, _) <- accept sock+ requestResult <- newEmptyMVar+ void . forkIO $+ request+ client+ "LIMIT.REQUEST"+ (BS.replicate 2048 _x)+ [withRequestTimeout 5]+ >>= putMVar requestResult+ timeout 50000 (takeMVar requestResult) `shouldReturn` Nothing++ void (completeHandshakeWithInfo secondConn (infoWithMaxPayload 1024))+ expectMVar "request did not revalidate on reconnect" requestResult+ `shouldReturn` Left (NatsPayloadTooLarge 2048 1024)+ pingResult <- newEmptyMVar+ void . forkIO $ ping client [] >>= putMVar pingResult+ expectNoClientBytesBefore+ secondConn+ (BS.isInfixOf "LIMIT.REQUEST")+ "PING\r\n"+ "oversized request reached the reconnect target"+ sendAll secondConn "PONG\r\n"+ expectMVar "ping failed after rejected reconnect request" pingResult+ `shouldReturn` Right ()++ close client []+ Network.Socket.close secondConn+ Network.Socket.close sock it "resubscribes with a queue group after reconnect" $ do (p, sock) <- openFreePort listen sock 2@@ -1238,6 +1776,104 @@ expectQueuedSub secondConn "JOBS" "WORKERS" sid close client [] Network.Socket.close secondConn+ Network.Socket.close sock+ it "resubscribes one-shot subscriptions with the server-side limit" $ do+ (p, sock) <- openFreePort+ listen sock 2+ clientVar <- newEmptyMVar+ void . forkIO $ do+ client <-+ newClientOrFail+ [("127.0.0.1", p)]+ [ withConnectionAttempts 2+ , withConnectName "one-shot-client"+ ]+ putMVar clientVar client+ (firstConn, _) <- accept sock+ void (completeHandshake firstConn)+ client <- expectMVar "one-shot client did not connect" clientVar+ delivered <- newIORef ([] :: [Message])+ Right subscription <- subscribeOnce client "ONCE.RECONNECT" [] $ \message ->+ atomicModifyIORef' delivered $ \messages -> (messages ++ [message], ())+ let sid = subscriptionSid subscription+ subscribeCommand = BS.concat ["SUB ONCE.RECONNECT ", sid, "\r\n"]+ limitCommand = BS.concat ["UNSUB ", sid, " 1\r\n"]+ firstWire <- expectClientBytes+ firstConn+ (\bytes -> BS.isInfixOf subscribeCommand bytes && BS.isInfixOf limitCommand bytes)+ "initial one-shot commands were not sent"+ unless (appearsBefore subscribeCommand limitCommand firstWire) $+ expectationFailure ("one-shot limit preceded SUB: " ++ show firstWire)++ Network.Socket.close firstConn+ (secondConn, _) <- accept sock+ void (completeHandshake secondConn)+ secondWire <- expectClientBytes+ secondConn+ (\bytes -> BS.isInfixOf subscribeCommand bytes && BS.isInfixOf limitCommand bytes)+ "resumed one-shot commands were not sent"+ unless (appearsBefore subscribeCommand limitCommand secondWire) $+ expectationFailure ("resumed one-shot limit preceded SUB: " ++ show secondWire)++ barrierDone <- newEmptyMVar+ Right barrier <- subscribe client "ONCE.BARRIER" [] (const (putMVar barrierDone ()))+ let barrierSid = subscriptionSid barrier+ expectClientCommand secondConn (BS.concat ["SUB ONCE.BARRIER ", barrierSid, "\r\n"])+ sendAll secondConn $+ BS.concat+ [ msgFrame "ONCE.RECONNECT" sid "first"+ , msgFrame "ONCE.RECONNECT" sid "second"+ , msgFrame "ONCE.BARRIER" barrierSid "done"+ ]+ expectMVar "one-shot callback barrier did not run" barrierDone+ messages <- readIORef delivered+ fmap payload messages `shouldBe` ["first"]++ close client []+ Network.Socket.close secondConn+ Network.Socket.close sock+ it "does not revive a one-shot subscription that expires during reconnect" $ do+ (p, sock) <- openFreePort+ listen sock 2+ clientVar <- newEmptyMVar+ void . forkIO $ do+ client <-+ newClientOrFail+ [("127.0.0.1", p)]+ [ withConnectionAttempts 2+ , withConnectName "expiring-reconnect-client"+ ]+ putMVar clientVar client+ (firstConn, _) <- accept sock+ void (completeHandshake firstConn)+ client <- expectMVar "expiring client did not connect" clientVar+ Right subscription <-+ subscribeOnce+ client+ "EXPIRY.RECONNECT"+ [withSubscriptionExpiry 0.01]+ (const (pure ()))+ let sid = subscriptionSid subscription+ expectClientCommand firstConn $+ BS.concat ["SUB EXPIRY.RECONNECT ", sid, "\r\nUNSUB ", sid, " 1\r\n"]++ Network.Socket.close firstConn+ (secondConn, _) <- accept sock+ threadDelay 1500000+ void (completeHandshake secondConn)+ pingResult <- newEmptyMVar+ void . forkIO $ ping client [] >>= putMVar pingResult+ expectNoClientBytesBefore+ secondConn+ (BS.isInfixOf "EXPIRY.RECONNECT")+ "PING\r\n"+ "expired one-shot subscription was revived"+ sendAll secondConn "PONG\r\n"+ expectMVar "ping failed after reconnect expiry" pingResult+ `shouldReturn` Right ()++ close client []+ Network.Socket.close secondConn Network.Socket.close sock it "resubscribes JetStream push consumers after reconnect" $ do (p, sock) <- openFreePort
test/PublicAPI/Main.hs view
@@ -24,6 +24,8 @@ , Client.withMessageLimit (1024 * 1024) , Client.withPendingDeliveryLimits 65536 (64 * 1024 * 1024) , Client.withErrorHandler (const (pure ()))+ , Client.withServerErrorHandler (const (pure ()))+ , Client.withConnectionEventHandler inspectConnectionEvent ] configuredServer :: Client.Server@@ -61,10 +63,39 @@ "service.echo" "request" [Nats.withRequestTimeout 1]- _ <- Nats.ping client []- _ <- Nats.flush client []+ _ <- Nats.ping client [Nats.withPingTimeout 1]+ _ <- Nats.flush client [Nats.withFlushTimeout 1]+ state <- Nats.connectionState client+ inspectConnectionState state Nats.close client [] +inspectConnectionState :: Nats.ConnectionState -> IO ()+inspectConnectionState state =+ case state of+ Nats.ConnectionConnecting -> pure ()+ Nats.ConnectionConnected -> pure ()+ Nats.ConnectionReconnecting -> pure ()+ Nats.ConnectionClosing _ -> pure ()+ Nats.ConnectionClosed _ -> pure ()++inspectServerError :: Client.ServerError -> IO ()+inspectServerError serverError = do+ Client.serverErrorReason serverError `seq` pure ()+ case Client.serverErrorKind serverError of+ Client.ServerErrorAuthentication -> pure ()+ Client.ServerErrorPermission -> pure ()+ Client.ServerErrorProtocol -> pure ()+ Client.ServerErrorResourceLimit -> pure ()+ Client.ServerErrorConnection -> pure ()+ Client.ServerErrorUnknown -> pure ()++inspectConnectionEvent :: Client.ConnectionEvent -> IO ()+inspectConnectionEvent event =+ case event of+ Client.ConnectionEventDisconnected -> pure ()+ Client.ConnectionEventReconnected -> pure ()+ Client.ConnectionEventClosed _ -> pure ()+ observeMessage :: Nats.Message -> IO () observeMessage message = do let _ = Nats.subject message@@ -94,13 +125,16 @@ ["orders.>"] [ Stream.withStorage Stream.FileStorage , Stream.withMaxMessages 1000+ , Stream.withMaxMessageSize 1024 ] [] case created of Left _ -> pure () Right handle -> do- _ <- Stream.getStreamInfo (JetStream.streams jetStream) handle []- pure ()+ info <- Stream.getStreamInfo (JetStream.streams jetStream) handle []+ case info of+ Left _ -> pure ()+ Right detail -> Stream.streamConfigMaxMessageSize (Stream.streamInfoConfig detail) `seq` pure () _ <- JetStreamPublish.publish (JetStream.publisher jetStream) "orders.created"@@ -113,13 +147,20 @@ Consumer.ConsumerCreate (Consumer.DurableConsumer "processor") Consumer.PullConsumer- [Consumer.withConsumerAckPolicy Consumer.AckExplicit]+ [ Consumer.withConsumerAckPolicy Consumer.AckExplicit+ , Consumer.withConsumerBackoff [1, 2]+ , Consumer.withConsumerMaxRequestBatch 64+ ] [] case consumer of Left _ -> pure () Right handle -> do- _ <- Consumer.getConsumerInfo (JetStream.consumers jetStream) handle []- pure ()+ info <- Consumer.getConsumerInfo (JetStream.consumers jetStream) handle []+ case info of+ Left _ -> pure ()+ Right detail -> do+ Consumer.consumerConfigBackoff (Consumer.consumerInfoConfig detail) `seq` pure ()+ Consumer.consumerConfigMaxRequestBatch (Consumer.consumerInfoConfig detail) `seq` pure () _ <- JetStreamMessage.fetch (JetStream.messages jetStream) "ORDERS"@@ -141,7 +182,14 @@ messageOperations messageAPI message = do JetStreamMessage.messageMetadata message `seq` pure () _ <- JetStreamMessage.ack messageAPI message []+ _ <- JetStreamMessage.ackSync messageAPI message [] _ <- JetStreamMessage.nak messageAPI message []+ case JetStreamMessage.nakDelay 1 of+ Nothing -> pure ()+ Just delay -> do+ _ <- JetStreamMessage.nakWithDelay messageAPI message delay []+ pure ()+ _ <- JetStreamMessage.termSync messageAPI message [] pure () inspectJetStreamApiError :: JetStream.JetStreamApiError -> IO ()@@ -155,6 +203,7 @@ connect `seq` serverBuilders `seq` coreOperations+ `seq` inspectServerError `seq` jetStreamContext `seq` jetStreamOperations `seq` messageOperations
test/Unit/ConnectionSpec.hs view
@@ -1,3 +1,4 @@+{-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} module ConnectionSpec (spec) where@@ -6,58 +7,163 @@ import qualified Client import Control.Concurrent import Control.Concurrent.STM+import Control.Exception (finally, throwIO)+import Control.Monad (replicateM_, void, when) import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LBS-import Data.IORef (newIORef, readIORef, writeIORef)+import Engine (runEngine) import Handshake.Nats- ( HandshakeError (HandshakeAuthError, HandshakeTLSError, HandshakeTimeout)+ ( HandshakeError (HandshakeAuthError, HandshakeProtocolError, HandshakeTLSError) , performHandshake ) import Lib.Logger- ( LogLevel (Debug)- , LoggerConfig (LoggerConfig)+ ( LogEntry (..)+ , LogLevel (Debug)+ , LoggerConfig (..) , newLogContext ) import Network.Connection (connectionApi) import Network.Connection.Core (Transport (..), pointTransport) import Network.ConnectionAPI- ( closeReader- , newConn- , readData- , reader+ ( ConnectionAPI (..)+ , ReaderAPI (..)+ , WriterAPI (..) ) import Parser.API ( ParseStep (DropPrefix, Emit, NeedMore, Reject)- , ParsedMessage (ParsedInfo, ParsedPing, ParsedPong)+ , ParsedMessage (ParsedErr, ParsedInfo, ParsedPing, ParsedPong) , ParserAPI (ParserAPI) ) import qualified Parser.Attoparsec as Attoparsec+import Pipeline.Broadcasting (broadcastingApi)+import Pipeline.Streaming (streamingApi) import qualified Queue.API as Queue-import Queue.TransactionalQueue (newQueue)+import Queue.TransactionalQueue (newQueue, newQueueWithCapacity) import Router.Nats ( RouteDirective (RouteContinue) , routeMessage ) import State.Store- ( newClientState- , pushPingAction+ ( PingResult (..)+ , PublishEnqueueResult (..)+ , ServerErrorEnqueueResult (..)+ , activateConnectionAttempt+ , beginConnectionAttempt+ , clearPendingPings+ , enqueueOnGenerationTracked+ , enqueuePublishOnConnected+ , enqueuePublishOnConnectedGenerationTracked+ , markClosed+ , markConnectionReady+ , newClientState+ , notifyServerError , queue+ , readConnectionGeneration , readServerInfo+ , readStatus+ , registerPingWaiter+ , registerPingWaiterAndEnqueue+ , setClosing+ , setReconnecting+ , setServerInfo+ , startCallbackWorker+ , stopCallbackWorker+ , waitForConnectionGenerationAfter+ , withSubscriptionGate )-import State.Types (ClientConfig (..))-import Subscription.Store (newSubscriptionStore)-import Subscription.Types (defaultPendingLimits)+import State.Types+ ( ClientConfig (..)+ , ClientExitReason (..)+ , ConnectionEvent (..)+ , ConnectionState (..)+ , ServerErrorKind (..)+ , serverErrorFromProtocol+ , serverErrorKind+ )+import Subscription.Store+ ( awaitCallbackDrain+ , enqueueControl+ , newSubscriptionStore+ , register+ , startWorkers+ )+import Subscription.Types+ ( SubscribeConfig (..)+ , SubscriptionKind (StandardSubscription)+ , SubscriptionMeta (..)+ , defaultPendingLimits+ ) import System.Timeout (timeout) import Test.Hspec import Transformers.Transformers (Transformer (transform)) import qualified Types.Connect as Connect-import Types.Info (Info (Info))+import qualified Types.Err as Err+import Types.Info (Info (..)) import Types.Ping (Ping (Ping)) import Types.Pong (Pong (Pong))+import qualified Types.Pub as Pub import Types.TLS (TLSConfig (..), defaultTLSConfig) spec :: Spec spec = do+ describe "Transactional queue" $ do+ it "wakes a full blocked enqueue when the queue closes" $ do+ messageQueue <- newQueueWithCapacity 1+ Queue.enqueue messageQueue (Queue.QueueItem Ping) `shouldReturn` Right ()+ started <- newEmptyMVar+ result <- newEmptyMVar+ _ <- forkIO $ do+ putMVar started ()+ Queue.enqueue messageQueue (Queue.QueueItem Ping) >>= putMVar result+ takeMVar started+ timeout 20000 (takeMVar result) `shouldReturn` Nothing++ Queue.close messageQueue++ timeout 1000000 (takeMVar result) `shouldReturn` Just (Left "Queue is closed")++ it "discards only connection-scoped commands at a connection boundary" $ do+ messageQueue <- newQueueWithCapacity 3+ Queue.enqueue messageQueue (Queue.QueueItem Ping) `shouldReturn` Right ()+ Queue.enqueue messageQueue (Queue.QueueConnectionScoped (Queue.QueueItem Ping))+ `shouldReturn` Right ()++ atomically (Queue.closeAndDiscardConnectionScoped messageQueue)++ Queue.dequeue messageQueue >>= \case+ Left err -> err `shouldBe` "Queue is closed"+ Right _ -> expectationFailure "expected closed queue"+ Queue.open messageQueue+ queued <- Queue.dequeue messageQueue+ case queued of+ Right item -> LBS.toStrict (transform item) `shouldBe` "PING\r\n"+ Left err -> expectationFailure err+ it "discards every buffered command on terminal close" $ do+ messageQueue <- newQueueWithCapacity 2+ Queue.enqueue messageQueue (Queue.QueueItem Ping) `shouldReturn` Right ()+ Queue.enqueue messageQueue+ (Queue.QueueConnectionScoped (Queue.QueueItem Ping)) `shouldReturn` Right ()++ atomically (Queue.closeAndDiscardAll messageQueue)++ Queue.dequeue messageQueue >>= \case+ Left err -> err `shouldBe` "Queue is closed"+ Right _ -> expectationFailure "expected closed queue"+ Queue.open messageQueue+ blocked <- timeout 20000 (Queue.dequeue messageQueue)+ case blocked of+ Nothing -> pure ()+ Just _ -> expectationFailure "terminal close retained a buffered command"+ Queue.close messageQueue++ it "offers a nonblocking enqueue for cancellation cleanup" $ do+ messageQueue <- newQueueWithCapacity 1+ Queue.enqueue messageQueue (Queue.QueueItem Ping) `shouldReturn` Right ()++ timeout 1000000+ (Queue.tryEnqueue messageQueue (Queue.QueueItem Ping))+ `shouldReturn` Just Queue.TryQueueFull+ describe "TLS configuration" $ do it "does not expose client private keys when rendered" $ do let config = defaultTLSConfig@@ -80,6 +186,7 @@ , transportWriteLazy = \_ -> pure () , transportFlush = pure () , transportClose = pure ()+ , transportAbort = pure () , transportUpgrade = Nothing } pointTransport conn transport@@ -89,10 +196,34 @@ closeReader (reader connectionApi) conn result <- timeout 1000000 (takeMVar resultVar) result `shouldBe` Just (Left "Read operation is blocked")++ it "force-aborts a blocked transport write" $ do+ conn <- newConn connectionApi+ started <- newEmptyMVar+ release <- newEmptyMVar+ let transport = Transport+ { transportRead = const (pure BS.empty)+ , transportWrite = \_ -> putMVar started () >> takeMVar release+ , transportWriteLazy = \_ -> putMVar started () >> takeMVar release+ , transportFlush = pure ()+ , transportClose = pure ()+ , transportAbort = void (tryPutMVar release ())+ , transportUpgrade = Nothing+ }+ pointTransport conn transport+ resultVar <- newEmptyMVar+ _ <- forkIO $+ writeData (writer connectionApi) conn "blocked" >>= putMVar resultVar+ takeMVar started++ abort connectionApi conn++ timeout 1000000 (takeMVar resultVar) `shouldReturn` Just (Right ()) describe "Router ping lifecycle" $ do it "responds to server PING by enqueueing PONG" $ do state <- newTestState store <- newSubscriptionStore defaultPendingLimits (pure ())+ markTestStateConnected state directive <- routeMessage state store (ParsedPing Ping) @@ -104,24 +235,576 @@ Left err -> expectationFailure ("expected queued PONG, got queue error: " ++ err) - it "runs exactly one pending ping action for each PONG" $ do+ it "keeps routing when the configured logger throws" $ do+ state <- newTestStateWithConfig $ \cfg ->+ cfg+ { loggerConfig =+ (loggerConfig cfg)+ { logFn = const (throwIO (userError "logger failed")) }+ }+ store <- newSubscriptionStore defaultPendingLimits (pure ())+ markTestStateConnected state++ routeMessage state store (ParsedPing Ping) `shouldReturn` RouteContinue+ queued <- Queue.dequeue (queue state)+ case queued of+ Right item -> LBS.toStrict (transform item) `shouldBe` "PONG\r\n"+ Left err -> expectationFailure err++ it "resolves exactly one pending ping waiter for each PONG" $ do state <- newTestState store <- newSubscriptionStore defaultPendingLimits (pure ())- firstActionRan <- newIORef False- secondActionRan <- newIORef False- pushPingAction state (writeIORef firstActionRan True)- pushPingAction state (writeIORef secondActionRan True)+ markTestStateConnected state+ firstWaiter <- newEmptyTMVarIO+ secondWaiter <- newEmptyTMVarIO+ registerPingWaiter state firstWaiter `shouldReturn` Right ()+ registerPingWaiter state secondWaiter `shouldReturn` Right () firstDirective <- routeMessage state store (ParsedPong Pong) firstDirective `shouldBe` RouteContinue- readIORef firstActionRan `shouldReturn` True- readIORef secondActionRan `shouldReturn` False+ timeout 1000000 (atomically (readTMVar firstWaiter))+ `shouldReturn` Just PingReceived+ atomically (tryReadTMVar secondWaiter) `shouldReturn` Nothing secondDirective <- routeMessage state store (ParsedPong Pong) secondDirective `shouldBe` RouteContinue- readIORef secondActionRan `shouldReturn` True+ timeout 1000000 (atomically (readTMVar secondWaiter))+ `shouldReturn` Just PingReceived++ it "releases pending ping waiters at a connection boundary" $ do+ state <- newTestState+ markTestStateConnected state+ waiter <- newEmptyTMVarIO+ registerPingWaiter state waiter `shouldReturn` Right ()++ clearPendingPings state++ atomically (readTMVar waiter) `shouldReturn` PingConnectionLost++ it "atomically invalidates ping waiters and scoped PINGs on reconnect" $ do+ state <- newTestState+ markTestStateConnected state+ waiter <- newEmptyTMVarIO+ registerPingWaiter state waiter `shouldReturn` Right ()+ Queue.enqueue (queue state) (Queue.QueueConnectionScoped (Queue.QueueItem Ping))+ `shouldReturn` Right ()++ setReconnecting state++ atomically (readTMVar waiter) `shouldReturn` PingConnectionLost+ Queue.dequeue (queue state) >>= \case+ Left err -> err `shouldBe` "Queue is closed"+ Right _ -> expectationFailure "expected scoped PING to be discarded"++ it "does not orphan a ping waiter when reset wins a full-queue race" $ do+ state <- newTestStateWithQueueCapacity 1+ markTestStateConnected state+ Queue.enqueue (queue state) (Queue.QueueItem Ping) `shouldReturn` Right ()+ waiter <- newEmptyTMVarIO+ started <- newEmptyMVar+ result <- newEmptyMVar+ _ <- forkIO $ do+ putMVar started ()+ registerPingWaiterAndEnqueue+ state+ waiter+ (Queue.QueueConnectionScoped (Queue.QueueItem Ping))+ >>= putMVar result+ takeMVar started+ timeout 20000 (takeMVar result) `shouldReturn` Nothing++ _ <- setReconnecting state++ timeout 1000000 (takeMVar result)+ `shouldReturn` Just (Left ExitResetRequested)+ atomically (tryReadTMVar waiter) `shouldReturn` Nothing+ Queue.dequeue (queue state) >>= \case+ Left err -> err `shouldBe` "Queue is closed"+ Right _ -> expectationFailure "expected closed queue"++ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ markConnectionReady state attempt `shouldReturn` True+ queued <- Queue.dequeue (queue state)+ case queued of+ Right item -> LBS.toStrict (transform item) `shouldBe` "PING\r\n"+ Left err -> expectationFailure err++ it "revalidates a standalone publish before a generation-changing commit" $ do+ state <- newTestStateWithQueueCapacity 1+ setServerInfo state (testInfo { max_payload = 4096 })+ markTestStateConnected state+ Queue.enqueue (queue state) (Queue.QueueItem Ping) `shouldReturn` Right ()+ started <- newEmptyMVar+ result <- newEmptyMVar+ _ <- forkIO $ do+ putMVar started ()+ enqueuePublishOnConnected state 2048 oversizedPublish >>= putMVar result+ takeMVar started+ timeout 20000 (takeMVar result) `shouldReturn` Nothing++ _ <- setReconnecting state+ setServerInfo state (testInfo { max_payload = 1024 })+ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ markConnectionReady state attempt `shouldReturn` True++ timeout 1000000 (takeMVar result)+ `shouldReturn` Just (PublishTooLarge 1024)++ it "revalidates a request publish before a generation-changing commit" $ do+ state <- newTestStateWithQueueCapacity 1+ setServerInfo state (testInfo { max_payload = 4096 })+ markTestStateConnected state+ Queue.enqueue (queue state) (Queue.QueueItem Ping) `shouldReturn` Right ()+ committedGeneration <- newEmptyTMVarIO+ started <- newEmptyMVar+ result <- newEmptyMVar+ _ <- forkIO $ do+ putMVar started ()+ enqueuePublishOnConnectedGenerationTracked+ state+ committedGeneration+ 2048+ (Queue.QueueConnectionScoped oversizedPublish)+ >>= putMVar result+ takeMVar started+ timeout 20000 (takeMVar result) `shouldReturn` Nothing++ _ <- setReconnecting state+ setServerInfo state (testInfo { max_payload = 1024 })+ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ markConnectionReady state attempt `shouldReturn` True++ timeout 1000000 (takeMVar result)+ `shouldReturn` Just (PublishTooLarge 1024)+ atomically (tryReadTMVar committedGeneration) `shouldReturn` Nothing++ it "drops durable publishes that exceed a reconnect target's limit" $ do+ state <- newTestStateWithQueueCapacity 2+ rejected <- newEmptyMVar+ setServerInfo state (testInfo { max_payload = 4096 })+ markTestStateConnected state+ Queue.enqueue (queue state) (oversizedPublishWith (putMVar rejected))+ `shouldReturn` Right ()+ Queue.enqueue (queue state) (Queue.QueueItem Ping) `shouldReturn` Right ()++ _ <- setReconnecting state+ setServerInfo state (testInfo { max_payload = 1024 })+ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ markConnectionReady state attempt `shouldReturn` True+ takeMVar rejected `shouldReturn` 1024++ Queue.dequeue (queue state) >>= \case+ Right item -> LBS.toStrict (transform item) `shouldBe` "PING\r\n"+ Left err -> expectationFailure err+ remaining <- timeout 20000 (Queue.dequeue (queue state))+ case remaining of+ Nothing -> pure ()+ Just _ -> expectationFailure "oversized durable publish remained queued"++ it "dispatches durable publish rejection without blocking readiness" $ do+ state <- newTestStateWithQueueCapacity 1+ store <- newSubscriptionStore defaultPendingLimits (pure ())+ stopping <- newTVarIO False+ callbackStarted <- newEmptyMVar+ releaseCallback <- newEmptyMVar+ _ <- startWorkers+ 1+ store+ (readTVar stopping >>= check)+ (const (pure ()))+ let handleRejection maximumSize = do+ putMVar callbackStarted maximumSize+ takeMVar releaseCallback++ setServerInfo state (testInfo { max_payload = 4096 })+ markTestStateConnected state+ Queue.enqueue+ (queue state)+ (oversizedPublishWith (enqueueControl store . handleRejection))+ `shouldReturn` Right ()++ _ <- setReconnecting state+ setServerInfo state (testInfo { max_payload = 1024 })+ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ ready <- newEmptyMVar+ void . forkIO $ markConnectionReady state attempt >>= putMVar ready++ takeMVar callbackStarted `shouldReturn` 1024+ timeout 100000 (takeMVar ready) `shouldReturn` Just True+ putMVar releaseCallback ()+ atomically (awaitCallbackDrain store)+ atomically (writeTVar stopping True)++ describe "Connection lifecycle" $ do+ it "keeps the outbound queue closed until the attempt is ready" $ do+ state <- newTestState+ Just attempt <- beginConnectionAttempt state++ activateConnectionAttempt state attempt `shouldReturn` True+ Queue.enqueue (queue state) (Queue.QueueItem Ping)+ `shouldReturn` Left "Queue is closed"++ markConnectionReady state attempt `shouldReturn` True+ Queue.enqueue (queue state) (Queue.QueueItem Ping) `shouldReturn` Right ()++ it "rejects readiness from an attempt invalidated by reset" $ do+ state <- newTestState+ Just attempt <- beginConnectionAttempt state++ setReconnecting state+ activateConnectionAttempt state attempt `shouldReturn` False+ markConnectionReady state attempt `shouldReturn` False++ readStatus state `shouldReturn` ConnectionReconnecting++ it "does not overwrite a closing state with readiness" $ do+ state <- newTestState+ Just attempt <- beginConnectionAttempt state++ setClosing state ExitClosedByUser+ markConnectionReady state attempt `shouldReturn` False++ readStatus state `shouldReturn` ConnectionClosing ExitClosedByUser++ it "does not reopen the outbound queue after closing starts" $ do+ state <- newTestState+ setClosing state ExitClosedByUser++ beginConnectionAttempt state `shouldReturn` Nothing+ Queue.enqueue (queue state) (Queue.QueueItem Ping)+ `shouldReturn` Left "Queue is closed"++ it "waits for a connected generation strictly after the invalidated one" $ do+ state <- newTestState+ markTestStateConnected state+ invalidated <- setReconnecting state+ Just secondAttempt <- beginConnectionAttempt state+ activateConnectionAttempt state secondAttempt `shouldReturn` True+ markConnectionReady state secondAttempt `shouldReturn` True+ _ <- setReconnecting state++ timeout 20000+ (atomically (waitForConnectionGenerationAfter state invalidated))+ `shouldReturn` Nothing++ Just thirdAttempt <- beginConnectionAttempt state+ activateConnectionAttempt state thirdAttempt `shouldReturn` True+ markConnectionReady state thirdAttempt `shouldReturn` True+ timeout 1000000+ (atomically (waitForConnectionGenerationAfter state invalidated))+ `shouldReturn` Just ()++ it "keeps subscription operations behind the readiness transition" $ do+ state <- newTestState+ markTestStateConnected state+ _ <- setReconnecting state+ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ callbackStarted <- newEmptyMVar+ callbackStatus <- newEmptyMVar++ withSubscriptionGate state $ do+ _ <- forkIO $ do+ putMVar callbackStarted ()+ withSubscriptionGate state (readStatus state >>= putMVar callbackStatus)+ takeMVar callbackStarted+ timeout 20000 (takeMVar callbackStatus) `shouldReturn` Nothing+ markConnectionReady state attempt `shouldReturn` True++ timeout 1000000 (takeMVar callbackStatus)+ `shouldReturn` Just ConnectionConnected++ it "rejects a stale subscription enqueue without holding readiness" $ do+ state <- newTestState+ markTestStateConnected state+ committedGeneration <- newEmptyTMVarIO+ readinessStarted <- newEmptyMVar+ readinessResult <- newEmptyMVar++ withSubscriptionGate state $ do+ generation <- readConnectionGeneration state+ _ <- setReconnecting state+ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ void . forkIO $ do+ putMVar readinessStarted ()+ withSubscriptionGate state (markConnectionReady state attempt)+ >>= putMVar readinessResult+ takeMVar readinessStarted+ timeout 20000 (takeMVar readinessResult) `shouldReturn` Nothing++ enqueueOnGenerationTracked+ state+ generation+ committedGeneration+ (Queue.QueueConnectionScoped (Queue.QueueItem Ping))+ `shouldReturn` Left ExitResetRequested+ atomically (tryReadTMVar committedGeneration) `shouldReturn` Nothing++ timeout 1000000 (takeMVar readinessResult) `shouldReturn` Just True+ readStatus state `shouldReturn` ConnectionConnected++ it "serializes a delayed pre-ready PONG with runtime writes" $ do+ preReadyPing <- newEmptyMVar+ loggerLock <- newTMVarIO ()+ let logger = LoggerConfig Debug (\entry ->+ when (leMessage entry == "routing pre-ready PING") $+ void (tryPutMVar preReadyPing ())) loggerLock+ cfg =+ (testConfig logger)+ { connectionAttempts = 1+ , connectTimeoutMicros = 2 * 1000000+ , connectOptions = [("fake", 4222)]+ }+ messageQueue <- newQueue+ ctx <- newLogContext+ conn <- newConn connectionApi+ state <- newClientState cfg messageQueue conn ctx+ store <- newSubscriptionStore defaultPendingLimits (pure ())+ markTestStateConnected state+ Queue.enqueue+ (queue state)+ (Queue.QueueItem (Pub.Pub "RUNTIME" Nothing Nothing (Just "x")))+ `shouldReturn` Right ()+ _ <- setReconnecting state+ register+ store+ "1"+ (SubscriptionMeta "READY" Nothing StandardSubscription)+ (SubscribeConfig Nothing Nothing)+ (const (pure ()))+ (pure ())+ reads <- newTQueueIO+ atomically $ do+ writeTQueue reads infoFrame+ writeTQueue reads "PONG\r\n"+ resubscribeStarted <- newEmptyMVar+ releaseResubscribe <- newEmptyMVar+ pongStarted <- newEmptyMVar+ releasePong <- newEmptyMVar+ runtimeWriteStarted <- newEmptyMVar+ let writeFrame bytes+ | "CONNECT " `BS.isInfixOf` bytes = pure ()+ | "SUB READY 1\r\n" `BS.isInfixOf` bytes = do+ atomically (writeTQueue reads "PING\r\n")+ putMVar resubscribeStarted ()+ takeMVar releaseResubscribe+ | bytes == "PONG\r\n" = do+ putMVar pongStarted ()+ takeMVar releasePong+ | "PUB RUNTIME " `BS.isInfixOf` bytes =+ putMVar runtimeWriteStarted ()+ | otherwise = pure ()+ transport = Transport+ { transportRead = const (atomically (readTQueue reads))+ , transportWrite = writeFrame+ , transportWriteLazy = writeFrame . LBS.toStrict+ , transportFlush = pure ()+ , transportClose = pure ()+ , transportAbort = do+ void (tryPutMVar releaseResubscribe ())+ void (tryPutMVar releasePong ())+ atomically (writeTQueue reads BS.empty)+ , transportUpgrade = Nothing+ }+ fakeConnectionApi = connectionApi+ { connectTcp = \target _ _ -> do+ pointTransport target transport+ pure (Right ())+ , configure = \_ _ -> pure (Right ())+ }+ engineDone <- newEmptyMVar+ _ <- forkIO $+ runEngine+ fakeConnectionApi+ streamingApi+ broadcastingApi+ Attoparsec.parserApi+ state+ store+ auth+ `finally` putMVar engineDone ()++ timeout 1000000 (takeMVar resubscribeStarted) `shouldReturn` Just ()+ timeout 1000000 (takeMVar preReadyPing) `shouldReturn` Just ()+ putMVar releaseResubscribe ()+ timeout 1000000 (takeMVar pongStarted) `shouldReturn` Just ()+ timeout 20000 (takeMVar runtimeWriteStarted) `shouldReturn` Nothing+ putMVar releasePong ()+ timeout 1000000 (takeMVar runtimeWriteStarted) `shouldReturn` Just ()++ setClosing state ExitClosedByUser+ abort fakeConnectionApi conn+ timeout 1000000 (takeMVar engineDone) `shouldReturn` Just ()++ it "applies the attempt deadline to stalled resubscription writes" $ do+ messageQueue <- newQueue+ ctx <- newLogContext+ conn <- newConn connectionApi+ logger <- newSilentLogger+ let cfg =+ (testConfig logger)+ { connectionAttempts = 1+ , connectTimeoutMicros = 3000000+ , connectOptions = [("fake", 4222)]+ }+ state <- newClientState cfg messageQueue conn ctx+ store <- newSubscriptionStore defaultPendingLimits (pure ())+ markTestStateConnected state+ _ <- setReconnecting state+ register+ store+ "1"+ (SubscriptionMeta "READY" Nothing StandardSubscription)+ (SubscribeConfig Nothing Nothing)+ (const (pure ()))+ (pure ())+ reads <- newTQueueIO+ atomically $ do+ writeTQueue reads infoFrame+ writeTQueue reads "PONG\r\n"+ resubscribeStarted <- newEmptyMVar+ releaseWrite <- newEmptyMVar+ let transport = Transport+ { transportRead = const (atomically (readTQueue reads))+ , transportWrite = const (pure ())+ , transportWriteLazy = \bytes ->+ when ("SUB READY 1\r\n" `BS.isInfixOf` LBS.toStrict bytes) $ do+ putMVar resubscribeStarted ()+ takeMVar releaseWrite+ , transportFlush = pure ()+ , transportClose = pure ()+ , transportAbort = do+ void (tryPutMVar releaseWrite ())+ atomically (writeTQueue reads BS.empty)+ , transportUpgrade = Nothing+ }+ fakeConnectionApi = connectionApi+ { connectTcp = \target _ _ -> do+ threadDelay 1500000+ pointTransport target transport+ pure (Right ())+ , configure = \_ _ -> pure (Right ())+ }+ engineDone <- newEmptyMVar+ engineThread <- forkIO $+ runEngine+ fakeConnectionApi+ streamingApi+ broadcastingApi+ Attoparsec.parserApi+ state+ store+ auth+ `finally` putMVar engineDone ()+ let stopFailedEngine = do+ void (tryPutMVar releaseWrite ())+ setClosing state ExitClosedByUser+ abort fakeConnectionApi conn+ killThread engineThread+ readMVar engineDone++ started <- timeout 5000000 (takeMVar resubscribeStarted)+ case started of+ Just () -> pure ()+ Nothing -> do+ status <- readStatus state+ stopFailedEngine+ expectationFailure+ ("resubscription did not start; status=" ++ show status)+ completed <- timeout 2250000 (readMVar engineDone)+ case completed of+ Just () -> pure ()+ Nothing -> do+ status <- readStatus state+ stopFailedEngine+ expectationFailure+ ("engine exceeded the shared attempt deadline; status=" ++ show status)+ readStatus state >>= \case+ ConnectionClosed _ -> pure ()+ status -> expectationFailure+ ("engine remained live after readiness deadline: " ++ show status)++ describe "Server errors" $ do+ it "classifies typed protocol errors without parsing their reason" $ do+ serverErrorKind (serverErrorFromProtocol (Err.ErrAuthViolation "secret"))+ `shouldBe` ServerErrorAuthentication+ serverErrorKind (serverErrorFromProtocol (Err.ErrPermViolation "private.subject"))+ `shouldBe` ServerErrorPermission+ serverErrorKind (serverErrorFromProtocol (Err.ErrInvalidProtocol "bad"))+ `shouldBe` ServerErrorProtocol+ serverErrorKind (serverErrorFromProtocol (Err.ErrErr "future server error"))+ `shouldBe` ServerErrorUnknown++ it "retains the legacy ServerError Show representation" $ do+ let serverError = serverErrorFromProtocol (Err.ErrPermViolation "private")+ show serverError+ `shouldBe` "ServerError \"private\""+ show (ExitServerError serverError)+ `shouldBe` "ExitServerError (ServerError \"private\")"+ show (ConnectionEventClosed (ExitServerError serverError))+ `shouldBe` "ConnectionEventClosed (ExitServerError (ServerError \"private\"))"++ it "bounds handler events without blocking the protocol producer" $ do+ state <- newTestState+ let serverError = serverErrorFromProtocol (Err.ErrPermViolation "private")++ replicateM_ 256+ (notifyServerError state serverError `shouldReturn` ServerErrorQueued)++ notifyServerError state serverError `shouldReturn` ServerErrorDroppedReport+ notifyServerError state serverError `shouldReturn` ServerErrorDropped++ it "bounds handler events by retained reason bytes" $ do+ state <- newTestState+ let largeReason = BS.replicate (1024 * 1024) 120+ serverError = serverErrorFromProtocol (Err.ErrPermViolation largeReason)++ notifyServerError state serverError `shouldReturn` ServerErrorQueued+ notifyServerError state serverError `shouldReturn` ServerErrorDroppedReport++ it "drops nonfatal errors instead of reconnecting when callbacks are full" $ do+ state <- newTestState+ store <- newSubscriptionStore defaultPendingLimits (pure ())+ let protocolError = Err.ErrPermViolation "private"+ serverError = serverErrorFromProtocol protocolError+ replicateM_ 256+ (notifyServerError state serverError `shouldReturn` ServerErrorQueued)++ routeMessage state store (ParsedErr protocolError)+ `shouldReturn` RouteContinue++ it "guarantees lifecycle callback order when server callbacks are full" $ do+ observed <- newTQueueIO+ state <- newTestStateWithConfig $ \cfg ->+ cfg { connectionEventHandler = atomically . writeTQueue observed }+ let serverError = serverErrorFromProtocol (Err.ErrPermViolation "private")+ replicateM_ 256+ (notifyServerError state serverError `shouldReturn` ServerErrorQueued)+ notifyServerError state serverError `shouldReturn` ServerErrorDroppedReport+ markTestStateConnected state+ setReconnecting state+ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ markConnectionReady state attempt `shouldReturn` True+ setClosing state ExitClosedByUser+ markClosed state ExitClosedByUser `shouldReturn` Just ExitClosedByUser+ eventWorker <- startCallbackWorker state++ atomically (readTQueue observed)+ `shouldReturn` ConnectionEventDisconnected+ atomically (readTQueue observed)+ `shouldReturn` ConnectionEventReconnected+ atomically (readTQueue observed)+ `shouldReturn` ConnectionEventClosed ExitClosedByUser+ stopCallbackWorker eventWorker describe "Handshake" $ do it "returns a typed error when no servers are configured" $ do result <- Client.newClient [] [Client.withConnectionAttempts 1]@@ -150,6 +833,42 @@ readTVarIO writes `shouldReturn` [LBS.toStrict (transform testConnect <> transform Ping)] + it "accepts normally chunked INFO and PONG frames" $ do+ state <- newTestState+ conn <- newConn connectionApi+ writes <- newTVarIO []+ let (infoStart, infoEnd) = BS.splitAt 40 infoFrame+ transport <- newScriptedTransport [infoStart, infoEnd, "PO", "NG\r\n"] writes+ pointTransport conn transport++ result <- performHandshake connectionApi Attoparsec.parserApi state auth conn "127.0.0.1"++ result `shouldBe` Right ()++ it "rejects an oversized incomplete INFO control frame" $ do+ state <- newTestState+ conn <- newConn connectionApi+ writes <- newTVarIO []+ transport <- newScriptedTransport oversizedHandshakeChunks writes+ pointTransport conn transport++ result <- performHandshake connectionApi stalledInfoParser state auth conn "127.0.0.1"++ result `shouldBe`+ Left (HandshakeProtocolError "INFO control frame exceeds 65536 byte limit")++ it "rejects an oversized incomplete control frame before PONG" $ do+ state <- newTestState+ conn <- newConn connectionApi+ writes <- newTVarIO []+ transport <- newScriptedTransport (infoFrame : oversizedHandshakeChunks) writes+ pointTransport conn transport++ result <- performHandshake connectionApi stalledPongParser state auth conn "127.0.0.1"++ result `shouldBe`+ Left (HandshakeProtocolError "PONG control frame exceeds 65536 byte limit")+ it "accepts a parser backend that drops an invalid prefix before INFO" $ do state <- newTestState conn <- newConn connectionApi@@ -187,13 +906,16 @@ atomically $ modifyTVar' writes (<> [LBS.toStrict bytes]) , transportFlush = pure () , transportClose = pure ()+ , transportAbort = pure () , transportUpgrade = Nothing } pointTransport conn transport - result <- performHandshake connectionApi incrementalInfoParser state auth conn "127.0.0.1"+ result <- timeout 20000 $+ performHandshake connectionApi incrementalInfoParser state auth conn "127.0.0.1" - result `shouldBe` Left HandshakeTimeout+ result `shouldBe` Nothing+ putMVar blocker BS.empty it "accepts +OK before the handshake PONG" $ do state <- newTestState@@ -206,6 +928,21 @@ result `shouldBe` Right () + it "preserves protocol bytes read after the handshake PONG" $ do+ state <- newTestState+ conn <- newConn connectionApi+ writes <- newTVarIO []+ transport <- newScriptedTransport+ [infoFrame, "PONG\r\nMSG READY 1 2\r\nok\r\n"] writes+ pointTransport conn transport++ result <- performHandshake+ connectionApi Attoparsec.parserApi state auth conn "127.0.0.1"++ result `shouldBe` Right ()+ readData (reader connectionApi) conn 4096+ `shouldReturn` Right "MSG READY 1 2\r\nok\r\n"+ it "returns authentication errors received before PONG" $ do state <- newTestState conn <- newConn connectionApi@@ -242,11 +979,23 @@ newTestStateWithConfig updateConfig = do queue <- newQueue+ newTestStateWithQueue updateConfig queue++newTestStateWithQueueCapacity capacity = do+ queue <- newQueueWithCapacity capacity+ newTestStateWithQueue id queue++newTestStateWithQueue updateConfig queue = do ctx <- newLogContext conn <- newConn connectionApi logger <- newSilentLogger newClientState (updateConfig (testConfig logger)) queue conn ctx +markTestStateConnected state = do+ Just attempt <- beginConnectionAttempt state+ activateConnectionAttempt state attempt `shouldReturn` True+ markConnectionReady state attempt `shouldReturn` True+ newSilentLogger :: IO LoggerConfig newSilentLogger = do lock <- newTMVarIO ()@@ -262,6 +1011,8 @@ , connectConfig = testConnect , loggerConfig = logger , tlsConfig = Nothing+ , serverErrorHandler = const (pure ())+ , connectionEventHandler = const (pure ()) , exitAction = const (pure ()) , connectOptions = [] }@@ -306,6 +1057,19 @@ Nothing Nothing +oversizedPublish :: Queue.QueueItem+oversizedPublish = oversizedPublishWith (const (pure ()))++oversizedPublishWith :: (Int -> IO ()) -> Queue.QueueItem+oversizedPublishWith reject =+ Queue.QueuePayloadBound 2048 reject . Queue.QueueItem $+ Pub.Pub+ { Pub.subject = "BOUNDARY.TOO_BIG"+ , Pub.payload = Just (BS.replicate 2048 120)+ , Pub.replyTo = Nothing+ , Pub.headers = Nothing+ }+ incrementalInfoParser :: ParserAPI ParsedMessage incrementalInfoParser = ParserAPI parseInfo where@@ -325,6 +1089,8 @@ parseInfo bytes | bytes == "XIN" = DropPrefix 1 "invalid prefix"+ | bytes == "IN" =+ NeedMore | bytes == infoFrame = Emit (ParsedInfo testInfo) "" | bytes == "PONG\r\n" =@@ -332,6 +1098,19 @@ | otherwise = Reject ("unexpected input: " ++ show bytes) +stalledInfoParser :: ParserAPI ParsedMessage+stalledInfoParser = ParserAPI (const NeedMore)++stalledPongParser :: ParserAPI ParsedMessage+stalledPongParser = ParserAPI parseFrame+ where+ parseFrame bytes+ | bytes == infoFrame = Emit (ParsedInfo testInfo) ""+ | otherwise = NeedMore++oversizedHandshakeChunks :: [BS.ByteString]+oversizedHandshakeChunks = replicate 17 (BS.replicate 4096 73)+ infoFrame :: BS.ByteString infoFrame = "INFO {\"server_id\":\"srv\",\"version\":\"1.0.0\",\"go\":\"go1\",\"host\":\"127.0.0.1\",\"port\":4222,\"max_payload\":1024,\"proto\":1}\r\n"@@ -356,5 +1135,6 @@ atomically $ modifyTVar' writes (<> [LBS.toStrict bytes]) , transportFlush = pure () , transportClose = pure ()+ , transportAbort = pure () , transportUpgrade = Nothing }
test/Unit/JetStream/ConsumerSpec.hs view
@@ -10,7 +10,12 @@ , (.=) ) import qualified Data.ByteString.Lazy as LBS+import Data.Int (Int64)+import Data.Ratio ((%)) import JetStream.Consumer.Types+import JetStream.Error+ ( JetStreamError (JetStreamDecodeError)+ ) import Test.Hspec spec :: Spec@@ -34,6 +39,57 @@ eitherDecode (encode request) `shouldBe` Right (object ["deliver_policy" .= ("last" :: String)]) + it "encodes backoff durations as exact integer nanoseconds" $ do+ let request = consumerConfigRequest+ [ withConsumerBackoff+ [ fromRational (1234567891 % 1000000000)+ , fromRational (1 % 1000000000)+ ]+ , withConsumerMaxRequestBatch maxBound+ ]+ validateConsumerConfigRequest request `shouldBe` Right ()+ eitherDecode (encode request) `shouldBe` Right (object+ [ "backoff" .= [1234567891 :: Integer, 1]+ , "max_batch" .= (maxBound :: Int)+ ])++ it "floors backoff at sub-nanosecond and Int64 boundaries" $ do+ let request = consumerConfigRequest+ [withConsumerBackoff [subNanosecond, oneNanosecond, maxInt64Nanoseconds]]+ validateConsumerConfigRequest request `shouldBe` Right ()+ eitherDecode (encode request) `shouldBe` Right (object+ ["backoff" .= [0 :: Integer, 1, toInteger (maxBound :: Int64)]])++ it "rejects negative and overflowing backoff durations" $ do+ validateConsumerConfigRequest+ (consumerConfigRequest [withConsumerBackoff [negativeSubNanosecond]])+ `shouldBe` Left (JetStreamDecodeError backoffRangeError)+ validateConsumerConfigRequest+ (consumerConfigRequest [withConsumerBackoff [overflowingNanoseconds]])+ `shouldBe` Left (JetStreamDecodeError backoffRangeError)++ it "rejects negative max request batch" $ do+ validateConsumerConfigRequest+ (consumerConfigRequest [withConsumerMaxRequestBatch (-1)])+ `shouldBe` Left (JetStreamDecodeError "consumer max request batch must be zero or greater")++ it "allows backoff no longer than positive max deliver" $ do+ validateConsumerConfigRequest (backoffWithMaxDeliver 2 [1]) `shouldBe` Right ()+ validateConsumerConfigRequest (backoffWithMaxDeliver 2 [1, 2]) `shouldBe` Right ()+ validateConsumerConfigRequest (backoffWithMaxDeliver 0 [1, 2, 3]) `shouldBe` Right ()+ validateConsumerConfigRequest (backoffWithMaxDeliver (-1) [1, 2, 3]) `shouldBe` Right ()+ validateConsumerConfigRequest (backoffWithMaxDeliver 2 [1, 2, 3])+ `shouldBe` Left (JetStreamDecodeError+ "consumer backoff length cannot exceed positive max deliver")++ it "omits empty backoff and zero max request batch values" $ do+ let request = consumerConfigRequest+ [ withConsumerBackoff []+ , withConsumerMaxRequestBatch 0+ ]+ validateConsumerConfigRequest request `shouldBe` Right ()+ eitherDecode (encode request) `shouldBe` Right (object [])+ describe "Consumer reset request JSON" $ do it "encodes an optional reset sequence" $ do eitherDecode (encode (consumerResetRequest [withConsumerResetSequence 7]))@@ -43,6 +99,15 @@ it "decodes server responses into concrete fields" $ do eitherDecode consumerInfoJSON `shouldBe` Right consumerInfoFixture + describe "ConsumerConfig response JSON" $ do+ it "normalizes absent backoff and zero max request batch to unset" $ do+ fmap configResourceLimits (eitherDecode zeroMaxBatchConsumerConfigJSON)+ `shouldBe` Right (Nothing, Nothing)++ it "round-trips backoff and max request batch" $ do+ eitherDecode (encode (consumerInfoConfig consumerInfoFixture))+ `shouldBe` Right (consumerInfoConfig consumerInfoFixture)+ describe "ConsumerResetResponse JSON" . it "decodes reset metadata with the updated consumer info" $ do fmap consumerResetResponseSequence (eitherDecode consumerResetJSON)@@ -67,6 +132,30 @@ , consumerNamesConsumers = ["orders-puller", "orders-worker"] } +configResourceLimits config =+ (consumerConfigBackoff config, consumerConfigMaxRequestBatch config)++backoffWithMaxDeliver maxDeliver backoff =+ consumerConfigRequest+ [ withConsumerMaxDeliver maxDeliver+ , withConsumerBackoff backoff+ ]++subNanosecond = fromRational (1 % 2000000000)++negativeSubNanosecond = fromRational ((-1) % 2000000000)++oneNanosecond = fromRational (1 % 1000000000)++maxInt64Nanoseconds =+ fromRational (toInteger (maxBound :: Int64) % 1000000000)++overflowingNanoseconds =+ fromRational ((toInteger (maxBound :: Int64) + 1) % 1000000000)++backoffRangeError =+ "consumer backoff durations must floor to nanoseconds between 0 and 9223372036854775807"+ durablePullConsumerConfigRequest :: ConsumerConfigRequest durablePullConsumerConfigRequest = consumerConfigRequest@@ -142,15 +231,19 @@ consumerInfoJSON :: LBS.ByteString consumerInfoJSON =- "{\"type\":\"io.nats.jetstream.api.v1.consumer_info_response\",\"stream_name\":\"ORDERS\",\"name\":\"orders-puller\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\"},\"delivered\":{\"consumer_seq\":2,\"stream_seq\":10},\"ack_floor\":{\"consumer_seq\":1,\"stream_seq\":9},\"num_ack_pending\":1,\"num_redelivered\":0,\"num_waiting\":0,\"num_pending\":3}"+ "{\"type\":\"io.nats.jetstream.api.v1.consumer_info_response\",\"stream_name\":\"ORDERS\",\"name\":\"orders-puller\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\",\"backoff\":[1000000001,2],\"max_batch\":128},\"delivered\":{\"consumer_seq\":2,\"stream_seq\":10},\"ack_floor\":{\"consumer_seq\":1,\"stream_seq\":9},\"num_ack_pending\":1,\"num_redelivered\":0,\"num_waiting\":0,\"num_pending\":3}" +zeroMaxBatchConsumerConfigJSON :: LBS.ByteString+zeroMaxBatchConsumerConfigJSON =+ "{\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\",\"max_batch\":0}"+ consumerResetJSON :: LBS.ByteString consumerResetJSON = "{\"type\":\"io.nats.jetstream.api.v1.consumer_reset_response\",\"stream_name\":\"ORDERS\",\"name\":\"orders-puller\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\"},\"delivered\":{\"consumer_seq\":1,\"stream_seq\":7},\"ack_floor\":{\"consumer_seq\":0,\"stream_seq\":0},\"num_ack_pending\":0,\"num_redelivered\":0,\"num_waiting\":0,\"num_pending\":3,\"reset_seq\":7}" consumerListJSON :: LBS.ByteString consumerListJSON =- "{\"type\":\"io.nats.jetstream.api.v1.consumer_list_response\",\"total\":1,\"offset\":0,\"limit\":1024,\"consumers\":[{\"type\":\"io.nats.jetstream.api.v1.consumer_info_response\",\"stream_name\":\"ORDERS\",\"name\":\"orders-puller\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\"},\"delivered\":{\"consumer_seq\":2,\"stream_seq\":10},\"ack_floor\":{\"consumer_seq\":1,\"stream_seq\":9},\"num_ack_pending\":1,\"num_redelivered\":0,\"num_waiting\":0,\"num_pending\":3}]}"+ "{\"type\":\"io.nats.jetstream.api.v1.consumer_list_response\",\"total\":1,\"offset\":0,\"limit\":1024,\"consumers\":[{\"type\":\"io.nats.jetstream.api.v1.consumer_info_response\",\"stream_name\":\"ORDERS\",\"name\":\"orders-puller\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\",\"backoff\":[1000000001,2],\"max_batch\":128},\"delivered\":{\"consumer_seq\":2,\"stream_seq\":10},\"ack_floor\":{\"consumer_seq\":1,\"stream_seq\":9},\"num_ack_pending\":1,\"num_redelivered\":0,\"num_waiting\":0,\"num_pending\":3}]}" consumerNamesJSON :: LBS.ByteString consumerNamesJSON =@@ -176,6 +269,8 @@ , consumerConfigMaxDeliver = Nothing , consumerConfigMaxWaiting = Nothing , consumerConfigMaxAckPending = Nothing+ , consumerConfigBackoff = Just [fromRational (1000000001 % 1000000000), fromRational (2 % 1000000000)]+ , consumerConfigMaxRequestBatch = Just 128 , consumerConfigInactiveThreshold = Nothing , consumerConfigIdleHeartbeat = Nothing , consumerConfigHeadersOnly = Nothing
test/Unit/JetStream/MessageSpec.hs view
@@ -3,11 +3,14 @@ module JetStream.MessageSpec (spec) where import qualified Client.API as Nats+import Data.Int (Int64) import Data.IORef+import Data.Ratio ((%)) import JetStream.Error (JetStreamError (..))-import JetStream.Message (fetchMessages)+import JetStream.Message (fetchMessages, messageAPI) import JetStream.Message.Types import JetStream.Options (newJetStreamContext)+import JetStream.Types (withRequestTimeout) import Publish (defaultPublishConfig) import Publish.Config (publishReplyTo) import Test.Hspec@@ -39,6 +42,112 @@ inProgressPayload `shouldBe` "+WPI" termPayload `shouldBe` "+TERM" + it "encodes delayed NAKs as exact integer nanoseconds" $ do+ fmap nakDelayPayload (nakDelay 1.234567891234)+ `shouldBe` Just "-NAK {\"delay\":1234567891}"++ it "rejects delays shorter than one nanosecond" $ do+ nakDelay 0 `shouldBe` Nothing+ nakDelay 0.000000000999 `shouldBe` Nothing+ nakDelay (-1) `shouldBe` Nothing++ it "accepts the exact signed 64-bit nanosecond boundaries" $ do+ let oneNanosecond = fromRational (1 % 1000000000)+ maxNanoseconds = toInteger (maxBound :: Int64)+ maxDelay = fromRational (maxNanoseconds % 1000000000)+ overflowDelay = fromRational ((maxNanoseconds + 1) % 1000000000)+ fmap nakDelayPayload (nakDelay oneNanosecond)+ `shouldBe` Just "-NAK {\"delay\":1}"+ fmap nakDelayPayload (nakDelay maxDelay)+ `shouldBe` Just "-NAK {\"delay\":9223372036854775807}"+ nakDelay overflowDelay `shouldBe` Nothing++ describe "message dispositions" $ do+ it "keeps empty-option acknowledgements asynchronous" $ do+ publishes <- newIORef []+ requests <- newIORef []+ let api = messageAPI+ (newJetStreamContext (ackFakeClient publishes requests ackReply) [])+ (error "unused consumer API")++ ack api ackMessage [] `shouldReturn` Right ()+ readIORef publishes `shouldReturn` [(ackSubject, "+ACK")]+ readIORef requests `shouldReturn` []++ it "uses request-reply when acknowledgement request options are supplied" $ do+ publishes <- newIORef []+ requests <- newIORef []+ let api = messageAPI+ (newJetStreamContext (ackFakeClient publishes requests ackReply) [])+ (error "unused consumer API")++ ack api ackMessage [withRequestTimeout 0.25] `shouldReturn` Right ()+ readIORef publishes `shouldReturn` []+ readIORef requests `shouldReturn` [(ackSubject, "+ACK")]++ it "always confirms ackSync and termSync" $ do+ publishes <- newIORef []+ requests <- newIORef []+ let api = messageAPI+ (newJetStreamContext (ackFakeClient publishes requests ackReply) [])+ (error "unused consumer API")++ ackSync api ackMessage [] `shouldReturn` Right ()+ termSync api ackMessage [] `shouldReturn` Right ()+ readIORef publishes `shouldReturn` []+ readIORef requests `shouldReturn`+ [ (ackSubject, "+ACK")+ , (ackSubject, "+TERM")+ ]++ it "sends delayed NAKs with the selected confirmation mode" $ do+ publishes <- newIORef []+ requests <- newIORef []+ let api = messageAPI+ (newJetStreamContext (ackFakeClient publishes requests ackReply) [])+ (error "unused consumer API")+ Just delay = nakDelay 2.5++ nakWithDelay api ackMessage delay [] `shouldReturn` Right ()+ nakWithDelay api ackMessage delay [withRequestTimeout 0.25]+ `shouldReturn` Right ()+ readIORef publishes `shouldReturn`+ [(ackSubject, "-NAK {\"delay\":2500000000}")]+ readIORef requests `shouldReturn`+ [(ackSubject, "-NAK {\"delay\":2500000000}")]++ it "maps confirmed acknowledgement timeouts" $ do+ publishes <- newIORef []+ requests <- newIORef []+ let api = messageAPI+ (newJetStreamContext+ (ackFakeClient publishes requests (Left Nats.NatsRequestTimedOut))+ [])+ (error "unused consumer API")++ ackSync api ackMessage [] `shouldReturn` Left JetStreamTimeout++ it "maps confirmed acknowledgement status errors" $ do+ publishes <- newIORef []+ requests <- newIORef []+ let response = Right (ackReplyMessage "" (Just [("Status", "409"), ("Description", "consumer deleted")]))+ api = messageAPI+ (newJetStreamContext (ackFakeClient publishes requests response) [])+ (error "unused consumer API")++ termSync api ackMessage []+ `shouldReturn` Left (JetStreamStatusError 409 (Just "consumer deleted"))++ it "accepts confirmed acknowledgement reply payloads" $ do+ publishes <- newIORef []+ requests <- newIORef []+ let response = Right (ackReplyMessage "server metadata" Nothing)+ api = messageAPI+ (newJetStreamContext (ackFakeClient publishes requests response) [])+ (error "unused consumer API")++ ackSync api ackMessage [] `shouldReturn` Right ()+ describe "message metadata" $ do it "parses v1 JetStream ack reply subjects" $ do let metadata = messageMetadata $@@ -148,6 +257,7 @@ , Nats.newInbox = pure "_INBOX.batch" , Nats.ping = \_ -> pure (Right ()) , Nats.flush = \_ -> pure (Right ())+ , Nats.connectionState = pure Nats.ConnectionConnected , Nats.reset = \_ -> pure () , Nats.close = \_ -> pure () }@@ -182,6 +292,50 @@ , Nats.newInbox = pure "_INBOX.timeout" , Nats.ping = \_ -> pure (Right ()) , Nats.flush = \_ -> pure (Right ())+ , Nats.connectionState = pure Nats.ConnectionConnected+ , Nats.reset = \_ -> pure ()+ , Nats.close = \_ -> pure ()+ }++ackSubject :: Subject+ackSubject = "$JS.ACK.ORDERS.WORKER.1.1.1"++ackMessage :: Message+ackMessage = Message "ORDERS.created" "payload" Nothing (Just ackSubject) Nothing++ackReply :: Either Nats.NatsError Nats.MsgView+ackReply = Right (ackReplyMessage "" Nothing)++ackReplyMessage :: Payload -> Maybe Nats.Headers -> Nats.Message+ackReplyMessage body responseHeaders =+ Nats.Message+ { Nats.subject = "_INBOX.ack"+ , Nats.sid = "sid-ack"+ , Nats.replyTo = Nothing+ , Nats.payload = body+ , Nats.headers = responseHeaders+ }++ackFakeClient+ :: IORef [(Subject, Payload)]+ -> IORef [(Subject, Payload)]+ -> Either Nats.NatsError Nats.MsgView+ -> Nats.Client+ackFakeClient publishCalls requestCalls response =+ Nats.Client+ { Nats.publish = \subject body _ -> do+ modifyIORef' publishCalls (++ [(subject, body)])+ pure (Right ())+ , Nats.subscribe = \_ _ _ -> pure (Right (Nats.Subscription "sid-ack"))+ , Nats.subscribeOnce = \_ _ _ -> pure (Right (Nats.Subscription "sid-ack-once"))+ , Nats.request = \subject body _ -> do+ modifyIORef' requestCalls (++ [(subject, body)])+ pure response+ , Nats.unsubscribe = \_ _ -> pure (Right ())+ , Nats.newInbox = pure "_INBOX.ack"+ , Nats.ping = \_ -> pure (Right ())+ , Nats.flush = \_ -> pure (Right ())+ , Nats.connectionState = pure Nats.ConnectionConnected , Nats.reset = \_ -> pure () , Nats.close = \_ -> pure () }
test/Unit/JetStream/ProtocolSpec.hs view
@@ -116,6 +116,7 @@ , Nats.newInbox = pure "_INBOX.TEST" , Nats.ping = \_ -> pure (Right ()) , Nats.flush = \_ -> pure (Right ())+ , Nats.connectionState = pure Nats.ConnectionConnected , Nats.reset = \_ -> pure () , Nats.close = \_ -> pure () }
+ test/Unit/JetStream/StreamSpec.hs view
@@ -0,0 +1,76 @@+{-# LANGUAGE OverloadedStrings #-}++module JetStream.StreamSpec (spec) where++import Data.Aeson (eitherDecode, encode, object, (.=))+import qualified Data.ByteString.Lazy as LBS+import Data.Int (Int32)+import JetStream.Error (JetStreamError (JetStreamDecodeError))+import JetStream.Stream.Types+import Test.Hspec++spec :: Spec+spec = do+ describe "StreamConfig request JSON" $ do+ it "omits max message size unless configured" $ do+ let request = streamConfigRequest "ORDERS" ["orders.>"] []+ eitherDecode (encode request) `shouldBe` Right (object+ [ "name" .= ("ORDERS" :: String)+ , "subjects" .= ["orders.>" :: String]+ ])++ it "encodes valid max message size boundaries" $ do+ let encodedMaxMessageSize value =+ eitherDecode . encode $+ streamConfigRequest "ORDERS" [] [withMaxMessageSize value]+ validatedMaxMessageSize value =+ validateStreamConfigRequest $+ streamConfigRequest "ORDERS" [] [withMaxMessageSize value]+ expected value = object+ [ "name" .= ("ORDERS" :: String)+ , "subjects" .= ([] :: [String])+ , "max_msg_size" .= value+ ]+ encodedMaxMessageSize (-1) `shouldBe` Right (expected (-1 :: Int32))+ encodedMaxMessageSize 0 `shouldBe` Right (expected (0 :: Int32))+ encodedMaxMessageSize maxBound `shouldBe` Right (expected (maxBound :: Int32))+ validatedMaxMessageSize (-1) `shouldBe` Right ()+ validatedMaxMessageSize 0 `shouldBe` Right ()+ validatedMaxMessageSize maxBound `shouldBe` Right ()++ it "rejects max message sizes below unlimited" $ do+ validateStreamConfigRequest+ (streamConfigRequest "ORDERS" [] [withMaxMessageSize (-2)])+ `shouldBe` Left (JetStreamDecodeError "stream max message size must be -1 or greater")+ validateStreamConfigRequest+ (streamConfigRequest "ORDERS" [] [withMaxMessageSize minBound])+ `shouldBe` Left (JetStreamDecodeError "stream max message size must be -1 or greater")++ describe "StreamConfig response JSON" $ do+ it "normalizes a missing max message size to unlimited" $ do+ fmap streamConfigMaxMessageSize (eitherDecode streamConfigWithoutMaxMessageSizeJSON)+ `shouldBe` Right (-1)++ it "round-trips max message size" $ do+ eitherDecode (encode streamConfigFixture) `shouldBe` Right streamConfigFixture++streamConfigFixture :: StreamConfig+streamConfigFixture =+ StreamConfig+ { streamConfigName = "ORDERS"+ , streamConfigSubjects = Just ["orders.>"]+ , streamConfigRetention = LimitsPolicy+ , streamConfigStorage = MemoryStorage+ , streamConfigDiscard = DiscardOld+ , streamConfigMaxMessages = -1+ , streamConfigMaxBytes = -1+ , streamConfigMaxAge = 0+ , streamConfigMaxMessageSize = maxBound+ , streamConfigReplicas = 1+ , streamConfigDuplicateWindow = Nothing+ , streamConfigAllowDirect = False+ }++streamConfigWithoutMaxMessageSizeJSON :: LBS.ByteString+streamConfigWithoutMaxMessageSizeJSON =+ "{\"name\":\"ORDERS\",\"subjects\":[\"orders.>\"],\"retention\":\"limits\",\"storage\":\"memory\",\"discard\":\"old\",\"max_msgs\":-1,\"max_bytes\":-1,\"max_age\":0,\"num_replicas\":1,\"allow_direct\":false}"
test/Unit/StreamingSpec.hs view
@@ -82,6 +82,7 @@ case result of Left err -> return $ Left (show err) Right bytes -> return $ Right bytes+ , bufferRead = \_ _ -> pure () , closeReader = hClose , openReader = \_ -> pure () }@@ -148,11 +149,18 @@ ctx <- newLogContext q <- newQueue output <- newTVarIO []+ done <- newEmptyMVar let pub = Pub.Pub "FOO" Nothing Nothing (Just "0123456789") let BroadcastingAPI runBroadcasting = broadcastingApi+ _ <- forkIO $+ runWithLogger dl ctx+ (runBroadcasting q captureWriterApi (CaptureWriter output) :: AppM ())+ `finally` putMVar done () enqueue q (QueueItem pub) `shouldReturn` Right ()+ atomically $+ assertTVarWithRetry output ["PUB FOO 10\r\n0123456789\r\n"] close q- runWithLogger dl ctx (runBroadcasting q captureWriterApi (CaptureWriter output) :: AppM ())+ takeMVar done readTVarIO output `shouldReturn` ["PUB FOO 10\r\n0123456789\r\n"] ensureTVarIsEmpty :: TVar ByteString -> STM ()
test/Unit/SubscriptionStoreSpec.hs view
@@ -14,6 +14,40 @@ spec :: Spec spec = do+ describe "terminal cleanup" $ do+ it "clears active and expiry state while preserving queued callbacks" $ do+ delivered <- newEmptyMVar+ stopping <- newTVarIO False+ store <- newSubscriptionStore defaultPendingLimits (pure ())+ register+ store+ "1"+ (SubscriptionMeta "A" Nothing StandardSubscription)+ (SubscribeConfig Nothing Nothing)+ (const (putMVar delivered ()))+ (pure ())+ register+ store+ "2"+ (SubscriptionMeta "B" Nothing OneShotSubscription)+ (SubscribeConfig (Just 60) Nothing)+ (const (pure ()))+ (pure ())+ dispatchMessage store (message "1" "one") `shouldReturn` DispatchQueued++ closeStore store++ active store `shouldReturn` []+ hasTrackedExpiries store `shouldReturn` False+ _ <- startWorkers+ 1+ store+ (readTVar stopping >>= check)+ (const (pure ()))+ takeMVar delivered+ atomically (awaitCallbackDrain store)+ atomically (writeTVar stopping True)+ describe "global pending delivery limits" $ do it "shares the message limit across subscriptions and coalesces notification" $ do slowEvents <- newIORef (0 :: Int)@@ -56,7 +90,7 @@ register store "1"- (SubscriptionMeta "A" Nothing False)+ (SubscriptionMeta "A" Nothing StandardSubscription) (SubscribeConfig Nothing Nothing) (\_ -> putMVar callbackStarted () >> takeMVar releaseCallback) (pure ())@@ -92,7 +126,7 @@ register store "1"- (SubscriptionMeta "A" Nothing False)+ (SubscriptionMeta "A" Nothing StandardSubscription) (SubscribeConfig Nothing Nothing) (const (throwIO (userError "callback failed"))) (pure ())@@ -118,7 +152,7 @@ register store "2"- (SubscriptionMeta "B" Nothing True)+ (SubscriptionMeta "B" Nothing OneShotSubscription) (SubscribeConfig Nothing Nothing) (const (pure ())) (pure ())@@ -130,6 +164,61 @@ dispatchMessage store (message "2" "two") `shouldReturn` DispatchMissing + it "accepts a dropped one-shot before reporting overflow" $ do+ accepted <- newTVarIO (0 :: Int)+ store <- newSubscriptionStore (PendingLimits 1 1024) (pure ())+ registerSubscription store "1" "A" (pure ())+ registerWithAcceptance+ store+ "2"+ (SubscriptionMeta "B" Nothing OneShotSubscription)+ (SubscribeConfig Nothing Nothing)+ (modifyTVar' accepted (+ 1))+ (const (pure ()))+ (pure ())++ dispatchMessage store (message "1" "one")+ `shouldReturn` DispatchQueued+ dispatchMessage store (message "2" "two")+ `shouldReturn` DispatchDropped True+ readTVarIO accepted `shouldReturn` 1+ dispatchMessage store (message "2" "duplicate")+ `shouldReturn` DispatchMissing+ readTVarIO accepted `shouldReturn` 1++ it "commits request acceptance and overflow atomically" $ do+ accepted <- newTVarIO False+ rejected <- newTVarIO False+ throwOnRejection <- newTVarIO True+ store <- newSubscriptionStore (PendingLimits 1 1024) (pure ())+ registerSubscription store "1" "A" (pure ())+ registerWithDispatchHooks+ store+ "2"+ (SubscriptionMeta "B" Nothing RequestReplySubscription)+ (SubscribeConfig Nothing Nothing)+ (writeTVar accepted True)+ (do+ shouldThrow <- readTVar throwOnRejection+ when shouldThrow (throwSTM (userError "rejection failed"))+ writeTVar rejected True)+ (const (pure ()))+ (pure ())++ dispatchMessage store (message "1" "one")+ `shouldReturn` DispatchQueued+ dispatchMessage store (message "2" "two")+ `shouldThrow` anyIOException+ atomically ((,) <$> readTVar accepted <*> readTVar rejected)+ `shouldReturn` (False, False)+ map fst <$> active store `shouldReturn` ["1", "2"]++ atomically (writeTVar throwOnRejection False)+ dispatchMessage store (message "2" "two")+ `shouldReturn` DispatchDropped True+ atomically ((,) <$> readTVar accepted <*> readTVar rejected)+ `shouldReturn` (True, True)+ registerSubscription :: SubscriptionStore -> SID@@ -140,7 +229,7 @@ register store sidValue- (SubscriptionMeta subjectValue Nothing False)+ (SubscriptionMeta subjectValue Nothing StandardSubscription) (SubscribeConfig Nothing Nothing) (const (pure ()))