diff --git a/hsnsq.cabal b/hsnsq.cabal --- a/hsnsq.cabal +++ b/hsnsq.cabal @@ -1,5 +1,5 @@ name: hsnsq -version: 0.1.0.0 +version: 0.1.1.0 synopsis: Haskell NSQ client. description: Currently a primitive NSQ client, it implements the very basics of an NSQ client for communicating with single NSQ servers. homepage: https://github.com/gamelost/hsnsq @@ -15,6 +15,10 @@ library exposed-modules: Network.NSQ + , Network.NSQ.Connection + , Network.NSQ.Identify + , Network.NSQ.Parser + , Network.NSQ.Types default-language: Haskell2010 ghc-options: -Wall -fno-warn-missing-signatures -O2 hs-source-dirs: src diff --git a/src/Network/NSQ/Connection.hs b/src/Network/NSQ/Connection.hs new file mode 100644 --- /dev/null +++ b/src/Network/NSQ/Connection.hs @@ -0,0 +1,167 @@ +{-# LANGUAGE OverloadedStrings #-} +module Network.NSQ.Connection + ( defaultConfig + , establish + + ) where + +import Control.Monad.Reader +import Control.Monad.State.Strict +import Data.List +import Data.Maybe +import Prelude hiding (log) +import Network.HostName + +import qualified Data.Text as T +import qualified Data.ByteString as BS + +import qualified Pipes.Network.TCP as PNT +import qualified Pipes.Attoparsec as PA +import qualified Pipes.Prelude as PP +import Pipes + +import Control.Applicative +import Control.Concurrent.Async +import Control.Concurrent.STM + +import System.Log.Logger (debugM, errorM, warningM, infoM) + +import Network.NSQ.Types (NSQConnection(..), Message, Command, server, port, logName, LogName, identConf) +import qualified Network.NSQ.Types as NSQ +import qualified Network.NSQ.Identify as NSQ +import qualified Network.NSQ.Parser as NSQ + + +-- +-- Default Config +-- +-- TODO: standardize on some sort of logger hierchary (nsq server/topic?) +-- NSQ.[subsystem].[topic].[connection] - message +-- NSQ.[subsystem].[custom] .... +-- +defaultConfig :: String -> IO NSQConnection +defaultConfig serverHost = do + localHost <- T.pack <$> getHostName + let localClientId = T.takeWhile (/= '.') localHost + let localIdent = NSQ.defaultIdentify localClientId localHost + + return $ NSQConnection serverHost 4150 "NSQ.Connection." localIdent + +-- +-- Establish a session with this server +-- +-- Detail: +-- * Support connecting to a particular nsqd and doing the needful to establish identification and so forth +-- * Auto-handle heartbeat and all related stuff +-- * A higher layer will handle the message reading/balancing between multiplex nsqd connection for a particular topic/channel +-- +establish :: NSQConnection -> TQueue Message -> TQueue Command -> IO () +establish conn topicQueue reply = PNT.withSocketsDo $ + -- Establish stocket + PNT.connect (server conn) (show $ port conn) (\(sock, _) -> do + + -- TODO: maybe consider PNT.fromSocketN so that we can adjust fetch size if needed downstream + let send = (log "send" $ logName conn) >-> PNT.toSocket sock + let recv = PNT.fromSocket sock 8192 >-> (log "recv" $ logName conn) + race_ + (handleNSQ conn recv send topicQueue) + (runEffect $ handleReply reply >-> showCommand >-> send) -- Handles user replies + ) + +-- +-- NSQ Reply handler +-- +handleReply :: (Monad m, MonadIO m) => TQueue Command -> Producer NSQ.Command m () +handleReply queue = forever $ do + cmd <- liftIO $ atomically $ readTQueue queue + yield cmd + +-- +-- The NSQ handler and protocol +-- +handleNSQ :: (Monad m, MonadIO m) => NSQConnection -> Producer BS.ByteString m () -> Consumer BS.ByteString m () -> TQueue Message -> m () +handleNSQ sc recv send topicQueue = do + -- Initial handshake to kick off the handshake + runEffect $ (initialHandshake $ identConf sc) >-> showCommand >-> send + + -- Rest of the handshake process (parsing and dealing with identification) + runEffect $ (nsqParserErrorLogging (logName sc) recv) >-> identReply + + -- Setup the topic/channel/rdy + runEffect $ setupTopic >-> showCommand >-> send + + -- Regular nsq streaming + runEffect $ (nsqParserErrorLogging (logName sc) recv) >-> (command (logName sc) topicQueue) >-> showCommand >-> send + + return () + + where + -- Initial Handshake + initialHandshake im = do + yield $ NSQ.Protocol + yield $ NSQ.Identify im + return () + + -- TODO; replace it with actual logic + setupTopic = do + yield $ NSQ.Sub "glc-gamestate" "netheril.elder.lan." False + yield $ NSQ.Rdy 1 + return () + + -- Process the ident reply + identReply = do + identR <- await + + -- TODO: do stuff with it + liftIO $ debugM (logName sc) ("IDENT: " ++ show identR) + + return () + +-- +-- Parses incoming nsq messages and emits any errors to a log and keep going +-- +nsqParserErrorLogging :: MonadIO m => LogName -> Producer BS.ByteString m () -> Producer NSQ.Message m () +nsqParserErrorLogging l producer = do + (result, rest) <- lift $ runStateT (PA.parse NSQ.message) producer + + case result of + Nothing -> liftIO $ errorM l "Pipe is exhausted for nsq parser\n" + Just y -> do + case y of + Right x -> (liftIO $ debugM l ("msg: " ++ show x)) >> yield x + Left x -> liftIO $ errorM l (show x) + nsqParserErrorLogging l rest + +-- +-- Format outbound NSQ Commands +-- +showCommand :: Monad m => Pipe NSQ.Command BS.ByteString m () +showCommand = PP.map NSQ.encode + +-- +-- Log anything that passes through this stream to a logfile +-- +log :: MonadIO m => String -> LogName -> Pipe BS.ByteString BS.ByteString m r +log w l = forever $ do + x <- await + liftIO $ debugM l (w ++ ": " ++ show x) -- TODO: need a better way to log raw protocol messages + yield x + +-- +-- Do something with the inbound message +-- +command :: (Monad m, MonadIO m) => LogName -> TQueue Message -> Pipe NSQ.Message NSQ.Command m () +command l topicQueue = forever $ do + msg <- await + + case msg of + -- TODO: currently no-op + NSQ.OK -> return () + NSQ.Heartbeat -> yield $ NSQ.NOP + -- TODO: Implement a way to close our connection gracefully + NSQ.CloseWait -> liftIO $ infoM l ("Error: Server closed queue") + -- TODO: should pass it onto the client or have a callback + NSQ.Error e -> liftIO $ errorM l ("Error: " ++ show e) + NSQ.Message _ _ _ _ -> liftIO $ atomically $ writeTQueue topicQueue msg + -- TODO: should pass it onto the client or have a callback + NSQ.CatchAllMessage f m -> liftIO $ warningM l ("Error: Frame - " ++ show f ++ " - Msg - " ++ show m) diff --git a/src/Network/NSQ/Identify.hs b/src/Network/NSQ/Identify.hs new file mode 100644 --- /dev/null +++ b/src/Network/NSQ/Identify.hs @@ -0,0 +1,96 @@ +{-# LANGUAGE OverloadedStrings #-} +module Network.NSQ.Identify + ( defaultIdentify + , defaultUserAgent + , encodeMetadata + + ) where + +import Prelude hiding (take) +import Data.Maybe + +import qualified Data.Map.Strict as Map + +import qualified Data.ByteString.Lazy as BL + +import qualified Data.Text as T + +import Data.Aeson ((.=)) +import qualified Data.Aeson as A +import qualified Data.Aeson.Types as A + +import Network.NSQ.Types + + +defaultIdentify :: T.Text -> T.Text -> IdentifyMetadata +defaultIdentify cid host = IdentifyMetadata + { ident = Identification cid host Nothing Nothing Nothing + , tls = Nothing + , compression = Nothing + , heartbeatInterval = Nothing + , outputBufferSize = Nothing + , outputBufferTimeout = Nothing + , sampleRate = Nothing + , custom = Nothing + , customNegotiation = False + } + +defaultUserAgent :: T.Text +defaultUserAgent = "hsnsq/0.1.0.0" + +featureNegotiation :: IdentifyMetadata -> [A.Pair] +featureNegotiation im = catMaybes + ( + tlsSettings (tls im) + ++ + [ optionalSettings "heartbeat_interval" (-1) $ heartbeatInterval im + , optionalSettings "output_buffer_size" (-1) $ outputBufferSize im + , optionalSettings "output_buffer_timeout" (-1) $ outputBufferTimeout im + , optionalSettings "sample_rate" 0 $ sampleRate im + ] + ++ + optionalCompression (compression im) + ) + +optionalSettings :: T.Text -> Int -> Maybe OptionalSetting -> Maybe A.Pair +optionalSettings _ _ Nothing = Nothing +optionalSettings name def (Just Disabled) = Just (name, A.toJSON def) +optionalSettings name _ (Just (Custom val)) = Just (name, A.toJSON val) + +optionalCompression :: Maybe Compression -> [Maybe A.Pair] +optionalCompression Nothing = [] +optionalCompression (Just NoCompression) = Just `fmap` [ "snappy" .= False, "deflate" .= False ] +optionalCompression (Just Snappy) = Just `fmap` [ "snappy" .= True, "deflate" .= False ] +optionalCompression (Just (Deflate l)) = Just `fmap` [ "snappy" .= False, "deflate" .= True, "deflate_level" .= l ] + +customMetadata :: Maybe (Map.Map T.Text T.Text) -> [A.Pair] +customMetadata Nothing = [] +customMetadata (Just val) = Map.foldrWithKey (\k v xs -> (k .= v):xs) [] val + +tlsSettings :: Maybe TLS -> [Maybe A.Pair] +tlsSettings Nothing = [] +tlsSettings (Just NoTLS) = [Just $ "tls_v1" .= False] +tlsSettings (Just TLSV1) = [Just $ "tls_v1" .= True] + +-- TODO: This is an Orphan instance because the type is in types.hs, need to fix this +instance A.ToJSON IdentifyMetadata where + toJSON im@(IdentifyMetadata{ident=i}) = A.object + ( + -- Identification section + [ "client_id" .= clientId i + , "hostname" .= hostname i + , "short_id" .= fromMaybe (clientId i) (shortId i) + , "long_id" .= fromMaybe (hostname i) (longId i) + , "user_agent" .= fromMaybe defaultUserAgent (userAgent i) + + -- Feature Negotiation section + , "feature_negotiation" .= (not (null $ featureNegotiation im) || customNegotiation im) + ] + ++ + featureNegotiation im + ++ + customMetadata (custom im) + ) + +encodeMetadata :: IdentifyMetadata -> BL.ByteString +encodeMetadata = A.encode diff --git a/src/Network/NSQ/Parser.hs b/src/Network/NSQ/Parser.hs new file mode 100644 --- /dev/null +++ b/src/Network/NSQ/Parser.hs @@ -0,0 +1,142 @@ +{-# LANGUAGE OverloadedStrings #-} +module Network.NSQ.Parser + ( message + , decode + , encode + ) where + +import Data.Word +import Data.Monoid +import Control.Applicative +import Data.Attoparsec.Binary +import Data.Attoparsec.ByteString hiding (count) +import Data.Int +import Prelude hiding (take) + +import qualified Data.List as DL + +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as C8 +import qualified Data.ByteString.Lazy as BL +import qualified Data.ByteString.Lazy.Builder as BL + +import qualified Data.Text.Encoding as T + +import Network.NSQ.Types +import Network.NSQ.Identify + + +decode :: BS.ByteString -> Maybe Message +decode str = case parseOnly message str of + Left _ -> Nothing + Right r -> Just r + +-- TODO convert "E_*" into Error +command :: BS.ByteString -> Message +command "_heartbeat_" = Heartbeat +command "OK" = OK +command "CLOSE_WAIT" = CloseWait +command x = CatchAllMessage (FTUnknown 99) x -- TODO: Better FrameType + +frameType :: Int32 -> FrameType +frameType 0 = FTResponse +frameType 1 = FTError +frameType 2 = FTMessage +frameType x = FTUnknown x + +-- TODO: do sanity check such as checking that the size is of a minimal +-- size, then parsing the frameType, then the remainder (fail "messg") +message :: Parser Message +message = do + size <- fromIntegral <$> anyWord32be + ft <- frameType <$> fromIntegral <$> anyWord32be + frame ft (size - 4) -- Taking in accord the frameType + +frame :: FrameType -> Int -> Parser Message +frame FTError size = Error <$> take size +frame FTResponse size = command <$> take size +frame FTMessage size = Message + <$> (fromIntegral <$> anyWord64be) + <*> anyWord16be + -- 16 bytes message id + <*> take 16 + -- Taking in accord timestamp/attempts/msgid + <*> take (size - 26) +frame ft size = CatchAllMessage ft <$> take size + + +-- TODO: this won't work for streaming the data... +-- Should provide two api, one for in memory (ie where we count up the length of the data manualy +-- And a "streaming" version in which we know the actual size before streaming (ie streaming from a file for ex) +sizedData :: BS.ByteString -> BL.Builder +sizedData dat = BL.word32BE (fromIntegral $ BS.length dat) <> BL.byteString dat + + +-- Body of a foldl to build up a sequence of concat sized data +concatSizedData :: (Word32, Word32, BL.Builder) -> BS.ByteString -> (Word32, Word32, BL.Builder) +concatSizedData (totalSize, count, xs) dat = ( + totalSize + 4 + fromIntegral (BS.length dat), -- Add 4 to accord for message size + count + 1, + xs <> sizedData dat + ) + +-- Reply +-- Note: +-- * One sub (topic/channel) per nsqd connection max, any more will get an E_INVALID +-- * Seems to be able to publish to any topic/channel without limitation +encode :: Command -> BS.ByteString +encode Protocol = " V2" +encode NOP = "NOP\n" +encode Cls = "CLS\n" +encode (Identify identify) = BL.toStrict $ BL.toLazyByteString ( + BL.byteString "IDENTIFY\n" <> + sizedData (BL.toStrict $ encodeMetadata identify) + ) +encode (Sub topic channel ephemeral) = BL.toStrict $ BL.toLazyByteString ( + BL.byteString "SUB " <> + BL.byteString (T.encodeUtf8 topic) <> + BL.byteString " " <> + BL.byteString (T.encodeUtf8 channel) <> + BL.byteString (if ephemeral then "#ephemeral" else "") <> + BL.byteString "\n" + ) +encode (Pub topic dat) = BL.toStrict $ BL.toLazyByteString ( + BL.byteString "PUB " <> + BL.byteString (T.encodeUtf8 topic) <> + BL.byteString "\n" <> + sizedData dat + ) +encode (MPub topic dx) = BL.toStrict $ BL.toLazyByteString ( + BL.byteString "MPUB " <> + BL.byteString (T.encodeUtf8 topic) <> + BL.byteString "\n" <> + BL.word32BE (totalSize + 4) <> -- Accord for message count + BL.word32BE totalCount <> + content + ) + where + (totalSize, totalCount, content) = DL.foldl' concatSizedData (0, 0, mempty) dx + +encode (Rdy count) = BL.toStrict $ BL.toLazyByteString ( + BL.byteString "RDY " <> + BL.byteString (C8.pack $ show count) <> + BL.byteString "\n" + ) +encode (Fin msg_id) = BL.toStrict $ BL.toLazyByteString ( + BL.byteString "FIN " <> + BL.byteString msg_id <> + BL.byteString "\n" + ) +encode (Touch msg_id) = BL.toStrict $ BL.toLazyByteString ( + BL.byteString "TOUCH " <> + BL.byteString msg_id <> + BL.byteString "\n" + ) +encode (Req msg_id timeout) = BL.toStrict $ BL.toLazyByteString ( + BL.byteString "REQ " <> + BL.byteString msg_id <> + BL.byteString " " <> + BL.byteString (C8.pack $ show timeout) <> + BL.byteString "\n" + ) +encode (Command m) = m diff --git a/src/Network/NSQ/Types.hs b/src/Network/NSQ/Types.hs new file mode 100644 --- /dev/null +++ b/src/Network/NSQ/Types.hs @@ -0,0 +1,149 @@ +module Network.NSQ.Types + ( MsgId + , Topic + , Channel + , LogName + + , Message(..) + , Command(..) + + -- TODO: probably don't want to export constructor here but want pattern matching + , FrameType(..) + , ErrorType(..) + + , OptionalSetting(..) + , TLS(..) + , Compression(..) + , Identification(..) + , IdentifyMetadata(..) + + -- Connection config/state + , NSQConnection(..) + + ) where + +import Data.Int +import Data.Word +import Network + +import qualified Data.ByteString as BS +import qualified Data.Map.Strict as Map +import qualified Data.Text as T + +-- High level arch: +-- * One queue per topic/channel +-- * This queue can be feed by multiple nsqd (load balanced/nsqlookup for ex) +-- * Probably will have one set of state/config per nsqd connection and per queue/topic/channel +-- * Can probably later on provide helpers for consuming the queue + + +-- | Per Connection configuration +-- * Per nsqd (connection) state (rdy, load balance, etc) +-- * Per topic state (channel related info and which nsqd connection) +-- * Global? state (do we have any atm? maybe configuration?) +-- TODO: consider using monad-journal logger for pure code tracing +data NSQConnection = NSQConnection + { server :: String + , port :: PortNumber + , logName :: LogName + + , identConf :: IdentifyMetadata + } + +-- | Logger Name for a connection (hslogger format) +type LogName = String + +-- | Message Id, it is a 16-byte hexdecimal string encoded as ASCII +type MsgId = BS.ByteString + +-- | NSQ Topic, the only allowed character in a topic is @\\.a-zA-Z0-9_-@ +type Topic = T.Text + +-- | NSQ Channel, the only allowed character in a channel is @\\.a-zA-Z0-9_-@ +-- A channel can be marked as ephemeral by toggling the 'Bool' value to +-- true in 'Sub' +type Channel = T.Text + +-- | NSQ Command +data Command = Protocol -- ^ The protocol version + | NOP -- ^ No-op, usually used in reply to a 'Heartbeat' request from the server. + | Identify IdentifyMetadata -- ^ Client Identification + possible features negotiation. + | Sub Topic Channel Bool -- ^ Subscribe to a specified 'Topic'/'Channel', use 'True' if its an ephemeral channel. + | Pub Topic BS.ByteString -- ^ Publish a message to the specified 'Topic'. + | MPub Topic [BS.ByteString] -- ^ Publish multiple messages to a specified 'Topic'. + | Rdy Word64 -- ^ Update @RDY@ state (ready to recieve messages). Number of message you can process at once. + | Fin MsgId -- ^ Finish a message. + | Req MsgId Word64 -- ^ Re-queue a message (failure to process), Timeout is in milliseconds. + | Touch MsgId -- ^ Reset the timeout for an in-flight message. + | Cls -- ^ Cleanly close the connection to the NSQ daemon. + | Command BS.ByteString -- ^ Catch-all command for future expansion/custom commands. + deriving Show + +-- | Frame Type of the incoming data from the NSQ daemon. +data FrameType = FTResponse -- ^ Response to a 'Command' from the server. + | FTError -- ^ An error in response to a 'Command'. + | FTMessage -- ^ Messages. + | FTUnknown Int32 -- ^ For future extension for handling new Frame Types. + deriving Show + +-- | Types of error that the server can return in response to an 'Command' +data ErrorType = Invalid + | BadBody + | BadTopic + | BadChannel + | BadMessage + | PubFailed + | MPubFailed + | FinFailed + | ReqFailed + | TouchFailed + deriving Show + +-- | The message and replies back from the server. +data Message = OK -- ^ Everything is allright. + | Heartbeat -- ^ Heartbeat, reply with the 'NOP' 'Command'. + | CloseWait -- ^ Server has closed the connection. + + -- TODO: BS.ByteString -> ErrorType + | Error BS.ByteString -- ^ The server sent back an error. + + | Message Int64 Word16 MsgId BS.ByteString -- ^ A message to be processed. The values are: Nanosecond Timestamp, number of attempts, Message Id, and the content of the message to be processed. + + | CatchAllMessage FrameType BS.ByteString -- ^ Catch-all message for future expansion. This currently includes the reply from 'Identify' if feature negotiation is set. + deriving Show + + +-- Feature and Identification +data OptionalSetting = Disabled | Custom Word64 + deriving Show + +data TLS = NoTLS | TLSV1 + deriving Show + +data Compression = NoCompression | Snappy | Deflate Word8 + deriving Show + +data Identification = Identification + { clientId :: T.Text + , hostname :: T.Text + , shortId :: Maybe T.Text -- Deprecated in favor of client_id + , longId :: Maybe T.Text -- Deprecated in favor of hostname + , userAgent :: Maybe T.Text -- Default (client_library_name/version) + } + deriving Show + +-- feature_negotiation - set automatically if anything is set +data IdentifyMetadata = IdentifyMetadata + { ident :: Identification + , tls :: Maybe TLS + , compression :: Maybe Compression + , heartbeatInterval :: Maybe OptionalSetting -- disabled = -1 + , outputBufferSize :: Maybe OptionalSetting -- disabled = -1 + , outputBufferTimeout :: Maybe OptionalSetting -- disabled = -1 + , sampleRate :: Maybe OptionalSetting -- disabled = 0 + + -- Map of possible json value for future compat + , custom :: Maybe (Map.Map T.Text T.Text) + , customNegotiation :: Bool + } + deriving Show