diff --git a/client/Client.hs b/client/Client.hs --- a/client/Client.hs +++ b/client/Client.hs @@ -48,6 +48,7 @@ import Control.Exception (SomeException, displayException) import Control.Monad (void, when) import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BC import Engine (closeClient, resetClient, runEngine) import Lib.CallOption (CallOption, applyCallOptions) import Lib.Logger @@ -75,12 +76,13 @@ , nextInbox , nextSid , pushPingAction + , readServerInfo , readStatus , runClient , setConnectName , waitForClosed + , waitForConnectionGenerationAfter , waitForNotRunning - , waitForServerInfo ) import State.Types ( ClientConfig (..) @@ -105,11 +107,13 @@ , SubscriptionMeta (SubscriptionMeta) ) 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 import qualified Types.Sub as Sub import qualified Types.Unsub as Unsub +import Validators.Validators (validate) data ClientAuth = ClientAuthNone | ClientAuthToken AuthTokenData @@ -187,7 +191,7 @@ configuredAuth atomically $ - waitForServerInfo clientState + waitForConnectionGenerationAfter clientState 0 `orElse` waitForClosed clientState pure Client @@ -359,37 +363,66 @@ runClient client $ logMessage Debug ("publishing to subject: " ++ show subject) replyTo <- case callback of - Nothing -> - pure configuredReplyTo - Just replyCallback -> do - inbox <- maybe (nextInbox client) pure configuredReplyTo - sid <- nextSid client - let meta = - SubscriptionMeta inbox Nothing True - register store sid meta defaultSubscribeConfig replyCallback - enqueue client $ - QueueItem - Sub.Sub - { Sub.subject = inbox - , Sub.queueGroup = Nothing - , Sub.sid = sid - } - enqueue client $ - QueueItem - Unsub.Unsub - { Unsub.sid = sid - , Unsub.maxMsg = Just 1 - } - pure (Just inbox) - enqueue client $ - QueueItem - Pub.Pub - { Pub.subject = subject - , Pub.payload = payload - , Pub.replyTo = replyTo - , Pub.headers = headers - } + Nothing -> pure configuredReplyTo + Just _ -> Just <$> maybe (nextInbox client) pure configuredReplyTo + let publishMessage = + Pub.Pub + { Pub.subject = subject + , Pub.payload = payload + , Pub.replyTo = replyTo + , Pub.headers = headers + } + shouldPublish <- canPublish client publishMessage + when shouldPublish $ do + case callback of + Nothing -> + pure () + Just replyCallback -> do + case replyTo of + Nothing -> + pure () + Just replyInbox -> do + sid <- nextSid client + let meta = + SubscriptionMeta replyInbox Nothing True + register store sid meta defaultSubscribeConfig replyCallback + enqueue client $ + QueueItem + Sub.Sub + { Sub.subject = replyInbox + , Sub.queueGroup = Nothing + , Sub.sid = sid + } + enqueue client $ + QueueItem + Unsub.Unsub + { Unsub.sid = sid + , Unsub.maxMsg = Just 1 + } + enqueue client (QueueItem publishMessage) +canPublish :: ClientState -> Pub.Pub -> IO Bool +canPublish client publishMessage = + case validate publishMessage of + Left reason -> do + runClient client $ + logMessage Error ("dropping invalid publish: " ++ BC.unpack reason) + pure False + Right () -> do + serverInfo <- readServerInfo client + case serverInfo of + Just info + | Pub.messageSize publishMessage > Info.max_payload info -> do + runClient client $ + logMessage Error + ("dropping publish: payload size " + ++ show (Pub.messageSize publishMessage) + ++ " exceeds server max_payload " + ++ show (Info.max_payload info)) + pure False + _ -> + pure True + subscribeClient :: ClientState -> SubscriptionStore -> Bool -> Msg.Subject -> SubscribeConfig -> (Maybe Msg.Msg -> IO ()) -> IO Msg.SID subscribeClient client store isReply subject cfg callback = do runClient client $ @@ -397,23 +430,30 @@ sid <- nextSid client let queueGroup = subscribeQueueGroup cfg - let meta = - SubscriptionMeta subject queueGroup isReply - register store sid meta cfg callback - enqueue client $ - QueueItem - Sub.Sub - { Sub.subject = subject - , Sub.queueGroup = queueGroup - , Sub.sid = sid - } - when isReply $ - enqueue client - (QueueItem - Unsub.Unsub - { Unsub.sid = sid - , Unsub.maxMsg = Just 1 - }) + subscriptionMessage = + Sub.Sub + { Sub.subject = subject + , Sub.queueGroup = queueGroup + , Sub.sid = sid + } + case validate subscriptionMessage of + Left reason -> + runClient client $ + logMessage Error ("dropping invalid subscription: " ++ BC.unpack reason) + Right () -> do + let meta = + SubscriptionMeta subject queueGroup isReply + register store sid meta cfg callback + enqueue client $ + QueueItem + subscriptionMessage + when isReply $ + enqueue client + (QueueItem + Unsub.Unsub + { Unsub.sid = sid + , Unsub.maxMsg = Just 1 + }) pure sid unsubscribeClient :: ClientState -> SubscriptionStore -> Msg.SID -> IO () 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 @@ -80,7 +80,9 @@ shouldContinue <- enqueueReadResult conn result case result of Left _ -> return () - Right _ -> when shouldContinue loop + Right bytes + | BS.null bytes -> return () + | otherwise -> when shouldContinue loop enqueueReadResult :: Conn -> Either SomeException ByteString -> IO Bool enqueueReadResult conn result = atomically $ @@ -213,6 +215,7 @@ case current of Nothing -> return () Just currentTransport' -> void $ try @SomeException (transportClose currentTransport') + waitForReadWorkerStopped conn openConn :: Conn -> IO () openConn conn = do @@ -227,3 +230,9 @@ bufferRead conn bytes = unless (BS.null bytes) (atomically $ modifyTVar' (readBuffer conn) (bytes <>)) + +waitForReadWorkerStopped :: Conn -> IO () +waitForReadWorkerStopped conn = + atomically $ do + running <- readTVar (readWorkerRunning conn) + when running retry diff --git a/internal/Plumbing/Parser/Attoparsec.hs b/internal/Plumbing/Parser/Attoparsec.hs --- a/internal/Plumbing/Parser/Attoparsec.hs +++ b/internal/Plumbing/Parser/Attoparsec.hs @@ -235,7 +235,11 @@ headerKey <- A.takeTill (== colon) _ <- A.word8 colon headerValue <- lineParser - headerPairsParser ((stripHorizontalSpace headerKey, stripHorizontalSpace headerValue) : reversedPairs) + let trimmedKey = stripHorizontalSpace headerKey + trimmedValue = stripHorizontalSpace headerValue + if BS.null trimmedKey + then fail "empty HMSG header key" + else headerPairsParser ((trimmedKey, trimmedValue) : reversedPairs) classifyErr :: BS.ByteString -> Either String Err classifyErr reason diff --git a/internal/Plumbing/Validators/Validators.hs b/internal/Plumbing/Validators/Validators.hs --- a/internal/Plumbing/Validators/Validators.hs +++ b/internal/Plumbing/Validators/Validators.hs @@ -1,7 +1,66 @@ +{-# LANGUAGE OverloadedStrings #-} + module Validators.Validators where -import Data.ByteString +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BC +import Data.Char (isSpace) class Validator a where validate :: a -> Either ByteString () +validateSubjectSyntax :: ByteString -> Either ByteString () +validateSubjectSyntax subject + | BC.any isSpace subject = + Left "subject contains whitespace" + | any BS.null tokens = + Left "subject contains empty token" + | otherwise = + validateWildcardTokens tokens + where + tokens = BC.split '.' subject + +validateQueueGroupSyntax :: ByteString -> Either ByteString () +validateQueueGroupSyntax queueGroup + | BC.any isSpace queueGroup = + Left "queue group contains whitespace" + | otherwise = + Right () + +validateHeaderKeySyntax :: ByteString -> Either ByteString () +validateHeaderKeySyntax key + | ":" `BS.isInfixOf` key = + Left "header key contains colon" + | hasLineBreak key = + Left "header key contains line break" + | otherwise = + Right () + +validateHeaderValueSyntax :: ByteString -> Either ByteString () +validateHeaderValueSyntax value + | hasLineBreak value = + Left "header value contains line break" + | otherwise = + Right () + +validateWildcardTokens :: [ByteString] -> Either ByteString () +validateWildcardTokens [] = + Right () +validateWildcardTokens [">"] = + Right () +validateWildcardTokens (">":_) = + Left "full wildcard must be the final token" +validateWildcardTokens (token:rest) + | token == "*" = + validateWildcardTokens rest + | ">" `BS.isInfixOf` token = + Left "full wildcard must be a standalone token" + | "*" `BS.isInfixOf` token = + Left "partial wildcard must be a standalone token" + | otherwise = + validateWildcardTokens rest + +hasLineBreak :: ByteString -> Bool +hasLineBreak bytes = + "\r" `BS.isInfixOf` bytes || "\n" `BS.isInfixOf` bytes diff --git a/internal/Policy/Engine.hs b/internal/Policy/Engine.hs --- a/internal/Policy/Engine.hs +++ b/internal/Policy/Engine.hs @@ -36,14 +36,17 @@ , incrementAttemptIndex , markClosed , markConnected + , markConnectionReady , openQueue , queue , readAttemptIndex + , readConnectionGeneration , readStatus , runClient , setClosing , setEndpoint , waitForClosed + , waitForConnectionGenerationAfter ) import State.Types ( ClientConfig (..) @@ -124,6 +127,7 @@ pure (show err) Right () -> do resubscribeIfNeeded + markConnectionReady state connectionResult <- runConnection conn close connectionApi conn case connectionResult of @@ -241,8 +245,16 @@ shutdownClient connectionApi state store ExitClosedByUser "closing client connection" resetClient :: ConnectionAPI -> ClientState -> SubscriptionStore -> IO () -resetClient connectionApi state store = - shutdownClient connectionApi state store ExitResetRequested "resetting client connection" +resetClient connectionApi state _store = do + generation <- readConnectionGeneration state + runClient state $ + logMessage Info "resetting client connection" + closeQueue state + closeReader (reader connectionApi) (connection state) + closeWriter (writer connectionApi) (connection state) + atomically $ + waitForConnectionGenerationAfter state generation + `orElse` waitForClosed state shutdownClient :: ConnectionAPI 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 @@ -27,11 +27,15 @@ , waitForNotRunning , waitForServerInfo , markConnected + , markConnectionReady + , readConnectionGeneration + , waitForConnectionGenerationAfter , readAttemptIndex , incrementAttemptIndex ) where import Control.Concurrent.STM +import Control.Monad (unless) import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as BC import Data.Maybe (fromMaybe) @@ -51,23 +55,25 @@ import Types.Msg (SID, Subject) data ClientState = ClientState - { clientConfig :: ClientConfig - , clientQueue :: Queue - , clientPings :: TQueue (IO ()) - , clientConnectedOnce :: TVar Bool - , clientSidCounter :: TVar SIDCounter - , clientInboxNuid :: TVar Nuid - , clientServerInfo :: TVar (Maybe Info) - , clientAttemptIndex :: TVar Int - , clientStatus :: TVar ClientStatus - , clientConnection :: Conn - , clientLogContext :: TVar LogContext + { clientConfig :: ClientConfig + , clientQueue :: Queue + , clientPings :: TQueue (IO ()) + , clientConnectedOnce :: TVar Bool + , clientConnectionGeneration :: TVar Int + , clientSidCounter :: TVar SIDCounter + , clientInboxNuid :: TVar Nuid + , clientServerInfo :: TVar (Maybe Info) + , clientAttemptIndex :: TVar Int + , clientStatus :: TVar ClientStatus + , clientConnection :: Conn + , clientLogContext :: TVar LogContext } newClientState :: ClientConfig -> Queue -> Conn -> TVar LogContext -> IO ClientState newClientState cfg queue conn logContext = do pings <- newTQueueIO connectedOnce <- newTVarIO False + connectionGeneration <- newTVarIO 0 sidCounter <- newTVarIO initialSIDCounter inboxNuid <- newTVarIO =<< newNuidIO serverInfo <- newTVarIO Nothing @@ -78,6 +84,7 @@ , clientQueue = queue , clientPings = pings , clientConnectedOnce = connectedOnce + , clientConnectionGeneration = connectionGeneration , clientSidCounter = sidCounter , clientInboxNuid = inboxNuid , clientServerInfo = serverInfo @@ -215,6 +222,20 @@ alreadyConnected <- readTVar (clientConnectedOnce client) writeTVar (clientConnectedOnce client) True pure alreadyConnected + +markConnectionReady :: ClientState -> IO () +markConnectionReady client = + atomically $ + modifyTVar' (clientConnectionGeneration client) (+ 1) + +readConnectionGeneration :: ClientState -> IO Int +readConnectionGeneration = + readTVarIO . clientConnectionGeneration + +waitForConnectionGenerationAfter :: ClientState -> Int -> STM () +waitForConnectionGenerationAfter client generation = do + current <- readTVar (clientConnectionGeneration client) + unless (current > generation) retry readAttemptIndex :: ClientState -> IO Int readAttemptIndex = readTVarIO . clientAttemptIndex diff --git a/internal/Types/Pub.hs b/internal/Types/Pub.hs --- a/internal/Types/Pub.hs +++ b/internal/Types/Pub.hs @@ -49,12 +49,12 @@ validateSubject :: Pub -> Either ByteString () validateSubject p | subject p == "" = Left "explicit empty subject" - | otherwise = Right () + | otherwise = validateSubjectSyntax (subject p) validateReplyTo :: Pub -> Either ByteString () validateReplyTo p | replyTo p == Just "" = Left "explicit empty replyTo" - | otherwise = Right () + | otherwise = maybe (Right ()) validateSubjectSyntax (replyTo p) validatePayload :: Pub -> Either ByteString () validatePayload p @@ -67,7 +67,17 @@ | headers p == Just [] = Left "explicit empty headers" | Prelude.any (\(k, _) -> k == "") (fromJust (headers p)) = Left "explicit empty header key" | Prelude.any (\(_, v) -> v == "") (fromJust (headers p)) = Left "explicit empty header value" - | otherwise = Right () + | otherwise = do + mapM_ (validateHeaderKeySyntax . fst) (fromJust (headers p)) + mapM_ (validateHeaderValueSyntax . snd) (fromJust (headers p)) + +messageSize :: Pub -> Int +messageSize pubMsg = + case headers pubMsg of + Nothing -> + length (payloadChunk (payload pubMsg)) + Just headerList -> + length (headerString headerList) + length (payloadChunk (payload pubMsg)) headerString :: Headers -> ByteString headerString hs = diff --git a/internal/Types/Sub.hs b/internal/Types/Sub.hs --- a/internal/Types/Sub.hs +++ b/internal/Types/Sub.hs @@ -30,12 +30,12 @@ validateSubject :: Sub -> Either ByteString () validateSubject s | subject s == "" = Left "explicit empty subject" - | otherwise = Right () + | otherwise = validateSubjectSyntax (subject s) validateQueueGroup :: Sub -> Either ByteString () validateQueueGroup s | queueGroup s == Just "" = Left "explicit empty queue group" - | otherwise = Right () + | otherwise = maybe (Right ()) validateQueueGroupSyntax (queueGroup s) validateSid :: Sub -> Either ByteString () validateSid s diff --git a/jetstream-client/JetStream/API.hs b/jetstream-client/JetStream/API.hs new file mode 100644 --- /dev/null +++ b/jetstream-client/JetStream/API.hs @@ -0,0 +1,53 @@ +-- | Capability record for the JetStream client surface. +module JetStream.API + ( + JetStream (..) + , AccountAPIStats (..) + , AccountInfo (..) + , AccountLimits (..) + , AccountTier (..) + , JetStreamApiError (..) + , JetStreamError (..) + , StreamName + , ConsumerName + , Subject + , Payload + , module JetStream.Consumer.API + , module JetStream.Message.API + , module JetStream.Publish.API + , module JetStream.Stream.API + ) where + +import JetStream.Consumer.API +import JetStream.Error + ( JetStreamApiError (..) + , JetStreamError (..) + ) +import JetStream.Message.API +import JetStream.Publish.API +import JetStream.Stream.API +import JetStream.Types + ( AccountAPIStats (..) + , AccountInfo (..) + , AccountLimits (..) + , AccountTier (..) + , ConsumerName + , Payload + , StreamName + , Subject + ) + +-- | JetStream capabilities for streams, consumers, publishing, and pull +-- message workflows. +data JetStream = JetStream + { streams :: StreamAPI + -- ^ Manage streams and inspect stream state. + , consumers :: ConsumerAPI + -- ^ Manage stream consumers. + , publisher :: PublishAPI + -- ^ Publish messages to JetStream subjects. + , messages :: MessageAPI + -- ^ Fetch and acknowledge pull-consumer messages. + , accountInfo :: IO (Either JetStreamError AccountInfo) + -- ^ Fetch account-level JetStream usage and limits. + } diff --git a/jetstream-client/JetStream/Client.hs b/jetstream-client/JetStream/Client.hs new file mode 100644 --- /dev/null +++ b/jetstream-client/JetStream/Client.hs @@ -0,0 +1,37 @@ +-- | High-level client implementation for JetStream. +module JetStream.Client + ( newJetStream + , JetStream (..) + , JetStreamOption + , withDomain + , withRequestTimeoutMicros + ) where + +import qualified API as Nats +import JetStream.API (JetStream (..)) +import qualified JetStream.Consumer as Consumer +import qualified JetStream.Message as Message +import JetStream.Options + ( JetStreamOption + , newJetStreamContext + , withDomain + , withRequestTimeoutMicros + ) +import qualified JetStream.Protocol.Request as Request +import qualified JetStream.Protocol.Subject as Subject +import qualified JetStream.Publish as Publish +import qualified JetStream.Stream as Stream + +-- | Build a JetStream capability record from an existing NATS client. +newJetStream :: Nats.Client -> [JetStreamOption] -> JetStream +newJetStream client options = + let ctx = newJetStreamContext client options + consumerAPI = Consumer.consumerAPI ctx + in JetStream + { streams = Stream.streamAPI ctx + , consumers = consumerAPI + , publisher = Publish.publishAPI ctx + , messages = Message.messageAPI ctx consumerAPI + , accountInfo = + Request.requestJSON ctx (Subject.accountInfoSubject ctx) Nothing + } diff --git a/jetstream/JetStream/Consumer.hs b/jetstream/JetStream/Consumer.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Consumer.hs @@ -0,0 +1,130 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Consumer + ( consumerAPI + , module JetStream.Consumer.API + ) where + +import Data.Aeson (Value, object, toJSON, (.=)) +import Data.Maybe (catMaybes) +import JetStream.Consumer.API (ConsumerAPI (..)) +import JetStream.Consumer.Types + ( ConsumerAction (..) + , ConsumerConfigOption + , ConsumerConfigRequest + , ConsumerKind + , ConsumerTarget (..) + , applyConsumerKind + , applyConsumerTarget + , consumerActionValue + , consumerConfigRequest + , consumerListRequest + , consumerNamesRequest + , consumerPauseRequest + , consumerResetRequest + ) +import JetStream.Error + ( JetStreamError (JetStreamDecodeError) + ) +import JetStream.Options (JetStreamContext) +import qualified JetStream.Protocol.Request as Request +import qualified JetStream.Protocol.Subject as Subject +import JetStream.Types + ( StreamName + , Subject + , byteStringToJSON + ) + +consumerAPI :: JetStreamContext -> ConsumerAPI +consumerAPI context = + ConsumerAPI + { putConsumer = \stream action target kind options -> + case consumerRequest context stream action target kind options of + Left err -> + pure (Left err) + Right (subject, actionValue, config) -> + Request.requestJSON context + subject + (Just (createConsumerRequest stream actionValue config)) + , consumerInfo = \stream consumer -> + Request.requestJSON context + (Subject.consumerInfoSubject context stream consumer) + Nothing + , pauseConsumer = \stream consumer pauseUntil -> + Request.requestJSON context + (Subject.consumerPauseSubject context stream consumer) + (Just (toJSON (consumerPauseRequest (Just pauseUntil)))) + , resumeConsumer = \stream consumer -> + Request.requestJSON context + (Subject.consumerPauseSubject context stream consumer) + (Just (toJSON (consumerPauseRequest Nothing))) + , resetConsumer = \stream consumer options -> + Request.requestJSON context + (Subject.consumerResetSubject context stream consumer) + (Just (toJSON (consumerResetRequest options))) + , deleteConsumer = \stream consumer -> + Request.requestJSON context + (Subject.consumerDeleteSubject context stream consumer) + Nothing + , listConsumers = \stream options -> + Request.requestJSON context + (Subject.consumerListSubject context stream) + (Just (toJSON (consumerListRequest options))) + , consumerNames = \stream options -> + Request.requestJSON context + (Subject.consumerNamesSubject context stream) + (Just (toJSON (consumerNamesRequest options))) + } + +consumerRequest + :: JetStreamContext + -> StreamName + -> ConsumerAction + -> ConsumerTarget + -> ConsumerKind + -> [ConsumerConfigOption] + -> Either JetStreamError (Subject, Maybe StreamName, ConsumerConfigRequest) +consumerRequest context stream action target kind options = + case consumerSubject context stream action target of + Left err -> + Left err + Right subject -> + Right + ( subject + , actionValue + , applyConsumerTarget target . applyConsumerKind kind $ consumerConfigRequest options + ) + where + actionValue = + case target of + EphemeralConsumer -> + Nothing + DurableConsumer _ | action == ConsumerCreate -> + Nothing + _ -> + consumerActionValue action + +consumerSubject + :: JetStreamContext + -> StreamName + -> ConsumerAction + -> ConsumerTarget + -> Either JetStreamError Subject +consumerSubject context stream ConsumerCreate EphemeralConsumer = + Right (Subject.consumerCreateSubject context stream) +consumerSubject context stream ConsumerCreate (DurableConsumer durable) = + Right (Subject.durableConsumerCreateSubject context stream durable) +consumerSubject context stream _ (NamedConsumer consumer) = + Right (Subject.consumerCreateNamedSubject context stream consumer) +consumerSubject context stream _ (DurableConsumer durable) = + Right (Subject.consumerCreateNamedSubject context stream durable) +consumerSubject _ _ _ EphemeralConsumer = + Left (JetStreamDecodeError "ephemeral consumers can only be created") + +createConsumerRequest :: StreamName -> Maybe StreamName -> ConsumerConfigRequest -> Value +createConsumerRequest stream action config = + object . catMaybes $ + [ Just ("stream_name" .= byteStringToJSON stream) + , Just ("config" .= toJSON config) + , fmap (("action" .=) . byteStringToJSON) action + ] diff --git a/jetstream/JetStream/Consumer/API.hs b/jetstream/JetStream/Consumer/API.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Consumer/API.hs @@ -0,0 +1,63 @@ +module JetStream.Consumer.API + ( ConsumerAPI (..) + , AckPolicy (..) + , ConsumerAction (..) + , ConsumerConfig (..) + , ConsumerConfigOption + , ConsumerFilter (..) + , ConsumerKind (..) + , ConsumerInfo (..) + , ConsumerListOption + , ConsumerListResponse (..) + , ConsumerNamesResponse (..) + , ConsumerPauseResponse (..) + , ConsumerResetOption + , ConsumerResetResponse (..) + , ConsumerSequenceInfo (..) + , ConsumerTarget (..) + , DeleteConsumerResponse (..) + , DeliverPolicy (..) + , ReplayPolicy (..) + , withConsumerAckPolicy + , withConsumerAckWait + , withConsumerDescription + , withConsumerDeliverGroup + , withConsumerDeliverPolicy + , withConsumerFilter + , withConsumerHeadersOnly + , withConsumerIdleHeartbeat + , withConsumerInactiveThreshold + , withConsumerListOffset + , withConsumerMaxAckPending + , withConsumerMaxDeliver + , withConsumerMaxWaiting + , withConsumerMemoryStorage + , withConsumerReplicas + , withConsumerReplayPolicy + , withConsumerResetSequence + ) where + +import Data.Time.Clock (UTCTime) +import JetStream.Consumer.Types +import JetStream.Error (JetStreamError) +import JetStream.Types (ConsumerName, StreamName) + +-- | Consumer management operations. +data ConsumerAPI = ConsumerAPI + { putConsumer :: StreamName -> ConsumerAction -> ConsumerTarget -> ConsumerKind -> [ConsumerConfigOption] -> IO (Either JetStreamError ConsumerInfo) + -- ^ Create, update, or create-or-update a stream consumer. + , consumerInfo :: StreamName -> ConsumerName -> IO (Either JetStreamError ConsumerInfo) + -- ^ Fetch consumer configuration and state. + , pauseConsumer :: StreamName -> ConsumerName -> UTCTime -> IO (Either JetStreamError ConsumerPauseResponse) + -- ^ Pause a consumer until the given time. + , resumeConsumer :: StreamName -> ConsumerName -> IO (Either JetStreamError ConsumerPauseResponse) + -- ^ Resume a paused consumer. + , resetConsumer :: StreamName -> ConsumerName -> [ConsumerResetOption] -> IO (Either JetStreamError ConsumerResetResponse) + -- ^ Reset a consumer, optionally to a stream sequence. + , deleteConsumer :: StreamName -> ConsumerName -> IO (Either JetStreamError DeleteConsumerResponse) + -- ^ Delete a consumer from a stream. + , listConsumers :: StreamName -> [ConsumerListOption] -> IO (Either JetStreamError ConsumerListResponse) + -- ^ List consumers for a stream. + , consumerNames :: StreamName -> [ConsumerListOption] -> IO (Either JetStreamError ConsumerNamesResponse) + -- ^ List consumer names for a stream. + } diff --git a/jetstream/JetStream/Consumer/Types.hs b/jetstream/JetStream/Consumer/Types.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Consumer/Types.hs @@ -0,0 +1,625 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Consumer.Types + ( AckPolicy (..) + , ConsumerConfig (..) + , ConsumerAction (..) + , ConsumerConfigOption + , ConsumerConfigRequest + , ConsumerFilter (..) + , ConsumerKind (..) + , ConsumerInfo (..) + , ConsumerListOption + , ConsumerListResponse (..) + , ConsumerNamesResponse (..) + , ConsumerPauseResponse (..) + , ConsumerResetOption + , ConsumerResetResponse (..) + , ConsumerSequenceInfo (..) + , ConsumerTarget (..) + , DeleteConsumerResponse (..) + , DeliverPolicy (..) + , ReplayPolicy (..) + , applyConsumerKind + , applyConsumerTarget + , consumerConfigRequest + , consumerActionValue + , consumerListRequest + , consumerNamesRequest + , consumerPauseRequest + , consumerResetRequest + , withConsumerAckPolicy + , withConsumerAckWait + , withConsumerDescription + , withConsumerDeliverGroup + , withConsumerDurableName + , withConsumerDeliverSubject + , withConsumerFilter + , withConsumerHeadersOnly + , withConsumerIdleHeartbeat + , withConsumerInactiveThreshold + , withConsumerListOffset + , withConsumerMaxAckPending + , withConsumerMaxDeliver + , withConsumerMaxWaiting + , withConsumerMemoryStorage + , withConsumerName + , withConsumerReplicas + , withConsumerDeliverPolicy + , withConsumerReplayPolicy + , withConsumerResetSequence + ) where + +import Data.Aeson +import Data.Aeson.Types (Pair, Parser) +import qualified Data.ByteString as BS +import Data.Maybe (catMaybes) +import qualified Data.Text as T +import Data.Time.Clock (NominalDiffTime, UTCTime) +import JetStream.Types + ( AckPolicy (..) + , CallOption + , ConsumerName + , DeliverPolicy (..) + , ReplayPolicy (..) + , StreamName + , Subject + , applyCallOptions + , byteStringToJSON + , diffTimeNanosToJSON + , parseByteString + ) + +data ConsumerConfigRequest = ConsumerConfigRequest + { consumerConfigRequestDurableName :: Maybe ConsumerName + , consumerConfigRequestName :: Maybe ConsumerName + , consumerConfigRequestDescription :: Maybe BS.ByteString + , consumerConfigRequestDeliverSubject :: Maybe Subject + , consumerConfigRequestDeliverGroup :: Maybe Subject + , consumerConfigRequestDeliverPolicy :: Maybe DeliverPolicy + , consumerConfigRequestAckPolicy :: Maybe AckPolicy + , consumerConfigRequestReplayPolicy :: Maybe ReplayPolicy + , consumerConfigRequestFilter :: Maybe ConsumerFilter + , consumerConfigRequestAckWait :: Maybe NominalDiffTime + , consumerConfigRequestMaxDeliver :: Maybe Int + , consumerConfigRequestMaxWaiting :: Maybe Int + , consumerConfigRequestMaxAckPending :: Maybe Int + , consumerConfigRequestInactiveThreshold :: Maybe NominalDiffTime + , consumerConfigRequestIdleHeartbeat :: Maybe NominalDiffTime + , consumerConfigRequestHeadersOnly :: Maybe Bool + , consumerConfigRequestReplicas :: Maybe Int + , consumerConfigRequestMemoryStorage :: Maybe Bool + } + deriving (Eq, Show) + +data ConsumerAction = ConsumerCreate | ConsumerCreateOrUpdate | ConsumerUpdate + deriving (Eq, Show) + +consumerActionValue :: ConsumerAction -> Maybe BS.ByteString +consumerActionValue ConsumerCreate = Just "create" +consumerActionValue ConsumerCreateOrUpdate = Nothing +consumerActionValue ConsumerUpdate = Just "update" + +data ConsumerTarget = EphemeralConsumer + | NamedConsumer ConsumerName + | DurableConsumer ConsumerName + deriving (Eq, Show) + +data ConsumerKind = PullConsumer + | PushConsumer Subject + deriving (Eq, Show) + +type ConsumerConfigOption = CallOption ConsumerConfigRequest + +consumerConfigRequest :: [ConsumerConfigOption] -> ConsumerConfigRequest +consumerConfigRequest options = + applyCallOptions options emptyConsumerConfigRequest + +emptyConsumerConfigRequest :: ConsumerConfigRequest +emptyConsumerConfigRequest = + ConsumerConfigRequest + { consumerConfigRequestDurableName = Nothing + , consumerConfigRequestName = Nothing + , consumerConfigRequestDescription = Nothing + , consumerConfigRequestDeliverSubject = Nothing + , consumerConfigRequestDeliverGroup = Nothing + , consumerConfigRequestDeliverPolicy = Nothing + , consumerConfigRequestAckPolicy = Nothing + , consumerConfigRequestReplayPolicy = Nothing + , consumerConfigRequestFilter = Nothing + , consumerConfigRequestAckWait = Nothing + , consumerConfigRequestMaxDeliver = Nothing + , consumerConfigRequestMaxWaiting = Nothing + , consumerConfigRequestMaxAckPending = Nothing + , consumerConfigRequestInactiveThreshold = Nothing + , consumerConfigRequestIdleHeartbeat = Nothing + , consumerConfigRequestHeadersOnly = Nothing + , consumerConfigRequestReplicas = Nothing + , consumerConfigRequestMemoryStorage = Nothing + } + +applyConsumerTarget :: ConsumerTarget -> ConsumerConfigRequest -> ConsumerConfigRequest +applyConsumerTarget EphemeralConsumer config = + config + { consumerConfigRequestDurableName = Nothing + , consumerConfigRequestName = Nothing + } +applyConsumerTarget (NamedConsumer name) config = + config + { consumerConfigRequestDurableName = Nothing + , consumerConfigRequestName = Just name + } +applyConsumerTarget (DurableConsumer durable) config = + config + { consumerConfigRequestDurableName = Just durable + , consumerConfigRequestName = Just durable + } + +applyConsumerKind :: ConsumerKind -> ConsumerConfigRequest -> ConsumerConfigRequest +applyConsumerKind PullConsumer config = + config + { consumerConfigRequestDeliverSubject = Nothing + , consumerConfigRequestDeliverGroup = Nothing + } +applyConsumerKind (PushConsumer deliverSubject) config = + config { consumerConfigRequestDeliverSubject = Just deliverSubject } + +withConsumerDurableName :: ConsumerName -> ConsumerConfigOption +withConsumerDurableName durable config = + config { consumerConfigRequestDurableName = Just durable } + +withConsumerName :: ConsumerName -> ConsumerConfigOption +withConsumerName name config = + config { consumerConfigRequestName = Just name } + +withConsumerDescription :: BS.ByteString -> ConsumerConfigOption +withConsumerDescription description config = + config { consumerConfigRequestDescription = Just description } + +withConsumerDeliverSubject :: Subject -> ConsumerConfigOption +withConsumerDeliverSubject subject config = + config { consumerConfigRequestDeliverSubject = Just subject } + +withConsumerDeliverGroup :: Subject -> ConsumerConfigOption +withConsumerDeliverGroup group config = + config { consumerConfigRequestDeliverGroup = Just group } + +withConsumerDeliverPolicy :: DeliverPolicy -> ConsumerConfigOption +withConsumerDeliverPolicy policy config = + config { consumerConfigRequestDeliverPolicy = Just policy } + +withConsumerAckPolicy :: AckPolicy -> ConsumerConfigOption +withConsumerAckPolicy policy config = + config { consumerConfigRequestAckPolicy = Just policy } + +withConsumerReplayPolicy :: ReplayPolicy -> ConsumerConfigOption +withConsumerReplayPolicy policy config = + config { consumerConfigRequestReplayPolicy = Just policy } + +data ConsumerFilter = ConsumerFilterSubject Subject + | ConsumerFilterSubjects [Subject] + deriving (Eq, Show) + +withConsumerFilter :: ConsumerFilter -> ConsumerConfigOption +withConsumerFilter consumerFilter config = + config { consumerConfigRequestFilter = Just consumerFilter } + +withConsumerAckWait :: NominalDiffTime -> ConsumerConfigOption +withConsumerAckWait ackWait config = + config { consumerConfigRequestAckWait = Just ackWait } + +withConsumerMaxDeliver :: Int -> ConsumerConfigOption +withConsumerMaxDeliver maxDeliver config = + config { consumerConfigRequestMaxDeliver = Just maxDeliver } + +withConsumerMaxWaiting :: Int -> ConsumerConfigOption +withConsumerMaxWaiting maxWaiting config = + config { consumerConfigRequestMaxWaiting = Just maxWaiting } + +withConsumerMaxAckPending :: Int -> ConsumerConfigOption +withConsumerMaxAckPending maxAckPending config = + config { consumerConfigRequestMaxAckPending = Just maxAckPending } + +withConsumerInactiveThreshold :: NominalDiffTime -> ConsumerConfigOption +withConsumerInactiveThreshold inactiveThreshold config = + config { consumerConfigRequestInactiveThreshold = Just inactiveThreshold } + +withConsumerIdleHeartbeat :: NominalDiffTime -> ConsumerConfigOption +withConsumerIdleHeartbeat idleHeartbeat config = + config { consumerConfigRequestIdleHeartbeat = Just idleHeartbeat } + +withConsumerHeadersOnly :: Bool -> ConsumerConfigOption +withConsumerHeadersOnly headersOnly config = + config { consumerConfigRequestHeadersOnly = Just headersOnly } + +withConsumerReplicas :: Int -> ConsumerConfigOption +withConsumerReplicas replicas config = + config { consumerConfigRequestReplicas = Just replicas } + +withConsumerMemoryStorage :: Bool -> ConsumerConfigOption +withConsumerMemoryStorage memoryStorage config = + config { consumerConfigRequestMemoryStorage = Just memoryStorage } + +data ConsumerConfig = ConsumerConfig + { consumerConfigDurableName :: Maybe ConsumerName + , consumerConfigName :: Maybe ConsumerName + , consumerConfigDescription :: Maybe BS.ByteString + , consumerConfigDeliverSubject :: Maybe Subject + , consumerConfigDeliverGroup :: Maybe Subject + , consumerConfigDeliverPolicy :: DeliverPolicy + , consumerConfigAckPolicy :: AckPolicy + , consumerConfigReplayPolicy :: ReplayPolicy + , consumerConfigFilterSubject :: Maybe Subject + , consumerConfigFilterSubjects :: Maybe [Subject] + , consumerConfigAckWait :: Maybe NominalDiffTime + , consumerConfigMaxDeliver :: Maybe Int + , consumerConfigMaxWaiting :: Maybe Int + , consumerConfigMaxAckPending :: Maybe Int + , consumerConfigInactiveThreshold :: Maybe NominalDiffTime + , consumerConfigIdleHeartbeat :: Maybe NominalDiffTime + , consumerConfigHeadersOnly :: Maybe Bool + , consumerConfigReplicas :: Maybe Int + , consumerConfigMemoryStorage :: Maybe Bool + } + deriving (Eq, Show) + +data ConsumerInfo = ConsumerInfo + { consumerInfoStreamName :: StreamName + , consumerInfoName :: ConsumerName + , consumerInfoCreated :: UTCTime + , consumerInfoConfig :: ConsumerConfig + , consumerInfoDelivered :: ConsumerSequenceInfo + , consumerInfoAckFloor :: ConsumerSequenceInfo + , consumerInfoNumAckPending :: Int + , consumerInfoNumRedelivered :: Int + , consumerInfoNumWaiting :: Int + , consumerInfoNumPending :: Integer + } + deriving (Eq, Show) + +data ConsumerSequenceInfo = ConsumerSequenceInfo + { consumerSequenceConsumer :: Integer + , consumerSequenceStream :: Integer + , consumerSequenceLast :: Maybe UTCTime + } + deriving (Eq, Show) + +newtype DeleteConsumerResponse = DeleteConsumerResponse { deleteConsumerSuccess :: Bool } + deriving (Eq, Show) + +newtype ConsumerPauseRequest = ConsumerPauseRequest { consumerPauseUntil :: Maybe UTCTime } + deriving (Eq, Show) + +consumerPauseRequest :: Maybe UTCTime -> ConsumerPauseRequest +consumerPauseRequest = + ConsumerPauseRequest + +data ConsumerPauseResponse = ConsumerPauseResponse + { consumerPausePaused :: Bool + , consumerPauseUntilTime :: Maybe UTCTime + , consumerPauseRemaining :: Maybe NominalDiffTime + } + deriving (Eq, Show) + +newtype ConsumerResetRequest = ConsumerResetRequest { consumerResetRequestSequence :: Maybe Integer } + deriving (Eq, Show) + +type ConsumerResetOption = CallOption ConsumerResetRequest + +consumerResetRequest :: [ConsumerResetOption] -> ConsumerResetRequest +consumerResetRequest options = + applyCallOptions options $ + ConsumerResetRequest + { consumerResetRequestSequence = Nothing + } + +withConsumerResetSequence :: Integer -> ConsumerResetOption +withConsumerResetSequence sequenceNumber request = + request { consumerResetRequestSequence = Just sequenceNumber } + +data ConsumerResetResponse = ConsumerResetResponse + { consumerResetInfo :: ConsumerInfo + , consumerResetResponseSequence :: Integer + } + deriving (Eq, Show) + +newtype ConsumerListRequest = ConsumerListRequest { consumerListRequestOffset :: Maybe Int } + deriving (Eq, Show) + +type ConsumerListOption = CallOption ConsumerListRequest + +consumerListRequest :: [ConsumerListOption] -> ConsumerListRequest +consumerListRequest options = + applyCallOptions options defaultConsumerListRequest + +consumerNamesRequest :: [ConsumerListOption] -> ConsumerListRequest +consumerNamesRequest = + consumerListRequest + +defaultConsumerListRequest :: ConsumerListRequest +defaultConsumerListRequest = + ConsumerListRequest + { consumerListRequestOffset = Nothing + } + +withConsumerListOffset :: Int -> ConsumerListOption +withConsumerListOffset offset request = + request { consumerListRequestOffset = Just offset } + +data ConsumerListResponse = ConsumerListResponse + { consumerListTotal :: Int + , consumerListOffset :: Int + , consumerListLimit :: Int + , consumerListConsumers :: [ConsumerInfo] + } + deriving (Eq, Show) + +data ConsumerNamesResponse = ConsumerNamesResponse + { consumerNamesTotal :: Int + , consumerNamesOffset :: Int + , consumerNamesLimit :: Int + , consumerNamesConsumers :: [ConsumerName] + } + deriving (Eq, Show) + +instance ToJSON ConsumerConfigRequest where + toJSON config = + object . catMaybes $ + [ maybeByteStringPair "durable_name" (consumerConfigRequestDurableName config) + , maybeByteStringPair "name" (consumerConfigRequestName config) + , maybeByteStringPair "description" (consumerConfigRequestDescription config) + , maybeByteStringPair "deliver_subject" (consumerConfigRequestDeliverSubject config) + , maybeByteStringPair "deliver_group" (consumerConfigRequestDeliverGroup config) + , maybePair "ack_policy" (consumerConfigRequestAckPolicy config) + , maybePair "replay_policy" (consumerConfigRequestReplayPolicy config) + , maybeJsonPair "ack_wait" (diffTimeNanosToJSON <$> consumerConfigRequestAckWait config) + , maybePair "max_deliver" (consumerConfigRequestMaxDeliver config) + , maybePair "max_waiting" (consumerConfigRequestMaxWaiting config) + , maybePair "max_ack_pending" (consumerConfigRequestMaxAckPending config) + , maybeJsonPair "inactive_threshold" (diffTimeNanosToJSON <$> consumerConfigRequestInactiveThreshold config) + , maybeJsonPair "idle_heartbeat" (diffTimeNanosToJSON <$> consumerConfigRequestIdleHeartbeat config) + , maybePair "headers_only" (consumerConfigRequestHeadersOnly config) + , maybePair "num_replicas" (consumerConfigRequestReplicas config) + , maybePair "mem_storage" (consumerConfigRequestMemoryStorage config) + ] ++ maybe [] deliverPolicyRequestPairs (consumerConfigRequestDeliverPolicy config) + ++ maybe [] consumerFilterPairs (consumerConfigRequestFilter config) + +instance ToJSON ConsumerConfig where + toJSON config = + object . catMaybes $ + [ maybeByteStringPair "durable_name" (consumerConfigDurableName config) + , maybeByteStringPair "name" (consumerConfigName config) + , maybeByteStringPair "description" (consumerConfigDescription config) + , maybeByteStringPair "deliver_subject" (consumerConfigDeliverSubject config) + , maybeByteStringPair "deliver_group" (consumerConfigDeliverGroup config) + , Just ("ack_policy" .= consumerConfigAckPolicy config) + , Just ("replay_policy" .= consumerConfigReplayPolicy config) + , maybeByteStringPair "filter_subject" (consumerConfigFilterSubject config) + , maybeByteStringListPair "filter_subjects" (consumerConfigFilterSubjects config) + , maybeJsonPair "ack_wait" (diffTimeNanosToJSON <$> consumerConfigAckWait config) + , maybePair "max_deliver" (consumerConfigMaxDeliver config) + , maybePair "max_waiting" (consumerConfigMaxWaiting config) + , maybePair "max_ack_pending" (consumerConfigMaxAckPending config) + , maybeJsonPair "inactive_threshold" (diffTimeNanosToJSON <$> consumerConfigInactiveThreshold config) + , maybeJsonPair "idle_heartbeat" (diffTimeNanosToJSON <$> consumerConfigIdleHeartbeat config) + , maybePair "headers_only" (consumerConfigHeadersOnly config) + , maybePair "num_replicas" (consumerConfigReplicas config) + , maybePair "mem_storage" (consumerConfigMemoryStorage config) + ] ++ fmap Just (deliverPolicyPairs (consumerConfigDeliverPolicy config)) + +instance FromJSON ConsumerConfig where + parseJSON = withObject "ConsumerConfig" $ \obj -> + ConsumerConfig + <$> parseOptionalByteStringField obj "durable_name" + <*> parseOptionalByteStringField obj "name" + <*> parseOptionalByteStringField obj "description" + <*> parseOptionalByteStringField obj "deliver_subject" + <*> parseOptionalByteStringField obj "deliver_group" + <*> parseDeliverPolicyField obj "deliver_policy" + <*> obj .: "ack_policy" + <*> obj .: "replay_policy" + <*> parseOptionalByteStringField obj "filter_subject" + <*> parseOptionalByteStringListField obj "filter_subjects" + <*> parseOptionalDurationField obj "ack_wait" + <*> obj .:? "max_deliver" + <*> obj .:? "max_waiting" + <*> obj .:? "max_ack_pending" + <*> parseOptionalDurationField obj "inactive_threshold" + <*> parseOptionalDurationField obj "idle_heartbeat" + <*> obj .:? "headers_only" + <*> obj .:? "num_replicas" + <*> obj .:? "mem_storage" + +instance FromJSON ConsumerInfo where + parseJSON = withObject "ConsumerInfo" $ \obj -> + ConsumerInfo + <$> parseByteStringField obj "stream_name" + <*> parseByteStringField obj "name" + <*> obj .: "created" + <*> obj .: "config" + <*> obj .:? "delivered" .!= emptyConsumerSequenceInfo + <*> obj .:? "ack_floor" .!= emptyConsumerSequenceInfo + <*> obj .:? "num_ack_pending" .!= 0 + <*> obj .:? "num_redelivered" .!= 0 + <*> obj .:? "num_waiting" .!= 0 + <*> obj .:? "num_pending" .!= 0 + +instance ToJSON ConsumerInfo where + toJSON info = + object + [ "stream_name" .= byteStringToJSON (consumerInfoStreamName info) + , "name" .= byteStringToJSON (consumerInfoName info) + , "created" .= consumerInfoCreated info + , "config" .= consumerInfoConfig info + , "delivered" .= consumerInfoDelivered info + , "ack_floor" .= consumerInfoAckFloor info + , "num_ack_pending" .= consumerInfoNumAckPending info + , "num_redelivered" .= consumerInfoNumRedelivered info + , "num_waiting" .= consumerInfoNumWaiting info + , "num_pending" .= consumerInfoNumPending info + ] + +instance FromJSON ConsumerSequenceInfo where + parseJSON = withObject "ConsumerSequenceInfo" $ \obj -> + ConsumerSequenceInfo + <$> obj .:? "consumer_seq" .!= 0 + <*> obj .:? "stream_seq" .!= 0 + <*> obj .:? "last_active" + +instance ToJSON ConsumerSequenceInfo where + toJSON info = + object . catMaybes $ + [ Just ("consumer_seq" .= consumerSequenceConsumer info) + , Just ("stream_seq" .= consumerSequenceStream info) + , maybePair "last_active" (consumerSequenceLast info) + ] + +instance FromJSON DeleteConsumerResponse where + parseJSON = withObject "DeleteConsumerResponse" $ \obj -> + DeleteConsumerResponse + <$> obj .:? "success" .!= False + +instance ToJSON ConsumerPauseRequest where + toJSON request = + object . catMaybes $ + [ maybePair "pause_until" (consumerPauseUntil request) + ] + +instance FromJSON ConsumerPauseResponse where + parseJSON = withObject "ConsumerPauseResponse" $ \obj -> + ConsumerPauseResponse + <$> obj .: "paused" + <*> obj .:? "pause_until" + <*> parseOptionalDurationField obj "pause_remaining" + +instance ToJSON ConsumerResetRequest where + toJSON request = + object . catMaybes $ + [ maybePair "seq" (consumerResetRequestSequence request) + ] + +instance FromJSON ConsumerResetResponse where + parseJSON value = + ConsumerResetResponse + <$> parseJSON value + <*> parseResetSequence value + where + parseResetSequence = + withObject "ConsumerResetResponse" $ \obj -> + obj .: "reset_seq" + +instance ToJSON ConsumerListRequest where + toJSON request = + object . catMaybes $ + [ maybePair "offset" (consumerListRequestOffset request) + ] + +instance FromJSON ConsumerListResponse where + parseJSON = withObject "ConsumerListResponse" $ \obj -> + ConsumerListResponse + <$> obj .: "total" + <*> obj .: "offset" + <*> obj .: "limit" + <*> obj .:? "consumers" .!= [] + +instance FromJSON ConsumerNamesResponse where + parseJSON = withObject "ConsumerNamesResponse" $ \obj -> + ConsumerNamesResponse + <$> obj .: "total" + <*> obj .: "offset" + <*> obj .: "limit" + <*> parseOptionalByteStringListField obj "consumers" .!= [] + +byteStringPair :: Key -> BS.ByteString -> Pair +byteStringPair key value = + key .= byteStringToJSON value + +maybeByteStringPair :: Key -> Maybe BS.ByteString -> Maybe Pair +maybeByteStringPair key = + fmap (byteStringPair key) + +maybeByteStringListPair :: Key -> Maybe [BS.ByteString] -> Maybe Pair +maybeByteStringListPair key = + fmap ((key .=) . map byteStringToJSON) + +maybePair :: ToJSON value => Key -> Maybe value -> Maybe Pair +maybePair key = + fmap (key .=) + +maybeJsonPair :: Key -> Maybe Value -> Maybe Pair +maybeJsonPair key = + fmap (key .=) + +consumerFilterPairs :: ConsumerFilter -> [Maybe Pair] +consumerFilterPairs consumerFilter = + case consumerFilter of + ConsumerFilterSubject subject -> + [maybeByteStringPair "filter_subject" (Just subject)] + ConsumerFilterSubjects subjects -> + [maybeByteStringListPair "filter_subjects" (Just subjects)] + +deliverPolicyRequestPairs :: DeliverPolicy -> [Maybe Pair] +deliverPolicyRequestPairs policy = + fmap Just (deliverPolicyPairs policy) + +deliverPolicyPairs :: DeliverPolicy -> [Pair] +deliverPolicyPairs policy = + "deliver_policy" .= policy : extraPairs policy + where + extraPairs selectedPolicy = + case selectedPolicy of + DeliverByStartSequence sequenceNumber -> + ["opt_start_seq" .= sequenceNumber] + DeliverByStartTime startTime -> + ["opt_start_time" .= startTime] + _ -> + [] + +parseDeliverPolicyField :: Object -> Key -> Parser DeliverPolicy +parseDeliverPolicyField obj key = do + policyName <- obj .: key :: Parser T.Text + case policyName of + "all" -> + pure DeliverAll + "last" -> + pure DeliverLast + "new" -> + pure DeliverNew + "by_start_sequence" -> + fmap DeliverByStartSequence (obj .: "opt_start_seq") + "by_start_time" -> + fmap DeliverByStartTime (obj .: "opt_start_time") + "last_per_subject" -> + pure DeliverLastPerSubject + value -> + fail ("unknown deliver policy: " ++ T.unpack value) + +parseByteStringField :: Object -> Key -> Parser BS.ByteString +parseByteStringField obj key = + obj .: key >>= parseByteString + +parseOptionalByteStringField :: Object -> Key -> Parser (Maybe BS.ByteString) +parseOptionalByteStringField obj key = do + value <- obj .:? key + traverse parseByteString value + +parseOptionalByteStringListField :: Object -> Key -> Parser (Maybe [BS.ByteString]) +parseOptionalByteStringListField obj key = do + value <- obj .:? key + traverse (traverse parseByteString) value + +parseOptionalDurationField :: Object -> Key -> Parser (Maybe NominalDiffTime) +parseOptionalDurationField obj key = + fmap nanosToDiffTime <$> obj .:? key + +nanosToDiffTime :: Integer -> NominalDiffTime +nanosToDiffTime nanoseconds = + fromRational (toRational nanoseconds / 1000000000) + +emptyConsumerSequenceInfo :: ConsumerSequenceInfo +emptyConsumerSequenceInfo = + ConsumerSequenceInfo + { consumerSequenceConsumer = 0 + , consumerSequenceStream = 0 + , consumerSequenceLast = Nothing + } diff --git a/jetstream/JetStream/Error.hs b/jetstream/JetStream/Error.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Error.hs @@ -0,0 +1,39 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Error + ( JetStreamApiError (..) + , JetStreamError (..) + ) where + +import Data.Aeson +import qualified Data.ByteString as BS +import qualified Data.Text.Encoding as E + +data JetStreamApiError = JetStreamApiError + { apiErrorCode :: Int + , apiErrorCodeDetail :: Int + , apiErrorDescription :: BS.ByteString + } + deriving (Eq, Show) + +data JetStreamError = JetStreamApiFailure JetStreamApiError + | JetStreamStatusError Int (Maybe BS.ByteString) + | JetStreamDecodeError String + | JetStreamTimeout + | JetStreamNoReply + deriving (Eq, Show) + +instance FromJSON JetStreamApiError where + parseJSON = withObject "JetStreamApiError" $ \obj -> + JetStreamApiError + <$> obj .: "code" + <*> obj .:? "err_code" .!= 0 + <*> fmap E.encodeUtf8 (obj .:? "description" .!= "") + +instance ToJSON JetStreamApiError where + toJSON err = + object + [ "code" .= apiErrorCode err + , "err_code" .= apiErrorCodeDetail err + , "description" .= E.decodeUtf8 (apiErrorDescription err) + ] diff --git a/jetstream/JetStream/Message.hs b/jetstream/JetStream/Message.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Message.hs @@ -0,0 +1,361 @@ +module JetStream.Message + ( messageAPI + , fetchMessages + , ackMessage + , nakMessage + , inProgressMessage + , termMessage + ) where + +import qualified API as Nats +import Control.Concurrent.STM +import Control.Exception (bracket) +import Control.Monad (unless, void) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BC +import Data.Char (isAlphaNum) +import Data.Maybe (catMaybes, fromMaybe) +import JetStream.Consumer.API (ConsumerAPI (..)) +import JetStream.Consumer.Types + ( ConsumerAction (ConsumerCreateOrUpdate) + , ConsumerConfigOption + , ConsumerInfo (consumerInfoName) + , ConsumerKind (PullConsumer) + , ConsumerTarget (NamedConsumer) + , withConsumerAckPolicy + , withConsumerDeliverPolicy + , withConsumerFilter + , withConsumerHeadersOnly + , withConsumerInactiveThreshold + , withConsumerMaxDeliver + , withConsumerMaxWaiting + , withConsumerMemoryStorage + , withConsumerReplayPolicy + , withConsumerReplicas + ) +import JetStream.Error + ( JetStreamError (JetStreamDecodeError, JetStreamNoReply, JetStreamTimeout) + ) +import JetStream.Message.API (MessageAPI (..)) +import JetStream.Message.Types +import JetStream.Options (JetStreamContext (..)) +import qualified JetStream.Protocol.Subject as Subject +import JetStream.Types + ( AckPolicy (AckNone) + , ConsumerName + , DeliverPolicy (DeliverByStartSequence) + , StreamName + , Subject + ) +import System.Timeout (timeout) + +-- | Build the message capability from the shared JetStream context. +messageAPI :: JetStreamContext -> ConsumerAPI -> MessageAPI +messageAPI context consumerAPI = + MessageAPI + { fetch = \stream consumer options -> + fetchMessages context stream consumer (pullRequest options) + , consumePush = \deliverSubject options handler -> + consumePushMessages context deliverSubject (pushConsumeConfig options) handler + , createOrderedConsumer = \stream options -> + createOrderedConsumerHandle context consumerAPI stream (orderedConsumerConfig options) + , ack = ackMessage (contextClient context) + , nak = nakMessage (contextClient context) + , inProgress = inProgressMessage (contextClient context) + , term = termMessage (contextClient context) + } + +consumePushMessages + :: JetStreamContext + -> Subject + -> PushConsumeConfig + -> (Message -> IO ()) + -> IO PushSubscription +consumePushMessages context deliverSubject config handler = do + sid <- Nats.subscribe natsClient deliverSubject subscribeOptions $ \msgView -> + case fmap fromMsgView msgView of + Nothing -> + pure () + Just message -> + unless (isStatusMessage message) (handler message) + pure PushSubscription + { stopPushSubscription = Nats.unsubscribe natsClient sid + } + where + natsClient = contextClient context + subscribeOptions = + maybe [] (\queueGroup -> [Nats.withQueueGroup queueGroup]) (pushConsumeQueueGroup config) + +data OrderedState = OrderedState + { orderedStateContext :: JetStreamContext + , orderedStateConsumers :: ConsumerAPI + , orderedStateStream :: StreamName + , orderedStateConfig :: OrderedConsumerConfig + , orderedStateNamePrefix :: ConsumerName + , orderedStateSerial :: TVar Int + , orderedStateNextSequence :: TVar (Maybe Integer) + , orderedStateCurrentName :: TVar (Maybe ConsumerName) + , orderedStateStopped :: TVar Bool + } + +createOrderedConsumerHandle + :: JetStreamContext + -> ConsumerAPI + -> StreamName + -> OrderedConsumerConfig + -> IO (Either JetStreamError OrderedConsumer) +createOrderedConsumerHandle context consumerAPI stream config = do + namePrefix <- orderedNamePrefix context config + state <- OrderedState context consumerAPI stream config namePrefix + <$> newTVarIO 0 + <*> newTVarIO Nothing + <*> newTVarIO Nothing + <*> newTVarIO False + created <- resetOrderedConsumer state + pure $ do + void created + Right OrderedConsumer + { orderedConsumerInfo = orderedInfo state + , fetchOrdered = fetchOrderedMessages state + , stopOrderedConsumer = stopOrdered state + } + +orderedInfo :: OrderedState -> IO (Either JetStreamError ConsumerInfo) +orderedInfo state = do + currentName <- readTVarIO (orderedStateCurrentName state) + case currentName of + Nothing -> + pure (Left JetStreamNoReply) + Just consumerName -> + consumerInfo + (orderedStateConsumers state) + (orderedStateStream state) + consumerName + +fetchOrderedMessages :: OrderedState -> [FetchOption] -> IO (Either JetStreamError PullResponse) +fetchOrderedMessages state options = do + resetResult <- resetOrderedConsumer state + case resetResult of + Left err -> + pure (Left err) + Right info -> do + responseResult <- fetchMessages + (orderedStateContext state) + (orderedStateStream state) + (consumerInfoName info) + (pullRequest options) + case responseResult of + Left err -> + pure (Left err) + Right response -> + case orderedResponseNextSequence response of + Left err -> + pure (Left err) + Right Nothing -> + pure (Right response) + Right (Just nextSequence) -> do + atomically $ + writeTVar (orderedStateNextSequence state) (Just nextSequence) + pure (Right response) + +resetOrderedConsumer :: OrderedState -> IO (Either JetStreamError ConsumerInfo) +resetOrderedConsumer state = do + stopped <- readTVarIO (orderedStateStopped state) + if stopped + then pure (Left JetStreamNoReply) + else do + previous <- readTVarIO (orderedStateCurrentName state) + case previous of + Nothing -> + pure () + Just consumerName -> + void $ + deleteConsumer + (orderedStateConsumers state) + (orderedStateStream state) + consumerName + (serial, nextSequence) <- atomically $ do + serial <- succ <$> readTVar (orderedStateSerial state) + writeTVar (orderedStateSerial state) serial + nextSequence <- readTVar (orderedStateNextSequence state) + pure (serial, nextSequence) + let consumerName = orderedConsumerName (orderedStateNamePrefix state) serial + options = orderedConsumerOptions (orderedStateConfig state) nextSequence + result <- putConsumer + (orderedStateConsumers state) + (orderedStateStream state) + ConsumerCreateOrUpdate + (NamedConsumer consumerName) + PullConsumer + options + case result of + Left err -> + pure (Left err) + Right info -> do + atomically $ + writeTVar (orderedStateCurrentName state) (Just consumerName) + pure (Right info) + +stopOrdered :: OrderedState -> IO () +stopOrdered state = do + atomically $ + writeTVar (orderedStateStopped state) True + currentName <- readTVarIO (orderedStateCurrentName state) + case currentName of + Nothing -> + pure () + Just consumerName -> + void $ + deleteConsumer + (orderedStateConsumers state) + (orderedStateStream state) + consumerName + +orderedConsumerOptions + :: OrderedConsumerConfig + -> Maybe Integer + -> [ConsumerConfigOption] +orderedConsumerOptions config nextSequence = + catMaybes + [ Just (withConsumerAckPolicy AckNone) + , Just (withConsumerDeliverPolicy deliverPolicy) + , Just (withConsumerMaxDeliver (-1)) + , Just (withConsumerMaxWaiting 512) + , Just (withConsumerInactiveThreshold inactiveThreshold) + , Just (withConsumerReplicas 1) + , Just (withConsumerMemoryStorage True) + , withConsumerFilter <$> orderedConsumerFilter config + , withConsumerReplayPolicy <$> orderedConsumerReplayPolicy config + , withConsumerHeadersOnly <$> orderedConsumerHeadersOnly config + ] + where + deliverPolicy = + case nextSequence of + Nothing -> + orderedConsumerDeliverPolicy config + Just sequenceNumber -> + DeliverByStartSequence (fromInteger sequenceNumber) + inactiveThreshold = + fromMaybe 300 (orderedConsumerInactiveThreshold config) + +orderedResponseNextSequence :: PullResponse -> Either JetStreamError (Maybe Integer) +orderedResponseNextSequence response = + case reverse (pullResponseMessages response) of + [] -> + Right Nothing + lastMessage:_ -> + case messageMetadata lastMessage of + Nothing -> + Left (JetStreamDecodeError "missing JetStream message metadata") + Just metadata -> + Right (Just (messageMetadataStreamSequence metadata + 1)) + +orderedNamePrefix :: JetStreamContext -> OrderedConsumerConfig -> IO ConsumerName +orderedNamePrefix context config = + case orderedConsumerNamePrefix config of + Just namePrefix -> + pure namePrefix + Nothing -> do + inbox <- Nats.newInbox (contextClient context) + pure (BS.append (BC.pack "ORDERED_") (sanitizeConsumerName inbox)) + +orderedConsumerName :: ConsumerName -> Int -> ConsumerName +orderedConsumerName namePrefix serial = + BS.concat [namePrefix, BC.pack "_", BC.pack (show serial)] + +sanitizeConsumerName :: Subject -> ConsumerName +sanitizeConsumerName = + BC.map replace + where + replace char + | isAlphaNum char = char + | otherwise = '_' + +-- | Fetch messages for a pull consumer. +fetchMessages :: JetStreamContext -> StreamName -> ConsumerName -> PullRequest -> IO (Either JetStreamError PullResponse) +fetchMessages context stream consumer request + | pullRequestBatch request <= 0 = pure (Right (PullResponse [] Nothing)) + | otherwise = do + responseQueue <- newTQueueIO + inbox <- Nats.newInbox natsClient + bracket + (Nats.subscribe natsClient inbox [] (atomically . writeTQueue responseQueue)) + (Nats.unsubscribe natsClient) + (\_ -> do + Nats.publish natsClient requestSubject + [ Nats.withPayload (pullRequestPayload (pullRequestBatch request) request) + , Nats.withReplyTo inbox + ] + collectResponses (responseTimeoutMicros request) responseQueue (pullRequestBatch request) []) + where + natsClient = contextClient context + requestSubject = Subject.consumerNextSubject context stream consumer + +ackMessage :: Nats.Client -> Message -> IO (Either JetStreamError ()) +ackMessage natsClient = publishAck natsClient Ack + +nakMessage :: Nats.Client -> Message -> IO (Either JetStreamError ()) +nakMessage natsClient = publishAck natsClient Nak + +inProgressMessage :: Nats.Client -> Message -> IO (Either JetStreamError ()) +inProgressMessage natsClient = publishAck natsClient InProgress + +termMessage :: Nats.Client -> Message -> IO (Either JetStreamError ()) +termMessage natsClient = publishAck natsClient Term + +data PullResult = PullResultMessage Message + | PullResultStatus PullStatus + | PullResultTimeout + | PullResultClosed + +collectResponses :: Int -> TQueue (Maybe Nats.MsgView) -> Int -> [Message] -> IO (Either JetStreamError PullResponse) +collectResponses _ _ 0 messages = + pure (Right (PullResponse (reverse messages) Nothing)) +collectResponses waitMicros responseQueue remaining messages = do + response <- timeout waitMicros (atomically (readTQueue responseQueue)) + case classifyPullResult response of + PullResultMessage message -> + collectResponses waitMicros responseQueue (remaining - 1) (message:messages) + PullResultStatus status -> + pure (Right (PullResponse (reverse messages) (Just status))) + PullResultTimeout -> + pure (Left JetStreamTimeout) + PullResultClosed -> + pure (Left JetStreamNoReply) + +classifyPullResult :: Maybe (Maybe Nats.MsgView) -> PullResult +classifyPullResult Nothing = PullResultTimeout +classifyPullResult (Just Nothing) = PullResultClosed +classifyPullResult (Just (Just msgView)) = + case messageStatus message of + Nothing -> PullResultMessage message + Just status -> PullResultStatus status + where + message = fromMsgView msgView + +fromMsgView :: Nats.MsgView -> Message +fromMsgView msgView = + Message + { messageSubject = Nats.subject msgView + , messagePayload = Nats.payload msgView + , messageHeaders = Nats.headers msgView + , messageReplyTo = Nats.replyTo msgView + , messageStatus = classifyStatusHeaders (Nats.headers msgView) + } + +publishAck :: Nats.Client -> AckVerb -> Message -> IO (Either JetStreamError ()) +publishAck natsClient verb message = + case messageReplyTo message of + Nothing -> + pure (Left JetStreamNoReply) + Just reply -> + Right <$> Nats.publish natsClient reply [Nats.withPayload (ackPayload verb)] + +responseTimeoutMicros :: PullRequest -> Int +responseTimeoutMicros request = + case pullRequestWait request of + FetchNoWaitMicros timeoutMicros -> + max 1 timeoutMicros + FetchExpiresMicros timeoutMicros + | timeoutMicros <= 0 -> 1 + | otherwise -> timeoutMicros + 100000 diff --git a/jetstream/JetStream/Message/API.hs b/jetstream/JetStream/Message/API.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Message/API.hs @@ -0,0 +1,64 @@ +module JetStream.Message.API + ( MessageAPI (..) + , FetchOption + , FetchWait (..) + , Headers + , Message (..) + , OrderedConsumer (..) + , OrderedConsumerOption + , PushConsumeOption + , PushSubscription (..) + , PullResponse (..) + , PullStatus (..) + , withFetchBatch + , withFetchWait + , withOrderedConsumerDeliverPolicy + , withOrderedConsumerFilter + , withOrderedConsumerHeadersOnly + , withOrderedConsumerInactiveThreshold + , withOrderedConsumerNamePrefix + , withOrderedConsumerReplayPolicy + , withPushQueueGroup + ) where + +import JetStream.Error (JetStreamError) +import JetStream.Message.Types + ( FetchOption + , FetchWait (..) + , Headers + , Message (..) + , OrderedConsumer (..) + , OrderedConsumerOption + , PullResponse (..) + , PullStatus (..) + , PushConsumeOption + , PushSubscription (..) + , withFetchBatch + , withFetchWait + , withOrderedConsumerDeliverPolicy + , withOrderedConsumerFilter + , withOrderedConsumerHeadersOnly + , withOrderedConsumerInactiveThreshold + , withOrderedConsumerNamePrefix + , withOrderedConsumerReplayPolicy + , withPushQueueGroup + ) +import JetStream.Types (ConsumerName, StreamName, Subject) + +-- | Pull-consumer message operations. +data MessageAPI = MessageAPI + { fetch :: StreamName -> ConsumerName -> [FetchOption] -> IO (Either JetStreamError PullResponse) + -- ^ Fetch messages for a pull consumer. + , consumePush :: Subject -> [PushConsumeOption] -> (Message -> IO ()) -> IO PushSubscription + -- ^ Subscribe to a push consumer deliver subject. + , createOrderedConsumer :: StreamName -> [OrderedConsumerOption] -> IO (Either JetStreamError OrderedConsumer) + -- ^ Create a client-managed ordered pull consumer. + , ack :: Message -> IO (Either JetStreamError ()) + -- ^ Acknowledge successful message processing. + , nak :: Message -> IO (Either JetStreamError ()) + -- ^ Negatively acknowledge a message. + , inProgress :: Message -> IO (Either JetStreamError ()) + -- ^ Tell the server that message processing is still in progress. + , term :: Message -> IO (Either JetStreamError ()) + -- ^ Terminate message redelivery. + } diff --git a/jetstream/JetStream/Message/Types.hs b/jetstream/JetStream/Message/Types.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Message/Types.hs @@ -0,0 +1,364 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Message.Types + ( AckVerb (..) + , FetchOption + , FetchWait (..) + , Headers + , Message (..) + , MessageMetadata (..) + , OrderedConsumer (..) + , OrderedConsumerConfig (..) + , OrderedConsumerOption + , PushConsumeConfig (..) + , PushConsumeOption + , PushSubscription (..) + , PullRequest (..) + , PullResponse (..) + , PullStatus (..) + , ackPayload + , classifyStatusHeaders + , defaultPullRequest + , descriptionHeader + , pullRequest + , inProgressPayload + , isStatusMessage + , messageMetadata + , nakPayload + , orderedConsumerConfig + , pullRequestPayload + , pushConsumeConfig + , statusHeader + , termPayload + , withFetchBatch + , withFetchWait + , withOrderedConsumerDeliverPolicy + , withOrderedConsumerFilter + , withOrderedConsumerHeadersOnly + , withOrderedConsumerInactiveThreshold + , withOrderedConsumerNamePrefix + , withOrderedConsumerReplayPolicy + , withPushQueueGroup + ) where + +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BC +import Data.Char (isDigit, toLower) +import Data.Maybe (isJust) +import Data.Time.Clock (NominalDiffTime, UTCTime) +import qualified Data.Time.Clock.POSIX as Time +import JetStream.Consumer.Types (ConsumerFilter, ConsumerInfo) +import JetStream.Error (JetStreamError) +import JetStream.Types + ( CallOption + , ConsumerName + , DeliverPolicy (..) + , Headers + , Payload + , ReplayPolicy + , StreamName + , Subject + , applyCallOptions + ) + +data PullRequest = PullRequest + { pullRequestBatch :: Int + , pullRequestWait :: FetchWait + } + deriving (Eq, Show) + +data FetchWait = FetchExpiresMicros Int + | FetchNoWaitMicros Int + deriving (Eq, Show) + +defaultPullRequest :: PullRequest +defaultPullRequest = + PullRequest + { pullRequestBatch = 1 + , pullRequestWait = FetchExpiresMicros 1000000 + } + +type FetchOption = CallOption PullRequest + +pullRequest :: [FetchOption] -> PullRequest +pullRequest options = + applyCallOptions options defaultPullRequest + +withFetchBatch :: Int -> FetchOption +withFetchBatch batch request = + request { pullRequestBatch = batch } + +withFetchWait :: FetchWait -> FetchOption +withFetchWait wait request = + request { pullRequestWait = wait } + +data PullResponse = PullResponse + { pullResponseMessages :: [Message] + , pullResponseStatus :: Maybe PullStatus + } + deriving (Eq, Show) + +data Message = Message + { messageSubject :: Subject + , messagePayload :: Maybe Payload + , messageHeaders :: Maybe Headers + , messageReplyTo :: Maybe Subject + , messageStatus :: Maybe PullStatus + } + deriving (Eq, Show) + +data MessageMetadata = MessageMetadata + { messageMetadataStream :: StreamName + , messageMetadataConsumer :: ConsumerName + , messageMetadataStreamSequence :: Integer + , messageMetadataConsumerSequence :: Integer + , messageMetadataNumDelivered :: Integer + , messageMetadataNumPending :: Integer + , messageMetadataTimestamp :: UTCTime + , messageMetadataDomain :: Maybe ByteString + } + deriving (Eq, Show) + +newtype PushSubscription = PushSubscription { stopPushSubscription :: IO () } + +newtype PushConsumeConfig = PushConsumeConfig { pushConsumeQueueGroup :: Maybe Subject } + +type PushConsumeOption = CallOption PushConsumeConfig + +pushConsumeConfig :: [PushConsumeOption] -> PushConsumeConfig +pushConsumeConfig options = + applyCallOptions options $ + PushConsumeConfig + { pushConsumeQueueGroup = Nothing + } + +withPushQueueGroup :: Subject -> PushConsumeOption +withPushQueueGroup queueGroup config = + config { pushConsumeQueueGroup = Just queueGroup } + +data OrderedConsumer = OrderedConsumer + { orderedConsumerInfo :: IO (Either JetStreamError ConsumerInfo) + , fetchOrdered :: [FetchOption] -> IO (Either JetStreamError PullResponse) + , stopOrderedConsumer :: IO () + } + +data OrderedConsumerConfig = OrderedConsumerConfig + { orderedConsumerNamePrefix :: Maybe ConsumerName + , orderedConsumerDeliverPolicy :: DeliverPolicy + , orderedConsumerFilter :: Maybe ConsumerFilter + , orderedConsumerReplayPolicy :: Maybe ReplayPolicy + , orderedConsumerInactiveThreshold :: Maybe NominalDiffTime + , orderedConsumerHeadersOnly :: Maybe Bool + } + +type OrderedConsumerOption = CallOption OrderedConsumerConfig + +orderedConsumerConfig :: [OrderedConsumerOption] -> OrderedConsumerConfig +orderedConsumerConfig options = + applyCallOptions options $ + OrderedConsumerConfig + { orderedConsumerNamePrefix = Nothing + , orderedConsumerDeliverPolicy = DeliverAll + , orderedConsumerFilter = Nothing + , orderedConsumerReplayPolicy = Nothing + , orderedConsumerInactiveThreshold = Nothing + , orderedConsumerHeadersOnly = Nothing + } + +withOrderedConsumerNamePrefix :: ConsumerName -> OrderedConsumerOption +withOrderedConsumerNamePrefix namePrefix config = + config { orderedConsumerNamePrefix = Just namePrefix } + +withOrderedConsumerDeliverPolicy :: DeliverPolicy -> OrderedConsumerOption +withOrderedConsumerDeliverPolicy policy config = + config { orderedConsumerDeliverPolicy = policy } + +withOrderedConsumerFilter :: ConsumerFilter -> OrderedConsumerOption +withOrderedConsumerFilter consumerFilter config = + config { orderedConsumerFilter = Just consumerFilter } + +withOrderedConsumerReplayPolicy :: ReplayPolicy -> OrderedConsumerOption +withOrderedConsumerReplayPolicy replayPolicy config = + config { orderedConsumerReplayPolicy = Just replayPolicy } + +withOrderedConsumerInactiveThreshold :: NominalDiffTime -> OrderedConsumerOption +withOrderedConsumerInactiveThreshold inactiveThreshold config = + config { orderedConsumerInactiveThreshold = Just inactiveThreshold } + +withOrderedConsumerHeadersOnly :: Bool -> OrderedConsumerOption +withOrderedConsumerHeadersOnly headersOnly config = + config { orderedConsumerHeadersOnly = Just headersOnly } + +data PullStatus = PullNoMessages (Maybe ByteString) + | PullRequestTimeout (Maybe ByteString) + | PullStatusError ByteString (Maybe ByteString) + deriving (Eq, Show) + +data AckVerb = Ack | Nak | InProgress | Term + deriving (Eq, Show) + +ackPayload :: AckVerb -> Payload +ackPayload Ack = "+ACK" +ackPayload Nak = "-NAK" +ackPayload InProgress = "+WPI" +ackPayload Term = "+TERM" + +nakPayload :: Payload +nakPayload = ackPayload Nak + +inProgressPayload :: Payload +inProgressPayload = ackPayload InProgress + +termPayload :: Payload +termPayload = ackPayload Term + +statusHeader :: Headers -> Maybe ByteString +statusHeader = lookupHeader "Status" + +descriptionHeader :: Headers -> Maybe ByteString +descriptionHeader = lookupHeader "Description" + +classifyStatusHeaders :: Maybe Headers -> Maybe PullStatus +classifyStatusHeaders Nothing = Nothing +classifyStatusHeaders (Just headers) = + case statusHeader headers of + Nothing -> Nothing + Just "404" -> Just (PullNoMessages description) + Just "408" -> Just (PullRequestTimeout description) + Just status -> Just (PullStatusError status description) + where + description = descriptionHeader headers + +isStatusMessage :: Message -> Bool +isStatusMessage = isJust . messageStatus + +messageMetadata :: Message -> Maybe MessageMetadata +messageMetadata message = + messageReplyTo message >>= parseMetadataSubject + +pullRequestPayload :: Int -> PullRequest -> Payload +pullRequestPayload requestedBatch request = + BS.concat + [ "{\"batch\":" + , intBytes (max 1 requestedBatch) + , noWaitField + , expiresField + , "}" + ] + where + noWaitField + | FetchNoWaitMicros _ <- pullRequestWait request = ",\"no_wait\":true" + | otherwise = "" + expiresField = + case pullRequestWait request of + FetchNoWaitMicros _ -> + "" + FetchExpiresMicros timeoutMicros + | timeoutMicros <= 0 -> "" + | otherwise -> + BS.concat + [ ",\"expires\":" + , intBytes (timeoutMicros * 1000) + ] + +lookupHeader :: ByteString -> Headers -> Maybe ByteString +lookupHeader needle = go + where + normalizedNeedle = normalizeHeaderName needle + go [] = Nothing + go ((name, value):rest) + | normalizeHeaderName name == normalizedNeedle = Just value + | otherwise = go rest + +normalizeHeaderName :: ByteString -> ByteString +normalizeHeaderName = BC.map toLower + +intBytes :: Int -> ByteString +intBytes = BC.pack . show + +parseMetadataSubject :: Subject -> Maybe MessageMetadata +parseMetadataSubject subject = do + tokens <- normalizeMetadataTokens (BC.split '.' subject) + stream <- tokenAt ackStreamTokenPos tokens + consumer <- tokenAt ackConsumerTokenPos tokens + delivered <- parseNum =<< tokenAt ackNumDeliveredTokenPos tokens + streamSeq <- parseNum =<< tokenAt ackStreamSeqTokenPos tokens + consumerSeq <- parseNum =<< tokenAt ackConsumerSeqTokenPos tokens + timestamp <- nanosToTime <$> (parseNum =<< tokenAt ackTimestampTokenPos tokens) + pending <- parseNum =<< tokenAt ackNumPendingTokenPos tokens + let domain = nonEmpty =<< tokenAt ackDomainTokenPos tokens + pure MessageMetadata + { messageMetadataStream = stream + , messageMetadataConsumer = consumer + , messageMetadataStreamSequence = streamSeq + , messageMetadataConsumerSequence = consumerSeq + , messageMetadataNumDelivered = delivered + , messageMetadataNumPending = pending + , messageMetadataTimestamp = timestamp + , messageMetadataDomain = domain + } + +normalizeMetadataTokens :: [ByteString] -> Maybe [ByteString] +normalizeMetadataTokens tokens + | length tokens < 9 || (length tokens > 9 && length tokens < 11) = + Nothing + | take 2 tokens /= ["$JS", "ACK"] = + Nothing + | length tokens == 9 = + Just (take ackDomainTokenPos tokens ++ ["", ""] ++ drop ackDomainTokenPos tokens) + | otherwise = + Just (normalizeDomain tokens) + where + normalizeDomain values = + case tokenAt ackDomainTokenPos values of + Just "_" -> + take ackDomainTokenPos values ++ [""] ++ drop (ackDomainTokenPos + 1) values + _ -> + values + +tokenAt :: Int -> [a] -> Maybe a +tokenAt index values = + if index < length values + then Just (values !! index) + else Nothing + +parseNum :: ByteString -> Maybe Integer +parseNum bytes + | BS.null bytes = Just 0 + | BC.all isDigit bytes = Just (read (BC.unpack bytes)) + | otherwise = Nothing + +nanosToTime :: Integer -> UTCTime +nanosToTime nanoseconds = + Time.posixSecondsToUTCTime $ + fromRational (toRational nanoseconds / 1000000000) + +nonEmpty :: ByteString -> Maybe ByteString +nonEmpty bytes + | BS.null bytes = Nothing + | otherwise = Just bytes + +ackDomainTokenPos :: Int +ackDomainTokenPos = 2 + +ackStreamTokenPos :: Int +ackStreamTokenPos = 4 + +ackConsumerTokenPos :: Int +ackConsumerTokenPos = 5 + +ackNumDeliveredTokenPos :: Int +ackNumDeliveredTokenPos = 6 + +ackStreamSeqTokenPos :: Int +ackStreamSeqTokenPos = 7 + +ackConsumerSeqTokenPos :: Int +ackConsumerSeqTokenPos = 8 + +ackTimestampTokenPos :: Int +ackTimestampTokenPos = 9 + +ackNumPendingTokenPos :: Int +ackNumPendingTokenPos = 10 diff --git a/jetstream/JetStream/Options.hs b/jetstream/JetStream/Options.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Options.hs @@ -0,0 +1,50 @@ +module JetStream.Options + ( JetStreamConfig (..) + , JetStreamContext (..) + , JetStreamOption + , defaultJetStreamConfig + , newJetStreamContext + , withDomain + , withRequestTimeoutMicros + ) where + +import qualified API as Nats +import qualified Data.ByteString as BS + +data JetStreamConfig = JetStreamConfig + { configDomain :: Maybe BS.ByteString + , configRequestTimeoutMicros :: Int + } + deriving (Eq, Show) + +data JetStreamContext = JetStreamContext + { contextClient :: Nats.Client + , contextDomain :: Maybe BS.ByteString + , contextRequestTimeoutMicros :: Int + } + +type JetStreamOption = JetStreamConfig -> JetStreamConfig + +defaultJetStreamConfig :: JetStreamConfig +defaultJetStreamConfig = + JetStreamConfig + { configDomain = Nothing + , configRequestTimeoutMicros = 5 * 1000000 + } + +newJetStreamContext :: Nats.Client -> [JetStreamOption] -> JetStreamContext +newJetStreamContext client options = + let config = foldl (flip ($)) defaultJetStreamConfig options + in JetStreamContext + { contextClient = client + , contextDomain = configDomain config + , contextRequestTimeoutMicros = configRequestTimeoutMicros config + } + +withDomain :: BS.ByteString -> JetStreamOption +withDomain domain config = + config { configDomain = Just domain } + +withRequestTimeoutMicros :: Int -> JetStreamOption +withRequestTimeoutMicros timeoutMicros config = + config { configRequestTimeoutMicros = max 1 timeoutMicros } diff --git a/jetstream/JetStream/Protocol/Headers.hs b/jetstream/JetStream/Protocol/Headers.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Protocol/Headers.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Protocol.Headers + ( statusHeader + , descriptionHeader + , lookupHeader + , statusCode + , statusDescription + , statusError + ) where + +import qualified API as Nats +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BC +import Data.Maybe (fromMaybe, listToMaybe) +import JetStream.Error (JetStreamError (JetStreamStatusError)) + +statusHeader :: BS.ByteString +statusHeader = "Status" + +descriptionHeader :: BS.ByteString +descriptionHeader = "Description" + +lookupHeader :: BS.ByteString -> Nats.MsgView -> Maybe BS.ByteString +lookupHeader name msg = + listToMaybe + [ value + | (key, value) <- fromMaybe [] (Nats.headers msg) + , key == name + ] + +statusCode :: Nats.MsgView -> Maybe Int +statusCode msg = do + value <- lookupHeader statusHeader msg + case BC.readInt value of + Just (code, rest) + | BS.null rest -> + Just code + _ -> + Nothing + +statusDescription :: Nats.MsgView -> Maybe BS.ByteString +statusDescription = + lookupHeader descriptionHeader + +statusError :: Nats.MsgView -> Maybe JetStreamError +statusError msg = + fmap (\code -> JetStreamStatusError code (statusDescription msg)) (statusCode msg) diff --git a/jetstream/JetStream/Protocol/Request.hs b/jetstream/JetStream/Protocol/Request.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Protocol/Request.hs @@ -0,0 +1,100 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Protocol.Request + ( requestMsg + , requestJSON + , requestJSONWithHeaders + , requestUnit + , decodeJetStreamResponse + ) where + +import qualified API as Nats +import Control.Concurrent.STM +import Control.Monad (void) +import Data.Aeson +import qualified Data.Aeson as Aeson +import Data.Aeson.Types (Parser, parseEither) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as LBS +import JetStream.Error +import JetStream.Options (JetStreamContext (..)) +import JetStream.Protocol.Headers (statusError) +import JetStream.Types (Headers, Payload, Subject) +import System.Timeout (timeout) + +requestMsg :: JetStreamContext -> Subject -> Maybe Payload -> Headers -> IO (Either JetStreamError Nats.MsgView) +requestMsg ctx subject payload headers = do + replyBox <- newEmptyTMVarIO + let publishOptions = + maybe [] (\body -> [Nats.withPayload body]) payload + ++ [Nats.withReplyCallback (atomically . putTMVar replyBox)] + ++ [Nats.withHeaders headers | not (null headers)] + Nats.publish (contextClient ctx) subject publishOptions + result <- timeout (contextRequestTimeoutMicros ctx) (atomically (readTMVar replyBox)) + case result of + Nothing -> + pure (Left JetStreamTimeout) + Just Nothing -> + pure (Left JetStreamNoReply) + Just (Just msg) -> + pure (Right msg) + +requestJSON :: FromJSON a => JetStreamContext -> Subject -> Maybe Value -> IO (Either JetStreamError a) +requestJSON ctx subject body = + requestJSONWithHeaders ctx subject body [] + +requestJSONWithHeaders :: FromJSON a => JetStreamContext -> Subject -> Maybe Value -> Headers -> IO (Either JetStreamError a) +requestJSONWithHeaders ctx subject body headers = do + msgResult <- requestMsg ctx subject (strictEncode <$> body) headers + pure $ do + msg <- msgResult + case statusError msg of + Just err -> + Left err + Nothing -> + case Nats.payload msg of + Nothing -> + Left JetStreamNoReply + Just payload -> + decodeJetStreamResponse payload + +requestUnit :: JetStreamContext -> Subject -> Maybe Value -> IO (Either JetStreamError ()) +requestUnit ctx subject body = do + response <- requestJSON ctx subject body + pure (void (response :: Either JetStreamError SuccessResponse)) + +decodeJetStreamResponse :: FromJSON a => BS.ByteString -> Either JetStreamError a +decodeJetStreamResponse payload = + case Aeson.eitherDecode (LBS.fromStrict payload) :: Either String Value of + Left err -> + Left (JetStreamDecodeError err) + Right value -> + case parseEither parseApiResult value of + Left err -> + Left (JetStreamDecodeError err) + Right result -> + result + +parseApiResult :: FromJSON a => Value -> Parser (Either JetStreamError a) +parseApiResult value = + case value of + Object obj -> do + apiError <- obj .:? "error" + case apiError of + Just err -> + pure (Left (JetStreamApiFailure err)) + Nothing -> + Right <$> parseJSON value + _ -> + Right <$> parseJSON value + +strictEncode :: Value -> BS.ByteString +strictEncode = + LBS.toStrict . Aeson.encode + +newtype SuccessResponse = SuccessResponse Bool + deriving (Eq, Show) + +instance FromJSON SuccessResponse where + parseJSON = withObject "SuccessResponse" $ \obj -> + SuccessResponse <$> obj .: "success" diff --git a/jetstream/JetStream/Protocol/Subject.hs b/jetstream/JetStream/Protocol/Subject.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Protocol/Subject.hs @@ -0,0 +1,120 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Protocol.Subject + ( apiPrefix + , apiSubject + , accountInfoSubject + , streamCreateSubject + , streamUpdateSubject + , streamInfoSubject + , streamDeleteSubject + , streamMessageGetSubject + , streamMessageDeleteSubject + , streamPurgeSubject + , streamListSubject + , streamNamesSubject + , consumerCreateSubject + , consumerCreateNamedSubject + , durableConsumerCreateSubject + , consumerInfoSubject + , consumerDeleteSubject + , consumerPauseSubject + , consumerResetSubject + , consumerListSubject + , consumerNamesSubject + , consumerNextSubject + ) where + +import qualified Data.ByteString as BS +import JetStream.Options (JetStreamContext (..)) +import JetStream.Types (ConsumerName, StreamName, Subject) + +apiPrefix :: JetStreamContext -> Subject +apiPrefix ctx = + case contextDomain ctx of + Nothing -> "$JS.API" + Just domain -> BS.concat ["$JS.", domain, ".API"] + +apiSubject :: JetStreamContext -> [BS.ByteString] -> Subject +apiSubject ctx segments = + BS.intercalate "." (apiPrefix ctx : segments) + +accountInfoSubject :: JetStreamContext -> Subject +accountInfoSubject ctx = + apiSubject ctx ["INFO"] + +streamCreateSubject :: JetStreamContext -> StreamName -> Subject +streamCreateSubject ctx stream = + apiSubject ctx ["STREAM", "CREATE", stream] + +streamUpdateSubject :: JetStreamContext -> StreamName -> Subject +streamUpdateSubject ctx stream = + apiSubject ctx ["STREAM", "UPDATE", stream] + +streamInfoSubject :: JetStreamContext -> StreamName -> Subject +streamInfoSubject ctx stream = + apiSubject ctx ["STREAM", "INFO", stream] + +streamDeleteSubject :: JetStreamContext -> StreamName -> Subject +streamDeleteSubject ctx stream = + apiSubject ctx ["STREAM", "DELETE", stream] + +streamMessageGetSubject :: JetStreamContext -> StreamName -> Subject +streamMessageGetSubject ctx stream = + apiSubject ctx ["STREAM", "MSG", "GET", stream] + +streamMessageDeleteSubject :: JetStreamContext -> StreamName -> Subject +streamMessageDeleteSubject ctx stream = + apiSubject ctx ["STREAM", "MSG", "DELETE", stream] + +streamPurgeSubject :: JetStreamContext -> StreamName -> Subject +streamPurgeSubject ctx stream = + apiSubject ctx ["STREAM", "PURGE", stream] + +streamListSubject :: JetStreamContext -> Subject +streamListSubject ctx = + apiSubject ctx ["STREAM", "LIST"] + +streamNamesSubject :: JetStreamContext -> Subject +streamNamesSubject ctx = + apiSubject ctx ["STREAM", "NAMES"] + +consumerCreateSubject :: JetStreamContext -> StreamName -> Subject +consumerCreateSubject ctx stream = + apiSubject ctx ["CONSUMER", "CREATE", stream] + +consumerCreateNamedSubject :: JetStreamContext -> StreamName -> ConsumerName -> Subject +consumerCreateNamedSubject ctx stream consumer = + apiSubject ctx ["CONSUMER", "CREATE", stream, consumer] + +durableConsumerCreateSubject :: JetStreamContext -> StreamName -> ConsumerName -> Subject +durableConsumerCreateSubject ctx stream consumer = + apiSubject ctx ["CONSUMER", "DURABLE", "CREATE", stream, consumer] + +consumerInfoSubject :: JetStreamContext -> StreamName -> ConsumerName -> Subject +consumerInfoSubject ctx stream consumer = + apiSubject ctx ["CONSUMER", "INFO", stream, consumer] + +consumerDeleteSubject :: JetStreamContext -> StreamName -> ConsumerName -> Subject +consumerDeleteSubject ctx stream consumer = + apiSubject ctx ["CONSUMER", "DELETE", stream, consumer] + +consumerPauseSubject :: JetStreamContext -> StreamName -> ConsumerName -> Subject +consumerPauseSubject ctx stream consumer = + apiSubject ctx ["CONSUMER", "PAUSE", stream, consumer] + +consumerResetSubject :: JetStreamContext -> StreamName -> ConsumerName -> Subject +consumerResetSubject ctx stream consumer = + apiSubject ctx ["CONSUMER", "RESET", stream, consumer] + +consumerListSubject :: JetStreamContext -> StreamName -> Subject +consumerListSubject ctx stream = + apiSubject ctx ["CONSUMER", "LIST", stream] + +consumerNamesSubject :: JetStreamContext -> StreamName -> Subject +consumerNamesSubject ctx stream = + apiSubject ctx ["CONSUMER", "NAMES", stream] + +consumerNextSubject :: JetStreamContext -> StreamName -> ConsumerName -> Subject +consumerNextSubject ctx stream consumer = + apiSubject ctx ["CONSUMER", "MSG", "NEXT", stream, consumer] diff --git a/jetstream/JetStream/Publish.hs b/jetstream/JetStream/Publish.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Publish.hs @@ -0,0 +1,42 @@ +module JetStream.Publish + ( publishAPI + , module JetStream.Publish.API + ) where + +import qualified API as Nats +import JetStream.Error (JetStreamError (JetStreamNoReply)) +import JetStream.Options (JetStreamContext) +import JetStream.Protocol.Headers (statusError) +import JetStream.Protocol.Request + ( decodeJetStreamResponse + , requestMsg + ) +import JetStream.Publish.API +import JetStream.Publish.Types (publishHeaders) +import JetStream.Types (Payload, Subject) + +publishAPI :: JetStreamContext -> PublishAPI +publishAPI context = + PublishAPI + { publish = publishMessage context + } + +publishMessage + :: JetStreamContext + -> Subject + -> Payload + -> [PublishOption] + -> IO (Either JetStreamError PublishAck) +publishMessage context subject payload options = do + msgResult <- requestMsg context subject (Just payload) (publishHeaders options) + pure $ do + msg <- msgResult + case statusError msg of + Just err -> + Left err + Nothing -> + case Nats.payload msg of + Nothing -> + Left JetStreamNoReply + Just ackPayload -> + decodeJetStreamResponse ackPayload diff --git a/jetstream/JetStream/Publish/API.hs b/jetstream/JetStream/Publish/API.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Publish/API.hs @@ -0,0 +1,28 @@ +module JetStream.Publish.API + ( PublishAPI (..) + , PublishAck (..) + , PublishExpectation (..) + , PublishOption + , withMsgId + , withExpectedStream + , withPublishExpectation + , withHeaders + ) where + +import JetStream.Error (JetStreamError) +import JetStream.Publish.Types + ( PublishAck (..) + , PublishExpectation (..) + , PublishOption + , withExpectedStream + , withHeaders + , withMsgId + , withPublishExpectation + ) +import JetStream.Types (Payload, Subject) + +-- | Publish operations for JetStream subjects. +-- +-- The 'publish' field publishes a message and waits for a JetStream publish +-- acknowledgement. +newtype PublishAPI = PublishAPI { publish :: Subject -> Payload -> [PublishOption] -> IO (Either JetStreamError PublishAck) } diff --git a/jetstream/JetStream/Publish/Types.hs b/jetstream/JetStream/Publish/Types.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Publish/Types.hs @@ -0,0 +1,114 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Publish.Types + ( PublishAck (..) + , PublishExpectation (..) + , PublishOption + , withMsgId + , withExpectedStream + , withPublishExpectation + , withHeaders + , publishHeaders + ) where + +import Data.Aeson (FromJSON (..), withObject, (.:), (.:?)) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BC +import Data.Text.Encoding (encodeUtf8) +import Data.Word (Word64) +import JetStream.Types (StreamName) + +data PublishAck = PublishAck + { publishAckStream :: StreamName + , publishAckSequence :: Word64 + , publishAckDuplicate :: Maybe Bool + , publishAckDomain :: Maybe BS.ByteString + } + deriving (Eq, Show) + +data PublishExpectation = ExpectedLastSequence Word64 + | ExpectedLastSubjectSequence Word64 + | ExpectedLastMsgId BS.ByteString + deriving (Eq, Show) + +instance FromJSON PublishAck where + parseJSON = withObject "PublishAck" $ \object -> do + streamName <- object .: "stream" + sequence <- object .: "seq" + duplicate <- object .:? "duplicate" + domainName <- object .:? "domain" + pure PublishAck + { publishAckStream = encodeUtf8 streamName + , publishAckSequence = sequence + , publishAckDuplicate = duplicate + , publishAckDomain = encodeUtf8 <$> domainName + } + +type PublishOption = PublishConfig -> PublishConfig + +newtype PublishConfig = PublishConfig { publishConfigHeaders :: [(BS.ByteString, BS.ByteString)] } + +defaultPublishConfig :: PublishConfig +defaultPublishConfig = + PublishConfig + { publishConfigHeaders = [] + } + +withMsgId :: BS.ByteString -> PublishOption +withMsgId = + putHeader "Nats-Msg-Id" + +withExpectedStream :: StreamName -> PublishOption +withExpectedStream = + putHeader "Nats-Expected-Stream" + +withPublishExpectation :: PublishExpectation -> PublishOption +withPublishExpectation expectation = + putExpectationHeader key value + where + (key, value) = + case expectation of + ExpectedLastSequence sequenceNumber -> + ("Nats-Expected-Last-Sequence", renderWord64 sequenceNumber) + ExpectedLastSubjectSequence sequenceNumber -> + ("Nats-Expected-Last-Subject-Sequence", renderWord64 sequenceNumber) + ExpectedLastMsgId msgId -> + ("Nats-Expected-Last-Msg-Id", msgId) + +withHeaders :: [(BS.ByteString, BS.ByteString)] -> PublishOption +withHeaders headers config = + config + { publishConfigHeaders = + filter (not . isPublishExpectationHeader . fst) headers + ++ publishConfigHeaders config + } + +publishHeaders :: [PublishOption] -> [(BS.ByteString, BS.ByteString)] +publishHeaders options = + publishConfigHeaders (foldr ($) defaultPublishConfig options) + +putHeader :: BS.ByteString -> BS.ByteString -> PublishOption +putHeader key value config = + config + { publishConfigHeaders = + (key, value) : filter ((/= key) . fst) (publishConfigHeaders config) + } + +putExpectationHeader :: BS.ByteString -> BS.ByteString -> PublishOption +putExpectationHeader key value config = + config + { publishConfigHeaders = + (key, value) : filter (not . isPublishExpectationHeader . fst) (publishConfigHeaders config) + } + +isPublishExpectationHeader :: BS.ByteString -> Bool +isPublishExpectationHeader key = + key `elem` + [ "Nats-Expected-Last-Sequence" + , "Nats-Expected-Last-Subject-Sequence" + , "Nats-Expected-Last-Msg-Id" + ] + +renderWord64 :: Word64 -> BS.ByteString +renderWord64 = + BC.pack . show diff --git a/jetstream/JetStream/Stream.hs b/jetstream/JetStream/Stream.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Stream.hs @@ -0,0 +1,79 @@ +module JetStream.Stream + ( streamAPI + , module JetStream.Stream.API + ) where + +import Data.Aeson (toJSON) +import JetStream.Error + ( JetStreamApiError (..) + , JetStreamError (..) + ) +import JetStream.Options (JetStreamContext) +import qualified JetStream.Protocol.Request as Request +import qualified JetStream.Protocol.Subject as Subject +import JetStream.Stream.API (StreamAPI (..)) +import JetStream.Stream.Types + ( purgeStreamRequest + , streamConfigRequest + , streamListRequest + , streamMessageDeleteRequest + , streamMessageGetRequest + , streamNamesRequest + ) + +streamAPI :: JetStreamContext -> StreamAPI +streamAPI context = + StreamAPI + { create = \name subjects options -> + let config = streamConfigRequest name subjects options + in Request.requestJSON context + (Subject.streamCreateSubject context name) + (Just (toJSON config)) + , createOrUpdate = \name subjects options -> do + updateResult <- update (streamAPI context) name subjects options + case updateResult of + Left err + | isStreamNotFound err -> + create (streamAPI context) name subjects options + other -> + pure other + , update = \name subjects options -> + let config = streamConfigRequest name subjects options + in Request.requestJSON context + (Subject.streamUpdateSubject context name) + (Just (toJSON config)) + , info = \streamName -> + Request.requestJSON context + (Subject.streamInfoSubject context streamName) + Nothing + , getMessage = \streamName selector -> + Request.requestJSON context + (Subject.streamMessageGetSubject context streamName) + (Just (toJSON (streamMessageGetRequest selector))) + , deleteMessage = \streamName sequenceNumber mode -> + Request.requestJSON context + (Subject.streamMessageDeleteSubject context streamName) + (Just (toJSON (streamMessageDeleteRequest sequenceNumber mode))) + , delete = \streamName -> + Request.requestJSON context + (Subject.streamDeleteSubject context streamName) + Nothing + , purge = \streamName options -> + Request.requestJSON context + (Subject.streamPurgeSubject context streamName) + (Just (toJSON (purgeStreamRequest options))) + , list = + Request.requestJSON context + (Subject.streamListSubject context) . + Just . toJSON . streamListRequest + , names = + Request.requestJSON context + (Subject.streamNamesSubject context) . + Just . toJSON . streamNamesRequest + } + +isStreamNotFound :: JetStreamError -> Bool +isStreamNotFound (JetStreamApiFailure err) = + apiErrorCodeDetail err == 10059 +isStreamNotFound _ = + False diff --git a/jetstream/JetStream/Stream/API.hs b/jetstream/JetStream/Stream/API.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Stream/API.hs @@ -0,0 +1,65 @@ +module JetStream.Stream.API + ( StreamAPI (..) + , RetentionPolicy (..) + , StorageType (..) + , DiscardPolicy (..) + , StreamConfig (..) + , StreamConfigOption + , withRetention + , withStorage + , withDiscard + , withMaxMessages + , withMaxBytes + , withMaxAge + , withReplicas + , withDuplicateWindow + , withAllowDirect + , PurgeStreamOption + , withPurgeSubject + , withPurgeSequence + , withPurgeKeep + , StreamListOption + , withStreamListOffset + , withStreamListSubject + , StreamMessage (..) + , StreamMessageSelector (..) + , StreamMessageDeleteMode (..) + , DeleteStreamMessageResponse (..) + , StreamInfo (..) + , StreamState (..) + , StreamCluster (..) + , StreamPeer (..) + , StreamSourceInfo (..) + , DeleteStreamResponse (..) + , PurgeStreamResponse (..) + , StreamListResponse (..) + , StreamNamesResponse (..) + ) where + +import JetStream.Error (JetStreamError) +import JetStream.Stream.Types +import JetStream.Types (StreamName, Subject) + +-- | Stream management operations. +data StreamAPI = StreamAPI + { create :: StreamName -> [Subject] -> [StreamConfigOption] -> IO (Either JetStreamError StreamInfo) + -- ^ Create a stream. + , createOrUpdate :: StreamName -> [Subject] -> [StreamConfigOption] -> IO (Either JetStreamError StreamInfo) + -- ^ Create a stream or update it if it already exists. + , update :: StreamName -> [Subject] -> [StreamConfigOption] -> IO (Either JetStreamError StreamInfo) + -- ^ Update an existing stream configuration. + , info :: StreamName -> IO (Either JetStreamError StreamInfo) + -- ^ Fetch stream configuration and state. + , getMessage :: StreamName -> StreamMessageSelector -> IO (Either JetStreamError StreamMessage) + -- ^ Fetch a stored stream message by selector. + , deleteMessage :: StreamName -> Integer -> StreamMessageDeleteMode -> IO (Either JetStreamError DeleteStreamMessageResponse) + -- ^ Delete or securely erase a stored stream message. + , delete :: StreamName -> IO (Either JetStreamError DeleteStreamResponse) + -- ^ Delete a stream. + , purge :: StreamName -> [PurgeStreamOption] -> IO (Either JetStreamError PurgeStreamResponse) + -- ^ Purge messages from a stream. + , list :: [StreamListOption] -> IO (Either JetStreamError StreamListResponse) + -- ^ List streams and their metadata. + , names :: [StreamListOption] -> IO (Either JetStreamError StreamNamesResponse) + -- ^ List stream names. + } diff --git a/jetstream/JetStream/Stream/Types.hs b/jetstream/JetStream/Stream/Types.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Stream/Types.hs @@ -0,0 +1,676 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Stream.Types + ( RetentionPolicy (..) + , StorageType (..) + , DiscardPolicy (..) + , StreamConfig (..) + , StreamConfigOption + , streamConfigRequest + , withRetention + , withStorage + , withDiscard + , withMaxMessages + , withMaxBytes + , withMaxAge + , withReplicas + , withDuplicateWindow + , withAllowDirect + , PurgeStreamOption + , purgeStreamRequest + , withPurgeSubject + , withPurgeSequence + , withPurgeKeep + , StreamListOption + , streamListRequest + , streamNamesRequest + , withStreamListOffset + , withStreamListSubject + , StreamMessage (..) + , StreamMessageSelector (..) + , streamMessageGetRequest + , StreamMessageDeleteMode (..) + , streamMessageDeleteRequest + , DeleteStreamMessageResponse (..) + , StreamInfo (..) + , StreamState (..) + , StreamCluster (..) + , StreamPeer (..) + , StreamSourceInfo (..) + , DeleteStreamResponse (..) + , PurgeStreamResponse (..) + , StreamListResponse (..) + , StreamNamesResponse (..) + , durationToNanoseconds + , nanosecondsToDuration + ) where + +import Data.Aeson +import Data.Aeson.Types (Pair, Parser) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Base64 as Base64 +import Data.Maybe (catMaybes) +import Data.Time.Clock (NominalDiffTime, UTCTime) +import JetStream.Types + ( CallOption + , DiscardPolicy (..) + , Payload + , RetentionPolicy (..) + , StorageType (..) + , StreamName + , Subject + , applyCallOptions + , byteStringToJSON + , parseByteString + ) + +data StreamConfigRequest = StreamConfigRequest + { streamConfigRequestName :: StreamName + , streamConfigRequestSubjects :: [Subject] + , streamConfigRequestRetention :: Maybe RetentionPolicy + , streamConfigRequestStorage :: Maybe StorageType + , streamConfigRequestDiscard :: Maybe DiscardPolicy + , streamConfigRequestMaxMessages :: Maybe Integer + , streamConfigRequestMaxBytes :: Maybe Integer + , streamConfigRequestMaxAge :: Maybe NominalDiffTime + , streamConfigRequestReplicas :: Maybe Int + , streamConfigRequestDuplicateWindow :: Maybe NominalDiffTime + , streamConfigRequestAllowDirect :: Maybe Bool + } + deriving (Eq, Show) + +type StreamConfigOption = CallOption StreamConfigRequest + +streamConfigRequest :: StreamName -> [Subject] -> [StreamConfigOption] -> StreamConfigRequest +streamConfigRequest name subjects options = + applyCallOptions options $ + StreamConfigRequest + { streamConfigRequestName = name + , streamConfigRequestSubjects = subjects + , streamConfigRequestRetention = Nothing + , streamConfigRequestStorage = Nothing + , streamConfigRequestDiscard = Nothing + , streamConfigRequestMaxMessages = Nothing + , streamConfigRequestMaxBytes = Nothing + , streamConfigRequestMaxAge = Nothing + , streamConfigRequestReplicas = Nothing + , streamConfigRequestDuplicateWindow = Nothing + , streamConfigRequestAllowDirect = Nothing + } + +withRetention :: RetentionPolicy -> StreamConfigOption +withRetention retention config = + config { streamConfigRequestRetention = Just retention } + +withStorage :: StorageType -> StreamConfigOption +withStorage storage config = + config { streamConfigRequestStorage = Just storage } + +withDiscard :: DiscardPolicy -> StreamConfigOption +withDiscard discard config = + config { streamConfigRequestDiscard = Just discard } + +withMaxMessages :: Integer -> StreamConfigOption +withMaxMessages maxMessages config = + config { streamConfigRequestMaxMessages = Just maxMessages } + +withMaxBytes :: Integer -> StreamConfigOption +withMaxBytes maxBytes config = + config { streamConfigRequestMaxBytes = Just maxBytes } + +withMaxAge :: NominalDiffTime -> StreamConfigOption +withMaxAge maxAge config = + config { streamConfigRequestMaxAge = Just maxAge } + +withReplicas :: Int -> StreamConfigOption +withReplicas replicas config = + config { streamConfigRequestReplicas = Just replicas } + +withDuplicateWindow :: NominalDiffTime -> StreamConfigOption +withDuplicateWindow window config = + config { streamConfigRequestDuplicateWindow = Just window } + +withAllowDirect :: Bool -> StreamConfigOption +withAllowDirect allowDirect config = + config { streamConfigRequestAllowDirect = Just allowDirect } + +data StreamConfig = StreamConfig + { streamConfigName :: StreamName + , streamConfigSubjects :: Maybe [Subject] + , streamConfigRetention :: RetentionPolicy + , streamConfigStorage :: StorageType + , streamConfigDiscard :: DiscardPolicy + , streamConfigMaxMessages :: Integer + , streamConfigMaxBytes :: Integer + , streamConfigMaxAge :: NominalDiffTime + , streamConfigReplicas :: Int + , streamConfigDuplicateWindow :: Maybe NominalDiffTime + , streamConfigAllowDirect :: Bool + } + deriving (Eq, Show) + +data PurgeStreamRequest = PurgeStreamRequest + { purgeStreamSubject :: Maybe Subject + , purgeStreamSequence :: Maybe Integer + , purgeStreamKeep :: Maybe Integer + } + deriving (Eq, Show) + +type PurgeStreamOption = CallOption PurgeStreamRequest + +purgeStreamRequest :: [PurgeStreamOption] -> PurgeStreamRequest +purgeStreamRequest options = + applyCallOptions options $ + PurgeStreamRequest + { purgeStreamSubject = Nothing + , purgeStreamSequence = Nothing + , purgeStreamKeep = Nothing + } + +withPurgeSubject :: Subject -> PurgeStreamOption +withPurgeSubject subject request = + request { purgeStreamSubject = Just subject } + +withPurgeSequence :: Integer -> PurgeStreamOption +withPurgeSequence sequenceNumber request = + request { purgeStreamSequence = Just sequenceNumber } + +withPurgeKeep :: Integer -> PurgeStreamOption +withPurgeKeep keep request = + request { purgeStreamKeep = Just keep } + +data StreamListRequest = StreamListRequest + { streamListRequestOffset :: Maybe Int + , streamListRequestSubject :: Maybe Subject + } + deriving (Eq, Show) + +type StreamListOption = CallOption StreamListRequest + +streamListRequest :: [StreamListOption] -> StreamListRequest +streamListRequest options = + applyCallOptions options defaultStreamListRequest + +streamNamesRequest :: [StreamListOption] -> StreamListRequest +streamNamesRequest = + streamListRequest + +defaultStreamListRequest :: StreamListRequest +defaultStreamListRequest = + StreamListRequest + { streamListRequestOffset = Nothing + , streamListRequestSubject = Nothing + } + +withStreamListOffset :: Int -> StreamListOption +withStreamListOffset offset request = + request { streamListRequestOffset = Just offset } + +withStreamListSubject :: Subject -> StreamListOption +withStreamListSubject subject request = + request { streamListRequestSubject = Just subject } + +data StreamInfo = StreamInfo + { streamInfoConfig :: StreamConfig + , streamInfoCreated :: UTCTime + , streamInfoState :: StreamState + , streamInfoCluster :: Maybe StreamCluster + , streamInfoMirror :: Maybe StreamSourceInfo + , streamInfoSources :: [StreamSourceInfo] + } + deriving (Eq, Show) + +data StreamState = StreamState + { streamStateMessages :: Integer + , streamStateBytes :: Integer + , streamStateFirstSequence :: Integer + , streamStateFirstTime :: UTCTime + , streamStateLastSequence :: Integer + , streamStateLastTime :: UTCTime + , streamStateConsumerCount :: Int + , streamStateDeleted :: [Integer] + , streamStateNumDeleted :: Integer + , streamStateNumSubjects :: Integer + } + deriving (Eq, Show) + +data StreamCluster = StreamCluster + { streamClusterName :: BS.ByteString + , streamClusterLeader :: BS.ByteString + , streamClusterReplicas :: [StreamPeer] + } + deriving (Eq, Show) + +data StreamPeer = StreamPeer + { streamPeerName :: BS.ByteString + , streamPeerCurrent :: Bool + , streamPeerOffline :: Bool + , streamPeerActive :: NominalDiffTime + , streamPeerLag :: Integer + } + deriving (Eq, Show) + +data StreamSourceInfo = StreamSourceInfo + { streamSourceInfoName :: StreamName + , streamSourceInfoFilterSubject :: Subject + , streamSourceInfoLag :: Integer + , streamSourceInfoActive :: NominalDiffTime + } + deriving (Eq, Show) + +newtype DeleteStreamResponse = DeleteStreamResponse { deleteStreamSuccess :: Bool } + deriving (Eq, Show) + +data PurgeStreamResponse = PurgeStreamResponse + { purgeStreamSuccess :: Bool + , purgeStreamPurged :: Integer + } + deriving (Eq, Show) + +data StreamListResponse = StreamListResponse + { streamListTotal :: Int + , streamListOffset :: Int + , streamListLimit :: Int + , streamListStreams :: [StreamInfo] + } + deriving (Eq, Show) + +data StreamNamesResponse = StreamNamesResponse + { streamNamesTotal :: Int + , streamNamesOffset :: Int + , streamNamesLimit :: Int + , streamNamesStreams :: [StreamName] + } + deriving (Eq, Show) + +data StreamMessage = StreamMessage + { streamMessageSubject :: Subject + , streamMessageSequence :: Integer + , streamMessageHeadersRaw :: Maybe BS.ByteString + , streamMessagePayload :: Maybe Payload + , streamMessageTime :: UTCTime + } + deriving (Eq, Show) + +data StreamMessageSelector = StreamMessageBySequence Integer + | LastStreamMessageForSubject Subject + | NextStreamMessageForSubject Subject + deriving (Eq, Show) + +newtype StreamMessageGetRequest = StreamMessageGetRequest { streamMessageGetSelector :: StreamMessageSelector } + deriving (Eq, Show) + +streamMessageGetRequest :: StreamMessageSelector -> StreamMessageGetRequest +streamMessageGetRequest = + StreamMessageGetRequest + +data StreamMessageDeleteMode = DeleteMessage | SecureDeleteMessage + deriving (Eq, Show) + +data StreamMessageDeleteRequest = StreamMessageDeleteRequest + { streamMessageDeleteSequence :: Integer + , streamMessageDeleteNoErase :: Maybe Bool + } + deriving (Eq, Show) + +streamMessageDeleteRequest :: Integer -> StreamMessageDeleteMode -> StreamMessageDeleteRequest +streamMessageDeleteRequest sequenceNumber mode = + StreamMessageDeleteRequest + { streamMessageDeleteSequence = sequenceNumber + , streamMessageDeleteNoErase = + case mode of + DeleteMessage -> Just True + SecureDeleteMessage -> Nothing + } + +newtype DeleteStreamMessageResponse = DeleteStreamMessageResponse { deleteStreamMessageSuccess :: Bool } + deriving (Eq, Show) + +instance ToJSON StreamConfigRequest where + toJSON config = + object $ + [ byteStringPair "name" (streamConfigRequestName config) + , byteStringListPair "subjects" (streamConfigRequestSubjects config) + ] ++ catMaybes + [ maybePair "retention" (streamConfigRequestRetention config) + , maybePair "storage" (streamConfigRequestStorage config) + , maybePair "discard" (streamConfigRequestDiscard config) + , maybePair "max_msgs" (streamConfigRequestMaxMessages config) + , maybePair "max_bytes" (streamConfigRequestMaxBytes config) + , maybeDurationPair "max_age" (streamConfigRequestMaxAge config) + , maybePair "num_replicas" (streamConfigRequestReplicas config) + , maybeDurationPair "duplicate_window" (streamConfigRequestDuplicateWindow config) + , maybePair "allow_direct" (streamConfigRequestAllowDirect config) + ] + +instance ToJSON StreamConfig where + toJSON config = + object . catMaybes $ + [ Just (byteStringPair "name" (streamConfigName config)) + , maybeByteStringListPair "subjects" (streamConfigSubjects config) + , Just ("retention" .= streamConfigRetention config) + , Just ("storage" .= streamConfigStorage config) + , Just ("discard" .= streamConfigDiscard config) + , Just ("max_msgs" .= streamConfigMaxMessages config) + , Just ("max_bytes" .= streamConfigMaxBytes config) + , Just ("max_age" .= durationToNanoseconds (streamConfigMaxAge config)) + , Just ("num_replicas" .= streamConfigReplicas config) + , maybeDurationPair "duplicate_window" (streamConfigDuplicateWindow config) + , Just ("allow_direct" .= streamConfigAllowDirect config) + ] + +instance FromJSON StreamConfig where + parseJSON = + withObject "StreamConfig" $ \value -> + StreamConfig + <$> parseByteStringField value "name" + <*> parseOptionalByteStringListField value "subjects" + <*> value .: "retention" + <*> value .: "storage" + <*> value .: "discard" + <*> value .: "max_msgs" + <*> value .: "max_bytes" + <*> parseDurationField value "max_age" + <*> value .: "num_replicas" + <*> parseOptionalDurationField value "duplicate_window" + <*> value .: "allow_direct" + +instance ToJSON PurgeStreamRequest where + toJSON request = + object . catMaybes $ + [ maybeByteStringPair "filter" (purgeStreamSubject request) + , maybePair "seq" (purgeStreamSequence request) + , maybePair "keep" (purgeStreamKeep request) + ] + +instance ToJSON StreamListRequest where + toJSON request = + object . catMaybes $ + [ maybePair "offset" (streamListRequestOffset request) + , maybeByteStringPair "subject" (streamListRequestSubject request) + ] + +instance ToJSON StreamMessageGetRequest where + toJSON request = + case streamMessageGetSelector request of + StreamMessageBySequence sequenceNumber -> + object ["seq" .= sequenceNumber] + LastStreamMessageForSubject subject -> + object [byteStringPair "last_by_subj" subject] + NextStreamMessageForSubject subject -> + object [byteStringPair "next_by_subj" subject] + +instance ToJSON StreamMessageDeleteRequest where + toJSON request = + object . catMaybes $ + [ Just ("seq" .= streamMessageDeleteSequence request) + , maybePair "no_erase" (streamMessageDeleteNoErase request) + ] + +instance ToJSON StreamInfo where + toJSON info = + object . catMaybes $ + [ Just ("config" .= streamInfoConfig info) + , Just ("created" .= streamInfoCreated info) + , Just ("state" .= streamInfoState info) + , maybePair "cluster" (streamInfoCluster info) + , maybePair "mirror" (streamInfoMirror info) + , Just ("sources" .= streamInfoSources info) + ] + +instance FromJSON StreamInfo where + parseJSON = + withObject "StreamInfo" $ \value -> + StreamInfo + <$> value .: "config" + <*> value .: "created" + <*> value .: "state" + <*> value .:? "cluster" + <*> value .:? "mirror" + <*> value .:? "sources" .!= [] + +instance ToJSON StreamState where + toJSON state = + object + [ "messages" .= streamStateMessages state + , "bytes" .= streamStateBytes state + , "first_seq" .= streamStateFirstSequence state + , "first_ts" .= streamStateFirstTime state + , "last_seq" .= streamStateLastSequence state + , "last_ts" .= streamStateLastTime state + , "consumer_count" .= streamStateConsumerCount state + , "deleted" .= streamStateDeleted state + , "num_deleted" .= streamStateNumDeleted state + , "num_subjects" .= streamStateNumSubjects state + ] + +instance FromJSON StreamState where + parseJSON = + withObject "StreamState" $ \value -> + StreamState + <$> value .: "messages" + <*> value .: "bytes" + <*> value .: "first_seq" + <*> value .: "first_ts" + <*> value .: "last_seq" + <*> value .: "last_ts" + <*> value .: "consumer_count" + <*> value .:? "deleted" .!= [] + <*> value .:? "num_deleted" .!= 0 + <*> value .:? "num_subjects" .!= 0 + +instance ToJSON StreamCluster where + toJSON cluster = + object + [ "name" .= byteStringToJSON (streamClusterName cluster) + , "leader" .= byteStringToJSON (streamClusterLeader cluster) + , "replicas" .= streamClusterReplicas cluster + ] + +instance FromJSON StreamCluster where + parseJSON = + withObject "StreamCluster" $ \value -> + StreamCluster + <$> parseOptionalByteStringField value "name" .!= "" + <*> parseOptionalByteStringField value "leader" .!= "" + <*> value .:? "replicas" .!= [] + +instance ToJSON StreamPeer where + toJSON peer = + object + [ "name" .= byteStringToJSON (streamPeerName peer) + , "current" .= streamPeerCurrent peer + , "offline" .= streamPeerOffline peer + , "active" .= durationToNanoseconds (streamPeerActive peer) + , "lag" .= streamPeerLag peer + ] + +instance FromJSON StreamPeer where + parseJSON = + withObject "StreamPeer" $ \value -> + StreamPeer + <$> parseByteStringField value "name" + <*> value .: "current" + <*> value .:? "offline" .!= False + <*> parseDurationField value "active" + <*> value .:? "lag" .!= 0 + +instance ToJSON StreamSourceInfo where + toJSON source = + object + [ "name" .= byteStringToJSON (streamSourceInfoName source) + , "filter_subject" .= byteStringToJSON (streamSourceInfoFilterSubject source) + , "lag" .= streamSourceInfoLag source + , "active" .= durationToNanoseconds (streamSourceInfoActive source) + ] + +instance FromJSON StreamSourceInfo where + parseJSON = + withObject "StreamSourceInfo" $ \value -> + StreamSourceInfo + <$> parseByteStringField value "name" + <*> parseOptionalByteStringField value "filter_subject" .!= "" + <*> value .: "lag" + <*> parseDurationField value "active" + +instance ToJSON DeleteStreamResponse where + toJSON response = + object + [ "success" .= deleteStreamSuccess response + ] + +instance FromJSON DeleteStreamResponse where + parseJSON = + withObject "DeleteStreamResponse" $ \value -> + DeleteStreamResponse + <$> value .:? "success" .!= False + +instance ToJSON PurgeStreamResponse where + toJSON response = + object + [ "success" .= purgeStreamSuccess response + , "purged" .= purgeStreamPurged response + ] + +instance FromJSON PurgeStreamResponse where + parseJSON = + withObject "PurgeStreamResponse" $ \value -> + PurgeStreamResponse + <$> value .:? "success" .!= False + <*> value .:? "purged" .!= 0 + +instance ToJSON StreamListResponse where + toJSON response = + object + [ "total" .= streamListTotal response + , "offset" .= streamListOffset response + , "limit" .= streamListLimit response + , "streams" .= streamListStreams response + ] + +instance FromJSON StreamListResponse where + parseJSON = + withObject "StreamListResponse" $ \value -> + StreamListResponse + <$> value .: "total" + <*> value .: "offset" + <*> value .: "limit" + <*> value .:? "streams" .!= [] + +instance ToJSON StreamNamesResponse where + toJSON response = + object + [ "total" .= streamNamesTotal response + , "offset" .= streamNamesOffset response + , "limit" .= streamNamesLimit response + , "streams" .= map byteStringToJSON (streamNamesStreams response) + ] + +instance FromJSON StreamNamesResponse where + parseJSON = + withObject "StreamNamesResponse" $ \value -> + StreamNamesResponse + <$> value .: "total" + <*> value .: "offset" + <*> value .: "limit" + <*> parseOptionalByteStringListField value "streams" .!= [] + +instance FromJSON StreamMessage where + parseJSON = + withObject "StreamMessageResponse" $ \value -> + value .: "message" >>= parseStoredMessage + +instance ToJSON StreamMessage where + toJSON message = + object . catMaybes $ + [ Just (byteStringPair "subject" (streamMessageSubject message)) + , Just ("seq" .= streamMessageSequence message) + , maybeBase64Pair "hdrs" (streamMessageHeadersRaw message) + , maybeBase64Pair "data" (streamMessagePayload message) + , Just ("time" .= streamMessageTime message) + ] + +instance FromJSON DeleteStreamMessageResponse where + parseJSON = + withObject "DeleteStreamMessageResponse" $ \value -> + DeleteStreamMessageResponse + <$> value .:? "success" .!= False + +instance ToJSON DeleteStreamMessageResponse where + toJSON response = + object + [ "success" .= deleteStreamMessageSuccess response + ] + +parseStoredMessage :: Value -> Parser StreamMessage +parseStoredMessage = + withObject "StreamMessage" $ \value -> + StreamMessage + <$> parseByteStringField value "subject" + <*> value .: "seq" + <*> parseOptionalBase64Field value "hdrs" + <*> parseOptionalBase64Field value "data" + <*> value .: "time" + +durationToNanoseconds :: NominalDiffTime -> Integer +durationToNanoseconds duration = + round (realToFrac duration * (1000000000 :: Double)) + +nanosecondsToDuration :: Integer -> NominalDiffTime +nanosecondsToDuration nanoseconds = + fromRational (toRational nanoseconds / 1000000000) + +byteStringPair :: Key -> BS.ByteString -> Pair +byteStringPair key value = key .= byteStringToJSON value + +byteStringListPair :: Key -> [BS.ByteString] -> Pair +byteStringListPair key = (key .=) . map byteStringToJSON + +maybeByteStringPair :: Key -> Maybe BS.ByteString -> Maybe Pair +maybeByteStringPair key = fmap (byteStringPair key) + +maybeBase64Pair :: Key -> Maybe BS.ByteString -> Maybe Pair +maybeBase64Pair key = + fmap ((key .=) . byteStringToJSON . Base64.encode) + +maybeByteStringListPair :: Key -> Maybe [BS.ByteString] -> Maybe Pair +maybeByteStringListPair key = fmap (byteStringListPair key) + +maybePair :: ToJSON value => Key -> Maybe value -> Maybe Pair +maybePair key = fmap (key .=) + +maybeDurationPair :: Key -> Maybe NominalDiffTime -> Maybe Pair +maybeDurationPair key = fmap ((key .=) . durationToNanoseconds) + +parseByteStringField :: Object -> Key -> Parser BS.ByteString +parseByteStringField value key = + value .: key >>= parseByteString + +parseOptionalByteStringField :: Object -> Key -> Parser (Maybe BS.ByteString) +parseOptionalByteStringField value key = do + byteString <- value .:? key + traverse parseByteString byteString + +parseOptionalBase64Field :: Object -> Key -> Parser (Maybe BS.ByteString) +parseOptionalBase64Field value key = do + encoded <- value .:? key + traverse parseBase64 encoded + where + parseBase64 jsonValue = do + bytes <- parseByteString jsonValue + case Base64.decode bytes of + Left err -> fail err + Right decoded -> pure decoded + +parseOptionalByteStringListField :: Object -> Key -> Parser (Maybe [BS.ByteString]) +parseOptionalByteStringListField value key = do + byteStrings <- value .:? key + traverse (traverse parseByteString) byteStrings + +parseDurationField :: Object -> Key -> Parser NominalDiffTime +parseDurationField value key = + nanosecondsToDuration <$> value .: key + +parseOptionalDurationField :: Object -> Key -> Parser (Maybe NominalDiffTime) +parseOptionalDurationField value key = + fmap nanosecondsToDuration <$> value .:? key diff --git a/jetstream/JetStream/Types.hs b/jetstream/JetStream/Types.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Types.hs @@ -0,0 +1,341 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.Types + ( AccountAPIStats (..) + , AccountInfo (..) + , AccountLimits (..) + , AccountTier (..) + , StreamName + , ConsumerName + , Subject + , Payload + , Headers + , CallOption + , AckPolicy (..) + , DeliverPolicy (..) + , DiscardPolicy (..) + , ReplayPolicy (..) + , RetentionPolicy (..) + , StorageType (..) + , byteStringToJSON + , parseByteString + , diffTimeNanosToJSON + , applyCallOptions + ) where + +import Data.Aeson +import qualified Data.Aeson.Key as Key +import qualified Data.Aeson.KeyMap as KeyMap +import Data.Aeson.Types (Pair, Parser, typeMismatch) +import qualified Data.ByteString as BS +import Data.Maybe (catMaybes) +import qualified Data.Text as T +import qualified Data.Text.Encoding as E +import Data.Time.Clock (NominalDiffTime) +import qualified Data.Time.Clock as Time + +type StreamName = BS.ByteString +type ConsumerName = BS.ByteString +type Subject = BS.ByteString +type Payload = BS.ByteString +type Headers = [(BS.ByteString, BS.ByteString)] +type CallOption a = a -> a + +data AccountInfo = AccountInfo + { accountInfoTier :: AccountTier + , accountInfoDomain :: Maybe BS.ByteString + , accountInfoAPI :: AccountAPIStats + , accountInfoTiers :: [(BS.ByteString, AccountTier)] + } + deriving (Eq, Show) + +data AccountTier = AccountTier + { accountTierMemory :: Integer + , accountTierStorage :: Integer + , accountTierReservedMemory :: Integer + , accountTierReservedStorage :: Integer + , accountTierStreams :: Int + , accountTierConsumers :: Int + , accountTierLimits :: AccountLimits + } + deriving (Eq, Show) + +data AccountAPIStats = AccountAPIStats + { accountAPILevel :: Int + , accountAPITotal :: Integer + , accountAPIErrors :: Integer + , accountAPIInflight :: Maybe Integer + } + deriving (Eq, Show) + +data AccountLimits = AccountLimits + { accountLimitsMaxMemory :: Integer + , accountLimitsMaxStorage :: Integer + , accountLimitsMaxStreams :: Int + , accountLimitsMaxConsumers :: Int + , accountLimitsMaxAckPending :: Int + , accountLimitsMemoryMaxStreamBytes :: Integer + , accountLimitsStorageMaxStreamBytes :: Integer + , accountLimitsMaxBytesRequired :: Bool + } + deriving (Eq, Show) + +applyCallOptions :: [CallOption a] -> a -> a +applyCallOptions options value = + foldr ($) value options + +data AckPolicy = AckNone | AckAll | AckExplicit + deriving (Eq, Show) + +data DeliverPolicy = DeliverAll + | DeliverLast + | DeliverNew + | DeliverByStartSequence Int + | DeliverByStartTime Time.UTCTime + | DeliverLastPerSubject + deriving (Eq, Show) + +data DiscardPolicy = DiscardOld | DiscardNew + deriving (Eq, Show) + +data ReplayPolicy = ReplayInstant | ReplayOriginal + deriving (Eq, Show) + +data RetentionPolicy = LimitsPolicy | InterestPolicy | WorkQueuePolicy + deriving (Eq, Show) + +data StorageType = FileStorage | MemoryStorage + deriving (Eq, Show) + +instance ToJSON AckPolicy where + toJSON policy = + String $ + case policy of + AckNone -> "none" + AckAll -> "all" + AckExplicit -> "explicit" + +instance FromJSON AckPolicy where + parseJSON = withText "AckPolicy" $ \value -> + case value of + "none" -> pure AckNone + "all" -> pure AckAll + "explicit" -> pure AckExplicit + _ -> fail ("unknown ack policy: " ++ T.unpack value) + +instance ToJSON DeliverPolicy where + toJSON policy = + String $ + case policy of + DeliverAll -> "all" + DeliverLast -> "last" + DeliverNew -> "new" + DeliverByStartSequence _ -> "by_start_sequence" + DeliverByStartTime _ -> "by_start_time" + DeliverLastPerSubject -> "last_per_subject" + +instance FromJSON DeliverPolicy where + parseJSON = withText "DeliverPolicy" $ \value -> + case value of + "all" -> pure DeliverAll + "last" -> pure DeliverLast + "new" -> pure DeliverNew + "by_start_sequence" -> fail "deliver policy by_start_sequence requires opt_start_seq" + "by_start_time" -> fail "deliver policy by_start_time requires opt_start_time" + "last_per_subject" -> pure DeliverLastPerSubject + _ -> fail ("unknown deliver policy: " ++ T.unpack value) + +instance ToJSON DiscardPolicy where + toJSON policy = + String $ + case policy of + DiscardOld -> "old" + DiscardNew -> "new" + +instance FromJSON DiscardPolicy where + parseJSON = withText "DiscardPolicy" $ \value -> + case value of + "old" -> pure DiscardOld + "new" -> pure DiscardNew + _ -> fail ("unknown discard policy: " ++ T.unpack value) + +instance ToJSON ReplayPolicy where + toJSON policy = + String $ + case policy of + ReplayInstant -> "instant" + ReplayOriginal -> "original" + +instance FromJSON ReplayPolicy where + parseJSON = withText "ReplayPolicy" $ \value -> + case value of + "instant" -> pure ReplayInstant + "original" -> pure ReplayOriginal + _ -> fail ("unknown replay policy: " ++ T.unpack value) + +instance ToJSON RetentionPolicy where + toJSON policy = + String $ + case policy of + LimitsPolicy -> "limits" + InterestPolicy -> "interest" + WorkQueuePolicy -> "workqueue" + +instance FromJSON RetentionPolicy where + parseJSON = withText "RetentionPolicy" $ \value -> + case value of + "limits" -> pure LimitsPolicy + "interest" -> pure InterestPolicy + "workqueue" -> pure WorkQueuePolicy + _ -> fail ("unknown retention policy: " ++ T.unpack value) + +instance ToJSON StorageType where + toJSON storage = + String $ + case storage of + FileStorage -> "file" + MemoryStorage -> "memory" + +instance FromJSON StorageType where + parseJSON = withText "StorageType" $ \value -> + case value of + "file" -> pure FileStorage + "memory" -> pure MemoryStorage + _ -> fail ("unknown storage type: " ++ T.unpack value) + +byteStringToJSON :: BS.ByteString -> Value +byteStringToJSON = + String . E.decodeUtf8 + +parseByteString :: Value -> Parser BS.ByteString +parseByteString = + withText "ByteString" (pure . E.encodeUtf8) + +diffTimeNanosToJSON :: NominalDiffTime -> Value +diffTimeNanosToJSON diffTime = + Number (fromInteger (floor (realToFrac diffTime * (1000000000 :: Double)))) + +instance FromJSON AccountInfo where + parseJSON value@(Object obj) = + AccountInfo + <$> parseJSON value + <*> parseOptionalByteStringField obj "domain" + <*> obj .: "api" + <*> parseAccountTiers obj + parseJSON invalid = + typeMismatch "AccountInfo" invalid + +instance ToJSON AccountInfo where + toJSON info = + object $ + accountTierPairs (accountInfoTier info) + ++ catMaybes + [ maybeByteStringPair "domain" (accountInfoDomain info) + , Just ("api" .= accountInfoAPI info) + , Just ("tiers" .= accountTiersValue (accountInfoTiers info)) + ] + +instance FromJSON AccountTier where + parseJSON = + withObject "AccountTier" $ \obj -> + AccountTier + <$> obj .: "memory" + <*> obj .: "storage" + <*> obj .: "reserved_memory" + <*> obj .: "reserved_storage" + <*> obj .: "streams" + <*> obj .: "consumers" + <*> obj .: "limits" + +instance ToJSON AccountTier where + toJSON = + object . accountTierPairs + +instance FromJSON AccountAPIStats where + parseJSON = + withObject "AccountAPIStats" $ \obj -> + AccountAPIStats + <$> obj .: "level" + <*> obj .: "total" + <*> obj .: "errors" + <*> obj .:? "inflight" + +instance ToJSON AccountAPIStats where + toJSON stats = + object . catMaybes $ + [ Just ("level" .= accountAPILevel stats) + , Just ("total" .= accountAPITotal stats) + , Just ("errors" .= accountAPIErrors stats) + , maybePair "inflight" (accountAPIInflight stats) + ] + +instance FromJSON AccountLimits where + parseJSON = + withObject "AccountLimits" $ \obj -> + AccountLimits + <$> obj .: "max_memory" + <*> obj .: "max_storage" + <*> obj .: "max_streams" + <*> obj .: "max_consumers" + <*> obj .: "max_ack_pending" + <*> obj .: "memory_max_stream_bytes" + <*> obj .: "storage_max_stream_bytes" + <*> obj .: "max_bytes_required" + +instance ToJSON AccountLimits where + toJSON limits = + object + [ "max_memory" .= accountLimitsMaxMemory limits + , "max_storage" .= accountLimitsMaxStorage limits + , "max_streams" .= accountLimitsMaxStreams limits + , "max_consumers" .= accountLimitsMaxConsumers limits + , "max_ack_pending" .= accountLimitsMaxAckPending limits + , "memory_max_stream_bytes" .= accountLimitsMemoryMaxStreamBytes limits + , "storage_max_stream_bytes" .= accountLimitsStorageMaxStreamBytes limits + , "max_bytes_required" .= accountLimitsMaxBytesRequired limits + ] + +accountTierPairs :: AccountTier -> [Pair] +accountTierPairs tier = + [ "memory" .= accountTierMemory tier + , "storage" .= accountTierStorage tier + , "reserved_memory" .= accountTierReservedMemory tier + , "reserved_storage" .= accountTierReservedStorage tier + , "streams" .= accountTierStreams tier + , "consumers" .= accountTierConsumers tier + , "limits" .= accountTierLimits tier + ] + +parseAccountTiers :: Object -> Parser [(BS.ByteString, AccountTier)] +parseAccountTiers obj = do + tiers <- obj .:? "tiers" + case tiers of + Nothing -> + pure [] + Just (Object tierMap) -> + traverse parseTier (KeyMap.toList tierMap) + Just invalid -> + typeMismatch "AccountInfo.tiers" invalid + where + parseTier (name, value) = + (,) (E.encodeUtf8 (Key.toText name)) <$> parseJSON value + +accountTiersValue :: [(BS.ByteString, AccountTier)] -> Value +accountTiersValue = + Object . KeyMap.fromList . fmap tierPair + where + tierPair (name, tier) = + (Key.fromText (E.decodeUtf8 name), toJSON tier) + +parseOptionalByteStringField :: Object -> Key.Key -> Parser (Maybe BS.ByteString) +parseOptionalByteStringField obj key = do + value <- obj .:? key + traverse parseByteString value + +maybeByteStringPair :: Key.Key -> Maybe BS.ByteString -> Maybe Pair +maybeByteStringPair key = + fmap ((key .=) . byteStringToJSON) + +maybePair :: ToJSON value => Key.Key -> Maybe value -> Maybe Pair +maybePair key = + fmap (key .=) 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: 0.1.0.0 +version: 0.3.0.0 synopsis: A NATS client library written in Haskell tested-with: GHC ==8.8, @@ -59,6 +59,54 @@ hs-source-dirs: client +library jetstream + import: shared + visibility: public + build-depends: + natskell, + jetstream-internal + + exposed-modules: + JetStream.API + JetStream.Client + + hs-source-dirs: + jetstream-client + +library jetstream-internal + import: shared + build-depends: + natskell, + aeson >= 2.1 && < 2.3, + base64-bytestring >= 1.2 && < 1.3, + bytestring >= 0.10 && < 0.13, + stm >= 2.5 && < 2.6, + text (>= 1.2 && < 1.3) || (>= 2.0 && < 2.2), + time >= 1.9 && < 1.16 + + exposed-modules: + JetStream.Consumer + JetStream.Consumer.API + JetStream.Consumer.Types + JetStream.Error + JetStream.Message + JetStream.Message.API + JetStream.Message.Types + JetStream.Options + JetStream.Protocol.Headers + JetStream.Protocol.Request + JetStream.Protocol.Subject + JetStream.Publish + JetStream.Publish.API + JetStream.Publish.Types + JetStream.Stream + JetStream.Stream.API + JetStream.Stream.Types + JetStream.Types + + hs-source-dirs: + jetstream + library natskell-internal import: shared build-depends: @@ -157,7 +205,9 @@ hs-source-dirs: test/Unit main-is: Spec.hs build-depends: + natskell, natskell-internal, + jetstream-internal, hspec >= 2.7 && < 2.12, bytestring, network, @@ -187,6 +237,10 @@ NuidSpec WorkerPoolSpec WaitGroupSpec + JetStream.ConsumerSpec + JetStream.MessageSpec + JetStream.ProtocolSpec + JetStream.PublishSpec test-suite fuzz-test import: shared @@ -223,6 +277,7 @@ hs-source-dirs: test/Integration build-depends: natskell, + natskell:jetstream, test-common, port-utils >= 0.1 && < 0.3, bytestring, diff --git a/test/Fuzz/ValidatorsSpec.hs b/test/Fuzz/ValidatorsSpec.hs --- a/test/Fuzz/ValidatorsSpec.hs +++ b/test/Fuzz/ValidatorsSpec.hs @@ -3,7 +3,8 @@ module ValidatorsSpec (spec) where import qualified Data.ByteString as BS -import Data.Maybe +import qualified Data.ByteString.Char8 as BC +import Data.Char (isSpace) import Test.Hspec import Test.Hspec.QuickCheck (modifyMaxSuccess) import Test.QuickCheck @@ -61,15 +62,11 @@ pubRules :: Pub -> Bool pubRules p - | Types.Pub.subject p == "" = False - | Types.Pub.replyTo p == Just "" = False + | not (subjectRules (Types.Pub.subject p)) = False + | not (maybe True subjectRules (Types.Pub.replyTo p)) = False | Types.Pub.payload p == Just "" = False - | headers == Just [] = False - | isJust headers && Prelude.any (\(k, _) -> k == "") (fromJust headers) = False - | isJust headers && Prelude.any (\(_, v) -> v == "") (fromJust headers) = False + | not (maybe True headerRules (Types.Pub.headers p)) = False | otherwise = True - where - headers = Types.Pub.headers p propSub :: Sub -> Bool propSub s = @@ -79,8 +76,8 @@ subRules :: Sub -> Bool subRules s - | Types.Sub.subject s == "" = False - | Types.Sub.queueGroup s == Just "" = False + | not (subjectRules (Types.Sub.subject s)) = False + | not (maybe True queueGroupRules (Types.Sub.queueGroup s)) = False | Types.Sub.sid s == "" = False | otherwise = True @@ -94,6 +91,52 @@ unsubRules u | Types.Unsub.sid u == "" = False | otherwise = True + +subjectRules :: BS.ByteString -> Bool +subjectRules value = + not (BS.null value) + && not (BC.any isSpace value) + && not (any BS.null tokens) + && wildcardRules tokens + where + tokens = BC.split '.' value + +queueGroupRules :: BS.ByteString -> Bool +queueGroupRules value = + not (BS.null value) && not (BC.any isSpace value) + +headerRules :: [(BS.ByteString, BS.ByteString)] -> Bool +headerRules values = + not (null values) && all headerPairRules values + +headerPairRules :: (BS.ByteString, BS.ByteString) -> Bool +headerPairRules (key, value) = + not (BS.null key) + && not (BS.null value) + && not (":" `BS.isInfixOf` key) + && not (containsLineBreak key) + && not (containsLineBreak value) + +wildcardRules :: [BS.ByteString] -> Bool +wildcardRules [] = + True +wildcardRules [">"] = + True +wildcardRules (">":_) = + False +wildcardRules (token:rest) + | token == "*" = + wildcardRules rest + | ">" `BS.isInfixOf` token = + False + | "*" `BS.isInfixOf` token = + False + | otherwise = + wildcardRules rest + +containsLineBreak :: BS.ByteString -> Bool +containsLineBreak value = + "\r" `BS.isInfixOf` value || "\n" `BS.isInfixOf` value instance Arbitrary Connect where arbitrary = Connect <$> arbitrary diff --git a/test/Integration/ClientSpec.hs b/test/Integration/ClientSpec.hs --- a/test/Integration/ClientSpec.hs +++ b/test/Integration/ClientSpec.hs @@ -5,8 +5,10 @@ import API ( Client (..) , MsgView (..) + , withHeaders , withPayload , withQueueGroup + , withReplyCallback , withReplyTo , withSubscriptionExpiry ) @@ -25,6 +27,11 @@ , writeIORef ) import Data.Word8 +import qualified JetStream.API as JetStream +import JetStream.Client + ( newJetStream + , withRequestTimeoutMicros + ) import Network.Socket (Socket, accept, listen) import qualified Network.Socket import Network.Socket.ByteString (recv, sendAll) @@ -77,6 +84,22 @@ , "\r\n" ] +data CapturedPublish = CapturedPublish + { capturedInbox :: BS.ByteString + , capturedSid :: BS.ByteString + , capturedSubject :: BS.ByteString + , capturedReplyTo :: BS.ByteString + , capturedPayload :: BS.ByteString + } + deriving (Eq, Show) + +data CapturedSubscription = CapturedSubscription + { capturedSubSubject :: BS.ByteString + , capturedSubQueueGroup :: Maybe BS.ByteString + , capturedSubSid :: BS.ByteString + } + deriving (Eq, Show) + recvUntil :: Socket -> (BS.ByteString -> Bool) -> IO BS.ByteString recvUntil sock done = go BS.empty @@ -89,6 +112,31 @@ then pure acc else go (acc <> chunk) +expectClientBytes :: Socket -> (BS.ByteString -> Bool) -> String -> IO BS.ByteString +expectClientBytes sock done failureMessage = do + result <- timeout 1000000 $ + recvUntil sock done + case result of + Nothing -> + expectationFailure failureMessage >> + fail failureMessage + Just bytes -> + if done bytes + then pure bytes + else + expectationFailure ("unexpected client bytes: " ++ show bytes) >> + fail failureMessage + +expectNoClientBytes :: Socket -> (BS.ByteString -> Bool) -> String -> IO () +expectNoClientBytes sock unwanted failureMessage = do + result <- timeout 300000 $ + recvUntil sock unwanted + case result of + Just bytes | unwanted bytes -> + expectationFailure (failureMessage ++ ": " ++ show bytes) + _ -> + pure () + expectClientCommand :: Socket -> BS.ByteString -> IO () expectClientCommand sock command = do result <- timeout 1000000 $ @@ -100,10 +148,196 @@ unless (BS.isInfixOf command bytes) $ expectationFailure ("unexpected client bytes: " ++ show bytes) +expectMVar :: String -> MVar a -> IO a +expectMVar failureMessage var = do + result <- timeout 1000000 (takeMVar var) + case result of + Nothing -> + expectationFailure failureMessage >> + fail failureMessage + Just value -> + pure value + expectQueuedSub :: Socket -> BS.ByteString -> BS.ByteString -> BS.ByteString -> IO () expectQueuedSub sock subject queueGroup sid = expectClientCommand sock (BS.concat ["SUB ", subject, " ", queueGroup, " ", sid, "\r\n"]) +captureSubscription :: Socket -> BS.ByteString -> IO CapturedSubscription +captureSubscription sock expectedSubject = do + bytes <- recvUntil sock hasExpectedSubscription + case parseCapturedSubscription expectedSubject bytes of + Nothing -> + expectationFailure ("could not parse subscription from bytes: " ++ show bytes) >> + fail "could not parse subscription" + Just captured -> + pure captured + where + hasExpectedSubscription bytes = + case parseCapturedSubscription expectedSubject bytes of + Nothing -> False + Just _ -> True + +parseCapturedSubscription :: BS.ByteString -> BS.ByteString -> Maybe CapturedSubscription +parseCapturedSubscription expectedSubject bytes = + findSubscription normalizedLines + where + wireLines = C.split '\n' bytes + stripCR = C.takeWhile (/= '\r') + normalizedLines = fmap stripCR wireLines + findSubscription [] = Nothing + findSubscription (line:rest) = + case C.words line of + [prefix, subject', sid] + | prefix == "SUB" && subject' == expectedSubject -> + Just (CapturedSubscription subject' Nothing sid) + [prefix, subject', queueGroup, sid] + | prefix == "SUB" && subject' == expectedSubject -> + Just (CapturedSubscription subject' (Just queueGroup) sid) + _ -> + findSubscription rest + +capturePublish :: Socket -> BS.ByteString -> IO CapturedPublish +capturePublish sock expectedSubject = do + bytes <- recvUntil sock hasExpectedPublish + case parseCapturedPublish bytes of + Nothing -> + expectationFailure ("could not parse publish from bytes: " ++ show bytes) >> + fail "could not parse publish" + Just captured -> do + capturedSubject captured `shouldBe` expectedSubject + pure captured + where + hasExpectedPublish bytes = + case parseCapturedPublish bytes of + Nothing -> + False + Just captured -> + capturedSubject captured == expectedSubject + +parseCapturedPublish :: BS.ByteString -> Maybe CapturedPublish +parseCapturedPublish bytes = do + subLine <- findLineWithPrefix "SUB " + pubLine <- findLineWithPrefix "PUB " + let subParts = C.words subLine + pubParts = C.words pubLine + inbox <- subParts `at` 1 + sid <- subParts `at` 2 + subject' <- pubParts `at` 1 + reply <- pubParts `at` 2 + lenBytes <- pubParts `at` 3 + len <- readMaybeInt lenBytes + payloadLine <- payloadAfter pubLine + let payloadBytes = BS.take len payloadLine + pure CapturedPublish + { capturedInbox = inbox + , capturedSid = sid + , capturedSubject = subject' + , capturedReplyTo = reply + , capturedPayload = payloadBytes + } + where + wireLines = C.split '\n' bytes + stripCR = C.takeWhile (/= '\r') + normalizedLines = fmap stripCR wireLines + findLineWithPrefix prefix = + case filter (BS.isPrefixOf prefix) normalizedLines of + [] -> Nothing + (line:_) -> Just line + payloadAfter pubLine = + case dropWhile (/= pubLine) normalizedLines of + (_:payloadLine:_) -> Just payloadLine + _ -> Nothing + +at :: [a] -> Int -> Maybe a +at values index = + if index < length values + then Just (values !! index) + else Nothing + +readMaybeInt :: BS.ByteString -> Maybe Int +readMaybeInt bytes = + case reads (C.unpack bytes) of + [(value, "")] -> Just value + _ -> Nothing + +appearsBefore :: BS.ByteString -> BS.ByteString -> BS.ByteString -> Bool +appearsBefore first second bytes = + case (C.breakSubstring first bytes, C.breakSubstring second bytes) of + ((firstPrefix, firstMatch), (secondPrefix, secondMatch)) -> + not (BS.null firstMatch) + && not (BS.null secondMatch) + && BS.length firstPrefix < BS.length secondPrefix + +replyToCapturedPublish :: Socket -> CapturedPublish -> BS.ByteString -> IO () +replyToCapturedPublish sock captured body = + sendAll sock (msgFrame (capturedInbox captured) (capturedSid captured) body) + +replyStatusToCapturedPublish + :: Socket + -> CapturedPublish + -> BS.ByteString + -> BS.ByteString + -> IO () +replyStatusToCapturedPublish sock captured status description = + sendAll sock $ + hmsgFrame + (capturedInbox captured) + (capturedSid captured) + [ ("Status", status) + , ("Description", description) + ] + "" + +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}}" + +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\"]}}" + +protoPushConsumerInfoResponse :: BS.ByteString +protoPushConsumerInfoResponse = + "{\"stream_name\":\"PROTO_STREAM\",\"name\":\"PROTO_PUSH\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"name\":\"PROTO_PUSH\",\"deliver_subject\":\"PROTO.DELIVER\",\"deliver_group\":\"workers\",\"deliver_policy\":\"all\",\"ack_policy\":\"explicit\",\"replay_policy\":\"instant\"}}" + +protoOrderedConsumerInfoResponse :: BS.ByteString +protoOrderedConsumerInfoResponse = + "{\"stream_name\":\"PROTO_STREAM\",\"name\":\"PROTO_ORDERED_1\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"name\":\"PROTO_ORDERED_1\",\"deliver_policy\":\"by_start_sequence\",\"opt_start_seq\":3,\"ack_policy\":\"none\",\"replay_policy\":\"instant\",\"max_deliver\":-1,\"max_waiting\":512,\"inactive_threshold\":300000000000,\"num_replicas\":1,\"mem_storage\":true}}" + +protoOrderedConsumerInfoResponseFor :: BS.ByteString -> BS.ByteString +protoOrderedConsumerInfoResponseFor consumer = + BS.concat + [ "{\"stream_name\":\"PROTO_STREAM\",\"name\":\"" + , consumer + , "\",\"created\":\"2024-01-01T00:00:00Z\",\"config\":{\"name\":\"" + , consumer + , "\",\"deliver_policy\":\"all\",\"ack_policy\":\"none\",\"replay_policy\":\"instant\",\"max_deliver\":-1,\"max_waiting\":512,\"inactive_threshold\":300000000000,\"num_replicas\":1,\"mem_storage\":true}}" + ] + +protoAccountInfoResponse :: BS.ByteString +protoAccountInfoResponse = + "{\"memory\":10,\"storage\":20,\"reserved_memory\":0,\"reserved_storage\":0,\"streams\":1,\"consumers\":1,\"limits\":{\"max_memory\":-1,\"max_storage\":-1,\"max_streams\":-1,\"max_consumers\":-1,\"max_ack_pending\":-1,\"memory_max_stream_bytes\":-1,\"storage_max_stream_bytes\":-1,\"max_bytes_required\":false},\"api\":{\"level\":1,\"total\":3,\"errors\":0},\"tiers\":{}}" + +protoStreamMessageResponse :: BS.ByteString +protoStreamMessageResponse = + "{\"message\":{\"subject\":\"PROTO.SUBJECT\",\"seq\":2,\"data\":\"YWRtaW4=\",\"time\":\"2024-01-01T00:00:00Z\"}}" + +protoDeleteStreamMessageResponse :: BS.ByteString +protoDeleteStreamMessageResponse = + "{\"success\":true}" + +protoDeleteConsumerResponse :: BS.ByteString +protoDeleteConsumerResponse = + "{\"success\":true}" + +protoConsumerPauseResponse :: BS.ByteString +protoConsumerPauseResponse = + "{\"paused\":true,\"pause_until\":\"2024-01-01T00:00:01Z\",\"pause_remaining\":1000000000}" + +protoConsumerResetResponse :: BS.ByteString +protoConsumerResetResponse = + "{\"stream_name\":\"PROTO_STREAM\",\"name\":\"PROTO_CONSUMER\",\"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}" + startClientWith extraOptions = do (p, sock) <- openFreePort listen sock 1 @@ -285,6 +519,545 @@ ) close client wait wg + + describe "core NATS fake-server coverage" $ do + around withClient $ do + it "ignores duplicate request replies after the first message" $ \(serv, client, _, _) -> do + replies <- newIORef ([] :: [Maybe MsgView]) + firstReply <- newEmptyMVar + sid <- request client "SERVICE.ONCE" [] $ \reply -> do + atomicModifyIORef' replies $ \seen -> (seen ++ [reply], ()) + void (tryPutMVar firstReply ()) + sendAll serv (msgFrame "SERVICE.ONCE" sid "first") + sendAll serv (msgFrame "SERVICE.ONCE" sid "second") + expectMVar "first request reply did not arrive" firstReply + threadDelay 100000 + seen <- readIORef replies + case seen of + [Just msg] -> + payload msg `shouldBe` Just "first" + other -> + expectationFailure ("unexpected request replies: " ++ show other) + + it "handles many simultaneous request replies" $ \(serv, client, _, _) -> do + requests <- forM [1 .. 32 :: Int] $ \index -> do + replyBox <- newEmptyMVar + let requestSubject = C.pack ("SERVICE.MANY." ++ show index) + sid <- request client requestSubject [] (putMVar replyBox) + pure (index, requestSubject, sid, replyBox) + for_ requests $ \(index, requestSubject, sid, _) -> + sendAll serv $ + msgFrame requestSubject sid (C.pack ("reply-" ++ show index)) + for_ requests $ \(index, _, _, replyBox) -> do + reply <- expectMVar ("request reply " ++ show index ++ " did not arrive") replyBox + case reply of + Nothing -> + expectationFailure ("request reply " ++ show index ++ " expired") + Just msg -> + payload msg `shouldBe` Just (C.pack ("reply-" ++ show index)) + + it "keeps an active request alive after garbled inbound bytes" $ \(serv, client, _, _) -> do + replyBox <- newEmptyMVar + sid <- request client "SERVICE.GARBLED" [] (putMVar replyBox) + sendAll serv $ + BS.concat + [ "bad-prefix" + , msgFrame "SERVICE.GARBLED" sid "ok" + ] + reply <- expectMVar "request reply after garbled bytes did not arrive" replyBox + case reply of + Nothing -> + expectationFailure "request expired after garbled bytes" + Just msg -> + payload msg `shouldBe` Just "ok" + + it "round trips headers across request publish and reply" $ \(serv, client, _, _) -> do + replyBox <- newEmptyMVar + publish client + "HEADERS.SERVICE" + [ withHeaders [("X-Request", "one")] + , withPayload "ping" + , withReplyTo "_INBOX.headers" + , withReplyCallback (putMVar replyBox) + ] + bytes <- expectClientBytes + serv + (\sent -> + BS.isInfixOf "SUB _INBOX.headers " sent + && BS.isInfixOf "HPUB HEADERS.SERVICE _INBOX.headers " sent + && BS.isInfixOf "X-Request:one\r\n" sent + && BS.isInfixOf "\r\n\r\nping\r\n" sent) + "client did not send header request publish" + captured <- case parseCapturedSubscription "_INBOX.headers" bytes of + Nothing -> + expectationFailure ("could not parse header reply subscription: " ++ show bytes) >> + fail "could not parse header reply subscription" + Just sub -> + pure sub + sendAll serv $ + hmsgFrame + "_INBOX.headers" + (capturedSubSid captured) + [("X-Reply", "ok")] + "pong" + reply <- expectMVar "header reply did not arrive" replyBox + case reply of + Nothing -> + expectationFailure "header reply expired" + Just msg -> do + subject msg `shouldBe` "_INBOX.headers" + payload msg `shouldBe` Just "pong" + headers msg `shouldBe` Just [("X-Reply", "ok")] + + it "flush sends queued publishes before its PING" $ \(serv, client, _, _) -> do + expectClientCommand serv "CONNECT " + doneVar <- newEmptyMVar + let publishFrame = "PUB ORDER.FLUSH 3\r\none\r\n" + publish client "ORDER.FLUSH" [withPayload "one"] + void . forkIO $ do + flush client + putMVar doneVar () + bytes <- expectClientBytes + serv + (\sent -> + BS.isInfixOf publishFrame sent + && BS.isInfixOf "PING\r\n" sent) + "client did not send publish and flush PING" + unless (appearsBefore publishFrame "PING\r\n" bytes) $ + expectationFailure ("flush PING was sent before queued publish: " ++ show bytes) + sendAll serv "PONG\r\n" + expectMVar "flush did not complete after PONG" doneVar + + it "exits when the server disconnects while a subscription is active" $ \(serv, client, _, exited) -> do + sid <- subscribe client "ACTIVE.DISCONNECT" [] (const (pure ())) + expectClientCommand serv (BS.concat ["SUB ACTIVE.DISCONNECT ", sid, "\r\n"]) + Network.Socket.close serv + result <- atomically $ readTMVar exited + case result of + ExitRetriesExhausted _ -> + pure () + other -> + expectationFailure ("Unexpected exit reason: " ++ show other) + + it "does not write publishes above negotiated max_payload" $ \(serv, client, _, _) -> do + expectClientCommand serv "CONNECT " + publish client "BOUNDARY.TOO_BIG" [withPayload (BS.replicate 2048 _x)] + expectNoClientBytes + serv + (BS.isInfixOf "PUB BOUNDARY.TOO_BIG") + "oversized publish reached server" + wg <- newWaitGroup 1 + ping client (done wg) + expectClientCommand serv "PING\r\n" + sendAll serv "PONG\r\n" + wait wg + + it "does not write invalid publish commands" $ \(serv, client, _, _) -> do + expectClientCommand serv "CONNECT " + publish client "" [withPayload "x"] + expectNoClientBytes + serv + (\sent -> BS.isInfixOf "PUB " sent || BS.isInfixOf "HPUB " sent) + "invalid publish reached server" + + it "cleans up no responders replies for expiring requests" $ do + bracket startClient stopClientSafely $ \(serv, client, _, _) -> do + replyBox <- newEmptyMVar + sid <- request client + "SERVICE.NO_RESPONDERS" + [withSubscriptionExpiry 2] + (putMVar replyBox) + sendAll serv $ + hmsgFrame + "SERVICE.NO_RESPONDERS" + sid + [ ("Status", "503") + , ("Description", "No Responders") + ] + "" + reply <- expectMVar "no responders reply callback did not run" replyBox + case reply of + Nothing -> + expectationFailure "no responders reply expired" + Just msg -> + headers msg `shouldBe` Just + [ ("Status", "503") + , ("Description", "No Responders") + ] + result <- timeout 500000 (close client) + case result of + Nothing -> + expectationFailure "close blocked after no responders reply" + Just () -> + pure () + + it "does not resubscribe in-flight reply inboxes after reconnect" $ do + (p, sock) <- openFreePort + listen sock 2 + clientVar <- newEmptyMVar + void . forkIO $ do + let configOptions = + [ withMinimumLogLevel Debug + , withConnectionAttempts 2 + , withConnectName "test-client" + ] + c <- newClient [("127.0.0.1", p)] configOptions + putMVar clientVar c + (firstConn, _) <- accept sock + sendAll firstConn defaultINFO + clientResult <- timeout 1000000 (takeMVar clientVar) + case clientResult of + Nothing -> + expectationFailure "client did not connect" + Just client -> do + replyBox <- newEmptyMVar + publish client + "SERVICE.RECONNECT" + [ withPayload "ping" + , withReplyCallback (putMVar replyBox) + ] + captured <- capturePublish firstConn "SERVICE.RECONNECT" + Network.Socket.close firstConn + (secondConn, _) <- accept sock + sendAll secondConn defaultINFO + let replySubCommand = BS.concat ["SUB ", capturedInbox captured] + resubscribe <- timeout 300000 $ + recvUntil secondConn (BS.isInfixOf replySubCommand) + case resubscribe of + Just bytes | BS.isInfixOf replySubCommand bytes -> + expectationFailure ("reply inbox was resubscribed: " ++ show bytes) + _ -> + pure () + wg <- newWaitGroup 1 + ping client (done wg) + expectClientCommand secondConn "PING\r\n" + sendAll secondConn "PONG\r\n" + wait wg + close client + Network.Socket.close secondConn + Network.Socket.close sock + + describe "jetstream protocol integration" $ do + around withClient $ do + it "sends stream create options as a JetStream API request" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + done <- newEmptyMVar + void . forkIO $ do + result <- JetStream.create (JetStream.streams jetStream) + "PROTO_STREAM" + ["PROTO.SUBJECT"] + [ JetStream.withStorage JetStream.MemoryStorage + , JetStream.withMaxMessages 1 + , 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 "\"discard\":\"new\"" + replyToCapturedPublish serv captured protoStreamInfoResponse + result <- timeout 1000000 (takeMVar done) + case result of + Nothing -> + expectationFailure "stream create did not complete" + Just (Left err) -> + expectationFailure ("stream create failed: " ++ show err) + Just (Right info) -> + JetStream.streamConfigName (JetStream.streamInfoConfig info) `shouldBe` "PROTO_STREAM" + + it "sends only the selected consumer filter representation" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + done <- newEmptyMVar + void . forkIO $ do + result <- JetStream.putConsumer + (JetStream.consumers jetStream) + "PROTO_STREAM" + JetStream.ConsumerCreate + (JetStream.DurableConsumer "PROTO_CONSUMER") + JetStream.PullConsumer + [ JetStream.withConsumerDeliverPolicy JetStream.DeliverAll + , JetStream.withConsumerAckPolicy JetStream.AckExplicit + , JetStream.withConsumerFilter + (JetStream.ConsumerFilterSubjects ["PROTO.A", "PROTO.B"]) + ] + 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\":" + replyToCapturedPublish serv captured protoConsumerInfoResponse + result <- timeout 1000000 (takeMVar done) + case result of + Nothing -> + expectationFailure "consumer create did not complete" + Just (Left err) -> + expectationFailure ("consumer create failed: " ++ show err) + Just (Right info) -> + JetStream.consumerInfoName info `shouldBe` "PROTO_CONSUMER" + + it "sends push consumer create requests with a delivery subject" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + done <- newEmptyMVar + void . forkIO $ do + result <- JetStream.putConsumer + (JetStream.consumers jetStream) + "PROTO_STREAM" + JetStream.ConsumerCreate + (JetStream.NamedConsumer "PROTO_PUSH") + (JetStream.PushConsumer "PROTO.DELIVER") + [ JetStream.withConsumerDeliverGroup "workers" + , JetStream.withConsumerAckPolicy JetStream.AckExplicit + ] + putMVar done result + captured <- capturePublish serv "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_PUSH" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"action\":\"create\"" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"deliver_subject\":\"PROTO.DELIVER\"" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"deliver_group\":\"workers\"" + replyToCapturedPublish serv captured protoPushConsumerInfoResponse + result <- timeout 1000000 (takeMVar done) + case result of + Nothing -> + expectationFailure "push consumer create did not complete" + Just (Left err) -> + expectationFailure ("push consumer create failed: " ++ show err) + Just (Right info) -> do + JetStream.consumerInfoName info `shouldBe` "PROTO_PUSH" + JetStream.consumerConfigDeliverSubject (JetStream.consumerInfoConfig info) + `shouldBe` Just "PROTO.DELIVER" + + it "subscribes to push consumer delivery subjects with a queue group" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + subscription <- JetStream.consumePush + (JetStream.messages jetStream) + "PROTO.DELIVER" + [JetStream.withPushQueueGroup "workers"] + (const (pure ())) + expectClientCommand serv "SUB PROTO.DELIVER workers " + JetStream.stopPushSubscription subscription + expectClientCommand serv "UNSUB " + + it "keeps push consumers alive after garbled inbound bytes" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + delivered <- newEmptyMVar + subscription <- JetStream.consumePush + (JetStream.messages jetStream) + "PROTO.DELIVER" + [] + (putMVar delivered) + captured <- captureSubscription serv "PROTO.DELIVER" + sendAll serv $ + BS.concat + [ "garbled-prefix" + , msgFrame "PROTO.DELIVER" (capturedSubSid captured) "payload" + ] + result <- timeout 1000000 (takeMVar delivered) + JetStream.stopPushSubscription subscription + case result of + Nothing -> + expectationFailure "push message was not delivered after garbled bytes" + Just message -> + JetStream.messagePayload message `shouldBe` Just "payload" + + it "creates ordered consumers with the ordered defaults" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + done <- newEmptyMVar + void . forkIO $ do + result <- JetStream.createOrderedConsumer + (JetStream.messages jetStream) + "PROTO_STREAM" + [ JetStream.withOrderedConsumerNamePrefix "PROTO_ORDERED" + , JetStream.withOrderedConsumerDeliverPolicy (JetStream.DeliverByStartSequence 3) + ] + putMVar done result + captured <- capturePublish serv "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_ORDERED_1" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"name\":\"PROTO_ORDERED_1\"" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"deliver_policy\":\"by_start_sequence\"" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"opt_start_seq\":3" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"ack_policy\":\"none\"" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"max_deliver\":-1" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"max_waiting\":512" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"mem_storage\":true" + replyToCapturedPublish serv captured protoOrderedConsumerInfoResponse + result <- timeout 1000000 (takeMVar done) + case result of + Nothing -> + expectationFailure "ordered consumer create did not complete" + Just (Left err) -> + expectationFailure ("ordered consumer create failed: " ++ show err) + Just (Right _) -> + pure () + + it "sends no-wait pull requests and decodes no-message status" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + done <- newEmptyMVar + void . forkIO $ do + response <- JetStream.fetch (JetStream.messages jetStream) + "PROTO_STREAM" + "PROTO_CONSUMER" + [ JetStream.withFetchBatch 1 + , JetStream.withFetchWait (JetStream.FetchNoWaitMicros 1000000) + ] + putMVar done response + captured <- capturePublish serv "$JS.API.CONSUMER.MSG.NEXT.PROTO_STREAM.PROTO_CONSUMER" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"no_wait\":true" + capturedPayload captured `shouldSatisfy` not . BS.isInfixOf "\"expires\":" + capturedReplyTo captured `shouldBe` capturedInbox captured + replyStatusToCapturedPublish serv captured "404" "No Messages" + result <- timeout 1000000 (takeMVar done) + case result of + Nothing -> + expectationFailure "fetch did not complete" + Just (Left err) -> + expectationFailure ("fetch failed: " ++ show err) + Just (Right response) -> do + JetStream.pullResponseMessages response `shouldBe` [] + JetStream.pullResponseStatus response + `shouldBe` Just (JetStream.PullNoMessages (Just "No Messages")) + + it "sends account info requests" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + done <- newEmptyMVar + void . forkIO $ do + result <- JetStream.accountInfo jetStream + putMVar done result + captured <- capturePublish serv "$JS.API.INFO" + capturedPayload captured `shouldBe` "" + replyToCapturedPublish serv captured protoAccountInfoResponse + result <- timeout 1000000 (takeMVar done) + case result of + Nothing -> + expectationFailure "account info did not complete" + Just (Left err) -> + expectationFailure ("account info failed: " ++ show err) + Just (Right info) -> + JetStream.accountTierStreams (JetStream.accountInfoTier info) `shouldBe` 1 + + it "returns a JetStream timeout when admin requests receive no response" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [withRequestTimeoutMicros 1000] + done <- newEmptyMVar + void . forkIO $ do + result <- JetStream.accountInfo jetStream + putMVar done result + captured <- capturePublish serv "$JS.API.INFO" + capturedPayload captured `shouldBe` "" + result <- timeout 1000000 (takeMVar done) + case result of + Nothing -> + expectationFailure "account info timeout did not complete" + Just (Left JetStream.JetStreamTimeout) -> + pure () + Just other -> + expectationFailure ("unexpected account info timeout result: " ++ show other) + + it "sends stream message get and delete requests" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + fetched <- newEmptyMVar + void . forkIO $ do + result <- JetStream.getMessage + (JetStream.streams jetStream) + "PROTO_STREAM" + (JetStream.StreamMessageBySequence 2) + putMVar fetched result + getRequest <- capturePublish serv "$JS.API.STREAM.MSG.GET.PROTO_STREAM" + capturedPayload getRequest `shouldSatisfy` BS.isInfixOf "\"seq\":2" + replyToCapturedPublish serv getRequest protoStreamMessageResponse + fetchedResult <- timeout 1000000 (takeMVar fetched) + case fetchedResult of + Nothing -> + expectationFailure "stream message get did not complete" + Just (Left err) -> + expectationFailure ("stream message get failed: " ++ show err) + Just (Right message) -> + JetStream.streamMessagePayload message `shouldBe` Just "admin" + + deleted <- newEmptyMVar + void . forkIO $ do + result <- JetStream.deleteMessage + (JetStream.streams jetStream) + "PROTO_STREAM" + 2 + JetStream.DeleteMessage + putMVar deleted result + deleteRequest <- capturePublish serv "$JS.API.STREAM.MSG.DELETE.PROTO_STREAM" + capturedPayload deleteRequest `shouldSatisfy` BS.isInfixOf "\"seq\":2" + capturedPayload deleteRequest `shouldSatisfy` BS.isInfixOf "\"no_erase\":true" + replyToCapturedPublish serv deleteRequest protoDeleteStreamMessageResponse + deletedResult <- timeout 1000000 (takeMVar deleted) + case deletedResult of + Nothing -> + expectationFailure "stream message delete did not complete" + Just (Left err) -> + expectationFailure ("stream message delete failed: " ++ show err) + Just (Right response) -> + JetStream.deleteStreamMessageSuccess response `shouldBe` True + + it "sends consumer update actions" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + done <- newEmptyMVar + void . forkIO $ do + result <- JetStream.putConsumer + (JetStream.consumers jetStream) + "PROTO_STREAM" + JetStream.ConsumerUpdate + (JetStream.DurableConsumer "PROTO_CONSUMER") + JetStream.PullConsumer + [ JetStream.withConsumerDescription "updated" + , JetStream.withConsumerAckPolicy JetStream.AckExplicit + ] + putMVar done result + captured <- capturePublish serv "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_CONSUMER" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"action\":\"update\"" + capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"description\":\"updated\"" + replyToCapturedPublish serv captured protoConsumerInfoResponse + result <- timeout 1000000 (takeMVar done) + case result of + Nothing -> + expectationFailure "consumer update did not complete" + Just (Left err) -> + expectationFailure ("consumer update failed: " ++ show err) + Just (Right info) -> + JetStream.consumerInfoName info `shouldBe` "PROTO_CONSUMER" + + it "sends consumer pause and reset requests" $ \(serv, client, _, _) -> do + let jetStream = newJetStream client [] + pauseUntil = read "2024-01-01 00:00:01 UTC" + paused <- newEmptyMVar + void . forkIO $ do + result <- JetStream.pauseConsumer + (JetStream.consumers jetStream) + "PROTO_STREAM" + "PROTO_CONSUMER" + pauseUntil + putMVar paused result + pauseRequest <- capturePublish serv "$JS.API.CONSUMER.PAUSE.PROTO_STREAM.PROTO_CONSUMER" + capturedPayload pauseRequest `shouldSatisfy` BS.isInfixOf "\"pause_until\":\"2024-01-01T00:00:01Z\"" + replyToCapturedPublish serv pauseRequest protoConsumerPauseResponse + pausedResult <- timeout 1000000 (takeMVar paused) + case pausedResult of + Nothing -> + expectationFailure "consumer pause did not complete" + Just (Left err) -> + expectationFailure ("consumer pause failed: " ++ show err) + Just (Right response) -> + JetStream.consumerPausePaused response `shouldBe` True + + reset <- newEmptyMVar + void . forkIO $ do + result <- JetStream.resetConsumer + (JetStream.consumers jetStream) + "PROTO_STREAM" + "PROTO_CONSUMER" + [JetStream.withConsumerResetSequence 7] + putMVar reset result + resetRequest <- capturePublish serv "$JS.API.CONSUMER.RESET.PROTO_STREAM.PROTO_CONSUMER" + capturedPayload resetRequest `shouldSatisfy` BS.isInfixOf "\"seq\":7" + replyToCapturedPublish serv resetRequest protoConsumerResetResponse + resetResult <- timeout 1000000 (takeMVar reset) + case resetResult of + Nothing -> + expectationFailure "consumer reset did not complete" + Just (Left err) -> + expectationFailure ("consumer reset failed: " ++ show err) + Just (Right response) -> + JetStream.consumerResetResponseSequence response `shouldBe` 7 it "manual unsubscribe does not block close for tracked expiries" $ do bracket startClient stopClientSafely $ \(_, client, _, _) -> do sid <- request client "foo" [withSubscriptionExpiry 30] (const (pure ())) @@ -362,6 +1135,134 @@ (secondConn, _) <- accept sock sendAll secondConn defaultINFO expectQueuedSub secondConn "JOBS" "WORKERS" sid + close client + Network.Socket.close secondConn + Network.Socket.close sock + it "resubscribes JetStream push consumers after reconnect" $ do + (p, sock) <- openFreePort + listen sock 2 + clientVar <- newEmptyMVar + void . forkIO $ do + let configOptions = + [ withMinimumLogLevel Debug + , withConnectionAttempts 2 + , withConnectName "test-client" + ] + c <- newClient [("127.0.0.1", p)] configOptions + putMVar clientVar c + (firstConn, _) <- accept sock + sendAll firstConn defaultINFO + clientResult <- timeout 1000000 (takeMVar clientVar) + case clientResult of + Nothing -> + expectationFailure "client did not connect" + Just client -> do + let jetStream = newJetStream client [] + delivered <- newEmptyMVar + subscription <- JetStream.consumePush + (JetStream.messages jetStream) + "PROTO.DELIVER" + [] + (putMVar delivered) + firstSub <- captureSubscription firstConn "PROTO.DELIVER" + Network.Socket.close firstConn + (secondConn, _) <- accept sock + sendAll secondConn defaultINFO + secondSub <- captureSubscription secondConn "PROTO.DELIVER" + capturedSubSid secondSub `shouldBe` capturedSubSid firstSub + sendAll secondConn $ + msgFrame "PROTO.DELIVER" (capturedSubSid secondSub) "after-reconnect" + result <- timeout 1000000 (takeMVar delivered) + case result of + Nothing -> + expectationFailure "push message was not delivered after reconnect" + Just message -> + JetStream.messagePayload message `shouldBe` Just "after-reconnect" + JetStream.stopPushSubscription subscription + close client + Network.Socket.close secondConn + Network.Socket.close sock + it "ordered fetch returns after disconnect and recovers on reconnect" $ do + (p, sock) <- openFreePort + listen sock 2 + clientVar <- newEmptyMVar + void . forkIO $ do + let configOptions = + [ withMinimumLogLevel Debug + , withConnectionAttempts 2 + , withConnectName "test-client" + ] + c <- newClient [("127.0.0.1", p)] configOptions + putMVar clientVar c + (firstConn, _) <- accept sock + sendAll firstConn defaultINFO + clientResult <- timeout 1000000 (takeMVar clientVar) + case clientResult of + Nothing -> + expectationFailure "client did not connect" + Just client -> do + let jetStream = newJetStream client [] + orderedVar <- newEmptyMVar + void . forkIO $ do + result <- JetStream.createOrderedConsumer + (JetStream.messages jetStream) + "PROTO_STREAM" + [JetStream.withOrderedConsumerNamePrefix "PROTO_ORDERED"] + putMVar orderedVar result + createInitial <- capturePublish firstConn "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_ORDERED_1" + replyToCapturedPublish firstConn createInitial $ + protoOrderedConsumerInfoResponseFor "PROTO_ORDERED_1" + orderedResult <- timeout 1000000 (takeMVar orderedVar) + ordered <- case orderedResult of + Nothing -> + expectationFailure "ordered consumer create did not complete" >> + fail "ordered consumer create did not complete" + Just (Left err) -> + expectationFailure ("ordered consumer create failed: " ++ show err) >> + fail "ordered consumer create failed" + Just (Right ordered) -> + pure ordered + + fetchVar <- newEmptyMVar + void . forkIO $ do + result <- JetStream.fetchOrdered ordered + [ JetStream.withFetchBatch 1 + , JetStream.withFetchWait (JetStream.FetchExpiresMicros 100000) + ] + putMVar fetchVar result + deleteCurrent <- capturePublish firstConn "$JS.API.CONSUMER.DELETE.PROTO_STREAM.PROTO_ORDERED_1" + replyToCapturedPublish firstConn deleteCurrent protoDeleteConsumerResponse + createNext <- capturePublish firstConn "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_ORDERED_2" + replyToCapturedPublish firstConn createNext $ + protoOrderedConsumerInfoResponseFor "PROTO_ORDERED_2" + _nextRequest <- capturePublish firstConn "$JS.API.CONSUMER.MSG.NEXT.PROTO_STREAM.PROTO_ORDERED_2" + Network.Socket.close firstConn + fetchResult <- timeout 1000000 (takeMVar fetchVar) + case fetchResult of + Nothing -> + expectationFailure "ordered fetch did not return after disconnect" + Just (Left JetStream.JetStreamTimeout) -> + pure () + Just other -> + expectationFailure ("unexpected ordered fetch result after disconnect: " ++ show other) + + (secondConn, _) <- accept sock + sendAll secondConn defaultINFO + infoVar <- newEmptyMVar + void . forkIO $ do + result <- JetStream.orderedConsumerInfo ordered + putMVar infoVar result + infoRequest <- capturePublish secondConn "$JS.API.CONSUMER.INFO.PROTO_STREAM.PROTO_ORDERED_2" + replyToCapturedPublish secondConn infoRequest $ + protoOrderedConsumerInfoResponseFor "PROTO_ORDERED_2" + infoResult <- timeout 1000000 (takeMVar infoVar) + case infoResult of + Nothing -> + expectationFailure "ordered consumer info did not complete after reconnect" + Just (Left err) -> + expectationFailure ("ordered consumer info failed after reconnect: " ++ show err) + Just (Right info) -> + JetStream.consumerInfoName info `shouldBe` "PROTO_ORDERED_2" close client Network.Socket.close secondConn Network.Socket.close sock diff --git a/test/Unit/ConnectionSpec.hs b/test/Unit/ConnectionSpec.hs --- a/test/Unit/ConnectionSpec.hs +++ b/test/Unit/ConnectionSpec.hs @@ -7,6 +7,7 @@ import Control.Concurrent.STM import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LBS +import Data.IORef (newIORef, readIORef, writeIORef) import Handshake.Nats (performHandshake) import Lib.Logger ( LogLevel (Debug) @@ -23,17 +24,30 @@ ) import Parser.API ( ParseStep (DropPrefix, Emit, NeedMore, Reject) - , ParsedMessage (ParsedInfo) + , ParsedMessage (ParsedInfo, ParsedPing, ParsedPong) , ParserAPI (ParserAPI) ) +import qualified Queue.API as Queue import Queue.TransactionalQueue (newQueue) -import State.Store (newClientState, readServerInfo) +import Router.Nats + ( RouteDirective (RouteContinue) + , routeMessage + ) +import State.Store + ( newClientState + , pushPingAction + , queue + , readServerInfo + ) import State.Types (ClientConfig (..)) +import Subscription.Store (newSubscriptionStore) import System.Timeout (timeout) import Test.Hspec import Transformers.Transformers (Transformer (transform)) import qualified Types.Connect as Connect import Types.Info (Info (Info)) +import Types.Ping (Ping (Ping)) +import Types.Pong (Pong (Pong)) spec :: Spec spec = do @@ -57,6 +71,31 @@ closeReader (reader connectionApi) conn result <- timeout 1000000 (takeMVar resultVar) result `shouldBe` Just (Left "Read operation is blocked") + describe "Router ping lifecycle" $ do + it "responds to server PING by enqueueing PONG" $ do + state <- newTestState + store <- newSubscriptionStore + + directive <- routeMessage state store (ParsedPing Ping) + + directive `shouldBe` RouteContinue + queued <- Queue.dequeue (queue state) + case queued of + Right item -> + LBS.toStrict (transform item) `shouldBe` "PONG\r\n" + Left err -> + expectationFailure ("expected queued PONG, got queue error: " ++ err) + + it "runs the next pending ping action when PONG is routed" $ do + state <- newTestState + store <- newSubscriptionStore + actionRan <- newIORef False + pushPingAction state (writeIORef actionRan True) + + directive <- routeMessage state store (ParsedPong Pong) + + directive `shouldBe` RouteContinue + readIORef actionRan `shouldReturn` True describe "Handshake" $ do it "accepts an incremental parser backend for the initial INFO frame" $ do state <- newTestState diff --git a/test/Unit/JetStream/ConsumerSpec.hs b/test/Unit/JetStream/ConsumerSpec.hs new file mode 100644 --- /dev/null +++ b/test/Unit/JetStream/ConsumerSpec.hs @@ -0,0 +1,193 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.ConsumerSpec (spec) where + +import Data.Aeson + ( Value + , eitherDecode + , encode + , object + , (.=) + ) +import qualified Data.ByteString.Lazy as LBS +import JetStream.Consumer.Types +import Test.Hspec + +spec :: Spec +spec = do + describe "ConsumerConfig request JSON" $ do + it "encodes only populated fields" $ do + eitherDecode (encode durablePullConsumerConfigRequest) `shouldBe` Right durablePullConsumerConfigValue + + it "encodes push and ordered consumer fields" $ do + eitherDecode (encode pushConsumerConfigRequest) `shouldBe` Right pushConsumerConfigValue + + it "lets consumer target and kind own mutually exclusive fields" $ do + eitherDecode (encode targetKindConsumerConfigRequest) + `shouldBe` Right targetKindConsumerConfigValue + + describe "Consumer reset request JSON" $ do + it "encodes an optional reset sequence" $ do + eitherDecode (encode (consumerResetRequest [withConsumerResetSequence 7])) + `shouldBe` Right (object ["seq" .= (7 :: Int)]) + + describe "ConsumerInfo JSON" . + it "decodes server responses into concrete fields" $ do + eitherDecode consumerInfoJSON `shouldBe` Right consumerInfo + + describe "ConsumerResetResponse JSON" . + it "decodes reset metadata with the updated consumer info" $ do + fmap consumerResetResponseSequence (eitherDecode consumerResetJSON) + `shouldBe` Right 7 + + describe "Consumer list JSON" $ do + it "decodes detailed consumer list responses" $ do + eitherDecode consumerListJSON + `shouldBe` Right ConsumerListResponse + { consumerListTotal = 1 + , consumerListOffset = 0 + , consumerListLimit = 1024 + , consumerListConsumers = [consumerInfo] + } + + it "decodes consumer name responses" $ do + eitherDecode consumerNamesJSON + `shouldBe` Right ConsumerNamesResponse + { consumerNamesTotal = 2 + , consumerNamesOffset = 0 + , consumerNamesLimit = 1024 + , consumerNamesConsumers = ["orders-puller", "orders-worker"] + } + +durablePullConsumerConfigRequest :: ConsumerConfigRequest +durablePullConsumerConfigRequest = + consumerConfigRequest + [ withConsumerDurableName "orders-puller" + , withConsumerName "orders-puller" + , withConsumerDeliverPolicy DeliverAll + , withConsumerAckPolicy AckExplicit + , withConsumerReplayPolicy ReplayInstant + , withConsumerFilter (ConsumerFilterSubjects ["orders.created", "orders.updated"]) + , withConsumerMaxDeliver 5 + , withConsumerAckWait 30 + , withConsumerMaxAckPending 256 + , withConsumerInactiveThreshold 60 + ] + +durablePullConsumerConfigValue :: Value +durablePullConsumerConfigValue = object + [ "durable_name" .= ("orders-puller" :: String) + , "name" .= ("orders-puller" :: String) + , "deliver_policy" .= ("all" :: String) + , "ack_policy" .= ("explicit" :: String) + , "replay_policy" .= ("instant" :: String) + , "filter_subjects" .= ["orders.created" :: String, "orders.updated"] + , "max_deliver" .= (5 :: Int) + , "ack_wait" .= (30000000000 :: Integer) + , "max_ack_pending" .= (256 :: Int) + , "inactive_threshold" .= (60000000000 :: Integer) + ] + +pushConsumerConfigRequest :: ConsumerConfigRequest +pushConsumerConfigRequest = + consumerConfigRequest + [ withConsumerName "push" + , withConsumerDeliverSubject "deliver" + , withConsumerDeliverGroup "workers" + , withConsumerAckPolicy AckExplicit + , withConsumerMaxWaiting 512 + , withConsumerIdleHeartbeat 1 + , withConsumerHeadersOnly True + , withConsumerReplicas 1 + , withConsumerMemoryStorage True + ] + +pushConsumerConfigValue :: Value +pushConsumerConfigValue = object + [ "name" .= ("push" :: String) + , "deliver_subject" .= ("deliver" :: String) + , "deliver_group" .= ("workers" :: String) + , "ack_policy" .= ("explicit" :: String) + , "max_waiting" .= (512 :: Int) + , "idle_heartbeat" .= (1000000000 :: Integer) + , "headers_only" .= True + , "num_replicas" .= (1 :: Int) + , "mem_storage" .= True + ] + +targetKindConsumerConfigRequest :: ConsumerConfigRequest +targetKindConsumerConfigRequest = + applyConsumerTarget (NamedConsumer "right") . + applyConsumerKind (PushConsumer "deliver") $ + consumerConfigRequest + [ withConsumerName "wrong" + , withConsumerDeliverSubject "wrong" + , withConsumerDeliverGroup "workers" + ] + +targetKindConsumerConfigValue :: Value +targetKindConsumerConfigValue = object + [ "name" .= ("right" :: String) + , "deliver_subject" .= ("deliver" :: String) + , "deliver_group" .= ("workers" :: String) + ] + +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}" + +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}]}" + +consumerNamesJSON :: LBS.ByteString +consumerNamesJSON = + "{\"type\":\"io.nats.jetstream.api.v1.consumer_names_response\",\"total\":2,\"offset\":0,\"limit\":1024,\"consumers\":[\"orders-puller\",\"orders-worker\"]}" + +consumerInfo :: ConsumerInfo +consumerInfo = ConsumerInfo + { consumerInfoStreamName = "ORDERS" + , consumerInfoName = "orders-puller" + , consumerInfoCreated = timestamp + , consumerInfoConfig = ConsumerConfig + { consumerConfigDurableName = Nothing + , consumerConfigName = Nothing + , consumerConfigDescription = Nothing + , consumerConfigDeliverSubject = Nothing + , consumerConfigDeliverGroup = Nothing + , consumerConfigDeliverPolicy = DeliverAll + , consumerConfigAckPolicy = AckExplicit + , consumerConfigReplayPolicy = ReplayInstant + , consumerConfigFilterSubject = Nothing + , consumerConfigFilterSubjects = Nothing + , consumerConfigAckWait = Nothing + , consumerConfigMaxDeliver = Nothing + , consumerConfigMaxWaiting = Nothing + , consumerConfigMaxAckPending = Nothing + , consumerConfigInactiveThreshold = Nothing + , consumerConfigIdleHeartbeat = Nothing + , consumerConfigHeadersOnly = Nothing + , consumerConfigReplicas = Nothing + , consumerConfigMemoryStorage = Nothing + } + , consumerInfoDelivered = ConsumerSequenceInfo + { consumerSequenceConsumer = 2 + , consumerSequenceStream = 10 + , consumerSequenceLast = Nothing + } + , consumerInfoAckFloor = ConsumerSequenceInfo + { consumerSequenceConsumer = 1 + , consumerSequenceStream = 9 + , consumerSequenceLast = Nothing + } + , consumerInfoNumAckPending = 1 + , consumerInfoNumRedelivered = 0 + , consumerInfoNumWaiting = 0 + , consumerInfoNumPending = 3 + } + +timestamp = read "2024-01-01 00:00:00 UTC" diff --git a/test/Unit/JetStream/MessageSpec.hs b/test/Unit/JetStream/MessageSpec.hs new file mode 100644 --- /dev/null +++ b/test/Unit/JetStream/MessageSpec.hs @@ -0,0 +1,180 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.MessageSpec (spec) where + +import qualified API as Nats +import Data.IORef +import JetStream.Error (JetStreamError (..)) +import JetStream.Message (fetchMessages) +import JetStream.Message.Types +import JetStream.Options (newJetStreamContext) +import Publish (defaultPublishConfig) +import Test.Hspec +import Types.Msg (Payload, SID, Subject) + +spec :: Spec +spec = do + describe "status classification" $ do + it "classifies no-message status headers" $ do + classifyStatusHeaders (Just [("Status", "404"), ("Description", "No Messages")]) + `shouldBe` Just (PullNoMessages (Just "No Messages")) + + it "classifies request timeout status headers case-insensitively" $ do + classifyStatusHeaders (Just [("status", "408"), ("description", "Request Timeout")]) + `shouldBe` Just (PullRequestTimeout (Just "Request Timeout")) + + it "preserves unrecognised status details" $ do + classifyStatusHeaders (Just [("Status", "503"), ("Description", "Service Unavailable")]) + `shouldBe` Just (PullStatusError "503" (Just "Service Unavailable")) + + it "ignores normal message headers" $ do + classifyStatusHeaders (Just [("Nats-Stream", "ORDERS")]) + `shouldBe` Nothing + + describe "ack payloads" $ do + it "encodes ack verbs" $ do + ackPayload Ack `shouldBe` "+ACK" + nakPayload `shouldBe` "-NAK" + inProgressPayload `shouldBe` "+WPI" + termPayload `shouldBe` "+TERM" + + describe "message metadata" $ do + it "parses v1 JetStream ack reply subjects" $ do + let metadata = messageMetadata $ + Message "ORDERS.created" (Just "payload") Nothing + (Just "$JS.ACK.ORDERS.WORKER.3.10.2.123000000000.4") + Nothing + fmap messageMetadataStream metadata `shouldBe` Just "ORDERS" + fmap messageMetadataConsumer metadata `shouldBe` Just "WORKER" + fmap messageMetadataNumDelivered metadata `shouldBe` Just 3 + fmap messageMetadataStreamSequence metadata `shouldBe` Just 10 + fmap messageMetadataConsumerSequence metadata `shouldBe` Just 2 + fmap messageMetadataTimestamp metadata `shouldBe` Just (read "1970-01-01 00:02:03 UTC") + fmap messageMetadataNumPending metadata `shouldBe` Just 4 + fmap messageMetadataDomain metadata `shouldBe` Just Nothing + + it "parses v2 JetStream ack reply subjects with domains" $ do + let metadata = messageMetadata $ + Message "ORDERS.created" (Just "payload") Nothing + (Just "$JS.ACK.HUB.ACCOUNT.ORDERS.WORKER.1.11.3.124000000000.0") + Nothing + fmap messageMetadataDomain metadata `shouldBe` Just (Just "HUB") + fmap messageMetadataStreamSequence metadata `shouldBe` Just 11 + + it "rejects non-JetStream reply subjects" $ do + messageMetadata (Message "ORDERS.created" Nothing Nothing (Just "_INBOX.reply") Nothing) + `shouldBe` Nothing + + describe "pull request payloads" $ do + it "encodes one-message requests with an expires timeout" $ do + pullRequestPayload 1 defaultPullRequest + `shouldBe` "{\"batch\":1,\"expires\":1000000000}" + + it "encodes no-wait requests without expires" $ do + let request = pullRequest [withFetchWait (FetchNoWaitMicros 1000000)] + pullRequestPayload 1 request + `shouldBe` "{\"batch\":1,\"no_wait\":true}" + + describe "pull fetch" $ do + it "uses one explicit reply inbox for a batch request" $ do + publishCalls <- newIORef [] + unsubscribeCalls <- newIORef [] + callbackRef <- newIORef Nothing + let client = fakeClient publishCalls unsubscribeCalls callbackRef + request = pullRequest [withFetchBatch 2] + result <- fetchMessages (newJetStreamContext client []) "ORDERS" "WORKER" request + + publishCalls' <- readIORef publishCalls + publishCalls' `shouldBe` + [ ( "$JS.API.CONSUMER.MSG.NEXT.ORDERS.WORKER" + , Just "{\"batch\":2,\"expires\":1000000000}" + , Just "_INBOX.batch" + ) + ] + unsubscribeCalls' <- readIORef unsubscribeCalls + unsubscribeCalls' `shouldBe` ["sid-1"] + case result of + Left err -> + expectationFailure ("fetch failed: " ++ show err) + Right response -> do + fmap messagePayload (pullResponseMessages response) `shouldBe` [Just "one", Just "two"] + pullResponseStatus response `shouldBe` Nothing + + it "returns a JetStream timeout when no pull response arrives" $ do + publishCalls <- newIORef [] + unsubscribeCalls <- newIORef [] + let client = silentFakeClient publishCalls unsubscribeCalls + request = pullRequest [withFetchWait (FetchNoWaitMicros 1000)] + result <- fetchMessages (newJetStreamContext client []) "ORDERS" "WORKER" request + + result `shouldBe` Left JetStreamTimeout + publishCalls' <- readIORef publishCalls + publishCalls' `shouldBe` + [ ( "$JS.API.CONSUMER.MSG.NEXT.ORDERS.WORKER" + , Just "{\"batch\":1,\"no_wait\":true}" + , Just "_INBOX.timeout" + ) + ] + unsubscribeCalls' <- readIORef unsubscribeCalls + unsubscribeCalls' `shouldBe` ["sid-timeout"] + +fakeClient + :: IORef [(Subject, Maybe Payload, Maybe Subject)] + -> IORef [SID] + -> IORef (Maybe (Maybe Nats.MsgView -> IO ())) + -> Nats.Client +fakeClient publishCalls unsubscribeCalls callbackRef = + Nats.Client + { Nats.publish = \subject options -> do + let (body, _, _, replyTo') = foldr ($) defaultPublishConfig options + modifyIORef' publishCalls (++ [(subject, body, replyTo')]) + callback <- readIORef callbackRef + case callback of + Nothing -> + pure () + Just deliver -> do + deliver (Just (fakeMsg "one")) + deliver (Just (fakeMsg "two")) + , Nats.subscribe = \_ _ callback -> do + writeIORef callbackRef (Just callback) + pure "sid-1" + , Nats.request = \_ _ _ -> pure "sid-request" + , Nats.unsubscribe = \sid -> + modifyIORef' unsubscribeCalls (++ [sid]) + , Nats.newInbox = pure "_INBOX.batch" + , Nats.ping = id + , Nats.flush = pure () + , Nats.reset = pure () + , Nats.close = pure () + } + +fakeMsg :: Payload -> Nats.MsgView +fakeMsg body = + Nats.MsgView + { Nats.subject = "ORDERS.created" + , Nats.sid = "sid-1" + , Nats.replyTo = Just "$JS.ACK.ORDERS.WORKER.1.1.1" + , Nats.payload = Just body + , Nats.headers = Nothing + } + +silentFakeClient + :: IORef [(Subject, Maybe Payload, Maybe Subject)] + -> IORef [SID] + -> Nats.Client +silentFakeClient publishCalls unsubscribeCalls = + Nats.Client + { Nats.publish = \subject options -> do + let (body, _, _, replyTo') = foldr ($) defaultPublishConfig options + modifyIORef' publishCalls (++ [(subject, body, replyTo')]) + , Nats.subscribe = \_ _ _ -> + pure "sid-timeout" + , Nats.request = \_ _ _ -> pure "sid-request" + , Nats.unsubscribe = \sid -> + modifyIORef' unsubscribeCalls (++ [sid]) + , Nats.newInbox = pure "_INBOX.timeout" + , Nats.ping = id + , Nats.flush = pure () + , Nats.reset = pure () + , Nats.close = pure () + } diff --git a/test/Unit/JetStream/ProtocolSpec.hs b/test/Unit/JetStream/ProtocolSpec.hs new file mode 100644 --- /dev/null +++ b/test/Unit/JetStream/ProtocolSpec.hs @@ -0,0 +1,107 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.ProtocolSpec (spec) where + +import qualified API as Nats +import Data.Aeson (eitherDecodeStrict) +import Data.Either (isLeft) +import JetStream.Error (JetStreamError (..)) +import JetStream.Options (newJetStreamContext, withDomain) +import JetStream.Protocol.Headers (statusCode, statusDescription) +import JetStream.Protocol.Request (decodeJetStreamResponse) +import JetStream.Protocol.Subject +import JetStream.Stream.Types (StreamMessage (..)) +import JetStream.Types + ( AccountInfo (..) + , AccountTier (..) + , DeliverPolicy + ) +import Test.Hspec + +spec :: Spec +spec = do + describe "JetStream subjects" $ do + it "builds default API subjects" $ do + streamInfoSubject testContext "ORDERS" `shouldBe` "$JS.API.STREAM.INFO.ORDERS" + + it "builds domain API subjects" $ do + let ctx = newJetStreamContext fakeClient [withDomain "HUB"] + consumerNextSubject ctx "ORDERS" "WORKER" `shouldBe` "$JS.HUB.API.CONSUMER.MSG.NEXT.ORDERS.WORKER" + + it "builds administrative API subjects" $ do + accountInfoSubject testContext `shouldBe` "$JS.API.INFO" + streamMessageGetSubject testContext "ORDERS" `shouldBe` "$JS.API.STREAM.MSG.GET.ORDERS" + streamMessageDeleteSubject testContext "ORDERS" `shouldBe` "$JS.API.STREAM.MSG.DELETE.ORDERS" + consumerPauseSubject testContext "ORDERS" "WORKER" `shouldBe` "$JS.API.CONSUMER.PAUSE.ORDERS.WORKER" + consumerResetSubject testContext "ORDERS" "WORKER" `shouldBe` "$JS.API.CONSUMER.RESET.ORDERS.WORKER" + + describe "JetStream status headers" $ do + it "extracts status and description from normal headers" $ do + let msg = fakeMsg { Nats.headers = Just [("Status", "404"), ("Description", "No Messages")] } + statusCode msg `shouldBe` Just 404 + statusDescription msg `shouldBe` Just "No Messages" + + describe "JetStream response decoding" $ do + it "decodes API error envelopes" $ do + case decodeJetStreamResponse "{\"error\":{\"code\":404,\"err_code\":10059,\"description\":\"stream not found\"}}" of + Left (JetStreamApiFailure err) -> + show err `shouldContain` "stream not found" + other -> + expectationFailure ("unexpected decode result: " ++ show (other :: Either JetStreamError ())) + + it "decodes stored stream messages from base64 JSON fields" $ do + case decodeJetStreamResponse storedMessageResponse of + Right message -> do + streamMessageSubject message `shouldBe` "ORDERS.created" + streamMessageSequence message `shouldBe` 2 + streamMessagePayload message `shouldBe` Just "payload" + other -> + expectationFailure ("unexpected stored message decode: " ++ show (other :: Either JetStreamError StreamMessage)) + + it "decodes account information" $ do + case decodeJetStreamResponse accountInfoResponse of + Right info -> do + accountTierStreams (accountInfoTier info) `shouldBe` 1 + accountInfoDomain info `shouldBe` Just "hub" + length (accountInfoTiers info) `shouldBe` 1 + other -> + expectationFailure ("unexpected account info decode: " ++ show (other :: Either JetStreamError AccountInfo)) + + it "rejects parameterised deliver policies without their companion fields" $ do + (eitherDecodeStrict "\"by_start_sequence\"" :: Either String DeliverPolicy) + `shouldSatisfy` isLeft + (eitherDecodeStrict "\"by_start_time\"" :: Either String DeliverPolicy) + `shouldSatisfy` isLeft + +testContext = + newJetStreamContext fakeClient [] + +storedMessageResponse = + "{\"message\":{\"subject\":\"ORDERS.created\",\"seq\":2,\"data\":\"cGF5bG9hZA==\",\"time\":\"2024-01-01T00:00:00Z\"}}" + +accountInfoResponse = + "{\"memory\":10,\"storage\":20,\"reserved_memory\":0,\"reserved_storage\":0,\"streams\":1,\"consumers\":2,\"limits\":{\"max_memory\":-1,\"max_storage\":-1,\"max_streams\":-1,\"max_consumers\":-1,\"max_ack_pending\":-1,\"memory_max_stream_bytes\":-1,\"storage_max_stream_bytes\":-1,\"max_bytes_required\":false},\"domain\":\"hub\",\"api\":{\"level\":1,\"total\":3,\"errors\":0},\"tiers\":{\"R1\":{\"memory\":10,\"storage\":20,\"reserved_memory\":0,\"reserved_storage\":0,\"streams\":1,\"consumers\":2,\"limits\":{\"max_memory\":-1,\"max_storage\":-1,\"max_streams\":-1,\"max_consumers\":-1,\"max_ack_pending\":-1,\"memory_max_stream_bytes\":-1,\"storage_max_stream_bytes\":-1,\"max_bytes_required\":false}}}}" + +fakeClient :: Nats.Client +fakeClient = + Nats.Client + { Nats.publish = \_ _ -> pure () + , Nats.subscribe = \_ _ _ -> pure "0" + , Nats.request = \_ _ _ -> pure "0" + , Nats.unsubscribe = \_ -> pure () + , Nats.newInbox = pure "_INBOX.TEST" + , Nats.ping = id + , Nats.flush = pure () + , Nats.reset = pure () + , Nats.close = pure () + } + +fakeMsg :: Nats.MsgView +fakeMsg = + Nats.MsgView + { Nats.subject = "subject" + , Nats.sid = "sid" + , Nats.replyTo = Nothing + , Nats.payload = Nothing + , Nats.headers = Nothing + } diff --git a/test/Unit/JetStream/PublishSpec.hs b/test/Unit/JetStream/PublishSpec.hs new file mode 100644 --- /dev/null +++ b/test/Unit/JetStream/PublishSpec.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE OverloadedStrings #-} + +module JetStream.PublishSpec (spec) where + +import Data.Aeson (eitherDecodeStrict) +import JetStream.Publish.Types +import Test.Hspec + +spec :: Spec +spec = + parallel $ do + describe "publishHeaders" $ do + it "renders JetStream publish expectation headers" $ do + publishHeaders + [ withMsgId "msg-1" + , withExpectedStream "ORDERS" + , withPublishExpectation (ExpectedLastSequence 12) + , withHeaders [("Nats-Expected-Test", "custom")] + ] + `shouldBe` [ ("Nats-Msg-Id", "msg-1") + , ("Nats-Expected-Stream", "ORDERS") + , ("Nats-Expected-Last-Sequence", "12") + , ("Nats-Expected-Test", "custom") + ] + + it "keeps publish expectations mutually exclusive" $ do + publishHeaders + [ withPublishExpectation (ExpectedLastSequence 12) + , withPublishExpectation (ExpectedLastSubjectSequence 7) + , withPublishExpectation (ExpectedLastMsgId "msg-0") + ] + `shouldBe` [("Nats-Expected-Last-Sequence", "12")] + + describe "PublishAck" $ do + it "decodes required and optional publish ack fields" $ do + eitherDecodeStrict + "{\"stream\":\"ORDERS\",\"seq\":42,\"duplicate\":true,\"domain\":\"hub\"}" + `shouldBe` Right PublishAck + { publishAckStream = "ORDERS" + , publishAckSequence = 42 + , publishAckDuplicate = Just True + , publishAckDomain = Just "hub" + } + + it "decodes publish ack without optional fields" $ do + eitherDecodeStrict "{\"stream\":\"ORDERS\",\"seq\":42}" + `shouldBe` Right PublishAck + { publishAckStream = "ORDERS" + , publishAckSequence = 42 + , publishAckDuplicate = Nothing + , publishAckDomain = Nothing + } diff --git a/test/Unit/ParserSpec.hs b/test/Unit/ParserSpec.hs --- a/test/Unit/ParserSpec.hs +++ b/test/Unit/ParserSpec.hs @@ -98,3 +98,22 @@ reason `shouldContain` "invalid HMSG status" other -> expectationFailure ("unexpected parse result: " ++ show other) + + it "rejects HMSG frames whose total size is smaller than the header size" $ do + let input = "HMSG FOO 13 24 23\r\n" + case parse parserApi input of + DropPrefix n reason -> do + n `shouldSatisfy` (> 0) + reason `shouldContain` "invalid HMSG sizes" + other -> + expectationFailure ("unexpected parse result: " ++ show other) + + it "rejects HMSG headers with an empty key after trimming" $ do + let input = + "HMSG FOO 13 23 23\r\nNATS/1.0\r\n : VALUE\r\n\r\n\r\n" + case parse parserApi input of + DropPrefix n reason -> do + n `shouldSatisfy` (> 0) + reason `shouldContain` "invalid HMSG headers" + other -> + expectationFailure ("unexpected parse result: " ++ show other) diff --git a/test/Unit/PubSpec.hs b/test/Unit/PubSpec.hs --- a/test/Unit/PubSpec.hs +++ b/test/Unit/PubSpec.hs @@ -5,6 +5,8 @@ import Control.Monad import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LBS +import Data.Either (isLeft) +import Fixtures (invalidSubjectCases) import Test.Hspec import Text.Printf import Transformers.Transformers @@ -15,6 +17,8 @@ spec = do transformerCases validateCases + validateSubjectCases + validateHeaderCases explicitTransformerCases :: [(Pub, BS.ByteString)] explicitTransformerCases = [ @@ -54,3 +58,27 @@ forM_ explicitValidatorCases $ \(input, want) -> it (printf "correctly validates %s" (show input)) $ do validate input `shouldBe` want + +validateSubjectCases = parallel $ do + describe "PUB subject validater" $ do + forM_ invalidSubjectCases $ \input -> + it (printf "rejects invalid publish subject %s" (show input)) $ do + validate (Pub input Nothing Nothing Nothing) `shouldSatisfy` isLeft + + forM_ invalidSubjectCases $ \input -> + it (printf "rejects invalid reply subject %s" (show input)) $ do + validate (Pub "FOO" (Just input) Nothing Nothing) `shouldSatisfy` isLeft + +validateHeaderCases = parallel $ do + describe "PUB header validater" $ do + it "rejects header keys containing a colon" $ do + validate (Pub "FOO" Nothing (Just [("bad:key", "value")]) Nothing) + `shouldSatisfy` isLeft + + it "rejects header keys containing CRLF" $ do + validate (Pub "FOO" Nothing (Just [("bad\r\nkey", "value")]) Nothing) + `shouldSatisfy` isLeft + + it "rejects header values containing CRLF" $ do + validate (Pub "FOO" Nothing (Just [("key", "bad\r\nvalue")]) Nothing) + `shouldSatisfy` isLeft diff --git a/test/Unit/SubSpec.hs b/test/Unit/SubSpec.hs --- a/test/Unit/SubSpec.hs +++ b/test/Unit/SubSpec.hs @@ -5,6 +5,8 @@ import Control.Monad import Data.ByteString import qualified Data.ByteString.Lazy as LBS +import Data.Either (isLeft) +import Fixtures (invalidSubjectCases) import Subscription.Types import Test.Hspec import Text.Printf @@ -17,6 +19,7 @@ subscribeOptionCases transformerCases validateCases + validateSubjectCases subscribeOptionCases = parallel $ do describe "subscribe config" $ do @@ -50,3 +53,12 @@ forM_ explicitValidaterCases $ \(input, want) -> it (printf "correctly validates %s" (show input)) $ do validate input `shouldBe` want + +validateSubjectCases = parallel $ do + describe "SUB subject validater" $ do + forM_ invalidSubjectCases $ \input -> + it (printf "rejects invalid subscription subject %s" (show input)) $ do + validate (Sub input Nothing "a") `shouldSatisfy` isLeft + + it "rejects queue groups containing whitespace" $ do + validate (Sub "FOO" (Just "bad queue") "a") `shouldSatisfy` isLeft