diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c)2014, Tyler Holien + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Tyler Holien nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Network/Kafka.hs b/Network/Kafka.hs new file mode 100644 --- /dev/null +++ b/Network/Kafka.hs @@ -0,0 +1,315 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} + +module Network.Kafka where + +import Control.Applicative +import Control.Exception (bracket) +import Control.Lens +import Control.Monad (liftM) +import Control.Monad.Trans (liftIO, lift) +import Control.Monad.Trans.Either +import Control.Monad.Trans.State +import Data.ByteString.Char8 (ByteString) +import Data.Monoid ((<>)) +import Data.Serialize.Get +import System.IO +import System.Random (getStdRandom, randomR) +import qualified Data.ByteString.Char8 as B +import qualified Data.Map as M +import qualified Network + +import Network.Kafka.Protocol + +data KafkaState = KafkaState { -- | Name to use as a client ID. + _stateName :: KafkaString + -- | An incrementing counter of requests. + , _stateCorrelationId :: CorrelationId + -- | How many acknowledgements are required for producing. + , _stateRequiredAcks :: RequiredAcks + -- | Time in milliseconds to wait for messages to be produced by broker. + , _stateRequestTimeout :: Timeout + -- | Minimum size of response bytes to block for. + , _stateWaitSize :: MinBytes + -- | Maximum size of response bytes to retrieve. + , _stateBufferSize :: MaxBytes + -- | Maximum time in milliseconds to wait for response. + , _stateWaitTime :: MaxWaitTime + -- | Broker cache + , _stateBrokers :: M.Map Leader Broker + } + +makeLenses ''KafkaState + +data KafkaConsumer = KafkaConsumer { _consumerState :: KafkaState + , _consumerHandle :: Handle + } + +makeLenses ''KafkaConsumer + +-- | The core Kafka monad. +type Kafka = StateT KafkaConsumer (EitherT KafkaClientError IO) + +type KafkaAddress = (Host, Port) +type KafkaClientId = KafkaString + +-- | Errors given from the Kafka monad. +data KafkaClientError = -- | A response did not contain an offset. + KafkaNoOffset + -- | Got a different form of a response than was requested. + | KafkaExpected KafkaExpectedResponse + -- | A value could not be deserialized correctly. + | KafkaDeserializationError String -- TODO: cereal is Stringly typed, should use tickle + -- | Could not find a cached broker for the found leader. + | KafkaInvalidBroker Leader + deriving (Eq, Show) + +-- | Type of response to expect, used for 'KafkaExpected' error. +data KafkaExpectedResponse = ExpectedMetadata + | ExpectedFetch + | ExpectedProduce + deriving (Eq, Show) + +-- | An abstract form of Kafka's time. Used for querying offsets. +data KafkaTime = -- | The latest time on the broker. + LatestTime + -- | The earliest time on the broker. + | EarliestTime + -- | A specific time. + | OtherTime Time + +data PartitionAndLeader = PartitionAndLeader { _palTopic :: TopicName + , _palPartition :: Partition + , _palLeader :: Leader + } + +makeLenses ''PartitionAndLeader + +data TopicAndPartition = TopicAndPartition { _tapTopic :: TopicName + , _tapPartition :: Partition + } + deriving (Eq, Ord, Show) + +-- | A topic with a serializable message. +data TopicAndMessage = TopicAndMessage { _tamTopic :: TopicName + , _tamMessage :: Message + } + deriving (Eq, Show) + +makeLenses ''TopicAndMessage + +-- | Get the bytes from the Kafka message, ignoring the topic. +tamPayload :: TopicAndMessage -> ByteString +tamPayload = foldOf (tamMessage . payload) + +-- * Configuration + +-- | Default: @0@ +defaultCorrelationId :: CorrelationId +defaultCorrelationId = 0 + +-- | Default: @1@ +defaultRequiredAcks :: RequiredAcks +defaultRequiredAcks = 1 + +-- | Default: @10000@ +defaultRequestTimeout :: Timeout +defaultRequestTimeout = 10000 + +-- | Default: @0@ +defaultMinBytes :: MinBytes +defaultMinBytes = MinBytes 0 + +-- | Default: @1024 * 1024@ +defaultMaxBytes :: MaxBytes +defaultMaxBytes = 1024 * 1024 + +-- | Default: @0@ +defaultMaxWaitTime :: MaxWaitTime +defaultMaxWaitTime = 0 + +-- | Create a consumer using default values. +defaultState :: KafkaClientId -> KafkaState +defaultState cid = + KafkaState cid + defaultCorrelationId + defaultRequiredAcks + defaultRequestTimeout + defaultMinBytes + defaultMaxBytes + defaultMaxWaitTime + M.empty + +-- | Run the underlying Kafka monad at the given leader address and initial state. +runKafka :: KafkaAddress -> KafkaState -> Kafka a -> IO (Either KafkaClientError a) +runKafka (h, p) s k = + bracket (Network.connectTo (h ^. hostString) (p ^. portId)) hClose $ runEitherT . evalStateT k . KafkaConsumer s + +-- | Make a request, incrementing the `_stateCorrelationId`. +makeRequest :: RequestMessage -> Kafka Request +makeRequest m = do + corid <- use (consumerState . stateCorrelationId) + consumerState . stateCorrelationId += 1 + conid <- use (consumerState . stateName) + return $ Request (corid, ClientId conid, m) + +-- | Perform a request and deserialize the response. +doRequest :: Request -> Kafka Response +doRequest r = mapStateT (bimapEitherT KafkaDeserializationError id) $ do + h <- use consumerHandle + dataLength <- lift . EitherT $ do + B.hPut h $ requestBytes r + hFlush h + rawLength <- B.hGet h 4 + return $ runGet (liftM fromIntegral getWord32be) rawLength + resp <- liftIO $ B.hGet h dataLength + lift . hoistEither $ runGet (getResponse dataLength) resp + +-- | Send a metadata request +metadata :: MetadataRequest -> Kafka MetadataResponse +metadata request = + makeRequest (MetadataRequest request) >>= doRequest >>= expectResponse ExpectedMetadata _MetadataResponse + +-- | Function to give an error when the response seems wrong. +expectResponse :: KafkaExpectedResponse -> Getting (Leftmost b) ResponseMessage b -> Response -> Kafka b +expectResponse e p = lift . maybe (left $ KafkaExpected e) return . firstOf (responseMessage . p) + +-- | Convert an abstract time to a serializable protocol value. +protocolTime :: KafkaTime -> Time +protocolTime LatestTime = Time (-1) +protocolTime EarliestTime = Time (-2) +protocolTime (OtherTime o) = o + +-- * Messages + +-- | Group messages together with the leader they should be sent to. +partitionAndCollate :: [TopicAndMessage] -> Kafka (M.Map Leader (M.Map TopicAndPartition [TopicAndMessage])) +partitionAndCollate ks = recurse ks M.empty + where recurse [] accum = return accum + recurse (x:xs) accum = do + topicPartitionsList <- brokerPartitionInfo $ _tamTopic x + pal <- getPartition topicPartitionsList + let leader = maybe (Leader Nothing) _palLeader pal + tp = TopicAndPartition <$> pal ^? folded . palTopic <*> pal ^? folded . palPartition + b = M.singleton leader $ maybe M.empty (`M.singleton` [x]) tp + accum' = M.unionWith (M.unionWith (<>)) accum b + recurse xs accum' + +getPartition :: [PartitionAndLeader] -> Kafka (Maybe PartitionAndLeader) +getPartition ps = + liftIO $ (ps' ^?) . element <$> getStdRandom (randomR (0, length ps' - 1)) + where ps' = ps ^.. folded . filtered (has $ palLeader . leaderId . _Just) + +-- | Create a protocol message set from a list of messages. +groupMessagesToSet :: [TopicAndMessage] -> MessageSet +groupMessagesToSet xs = MessageSet $ uncurry msm <$> zip [0..] xs + where msm n = MessageSetMember (Offset n) . _tamMessage + +-- | Find a leader and partition for the topic. +brokerPartitionInfo :: TopicName -> Kafka [PartitionAndLeader] +brokerPartitionInfo t = do + md <- metadata $ MetadataReq [t] + let brokers = md ^.. metadataResponseFields . _1 . folded + consumerState . stateBrokers .= foldr addBroker M.empty brokers + return $ pal <$> md ^.. topicsMetadata . folded . partitionsMetadata . folded + where pal d = PartitionAndLeader t (d ^. partitionId) (d ^. partitionMetadataLeader) + addBroker b = M.insert (Leader . Just $ b ^. brokerFields . _1 . nodeId) b + +-- | Default: @1@ +defaultMessageCrc :: Crc +defaultMessageCrc = 1 + +-- | Default: @0@ +defaultMessageMagicByte :: MagicByte +defaultMessageMagicByte = 0 + +-- | Default: @Nothing@ +defaultMessageKey :: Key +defaultMessageKey = Key Nothing + +-- | Default: @0@ +defaultMessageAttributes :: Attributes +defaultMessageAttributes = 0 + +-- | Construct a message from a string of bytes using default attributes. +makeMessage :: ByteString -> Message +makeMessage m = Message (defaultMessageCrc, defaultMessageMagicByte, defaultMessageAttributes, defaultMessageKey, Value (Just (KBytes m))) + +-- * Fetching + +-- | Default: @-1@ +ordinaryConsumerId :: ReplicaId +ordinaryConsumerId = ReplicaId (-1) + +-- | Construct a fetch request from the values in the state. +fetchRequest :: Offset -> Partition -> TopicName -> Kafka FetchRequest +fetchRequest o p topic = do + wt <- use (consumerState . stateWaitTime) + ws <- use (consumerState . stateWaitSize) + bs <- use (consumerState . stateBufferSize) + return $ FetchReq (ordinaryConsumerId, wt, ws, [(topic, [(p, o, bs)])]) + +-- | Execute a fetch request and get the raw fetch response. +fetch :: FetchRequest -> Kafka FetchResponse +fetch request = + makeRequest (FetchRequest request) >>= doRequest >>= expectResponse ExpectedFetch _FetchResponse + +-- | Extract out messages with their topics from a fetch response. +fetchMessages :: FetchResponse -> [TopicAndMessage] +fetchMessages fr = (fr ^.. fetchResponseFields . folded) >>= tam + where tam a = TopicAndMessage (a ^. _1) <$> a ^.. _2 . folded . _4 . messageSetMembers . folded . setMessage + +-- * Producing + +-- | Execute a produce request and get the raw preduce response. +produce :: ProduceRequest -> Kafka ProduceResponse +produce request = + makeRequest (ProduceRequest request) >>= doRequest >>= expectResponse ExpectedProduce _ProduceResponse + +-- | Construct a produce request with explicit arguments. +produceRequest :: RequiredAcks -> Timeout -> [(TopicAndPartition, MessageSet)] -> ProduceRequest +produceRequest ra ti ts = + ProduceReq (ra, ti, M.toList . M.unionsWith (<>) $ fmap f ts) + where f (TopicAndPartition t p, i) = M.singleton t [(p, i)] + +-- | Send messages to partition calculated by 'partitionAndCollate'. +produceMessages :: [TopicAndMessage] -> Kafka [ProduceResponse] +produceMessages tams = do + m <- fmap (fmap groupMessagesToSet) <$> partitionAndCollate tams + mapM (uncurry send) $ fmap M.toList <$> M.toList m + +-- | Execute a produce request using the values in the state. +send :: Leader -> [(TopicAndPartition, MessageSet)] -> Kafka ProduceResponse +send l ts = do + foundBroker <- use (consumerState . stateBrokers . at l) + broker <- lift $ maybe (left $ KafkaInvalidBroker l) right foundBroker + requiredAcks <- use (consumerState . stateRequiredAcks) + requestTimeout <- use (consumerState . stateRequestTimeout) + let h' = broker ^. brokerFields . _2 + p' = broker ^. brokerFields . _3 + cstate <- use consumerState + r <- liftIO . runKafka (h', p') cstate . produce $ produceRequest requiredAcks requestTimeout ts + lift $ either left right r + +-- * Offsets + +-- | Fields to construct an offset request, per topic and partition. +data PartitionOffsetRequestInfo = + PartitionOffsetRequestInfo { -- | Time to find an offset for. + _kafkaTime :: KafkaTime + -- | Number of offsets to retrieve. + , _maxNumOffsets :: MaxNumberOfOffsets + } + +-- TODO: Properly look up the offset via the partition. +-- | Get the first found offset. +getLastOffset :: KafkaTime -> Partition -> TopicName -> Kafka Offset +getLastOffset m p t = + makeRequest (offsetRequest [(TopicAndPartition t p, PartitionOffsetRequestInfo m 1)]) >>= doRequest >>= maybe (StateT . const $ left KafkaNoOffset) return . firstOf (responseMessage . _OffsetResponse . offsetResponseOffset p) + +-- | Create an offset request. +offsetRequest :: [(TopicAndPartition, PartitionOffsetRequestInfo)] -> RequestMessage +offsetRequest ts = + OffsetRequest $ OffsetReq (ReplicaId (-1), M.toList . M.unionsWith (<>) $ fmap f ts) + where f (TopicAndPartition t p, i) = M.singleton t [g p i] + g p (PartitionOffsetRequestInfo kt mno) = (p, protocolTime kt, mno) diff --git a/Network/Kafka/Protocol.hs b/Network/Kafka/Protocol.hs new file mode 100644 --- /dev/null +++ b/Network/Kafka/Protocol.hs @@ -0,0 +1,519 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE Rank2Types #-} + +module Network.Kafka.Protocol where + +import Control.Applicative (Applicative(..), Alternative(..), (<$>), (<*>)) +import Control.Category (Category(..)) +import Control.Lens +import Control.Monad (replicateM, liftM, liftM2, liftM3, liftM4, liftM5) +import Data.ByteString.Char8 (ByteString) +import Data.ByteString.Lens (unpackedChars) +import Data.Digest.CRC32 +import Data.Int +import Data.Serialize.Get +import Data.Serialize.Put +import GHC.Exts (IsString(..)) +import Numeric.Lens +import Prelude hiding ((.), id) +import qualified Data.ByteString.Char8 as B +import qualified Network + +class Serializable a where + serialize :: a -> Put + +class Deserializable a where + deserialize :: Get a + +data Response = Response { _responseCorrelationId :: CorrelationId, _responseMessage :: ResponseMessage } deriving (Show, Eq) + +getResponse :: Int -> Get Response +getResponse l = Response <$> deserialize <*> getResponseMessage (l - 4) + +newtype ConsumerMetadataResponse = ConsumerMetadataResp (KafkaError, Broker) deriving (Show, Eq, Deserializable) -- CoordinatorId CoordinatorHost CoordinatorPort +-- newtype ErrorCode = ErrorCode int16 +-- newtype CoordinatorId = CoordinatorId int32 +-- newtype CoordinatorHost = CoordinatorHost string +-- newtype CoordinatorPort = CoordinatorPort int32 + +getResponseMessage :: Int -> Get ResponseMessage +getResponseMessage l = liftM MetadataResponse (isolate l deserialize) + <|> liftM OffsetResponse (isolate l deserialize) + <|> liftM ProduceResponse (isolate l deserialize) + <|> liftM OffsetCommitResponse (isolate l deserialize) + <|> liftM OffsetFetchResponse (isolate l deserialize) + <|> liftM ConsumerMetadataResponse (isolate l deserialize) + -- MUST try FetchResponse last! + -- + -- As an optimization, Kafka might return a partial message + -- at the end of a MessageSet, so this will consume the rest + -- of the message at the end of the input. + -- + -- Strictly speaking, this might not actually be necessary. + -- Parsing a MessageSet is isolated to the byte count that's + -- at the beginning of a MessageSet. I don't want to spend + -- the time right now to prove that will always be safe, but + -- I'd like to at some point. + <|> liftM FetchResponse (isolate l deserialize) + +newtype ApiKey = ApiKey Int16 deriving (Show, Eq, Deserializable, Serializable, Num) -- numeric ID for API (i.e. metadata req, produce req, etc.) +newtype ApiVersion = ApiVersion Int16 deriving (Show, Eq, Deserializable, Serializable, Num) +newtype CorrelationId = CorrelationId Int32 deriving (Show, Eq, Deserializable, Serializable, Num, Enum) +newtype ClientId = ClientId KafkaString deriving (Show, Eq, Deserializable, Serializable, IsString) + +data RequestMessage = MetadataRequest MetadataRequest + | ProduceRequest ProduceRequest + | FetchRequest FetchRequest + | OffsetRequest OffsetRequest + | OffsetCommitRequest OffsetCommitRequest + | OffsetFetchRequest OffsetFetchRequest + | ConsumerMetadataRequest ConsumerMetadataRequest + deriving (Show, Eq) + +newtype MetadataRequest = MetadataReq [TopicName] deriving (Show, Eq, Serializable, Deserializable) +newtype TopicName = TName { _tName :: KafkaString } deriving (Show, Eq, Ord, Deserializable, Serializable, IsString) + +newtype KafkaBytes = KBytes { _kafkaByteString :: ByteString } deriving (Show, Eq, IsString) +newtype KafkaString = KString { _kString :: ByteString } deriving (Show, Eq, Ord, IsString) + +newtype ProduceResponse = + ProduceResp { _produceResponseFields :: [(TopicName, [(Partition, KafkaError, Offset)])] } + deriving (Show, Eq, Deserializable, Serializable) + +newtype OffsetResponse = + OffsetResp { _offsetResponseFields :: [(TopicName, [PartitionOffsets])] } + deriving (Show, Eq, Deserializable) + +newtype PartitionOffsets = + PartitionOffsets { _partitionOffsetsFields :: (Partition, KafkaError, [Offset]) } + deriving (Show, Eq, Deserializable) + +newtype FetchResponse = + FetchResp { _fetchResponseFields :: [(TopicName, [(Partition, KafkaError, Offset, MessageSet)])] } + deriving (Show, Eq, Serializable, Deserializable) + +newtype MetadataResponse = MetadataResp { _metadataResponseFields :: ([Broker], [TopicMetadata]) } deriving (Show, Eq, Deserializable) +newtype Broker = Broker { _brokerFields :: (NodeId, Host, Port) } deriving (Show, Eq, Deserializable) +newtype NodeId = NodeId { _nodeId :: Int32 } deriving (Show, Eq, Ord, Deserializable, Num) +newtype Host = Host { _hostKString :: KafkaString } deriving (Show, Eq, Deserializable, IsString) +newtype Port = Port { _portInt :: Int32 } deriving (Show, Eq, Deserializable, Num) +newtype TopicMetadata = TopicMetadata { _topicMetadataFields :: (KafkaError, TopicName, [PartitionMetadata]) } deriving (Show, Eq, Deserializable) +newtype PartitionMetadata = PartitionMetadata { _partitionMetadataFields :: (KafkaError, Partition, Leader, Replicas, Isr) } deriving (Show, Eq, Deserializable) +newtype Leader = Leader { _leaderId :: Maybe Int32 } deriving (Show, Eq, Ord) + +newtype Replicas = Replicas [Int32] deriving (Show, Eq, Serializable, Deserializable) +newtype Isr = Isr [Int32] deriving (Show, Eq, Deserializable) + +newtype OffsetCommitResponse = OffsetCommitResp [(TopicName, [(Partition, KafkaError)])] deriving (Show, Eq, Deserializable) +newtype OffsetFetchResponse = OffsetFetchResp [(TopicName, [(Partition, Offset, Metadata, KafkaError)])] deriving (Show, Eq, Deserializable) + +newtype OffsetRequest = OffsetReq (ReplicaId, [(TopicName, [(Partition, Time, MaxNumberOfOffsets)])]) deriving (Show, Eq, Serializable) +newtype Time = Time { _timeInt :: Int64 } deriving (Show, Eq, Serializable, Num, Bounded) +newtype MaxNumberOfOffsets = MaxNumberOfOffsets Int32 deriving (Show, Eq, Serializable, Num) + +newtype FetchRequest = + FetchReq (ReplicaId, MaxWaitTime, MinBytes, + [(TopicName, [(Partition, Offset, MaxBytes)])]) + deriving (Show, Eq, Deserializable, Serializable) + +newtype ReplicaId = ReplicaId Int32 deriving (Show, Eq, Num, Serializable, Deserializable) +newtype MaxWaitTime = MaxWaitTime Int32 deriving (Show, Eq, Num, Serializable, Deserializable) +newtype MinBytes = MinBytes Int32 deriving (Show, Eq, Num, Serializable, Deserializable) +newtype MaxBytes = MaxBytes Int32 deriving (Show, Eq, Num, Serializable, Deserializable) + +newtype ProduceRequest = + ProduceReq (RequiredAcks, Timeout, + [(TopicName, [(Partition, MessageSet)])]) + deriving (Show, Eq, Serializable) + +newtype RequiredAcks = + RequiredAcks Int16 deriving (Show, Eq, Serializable, Deserializable, Num) +newtype Timeout = + Timeout Int32 deriving (Show, Eq, Serializable, Deserializable, Num) +newtype Partition = + Partition Int32 deriving (Show, Ord, Eq, Serializable, Deserializable, Num) + +newtype MessageSet = + MessageSet { _messageSetMembers :: [MessageSetMember] } deriving (Show, Eq) +data MessageSetMember = + MessageSetMember { _setOffset :: Offset, _setMessage :: Message } deriving (Show, Eq) + +newtype Offset = Offset Int64 deriving (Show, Eq, Serializable, Deserializable, Num) + +newtype Message = + Message { _messageFields :: (Crc, MagicByte, Attributes, Key, Value) } + deriving (Show, Eq, Deserializable) + +newtype Crc = Crc Int32 deriving (Show, Eq, Serializable, Deserializable, Num) +newtype MagicByte = MagicByte Int8 deriving (Show, Eq, Serializable, Deserializable, Num) +newtype Attributes = Attributes Int8 deriving (Show, Eq, Serializable, Deserializable, Num) + +newtype Key = Key { _keyBytes :: Maybe KafkaBytes } deriving (Show, Eq) +newtype Value = Value { _valueBytes :: Maybe KafkaBytes } deriving (Show, Eq) + +data ResponseMessage = MetadataResponse MetadataResponse + | ProduceResponse ProduceResponse + | FetchResponse FetchResponse + | OffsetResponse OffsetResponse + | OffsetCommitResponse OffsetCommitResponse + | OffsetFetchResponse OffsetFetchResponse + | ConsumerMetadataResponse ConsumerMetadataResponse + deriving (Show, Eq) + +newtype ConsumerMetadataRequest = ConsumerMetadataReq ConsumerGroup deriving (Show, Eq, Serializable) + +newtype OffsetCommitRequest = OffsetCommitReq (ConsumerGroup, [(TopicName, [(Partition, Offset, Time, Metadata)])]) deriving (Show, Eq, Serializable) +newtype OffsetFetchRequest = OffsetFetchReq (ConsumerGroup, [(TopicName, [Partition])]) deriving (Show, Eq, Serializable) +newtype ConsumerGroup = ConsumerGroup KafkaString deriving (Show, Eq, Serializable, Deserializable, IsString) +newtype Metadata = Metadata KafkaString deriving (Show, Eq, Serializable, Deserializable, IsString) + +errorKafka :: KafkaError -> Int16 +errorKafka NoError = 0 +errorKafka Unknown = -1 +errorKafka OffsetOutOfRange = 1 +errorKafka InvalidMessage = 2 +errorKafka UnknownTopicOrPartition = 3 +errorKafka InvalidMessageSize = 4 +errorKafka LeaderNotAvailable = 5 +errorKafka NotLeaderForPartition = 6 +errorKafka RequestTimedOut = 7 +errorKafka BrokerNotAvailable = 8 +errorKafka ReplicaNotAvailable = 9 +errorKafka MessageSizeTooLarge = 10 +errorKafka StaleControllerEpochCode = 11 +errorKafka OffsetMetadataTooLargeCode = 12 +errorKafka OffsetsLoadInProgressCode = 14 +errorKafka ConsumerCoordinatorNotAvailableCode = 15 +errorKafka NotCoordinatorForConsumerCode = 16 + +data KafkaError = NoError -- ^ @0@ No error--it worked! + | Unknown -- ^ @-1@ An unexpected server error + | OffsetOutOfRange -- ^ @1@ The requested offset is outside the range of offsets maintained by the server for the given topic/partition. + | InvalidMessage -- ^ @2@ This indicates that a message contents does not match its CRC + | UnknownTopicOrPartition -- ^ @3@ This request is for a topic or partition that does not exist on this broker. + | InvalidMessageSize -- ^ @4@ The message has a negative size + | LeaderNotAvailable -- ^ @5@ This error is thrown if we are in the middle of a leadership election and there is currently no leader for this partition and hence it is unavailable for writes. + | NotLeaderForPartition -- ^ @6@ This error is thrown if the client attempts to send messages to a replica that is not the leader for some partition. It indicates that the clients metadata is out of date. + | RequestTimedOut -- ^ @7@ This error is thrown if the request exceeds the user-specified time limit in the request. + | BrokerNotAvailable -- ^ @8@ This is not a client facing error and is used mostly by tools when a broker is not alive. + | ReplicaNotAvailable -- ^ @9@ If replica is expected on a broker, but is not. + | MessageSizeTooLarge -- ^ @10@ The server has a configurable maximum message size to avoid unbounded memory allocation. This error is thrown if the client attempt to produce a message larger than this maximum. + | StaleControllerEpochCode -- ^ @11@ Internal error code for broker-to-broker communication. + | OffsetMetadataTooLargeCode -- ^ @12@ If you specify a string larger than configured maximum for offset metadata + | OffsetsLoadInProgressCode -- ^ @14@ The broker returns this error code for an offset fetch request if it is still loading offsets (after a leader change for that offsets topic partition). + | ConsumerCoordinatorNotAvailableCode -- ^ @15@ The broker returns this error code for consumer metadata requests or offset commit requests if the offsets topic has not yet been created. + | NotCoordinatorForConsumerCode -- ^ @16@ The broker returns this error code if it receives an offset fetch or commit request for a consumer group that it is not a coordinator for. + deriving (Eq, Show) + +instance Serializable KafkaError where + serialize = serialize . errorKafka + +instance Deserializable KafkaError where + deserialize = do + x <- deserialize :: Get Int16 + case x of + 0 -> return NoError + (-1) -> return Unknown + 1 -> return OffsetOutOfRange + 2 -> return InvalidMessage + 3 -> return UnknownTopicOrPartition + 4 -> return InvalidMessageSize + 5 -> return LeaderNotAvailable + 6 -> return NotLeaderForPartition + 7 -> return RequestTimedOut + 8 -> return BrokerNotAvailable + 9 -> return ReplicaNotAvailable + 10 -> return MessageSizeTooLarge + 11 -> return StaleControllerEpochCode + 12 -> return OffsetMetadataTooLargeCode + 14 -> return OffsetsLoadInProgressCode + 15 -> return ConsumerCoordinatorNotAvailableCode + 16 -> return NotCoordinatorForConsumerCode + _ -> fail $ "invalid error code: " ++ show x + +newtype Request = Request (CorrelationId, ClientId, RequestMessage) deriving (Show, Eq) + +instance Serializable Request where + serialize (Request (correlationId, clientId, r)) = do + serialize (apiKey r) + serialize (apiVersion r) + serialize correlationId + serialize clientId + serialize r + +requestBytes :: Request -> ByteString +requestBytes x = runPut $ do + putWord32be . fromIntegral $ B.length mr + putByteString mr + where mr = runPut $ serialize x + +apiVersion :: RequestMessage -> ApiVersion +apiVersion _ = ApiVersion 0 -- everything is at version 0 right now + +apiKey :: RequestMessage -> ApiKey +apiKey (ProduceRequest{}) = ApiKey 0 +apiKey (FetchRequest{}) = ApiKey 1 +apiKey (OffsetRequest{}) = ApiKey 2 +apiKey (MetadataRequest{}) = ApiKey 3 +apiKey (OffsetCommitRequest{}) = ApiKey 8 +apiKey (OffsetFetchRequest{}) = ApiKey 9 +apiKey (ConsumerMetadataRequest{}) = ApiKey 10 + +instance Serializable RequestMessage where + serialize (ProduceRequest r) = serialize r + serialize (FetchRequest r) = serialize r + serialize (OffsetRequest r) = serialize r + serialize (MetadataRequest r) = serialize r + serialize (OffsetCommitRequest r) = serialize r + serialize (OffsetFetchRequest r) = serialize r + serialize (ConsumerMetadataRequest r) = serialize r + +instance Serializable Int64 where serialize = putWord64be . fromIntegral +instance Serializable Int32 where serialize = putWord32be . fromIntegral +instance Serializable Int16 where serialize = putWord16be . fromIntegral +instance Serializable Int8 where serialize = putWord8 . fromIntegral + +instance Serializable Key where + serialize (Key (Just bs)) = serialize bs + serialize (Key Nothing) = serialize (-1 :: Int32) + +instance Serializable Value where + serialize (Value (Just bs)) = serialize bs + serialize (Value Nothing) = serialize (-1 :: Int32) + +instance Serializable KafkaString where + serialize (KString bs) = do + let l = fromIntegral (B.length bs) :: Int16 + serialize l + putByteString bs + +instance Serializable MessageSet where + serialize (MessageSet ms) = do + let bytes = runPut $ mapM_ serialize ms + l = fromIntegral (B.length bytes) :: Int32 + serialize l + putByteString bytes + +instance Serializable KafkaBytes where + serialize (KBytes bs) = do + let l = fromIntegral (B.length bs) :: Int32 + serialize l + putByteString bs + +instance Serializable MessageSetMember where + serialize (MessageSetMember offset msg) = do + serialize offset + serialize msize + serialize msg + where msize = fromIntegral $ B.length $ runPut $ serialize msg :: Int32 + +instance Serializable Message where + serialize (Message (_, magic, attrs, k, v)) = do + let m = runPut $ serialize magic >> serialize attrs >> serialize k >> serialize v + putWord32be (crc32 m) + putByteString m + +instance (Serializable a) => Serializable [a] where + serialize xs = do + let l = fromIntegral (length xs) :: Int32 + serialize l + mapM_ serialize xs + +instance (Serializable a, Serializable b) => Serializable ((,) a b) where + serialize (x, y) = serialize x >> serialize y +instance (Serializable a, Serializable b, Serializable c) => Serializable ((,,) a b c) where + serialize (x, y, z) = serialize x >> serialize y >> serialize z +instance (Serializable a, Serializable b, Serializable c, Serializable d) => Serializable ((,,,) a b c d) where + serialize (w, x, y, z) = serialize w >> serialize x >> serialize y >> serialize z +instance (Serializable a, Serializable b, Serializable c, Serializable d, Serializable e) => Serializable ((,,,,) a b c d e) where + serialize (v, w, x, y, z) = serialize v >> serialize w >> serialize x >> serialize y >> serialize z + +instance Deserializable MessageSet where + deserialize = do + l <- deserialize :: Get Int32 + ms <- isolate (fromIntegral l) getMembers + return $ MessageSet ms + where getMembers :: Get [MessageSetMember] + getMembers = do + wasEmpty <- isEmpty + if wasEmpty + then return [] + else liftM2 (:) deserialize getMembers <|> (remaining >>= getBytes >> return []) + +instance Deserializable MessageSetMember where + deserialize = do + o <- deserialize + l <- deserialize :: Get Int32 + m <- isolate (fromIntegral l) deserialize + return $ MessageSetMember o m + +instance Deserializable Leader where + deserialize = do + x <- deserialize :: Get Int32 + let l = Leader $ if x == -1 then Nothing else Just x + return l + +instance Deserializable KafkaBytes where + deserialize = do + l <- deserialize :: Get Int32 + bs <- getByteString $ fromIntegral l + return $ KBytes bs + +instance Deserializable KafkaString where + deserialize = do + l <- deserialize :: Get Int16 + bs <- getByteString $ fromIntegral l + return $ KString bs + +instance Deserializable Key where + deserialize = do + l <- deserialize :: Get Int32 + case l of + -1 -> return (Key Nothing) + _ -> do + bs <- getByteString $ fromIntegral l + return $ Key (Just (KBytes bs)) + +instance Deserializable Value where + deserialize = do + l <- deserialize :: Get Int32 + case l of + -1 -> return (Value Nothing) + _ -> do + bs <- getByteString $ fromIntegral l + return $ Value (Just (KBytes bs)) + +instance (Deserializable a) => Deserializable [a] where + deserialize = do + l <- deserialize :: Get Int32 + replicateM (fromIntegral l) deserialize + +instance (Deserializable a, Deserializable b) => Deserializable ((,) a b) where + deserialize = liftM2 (,) deserialize deserialize +instance (Deserializable a, Deserializable b, Deserializable c) => Deserializable ((,,) a b c) where + deserialize = liftM3 (,,) deserialize deserialize deserialize +instance (Deserializable a, Deserializable b, Deserializable c, Deserializable d) => Deserializable ((,,,) a b c d) where + deserialize = liftM4 (,,,) deserialize deserialize deserialize deserialize +instance (Deserializable a, Deserializable b, Deserializable c, Deserializable d, Deserializable e) => Deserializable ((,,,,) a b c d e) where + deserialize = liftM5 (,,,,) deserialize deserialize deserialize deserialize deserialize + +instance Deserializable Int64 where deserialize = liftM fromIntegral getWord64be +instance Deserializable Int32 where deserialize = liftM fromIntegral getWord32be +instance Deserializable Int16 where deserialize = liftM fromIntegral getWord16be +instance Deserializable Int8 where deserialize = liftM fromIntegral getWord8 + +-- * Generated lenses + +makeLenses ''Response + +makeLenses ''TopicName + +makeLenses ''KafkaBytes +makeLenses ''KafkaString + +makeLenses ''ProduceResponse + +makeLenses ''OffsetResponse +makeLenses ''PartitionOffsets + +makeLenses ''FetchResponse + +makeLenses ''MetadataResponse +makeLenses ''Broker +makeLenses ''NodeId +makeLenses ''Host +makeLenses ''Port +makeLenses ''TopicMetadata +makeLenses ''PartitionMetadata +makeLenses ''Leader + +makeLenses ''Time + +makeLenses ''Partition + +makeLenses ''MessageSet +makeLenses ''MessageSetMember +makeLenses ''Offset + +makeLenses ''Message + +makeLenses ''Key +makeLenses ''Value + +makePrisms ''ResponseMessage + +-- * Composed lenses + +keyed :: (Field1 a a b b, Choice p, Applicative f, Eq b) => b -> Optic' p f a a +keyed k = filtered (view $ _1 . to (== k)) + +metadataResponseBrokers :: Lens' MetadataResponse [Broker] +metadataResponseBrokers = metadataResponseFields . _1 + +topicsMetadata :: Lens' MetadataResponse [TopicMetadata] +topicsMetadata = metadataResponseFields . _2 + +topicMetadataKafkaError :: Lens' TopicMetadata KafkaError +topicMetadataKafkaError = topicMetadataFields . _1 + +topicMetadataName :: Lens' TopicMetadata TopicName +topicMetadataName = topicMetadataFields . _2 + +partitionsMetadata :: Lens' TopicMetadata [PartitionMetadata] +partitionsMetadata = topicMetadataFields . _3 + +partitionId :: Lens' PartitionMetadata Partition +partitionId = partitionMetadataFields . _2 + +partitionMetadataLeader :: Lens' PartitionMetadata Leader +partitionMetadataLeader = partitionMetadataFields . _3 + +brokerNode :: Lens' Broker NodeId +brokerNode = brokerFields . _1 + +brokerHost :: Lens' Broker Host +brokerHost = brokerFields . _2 + +brokerPort :: Lens' Broker Port +brokerPort = brokerFields . _3 + +fetchResponseMessages :: Fold FetchResponse MessageSet +fetchResponseMessages = fetchResponseFields . folded . _2 . folded . _4 + +fetchResponseByTopic :: TopicName -> Fold FetchResponse (Partition, KafkaError, Offset, MessageSet) +fetchResponseByTopic t = fetchResponseFields . folded . keyed t . _2 . folded + +messageSetByPartition :: Partition -> Fold (Partition, KafkaError, Offset, MessageSet) MessageSetMember +messageSetByPartition p = keyed p . _4 . messageSetMembers . folded + +fetchResponseMessageMembers :: Fold FetchResponse MessageSetMember +fetchResponseMessageMembers = fetchResponseMessages . messageSetMembers . folded + +messageValue :: Lens' Message Value +messageValue = messageFields . _5 + +payload :: Fold Message ByteString +payload = messageValue . valueBytes . folded . kafkaByteString + +offsetResponseOffset :: Partition -> Fold OffsetResponse Offset +offsetResponseOffset p = offsetResponseFields . folded . _2 . folded . partitionOffsetsFields . keyed p . _3 . folded + +messageSet :: Partition -> TopicName -> Fold FetchResponse MessageSetMember +messageSet p t = fetchResponseByTopic t . messageSetByPartition p + +nextOffset :: Lens' MessageSetMember Offset +nextOffset = setOffset . adding 1 + +findPartitionMetadata :: Applicative f => TopicName -> LensLike' f TopicMetadata [PartitionMetadata] +findPartitionMetadata t = filtered (view $ topicMetadataName . to (== t)) . partitionsMetadata + +findPartition :: Partition -> Prism' PartitionMetadata PartitionMetadata +findPartition p = filtered (view $ partitionId . to (== p)) + +hostString :: Lens' Host String +hostString = hostKString . kString . unpackedChars + +portId :: IndexPreservingGetter Port Network.PortID +portId = portInt . to fromIntegral . to Network.PortNumber diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/milena.cabal b/milena.cabal new file mode 100644 --- /dev/null +++ b/milena.cabal @@ -0,0 +1,57 @@ +name: milena + +-- http://www.haskell.org/haskellwiki/Package_versioning_policy +-- PVP summary: +-+------- breaking API changes +-- | | +----- non-breaking API additions +-- | | | +--- code changes with no API change +version: 0.1.0.0 +synopsis: A Kafka client for Haskell. +description: Kafka client for Haskell (not recommended for production use). +license: BSD3 +license-file: LICENSE +author: Tyler Holien +maintainer: tyler.holien@gmail.com +copyright: 2014, Tyler Holien +category: Network +build-type: Simple +stability: alpha +cabal-version: >=1.10 + +source-repository head + type: git + location: https://github.com/tylerholien/milena.git + +source-repository this + type: git + location: https://github.com/tylerholien/milena.git + tag: 0.1.0.0 + +library + default-language: Haskell2010 + ghc-options: -Wall + exposed-modules: Network.Kafka, + Network.Kafka.Protocol + build-depends: base >=4.7 && <5, + mtl >=2.1 && <2.2, + cereal >=0.4 && <0.5, + bytestring >=0.10 && <0.12, + network >=2.4 && <2.7, + digest >=0.0.1.0 && <0.0.2, + containers >=0.5 && <0.6, + either >=4.3 && <4.4, + random >=1.0 && <1.2, + transformers >=0.3 && <0.4, + lens >=4.4 && <4.5 + +test-suite test + default-language: Haskell2010 + ghc-options: -Wall -threaded + build-depends: base, + milena, + network, + QuickCheck, + bytestring, + hspec + hs-source-dirs: test + main-is: tests.hs + type: exitcode-stdio-1.0 diff --git a/test/tests.hs b/test/tests.hs new file mode 100644 --- /dev/null +++ b/test/tests.hs @@ -0,0 +1,38 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Main where + +import Data.Functor +import Data.Either (isRight) +import Network.Kafka +import Network.Kafka.Protocol (Leader(..)) +import Test.Hspec +import Test.Hspec.QuickCheck +import qualified Data.ByteString.Char8 as B + +main :: IO () +main = hspec $ do + let topic = "milena-test" + run = runKafka ("localhost", 9092) $ defaultState "milena-test-client" + byteMessages = fmap (TopicAndMessage topic . makeMessage . B.pack) + + describe "can talk to local Kafka server" $ do + prop "can produce messages" $ \ms -> do + result <- run . produceMessages $ byteMessages ms + result `shouldSatisfy` isRight + + prop "can fetch messages" $ do + result <- run $ do + offset <- getLastOffset EarliestTime 0 topic + fetch =<< fetchRequest offset 0 topic + result `shouldSatisfy` isRight + + prop "can roundtrip messages" $ \ms -> do + let messages = byteMessages ms + result <- run $ do + info <- brokerPartitionInfo topic + leader <- maybe (Leader Nothing) _palLeader <$> getPartition info + offset <- getLastOffset LatestTime 0 topic + void $ send leader [(TopicAndPartition topic 0, groupMessagesToSet messages)] + fmap tamPayload . fetchMessages <$> (fetch =<< fetchRequest offset 0 topic) + result `shouldBe` Right (tamPayload <$> messages)