diff --git a/client/API.hs b/client/API.hs --- a/client/API.hs +++ b/client/API.hs @@ -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 diff --git a/client/Client.hs b/client/Client.hs --- a/client/Client.hs +++ b/client/Client.hs @@ -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 (..) diff --git a/client/Client/Implementation.hs b/client/Client/Implementation.hs --- a/client/Client/Implementation.hs +++ b/client/Client/Implementation.hs @@ -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 diff --git a/internal/Client/API.hs b/internal/Client/API.hs --- a/internal/Client/API.hs +++ b/internal/Client/API.hs @@ -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) diff --git a/internal/Lib/Exception.hs b/internal/Lib/Exception.hs new file mode 100644 --- /dev/null +++ b/internal/Lib/Exception.hs @@ -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) diff --git a/internal/Lib/Logger.hs b/internal/Lib/Logger.hs --- a/internal/Lib/Logger.hs +++ b/internal/Lib/Logger.hs @@ -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 = diff --git a/internal/Lib/WorkerPool.hs b/internal/Lib/WorkerPool.hs --- a/internal/Lib/WorkerPool.hs +++ b/internal/Lib/WorkerPool.hs @@ -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 diff --git a/internal/Plumbing/Network/Connection.hs b/internal/Plumbing/Network/Connection.hs --- a/internal/Plumbing/Network/Connection.hs +++ b/internal/Plumbing/Network/Connection.hs @@ -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 } diff --git a/internal/Plumbing/Network/Connection/Core.hs b/internal/Plumbing/Network/Connection/Core.hs --- a/internal/Plumbing/Network/Connection/Core.hs +++ b/internal/Plumbing/Network/Connection/Core.hs @@ -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 diff --git a/internal/Plumbing/Network/Connection/Tcp.hs b/internal/Plumbing/Network/Connection/Tcp.hs --- a/internal/Plumbing/Network/Connection/Tcp.hs +++ b/internal/Plumbing/Network/Connection/Tcp.hs @@ -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 () diff --git a/internal/Plumbing/Network/Connection/Tls.hs b/internal/Plumbing/Network/Connection/Tls.hs --- a/internal/Plumbing/Network/Connection/Tls.hs +++ b/internal/Plumbing/Network/Connection/Tls.hs @@ -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) diff --git a/internal/Plumbing/Network/Connection/Types.hs b/internal/Plumbing/Network/Connection/Types.hs --- a/internal/Plumbing/Network/Connection/Types.hs +++ b/internal/Plumbing/Network/Connection/Types.hs @@ -21,6 +21,7 @@ , transportWriteLazy :: LBS.ByteString -> IO () , transportFlush :: IO () , transportClose :: IO () + , transportAbort :: IO () , transportUpgrade :: Maybe (TLS.ClientParams -> IO (Either String Transport)) } diff --git a/internal/Plumbing/Network/ConnectionAPI.hs b/internal/Plumbing/Network/ConnectionAPI.hs --- a/internal/Plumbing/Network/ConnectionAPI.hs +++ b/internal/Plumbing/Network/ConnectionAPI.hs @@ -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 ()) } diff --git a/internal/Plumbing/Pipeline/Streaming/Parser.hs b/internal/Plumbing/Pipeline/Streaming/Parser.hs --- a/internal/Plumbing/Pipeline/Streaming/Parser.hs +++ b/internal/Plumbing/Pipeline/Streaming/Parser.hs @@ -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" diff --git a/internal/Plumbing/Queue/API.hs b/internal/Plumbing/Queue/API.hs --- a/internal/Plumbing/Queue/API.hs +++ b/internal/Plumbing/Queue/API.hs @@ -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 ()] } diff --git a/internal/Plumbing/Queue/TransactionalQueue.hs b/internal/Plumbing/Queue/TransactionalQueue.hs --- a/internal/Plumbing/Queue/TransactionalQueue.hs +++ b/internal/Plumbing/Queue/TransactionalQueue.hs @@ -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) diff --git a/internal/Policy/Engine.hs b/internal/Policy/Engine.hs --- a/internal/Policy/Engine.hs +++ b/internal/Policy/Engine.hs @@ -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 diff --git a/internal/Policy/Handshake/Nats.hs b/internal/Policy/Handshake/Nats.hs --- a/internal/Policy/Handshake/Nats.hs +++ b/internal/Policy/Handshake/Nats.hs @@ -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" diff --git a/internal/Policy/Router/Nats.hs b/internal/Policy/Router/Nats.hs --- a/internal/Policy/Router/Nats.hs +++ b/internal/Policy/Router/Nats.hs @@ -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 diff --git a/internal/Policy/State/Store.hs b/internal/Policy/State/Store.hs --- a/internal/Policy/State/Store.hs +++ b/internal/Policy/State/Store.hs @@ -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 diff --git a/internal/Policy/State/Types.hs b/internal/Policy/State/Types.hs --- a/internal/Policy/State/Types.hs +++ b/internal/Policy/State/Types.hs @@ -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) diff --git a/internal/Policy/Subscription/Store.hs b/internal/Policy/Subscription/Store.hs --- a/internal/Policy/Subscription/Store.hs +++ b/internal/Policy/Subscription/Store.hs @@ -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 diff --git a/internal/Policy/Subscription/Types.hs b/internal/Policy/Subscription/Types.hs --- a/internal/Policy/Subscription/Types.hs +++ b/internal/Policy/Subscription/Types.hs @@ -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 diff --git a/jetstream/JetStream/Consumer.hs b/jetstream/JetStream/Consumer.hs --- a/jetstream/JetStream/Consumer.hs +++ b/jetstream/JetStream/Consumer.hs @@ -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 = diff --git a/jetstream/JetStream/Consumer/API.hs b/jetstream/JetStream/Consumer/API.hs --- a/jetstream/JetStream/Consumer/API.hs +++ b/jetstream/JetStream/Consumer/API.hs @@ -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 diff --git a/jetstream/JetStream/Consumer/Types.hs b/jetstream/JetStream/Consumer/Types.hs --- a/jetstream/JetStream/Consumer/Types.hs +++ b/jetstream/JetStream/Consumer/Types.hs @@ -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 = diff --git a/jetstream/JetStream/Message.hs b/jetstream/JetStream/Message.hs --- a/jetstream/JetStream/Message.hs +++ b/jetstream/JetStream/Message.hs @@ -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 = diff --git a/jetstream/JetStream/Message/API.hs b/jetstream/JetStream/Message/API.hs --- a/jetstream/JetStream/Message/API.hs +++ b/jetstream/JetStream/Message/API.hs @@ -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 diff --git a/jetstream/JetStream/Message/Types.hs b/jetstream/JetStream/Message/Types.hs --- a/jetstream/JetStream/Message/Types.hs +++ b/jetstream/JetStream/Message/Types.hs @@ -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 diff --git a/jetstream/JetStream/Stream.hs b/jetstream/JetStream/Stream.hs --- a/jetstream/JetStream/Stream.hs +++ b/jetstream/JetStream/Stream.hs @@ -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) = diff --git a/jetstream/JetStream/Stream/API.hs b/jetstream/JetStream/Stream/API.hs --- a/jetstream/JetStream/Stream/API.hs +++ b/jetstream/JetStream/Stream/API.hs @@ -27,6 +27,7 @@ , streamConfigMaxMessages , streamConfigMaxBytes , streamConfigMaxAge + , streamConfigMaxMessageSize , streamConfigReplicas , streamConfigDuplicateWindow , streamConfigAllowDirect @@ -37,6 +38,7 @@ , withMaxMessages , withMaxBytes , withMaxAge + , withMaxMessageSize , withReplicas , withDuplicateWindow , withAllowDirect diff --git a/jetstream/JetStream/Stream/Types.hs b/jetstream/JetStream/Stream/Types.hs --- a/jetstream/JetStream/Stream/Types.hs +++ b/jetstream/JetStream/Stream/Types.hs @@ -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 = diff --git a/jetstream/JetStream/Types.hs b/jetstream/JetStream/Types.hs --- a/jetstream/JetStream/Types.hs +++ b/jetstream/JetStream/Types.hs @@ -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) = diff --git a/natskell.cabal b/natskell.cabal --- a/natskell.cabal +++ b/natskell.cabal @@ -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 diff --git a/test/Integration/ClientSpec.hs b/test/Integration/ClientSpec.hs --- a/test/Integration/ClientSpec.hs +++ b/test/Integration/ClientSpec.hs @@ -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 diff --git a/test/PublicAPI/Main.hs b/test/PublicAPI/Main.hs --- a/test/PublicAPI/Main.hs +++ b/test/PublicAPI/Main.hs @@ -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 diff --git a/test/Unit/ConnectionSpec.hs b/test/Unit/ConnectionSpec.hs --- a/test/Unit/ConnectionSpec.hs +++ b/test/Unit/ConnectionSpec.hs @@ -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 } diff --git a/test/Unit/JetStream/ConsumerSpec.hs b/test/Unit/JetStream/ConsumerSpec.hs --- a/test/Unit/JetStream/ConsumerSpec.hs +++ b/test/Unit/JetStream/ConsumerSpec.hs @@ -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 diff --git a/test/Unit/JetStream/MessageSpec.hs b/test/Unit/JetStream/MessageSpec.hs --- a/test/Unit/JetStream/MessageSpec.hs +++ b/test/Unit/JetStream/MessageSpec.hs @@ -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 () } diff --git a/test/Unit/JetStream/ProtocolSpec.hs b/test/Unit/JetStream/ProtocolSpec.hs --- a/test/Unit/JetStream/ProtocolSpec.hs +++ b/test/Unit/JetStream/ProtocolSpec.hs @@ -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 () } diff --git a/test/Unit/JetStream/StreamSpec.hs b/test/Unit/JetStream/StreamSpec.hs new file mode 100644 --- /dev/null +++ b/test/Unit/JetStream/StreamSpec.hs @@ -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}" diff --git a/test/Unit/StreamingSpec.hs b/test/Unit/StreamingSpec.hs --- a/test/Unit/StreamingSpec.hs +++ b/test/Unit/StreamingSpec.hs @@ -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 () diff --git a/test/Unit/SubscriptionStoreSpec.hs b/test/Unit/SubscriptionStoreSpec.hs --- a/test/Unit/SubscriptionStoreSpec.hs +++ b/test/Unit/SubscriptionStoreSpec.hs @@ -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 ()))