diff --git a/Network/Kafka.hs b/Network/Kafka.hs --- a/Network/Kafka.hs +++ b/Network/Kafka.hs @@ -12,9 +12,9 @@ import Control.Monad.Trans.State import Data.ByteString.Char8 (ByteString) import Data.Monoid ((<>)) +import qualified Data.Pool as Pool 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 @@ -23,8 +23,6 @@ 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. @@ -35,20 +33,26 @@ , _stateBufferSize :: MaxBytes -- | Maximum time in milliseconds to wait for response. , _stateWaitTime :: MaxWaitTime + -- | An incrementing counter of requests. + , _stateCorrelationId :: CorrelationId -- | Broker cache , _stateBrokers :: M.Map Leader Broker + -- | Connection cache + , _stateConnections :: M.Map Broker (Pool.Pool Handle) + -- | Topic metadata cache + , _stateTopicMetadata :: M.Map TopicName TopicMetadata } makeLenses ''KafkaState -data KafkaConsumer = KafkaConsumer { _consumerState :: KafkaState - , _consumerHandle :: Handle - } +data KafkaClient = KafkaClient { _kafkaClientState :: KafkaState + , _kafkaClientHandle :: Handle + } -makeLenses ''KafkaConsumer +makeLenses ''KafkaClient -- | The core Kafka monad. -type Kafka = StateT KafkaConsumer (EitherT KafkaClientError IO) +type Kafka = StateT KafkaClient (EitherT KafkaClientError IO) type KafkaAddress = (Host, Port) type KafkaClientId = KafkaString @@ -62,6 +66,7 @@ | KafkaDeserializationError String -- TODO: cereal is Stringly typed, should use tickle -- | Could not find a cached broker for the found leader. | KafkaInvalidBroker Leader + | KafkaFailedToFetchMetadata deriving (Eq, Show) -- | Type of response to expect, used for 'KafkaExpected' error. @@ -82,6 +87,7 @@ , _palPartition :: Partition , _palLeader :: Leader } + deriving (Show) makeLenses ''PartitionAndLeader @@ -132,31 +138,37 @@ defaultState :: KafkaClientId -> KafkaState defaultState cid = KafkaState cid - defaultCorrelationId defaultRequiredAcks defaultRequestTimeout defaultMinBytes defaultMaxBytes defaultMaxWaitTime + defaultCorrelationId M.empty + M.empty + 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 + bracket (Network.connectTo (h ^. hostString) (p ^. portId)) hClose $ runEitherT . evalStateT k . KafkaClient 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) + corid <- use (kafkaClientState . stateCorrelationId) + kafkaClientState . stateCorrelationId += 1 + conid <- use (kafkaClientState . 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 +doRequest r = do + h <- use kafkaClientHandle + doRequest' h r + +doRequest' :: Handle -> Request -> Kafka Response +doRequest' h r = mapStateT (bimapEitherT KafkaDeserializationError id) $ do dataLength <- lift . EitherT $ do B.hPut h $ requestBytes r hFlush h @@ -167,9 +179,15 @@ -- | Send a metadata request metadata :: MetadataRequest -> Kafka MetadataResponse -metadata request = - makeRequest (MetadataRequest request) >>= doRequest >>= expectResponse ExpectedMetadata _MetadataResponse +metadata request = do + h <- use kafkaClientHandle + metadata' h request +-- | Send a metadata request +metadata' :: Handle -> MetadataRequest -> Kafka MetadataResponse +metadata' handle request = + makeRequest (MetadataRequest request) >>= doRequest' handle >>= 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) @@ -180,61 +198,6 @@ 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@ @@ -244,9 +207,9 @@ -- | 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) + wt <- use (kafkaClientState . stateWaitTime) + ws <- use (kafkaClientState . stateWaitSize) + bs <- use (kafkaClientState . stateBufferSize) return $ FetchReq (ordinaryConsumerId, wt, ws, [(topic, [(p, o, bs)])]) -- | Execute a fetch request and get the raw fetch response. @@ -259,37 +222,44 @@ fetchMessages fr = (fr ^.. fetchResponseFields . folded) >>= tam where tam a = TopicAndMessage (a ^. _1) <$> a ^.. _2 . folded . _4 . messageSetMembers . folded . setMessage --- * Producing +updateMetadatas :: [TopicName] -> Kafka () +updateMetadatas ts = do + md <- metadata $ MetadataReq ts + let (brokers, tmds) = (md ^.. metadataResponseBrokers . folded, md ^.. topicsMetadata . folded) + kafkaClientState . stateBrokers %= \m -> foldr addBroker m brokers + kafkaClientState . stateTopicMetadata %= \m -> foldr addTopicMetadata m tmds + return () + where addBroker :: Broker -> M.Map Leader Broker -> M.Map Leader Broker + addBroker b = M.insert (Leader . Just $ b ^. brokerNode . nodeId) b + addTopicMetadata :: TopicMetadata -> M.Map TopicName TopicMetadata -> M.Map TopicName TopicMetadata + addTopicMetadata tm = M.insert (tm ^. topicMetadataName) tm --- | Execute a produce request and get the raw preduce response. -produce :: ProduceRequest -> Kafka ProduceResponse -produce request = - makeRequest (ProduceRequest request) >>= doRequest >>= expectResponse ExpectedProduce _ProduceResponse +updateMetadata :: TopicName -> Kafka () +updateMetadata t = updateMetadatas [t] --- | 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)] +updateAllMetadata :: Kafka () +updateAllMetadata = updateMetadatas [] --- | 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 handler action, creating a new Pool and updating the connections Map if needed. +withBrokerHandle :: Broker -> (Handle -> Kafka a) -> Kafka a +withBrokerHandle broker f = do + conns <- use (kafkaClientState . stateConnections) + let foundPool = conns ^. at broker + pool <- case foundPool of + Nothing -> do + newPool <- liftIO $ mkPool broker + kafkaClientState . stateConnections .= (at broker ?~ newPool $ conns) + return newPool + Just p -> return p + Pool.withResource pool f + where mkPool :: Broker -> IO (Pool.Pool Handle) + mkPool b = Pool.createPool (createHandle b) hClose 1 10 1 + createHandle b = do + let h = b ^. brokerHost ^. hostString + p = b ^. brokerPort ^. portId + Network.connectTo h p --- | 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 diff --git a/Network/Kafka/Producer.hs b/Network/Kafka/Producer.hs new file mode 100644 --- /dev/null +++ b/Network/Kafka/Producer.hs @@ -0,0 +1,108 @@ +module Network.Kafka.Producer where + +import Control.Applicative +import Control.Lens +import Control.Monad.Trans (liftIO, lift) +import Control.Monad.Trans.Either +import Data.ByteString.Char8 (ByteString) +import Data.Monoid ((<>)) +import System.IO +import qualified Data.Map as M +import System.Random (getStdRandom, randomR) + +import Network.Kafka +import Network.Kafka.Protocol + +-- * Producing + +-- | Execute a produce request and get the raw preduce response. +produce :: Handle -> ProduceRequest -> Kafka ProduceResponse +produce handle request = + makeRequest (ProduceRequest request) >>= doRequest' handle >>= 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 + +-- | Create a protocol message set from a list of messages. +groupMessagesToSet :: [TopicAndMessage] -> MessageSet +groupMessagesToSet xs = MessageSet $ msm <$> xs + where msm = MessageSetMember (Offset (-1)) . _tamMessage + +-- | 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' + +-- | Execute a produce request using the values in the state. +send :: Leader -> [(TopicAndPartition, MessageSet)] -> Kafka ProduceResponse +send l ts = do + let s = kafkaClientState . stateBrokers . at l + topicNames = map (_tapTopic . fst) ts + broker <- findMetadataOrElse topicNames s (KafkaInvalidBroker l) + requiredAcks <- use (kafkaClientState . stateRequiredAcks) + requestTimeout <- use (kafkaClientState . stateRequestTimeout) + withBrokerHandle broker $ \handle -> produce handle $ produceRequest requiredAcks requestTimeout ts + +-- | Find a leader and partition for the topic. +brokerPartitionInfo :: TopicName -> Kafka [PartitionAndLeader] +brokerPartitionInfo t = do + let s = kafkaClientState . stateTopicMetadata . at t + tmd <- findMetadataOrElse [t] s KafkaFailedToFetchMetadata + return $ pal <$> tmd ^. partitionsMetadata + where pal d = PartitionAndLeader t (d ^. partitionId) (d ^. partitionMetadataLeader) + +findMetadataOrElse :: [TopicName] -> Getting (Maybe a) KafkaClient (Maybe a) -> KafkaClientError -> Kafka a +findMetadataOrElse ts s err = do + maybeFound <- use s + case maybeFound of + Just x -> return x + Nothing -> do + updateMetadatas ts + maybeFound' <- use s + case maybeFound' of + Just x -> return x + Nothing -> lift $ left $ err + +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) + +-- * Messages + +-- | 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))) diff --git a/Network/Kafka/Protocol.hs b/Network/Kafka/Protocol.hs --- a/Network/Kafka/Protocol.hs +++ b/Network/Kafka/Protocol.hs @@ -95,10 +95,10 @@ 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 Broker = Broker { _brokerFields :: (NodeId, Host, Port) } deriving (Show, Eq, Ord, 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 Host = Host { _hostKString :: KafkaString } deriving (Show, Eq, Ord, Deserializable, IsString) +newtype Port = Port { _portInt :: Int32 } deriving (Show, Eq, Ord, 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) diff --git a/milena.cabal b/milena.cabal --- a/milena.cabal +++ b/milena.cabal @@ -4,7 +4,7 @@ -- PVP summary: +-+------- breaking API changes -- | | +----- non-breaking API additions -- | | | +--- code changes with no API change -version: 0.1.0.0 +version: 0.2.0.0 synopsis: A Kafka client for Haskell. description: Kafka client for Haskell (not recommended for production use). license: BSD3 @@ -24,24 +24,26 @@ source-repository this type: git location: https://github.com/tylerholien/milena.git - tag: 0.1.0.0 + tag: 0.2.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 + Network.Kafka.Producer + build-depends: base >=4.7 && <5, + mtl >=2.1 && <2.3, + cereal >=0.4 && <0.5, + bytestring >=0.10 && <0.11, + network >=2.4 && <2.7, + digest >=0.0.1.0 && <0.1, + containers >=0.5 && <0.6, + either >=4.3 && <4.4, + random >=1.0 && <1.2, + transformers >=0.3 && <0.5, + lens >=4.4 && <4.8, + resource-pool >=0.2.3.2 && <0.3 test-suite test default-language: Haskell2010 diff --git a/test/tests.hs b/test/tests.hs --- a/test/tests.hs +++ b/test/tests.hs @@ -5,6 +5,7 @@ import Data.Functor import Data.Either (isRight) import Network.Kafka +import Network.Kafka.Producer import Network.Kafka.Protocol (Leader(..)) import Test.Hspec import Test.Hspec.QuickCheck @@ -13,7 +14,7 @@ main :: IO () main = hspec $ do let topic = "milena-test" - run = runKafka ("localhost", 9092) $ defaultState "milena-test-client" + run = runKafka ("kafka1.dev", 9093) $ defaultState "milena-test-client" byteMessages = fmap (TopicAndMessage topic . makeMessage . B.pack) describe "can talk to local Kafka server" $ do