diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# hw-kafka-client +# hw-kafka-client [![CircleCI](https://circleci.com/gh/haskell-works/hw-kafka-client.svg?style=svg&circle-token=5f3ada2650dd600bc0fd4787143024867b2afc4e)](https://circleci.com/gh/haskell-works/hw-kafka-client) Kafka bindings for Haskell backed by the @@ -12,7 +12,7 @@ HaskellWorks Kafka ecosystem is described here: https://github.com/haskell-works/hw-kafka # Consumer -High level consumers are supported by `librdkafka` starting from version 0.9. +High level consumers are supported by `librdkafka` starting from version 0.9. High-level consumers provide an abstraction for consuming messages from multiple partitions and topics. They are also address scalability (up to a number of partitions) by providing automatic rebalancing functionality. When a new consumer joins a consumer @@ -30,7 +30,7 @@ ``` A working consumer example can be found here: [ConsumerExample.hs](example/ConsumerExample.hs)
-To run an example please compile with the `examples` flag: +To run an example please compile with the `examples` flag. ```Haskell import Data.Monoid ((<>)) @@ -39,10 +39,10 @@ -- Global consumer properties consumerProps :: ConsumerProperties -consumerProps = consumerBrokersList [BrokerAddress "localhost:9092"] +consumerProps = brokersList [BrokerAddress "localhost:9092"] <> groupId (ConsumerGroupId "consumer_example_group") <> noAutoCommit - <> consumerDebug [DebugAll] + <> consumerLogLevel KafkaLogInfo -- Subscription to topics consumerSub :: Subscription @@ -82,7 +82,8 @@ -- Global producer properties producerProps :: ProducerProperties -producerProps = producerBrokersList [BrokerAddress "localhost:9092"] +producerProps = brokersList [BrokerAddress "localhost:9092"] + <> logLevel KafkaLogDebug -- Topic to send messages to targetTopic :: TopicName @@ -96,23 +97,21 @@ sendMessages :: KafkaProducer -> IO (Either KafkaError ()) sendMessages prod = do - err1 <- produceMessage prod ProducerRecord - { prTopic = targetTopic - , prPartition = UnassignedPartition - , prKey = Nothing - , prValue = Just "test from producer" - } + err1 <- produceMessage prod (mkMessage Nothing (Just "test from producer") ) forM_ err1 print - err2 <- produceMessage prod ProducerRecord - { prTopic = targetTopic - , prPartition = UnassignedPartition - , prKey = Just "key" - , prValue = Just "test from producer (with key)" - } + err2 <- produceMessage prod (mkMessage (Just "key") (Just "test from producer (with key)")) forM_ err2 print return $ Right () + +mkMessage :: Maybe ByteString -> Maybe ByteString -> ProducerRecord +mkMessage k v = ProducerRecord + { prTopic = targetTopic + , prPartition = UnassignedPartition + , prKey = k + , prValue = v + } ``` # Installation diff --git a/example/ConsumerExample.hs b/example/ConsumerExample.hs --- a/example/ConsumerExample.hs +++ b/example/ConsumerExample.hs @@ -4,16 +4,16 @@ where import Control.Arrow ((&&&)) -import Data.Monoid ((<>)) +import Data.Monoid ((<>)) import Kafka.Consumer -- Global consumer properties consumerProps :: ConsumerProperties -consumerProps = consumerBrokersList [BrokerAddress "localhost:9092"] +consumerProps = brokersList [BrokerAddress "localhost:9092"] <> groupId (ConsumerGroupId "consumer_example_group") <> noAutoCommit - <> reballanceCallback (ReballanceCallback printingRebalanceCallback) - <> offsetsCommitCallback (OffsetsCommitCallback printingOffsetCallback) + <> setCallback (rebalanceCallback printingRebalanceCallback) + <> setCallback (offsetCommitCallback printingOffsetCallback) <> consumerLogLevel KafkaLogInfo -- Subscription to topics @@ -21,6 +21,7 @@ consumerSub = topics [TopicName "kafka-client-example-topic"] <> offsetReset Earliest +-- Running an example runConsumerExample :: IO () runConsumerExample = do print $ cpLogLevel consumerProps diff --git a/example/ProducerExample.hs b/example/ProducerExample.hs --- a/example/ProducerExample.hs +++ b/example/ProducerExample.hs @@ -3,15 +3,15 @@ module ProducerExample where -import Control.Monad (forM_) +import Control.Monad (forM_) +import Data.ByteString (ByteString) import Data.Monoid import Kafka.Producer -import Data.ByteString (ByteString) -- Global producer properties producerProps :: ProducerProperties -producerProps = producerBrokersList [BrokerAddress "localhost:9092"] - <> producerLogLevel KafkaLogDebug +producerProps = brokersList [BrokerAddress "localhost:9092"] + <> logLevel KafkaLogDebug -- Topic to send messages to targetTopic :: TopicName diff --git a/hw-kafka-client.cabal b/hw-kafka-client.cabal --- a/hw-kafka-client.cabal +++ b/hw-kafka-client.cabal @@ -1,5 +1,5 @@ name: hw-kafka-client -version: 1.1.4 +version: 2.0.0 homepage: https://github.com/haskell-works/hw-kafka-client bug-reports: https://github.com/haskell-works/hw-kafka-client/issues license: MIT @@ -62,22 +62,23 @@ , transformers , unix exposed-modules: - Kafka Kafka.Types Kafka.Consumer.ConsumerProperties Kafka.Consumer.Subscription Kafka.Consumer.Types Kafka.Consumer + Kafka.Metadata Kafka.Producer Kafka.Producer.ProducerProperties Kafka.Producer.Types other-modules: + Kafka.Callbacks Kafka.Consumer.Convert + Kafka.Consumer.Callbacks Kafka.Internal.Shared Kafka.Internal.Setup Kafka.Producer.Convert Kafka.Internal.RdKafka - Kafka.Internal.RdKafkaEnum Kafka.Internal.Dump hs-source-dirs: src default-language: Haskell2010 @@ -93,9 +94,26 @@ type: exitcode-stdio-1.0 default-language: Haskell2010 hs-source-dirs: tests - other-modules: Kafka.IntegrationSpec - , Kafka.Consumer.ConsumerRecordMapSpec + other-modules: Kafka.Consumer.ConsumerRecordMapSpec , Kafka.Consumer.ConsumerRecordTraverseSpec + main-is: Spec.hs + ghc-options: -Wall -threaded + build-depends: base >=4.6 && < 5 + , bifunctors + , bytestring + , containers + , hw-kafka-client + , monad-loops + , hspec + , regex-posix + , either + +test-suite integration-tests + type: exitcode-stdio-1.0 + default-language: Haskell2010 + hs-source-dirs: tests-it + other-modules: Kafka.IntegrationSpec + , Kafka.TestEnv main-is: Spec.hs ghc-options: -Wall -threaded build-depends: base >=4.6 && < 5 diff --git a/src/Kafka.hs b/src/Kafka.hs deleted file mode 100644 --- a/src/Kafka.hs +++ /dev/null @@ -1,6 +0,0 @@ -module Kafka -( module X -) where - -import Kafka.Consumer as X -import Kafka.Producer as X diff --git a/src/Kafka/Callbacks.hs b/src/Kafka/Callbacks.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Callbacks.hs @@ -0,0 +1,20 @@ +module Kafka.Callbacks +( errorCallback +, logCallback +) +where + +import Kafka.Internal.RdKafka +import Kafka.Types + +errorCallback :: HasKafkaConf k => (KafkaError -> String -> IO ()) -> k -> IO () +errorCallback callback k = + let (KafkaConf c) = getKafkaConf k + realCb _ err = callback (KafkaResponseError err) + in rdKafkaConfSetErrorCb c realCb + +logCallback :: HasKafkaConf k => (Int -> String -> String -> IO ()) -> k -> IO () +logCallback callback k = + let (KafkaConf c) = getKafkaConf k + realCb _ = callback + in rdKafkaConfSetLogCb c realCb diff --git a/src/Kafka/Consumer.hs b/src/Kafka/Consumer.hs --- a/src/Kafka/Consumer.hs +++ b/src/Kafka/Consumer.hs @@ -3,11 +3,10 @@ ( module X , runConsumer , newConsumer -, assign +, assign, assignment, subscription +, seek , pollMessage -, commitOffsetMessage -, commitAllOffsets -, commitPartitionsOffsets +, commitOffsetMessage, commitAllOffsets, commitPartitionsOffsets , closeConsumer -- ReExport Types @@ -16,29 +15,30 @@ , CIT.PartitionOffset (..) , CIT.TopicPartition (..) , CIT.ConsumerRecord (..) -, RDE.RdKafkaRespErrT (..) +, RdKafkaRespErrT (..) ) where +import Control.Arrow import Control.Exception -import Control.Monad (forM_) +import Control.Monad (forM_) import Control.Monad.IO.Class -import qualified Data.ByteString as BS -import qualified Data.Map as M -import Foreign hiding (void) +import Control.Monad.Trans.Except +import Data.Bifunctor +import qualified Data.ByteString as BS +import qualified Data.Map as M +import Foreign hiding (void) import Kafka.Consumer.Convert import Kafka.Internal.RdKafka -import Kafka.Internal.RdKafkaEnum import Kafka.Internal.Setup import Kafka.Internal.Shared -import qualified Kafka.Consumer.Types as CIT -import qualified Kafka.Internal.RdKafkaEnum as RDE +import qualified Kafka.Consumer.Types as CIT -import Kafka.Consumer.ConsumerProperties as X -import Kafka.Consumer.Subscription as X -import Kafka.Consumer.Types as X -import Kafka.Types as X +import Kafka.Consumer.ConsumerProperties as X +import Kafka.Consumer.Subscription as X +import Kafka.Consumer.Types as X +import Kafka.Types as X -- | Runs high-level kafka consumer. -- @@ -65,10 +65,10 @@ -> Subscription -> m (Either KafkaError KafkaConsumer) newConsumer cp (Subscription ts tp) = liftIO $ do - (KafkaConf kc) <- newConsumerConf cp + kc@(KafkaConf kc') <- newConsumerConf cp tp' <- topicConf (TopicProps $ M.toList tp) _ <- setDefaultTopicConf kc tp' - rdk <- mapLeft KafkaError <$> newRdKafkaT RdKafkaConsumer kc + rdk <- bimap KafkaError Kafka <$> newRdKafkaT RdKafkaConsumer kc' case flip KafkaConsumer kc <$> rdk of Left err -> return $ Left err Right kafka -> do @@ -88,7 +88,7 @@ => KafkaConsumer -> Timeout -- ^ the timeout, in milliseconds -> m (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))) -- ^ Left on error or timeout, right for success -pollMessage (KafkaConsumer k _) (Timeout ms) = +pollMessage (KafkaConsumer (Kafka k) _) (Timeout ms) = liftIO $ rdKafkaConsumerPoll k (fromIntegral ms) >>= fromMessagePtr -- | Commit message's offset on broker for the message's partition. @@ -120,70 +120,63 @@ -- | Assigns specified partitions to a current consumer. -- Assigning an empty list means unassigning from all partitions that are currently assigned. assign :: MonadIO m => KafkaConsumer -> [TopicPartition] -> m KafkaError -assign (KafkaConsumer k _) ps = +assign (KafkaConsumer (Kafka k) _) ps = let pl = if null ps then newForeignPtr_ nullPtr else toNativeTopicPartitionList ps in liftIO $ KafkaResponseError <$> (pl >>= rdKafkaAssign k) +-- | Returns current consumer's assignment +assignment :: MonadIO m => KafkaConsumer -> m (Either KafkaError (M.Map TopicName [PartitionId])) +assignment (KafkaConsumer (Kafka k) _) = liftIO $ do + tpl <- rdKafkaAssignment k + tps <- traverse fromNativeTopicPartitionList'' (left KafkaResponseError tpl) + return $ tpMap <$> tps + where + tpMap ts = toMap $ (tpTopicName &&& tpPartition) <$> ts +-- | Returns current consumer's subscription +subscription :: MonadIO m => KafkaConsumer -> m (Either KafkaError [(TopicName, SubscribedPartitions)]) +subscription (KafkaConsumer (Kafka k) _) = liftIO $ do + tpl <- rdKafkaSubscription k + tps <- traverse fromNativeTopicPartitionList'' (left KafkaResponseError tpl) + return $ toSub <$> tps + where + toSub ts = M.toList $ subParts <$> tpMap ts + tpMap ts = toMap $ (tpTopicName &&& tpPartition) <$> ts + subParts [PartitionId (-1)] = SubscribedPartitionsAll + subParts ps = SubscribedPartitions ps + +seek :: MonadIO m => KafkaConsumer -> Timeout -> [TopicPartition] -> m (Maybe KafkaError) +seek (KafkaConsumer (Kafka k) _) (Timeout timeout) tps = liftIO $ + either Just (const Nothing) <$> seekAll + where + seekAll = runExceptT $ do + tr <- traverse (ExceptT . topicPair) tps + _ <- traverse (\(kt, p, o) -> ExceptT (rdSeek kt p o)) tr + return () + + rdSeek kt (PartitionId p) o = do + res <- rdKafkaSeek kt (fromIntegral p) (offsetToInt64 o) timeout + return $ rdKafkaErrorToEither res + + topicPair tp = do + let (TopicName tn) = tpTopicName tp + nt <- newRdKafkaTopicT k tn Nothing + return $ bimap KafkaError (,tpPartition tp, tpOffset tp) nt + -- | Closes the consumer. closeConsumer :: MonadIO m => KafkaConsumer -> m (Maybe KafkaError) -closeConsumer (KafkaConsumer k _) = +closeConsumer (KafkaConsumer (Kafka k) _) = liftIO $ (kafkaErrorToMaybe . KafkaResponseError) <$> rdKafkaConsumerClose k ----------------------------------------------------------------------------- newConsumerConf :: ConsumerProperties -> IO KafkaConf -newConsumerConf (ConsumerProperties m rcb ccb _) = do +newConsumerConf (ConsumerProperties m _ cbs) = do conf <- kafkaConf (KafkaProps $ M.toList m) - forM_ rcb (\(ReballanceCallback cb) -> setRebalanceCallback conf cb) - forM_ ccb (\(OffsetsCommitCallback cb) -> setOffsetCommitCallback conf cb) + forM_ cbs (\setCb -> setCb conf) return conf --- | Sets a callback that is called when rebalance is needed. --- --- Callback implementations suppose to watch for 'KafkaResponseError' 'RdKafkaRespErrAssignPartitions' and --- for 'KafkaResponseError' 'RdKafkaRespErrRevokePartitions'. Other error codes are not expected and would indicate --- something really bad happening in a system, or bugs in @librdkafka@ itself. --- --- A callback is expected to call 'assign' according to the error code it receives. --- --- * When 'RdKafkaRespErrAssignPartitions' happens 'assign' should be called with all the partitions it was called with. --- It is OK to alter partitions offsets before calling 'assign'. --- --- * When 'RdKafkaRespErrRevokePartitions' happens 'assign' should be called with an empty list of partitions. -setRebalanceCallback :: KafkaConf - -> (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) - -> IO () -setRebalanceCallback (KafkaConf conf) callback = rdKafkaConfSetRebalanceCb conf realCb - where - realCb :: Ptr RdKafkaT -> RdKafkaRespErrT -> Ptr RdKafkaTopicPartitionListT -> Ptr Word8 -> IO () - realCb rk err pl _ = do - rk' <- newForeignPtr_ rk - ps <- fromNativeTopicPartitionList' pl - callback (KafkaConsumer rk' conf) (KafkaResponseError err) ps - --- | Sets a callback that is called when rebalance is needed. --- --- The results of automatic or manual offset commits will be scheduled --- for this callback and is served by `pollMessage`. --- --- A callback is expected to call 'assign' according to the error code it receives. --- --- If no partitions had valid offsets to commit this callback will be called --- with `KafkaError` == `KafkaResponseError` `RdKafkaRespErrNoOffset` which is not to be considered --- an error. -setOffsetCommitCallback :: KafkaConf - -> (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) - -> IO () -setOffsetCommitCallback (KafkaConf conf) callback = rdKafkaConfSetOffsetCommitCb conf realCb - where - realCb :: Ptr RdKafkaT -> RdKafkaRespErrT -> Ptr RdKafkaTopicPartitionListT -> Ptr Word8 -> IO () - realCb rk err pl _ = do - rk' <- newForeignPtr_ rk - ps <- fromNativeTopicPartitionList' pl - callback (KafkaConsumer rk' conf) (KafkaResponseError err) ps - -- | Subscribes to a given list of topics. -- -- Wildcard (regex) topics are supported by the librdkafka assignor: @@ -191,24 +184,24 @@ -- be regex-matched to the full list of topics in the cluster and matching -- topics will be added to the subscription list. subscribe :: KafkaConsumer -> [TopicName] -> IO (Maybe KafkaError) -subscribe (KafkaConsumer k _) ts = do +subscribe (KafkaConsumer (Kafka k) _) ts = do pl <- newRdKafkaTopicPartitionListT (length ts) mapM_ (\(TopicName t) -> rdKafkaTopicPartitionListAdd pl t (-1)) ts res <- KafkaResponseError <$> rdKafkaSubscribe k pl return $ kafkaErrorToMaybe res -setDefaultTopicConf :: RdKafkaConfTPtr -> TopicConf -> IO () -setDefaultTopicConf kc (TopicConf tc) = +setDefaultTopicConf :: KafkaConf -> TopicConf -> IO () +setDefaultTopicConf (KafkaConf kc) (TopicConf tc) = rdKafkaTopicConfDup tc >>= rdKafkaConfSetDefaultTopicConf kc commitOffsets :: OffsetCommit -> KafkaConsumer -> RdKafkaTopicPartitionListTPtr -> IO (Maybe KafkaError) -commitOffsets o (KafkaConsumer k _) pl = +commitOffsets o (KafkaConsumer (Kafka k) _) pl = (kafkaErrorToMaybe . KafkaResponseError) <$> rdKafkaCommit k pl (offsetCommitToBool o) setConsumerLogLevel :: KafkaConsumer -> KafkaLogLevel -> IO () -setConsumerLogLevel (KafkaConsumer k _) level = +setConsumerLogLevel (KafkaConsumer (Kafka k) _) level = liftIO $ rdKafkaSetLogLevel k (fromEnum level) redirectCallbacksPoll :: KafkaConsumer -> IO (Maybe KafkaError) -redirectCallbacksPoll (KafkaConsumer k _) = +redirectCallbacksPoll (KafkaConsumer (Kafka k) _) = (kafkaErrorToMaybe . KafkaResponseError) <$> rdKafkaPollSetConsumer k diff --git a/src/Kafka/Consumer/Callbacks.hs b/src/Kafka/Consumer/Callbacks.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Consumer/Callbacks.hs @@ -0,0 +1,52 @@ +module Kafka.Consumer.Callbacks +( rebalanceCallback +, offsetCommitCallback +, module X +) +where + +import Foreign +import Kafka.Callbacks as X +import Kafka.Consumer.Convert +import Kafka.Consumer.Types +import Kafka.Internal.RdKafka +import Kafka.Types + +-- | Sets a callback that is called when rebalance is needed. +-- +-- Callback implementations suppose to watch for 'KafkaResponseError' 'RdKafkaRespErrAssignPartitions' and +-- for 'KafkaResponseError' 'RdKafkaRespErrRevokePartitions'. Other error codes are not expected and would indicate +-- something really bad happening in a system, or bugs in @librdkafka@ itself. +-- +-- A callback is expected to call 'assign' according to the error code it receives. +-- +-- * When 'RdKafkaRespErrAssignPartitions' happens 'assign' should be called with all the partitions it was called with. +-- It is OK to alter partitions offsets before calling 'assign'. +-- +-- * When 'RdKafkaRespErrRevokePartitions' happens 'assign' should be called with an empty list of partitions. +rebalanceCallback :: (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) -> KafkaConf -> IO () +rebalanceCallback callback (KafkaConf conf) = rdKafkaConfSetRebalanceCb conf realCb + where + realCb k err pl = do + k' <- newForeignPtr_ k + pls <- fromNativeTopicPartitionList' pl + callback (KafkaConsumer (Kafka k') (KafkaConf conf)) (KafkaResponseError err) pls + +-- | Sets a callback that is called when rebalance is needed. +-- +-- The results of automatic or manual offset commits will be scheduled +-- for this callback and is served by `pollMessage`. +-- +-- A callback is expected to call 'assign' according to the error code it receives. +-- +-- If no partitions had valid offsets to commit this callback will be called +-- with `KafkaError` == `KafkaResponseError` `RdKafkaRespErrNoOffset` which is not to be considered +-- an error. +offsetCommitCallback :: (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) -> KafkaConf -> IO () +offsetCommitCallback callback (KafkaConf conf) = rdKafkaConfSetOffsetCommitCb conf realCb + where + realCb k err pl = do + k' <- newForeignPtr_ k + pls <- fromNativeTopicPartitionList' pl + callback (KafkaConsumer (Kafka k') (KafkaConf conf)) (KafkaResponseError err) pls + diff --git a/src/Kafka/Consumer/ConsumerProperties.hs b/src/Kafka/Consumer/ConsumerProperties.hs --- a/src/Kafka/Consumer/ConsumerProperties.hs +++ b/src/Kafka/Consumer/ConsumerProperties.hs @@ -1,113 +1,92 @@ module Kafka.Consumer.ConsumerProperties +( module Kafka.Consumer.ConsumerProperties +, module Kafka.Consumer.Callbacks +) where -- -import Control.Monad -import Data.Map (Map) -import Kafka.Types -import Kafka.Consumer.Types -import qualified Data.Map as M -import qualified Data.List as L +import Control.Monad +import qualified Data.List as L +import Data.Map (Map) +import qualified Data.Map as M +import Data.Monoid ((<>)) +import Kafka.Consumer.Callbacks +import Kafka.Consumer.Types +import Kafka.Types -- | Properties to create 'KafkaConsumer'. data ConsumerProperties = ConsumerProperties - { cpProps :: Map String String - , cpRebalanceCallback :: Maybe ReballanceCallback - , cpOffsetsCallback :: Maybe OffsetsCommitCallback - , cpLogLevel :: Maybe KafkaLogLevel + { cpProps :: Map String String + , cpLogLevel :: Maybe KafkaLogLevel + , callbacks :: [KafkaConf -> IO ()] } instance Monoid ConsumerProperties where - mempty = ConsumerProperties M.empty Nothing Nothing Nothing - mappend (ConsumerProperties m1 rb1 oc1 ll1) (ConsumerProperties m2 rb2 oc2 ll2) = - ConsumerProperties (M.union m1 m2) (rb2 `mplus` rb1) (oc2 `mplus` oc1) (ll2 `mplus` ll1) + mempty = ConsumerProperties M.empty Nothing [] + mappend (ConsumerProperties m1 ll1 cb1) (ConsumerProperties m2 ll2 cb2) = + ConsumerProperties (M.union m1 m2) (ll2 `mplus` ll1) (cb2 <> cb1) -consumerBrokersList :: [BrokerAddress] -> ConsumerProperties -consumerBrokersList bs = +brokersList :: [BrokerAddress] -> ConsumerProperties +brokersList bs = let bs' = L.intercalate "," ((\(BrokerAddress x) -> x) <$> bs) - in extraConsumerProps $ M.fromList [("bootstrap.servers", bs')] + in extraProps $ M.fromList [("bootstrap.servers", bs')] -- | Disables auto commit for the consumer noAutoCommit :: ConsumerProperties noAutoCommit = - extraConsumerProps $ M.fromList [("enable.auto.commit", "false")] + extraProps $ M.fromList [("enable.auto.commit", "false")] -- | Consumer group id groupId :: ConsumerGroupId -> ConsumerProperties groupId (ConsumerGroupId cid) = - extraConsumerProps $ M.fromList [("group.id", cid)] + extraProps $ M.fromList [("group.id", cid)] clientId :: ClientId -> ConsumerProperties clientId (ClientId cid) = - extraConsumerProps $ M.fromList [("client.id", cid)] - --- | Sets a callback that is called when rebalance is needed. --- --- Callback implementations suppose to watch for 'KafkaResponseError' 'RdKafkaRespErrAssignPartitions' and --- for 'KafkaResponseError' 'RdKafkaRespErrRevokePartitions'. Other error codes are not expected and would indicate --- something really bad happening in a system, or bugs in @librdkafka@ itself. --- --- A callback is expected to call 'assign' according to the error code it receives. --- --- * When 'RdKafkaRespErrAssignPartitions' happens 'assign' should be called with all the partitions it was called with. --- It is OK to alter partitions offsets before calling 'assign'. --- --- * When 'RdKafkaRespErrRevokePartitions' happens 'assign' should be called with an empty list of partitions. -reballanceCallback :: ReballanceCallback -> ConsumerProperties -reballanceCallback cb = ConsumerProperties M.empty (Just cb) Nothing Nothing + extraProps $ M.fromList [("client.id", cid)] --- | Sets offset commit callback for use with consumer groups. --- --- The results of automatic or manual offset commits will be scheduled --- for this callback and is served by `pollMessage`. --- --- A callback is expected to call 'assign' according to the error code it receives. --- --- If no partitions had valid offsets to commit this callback will be called --- with `KafkaError` == `KafkaResponseError` `RdKafkaRespErrNoOffset` which is not to be considered --- an error. -offsetsCommitCallback :: OffsetsCommitCallback -> ConsumerProperties -offsetsCommitCallback cb = ConsumerProperties M.empty Nothing (Just cb) Nothing +setCallback :: (KafkaConf -> IO ()) -> ConsumerProperties +setCallback cb = ConsumerProperties M.empty Nothing [cb] -- | Sets the logging level. --- Usually is used with 'consumerDebug' to configure which logs are needed. +-- Usually is used with 'debugOptions' to configure which logs are needed. consumerLogLevel :: KafkaLogLevel -> ConsumerProperties -consumerLogLevel ll = ConsumerProperties M.empty Nothing Nothing (Just ll) +consumerLogLevel ll = ConsumerProperties M.empty (Just ll) [] -- | Sets the compression codec for the consumer. -consumerCompression :: KafkaCompressionCodec -> ConsumerProperties -consumerCompression c = - extraConsumerProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) +compression :: KafkaCompressionCodec -> ConsumerProperties +compression c = + extraProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) -- | Suppresses consumer disconnects logs. -- -- It might be useful to turn this off when interacting with brokers -- with an aggressive connection.max.idle.ms value. -consumerSuppressDisconnectLogs :: ConsumerProperties -consumerSuppressDisconnectLogs = - extraConsumerProps $ M.fromList [("log.connection.close", "false")] +suppressDisconnectLogs :: ConsumerProperties +suppressDisconnectLogs = + extraProps $ M.fromList [("log.connection.close", "false")] -- | Any configuration options that are supported by /librdkafka/. -- The full list can be found -extraConsumerProps :: Map String String -> ConsumerProperties -extraConsumerProps m = ConsumerProperties m Nothing Nothing Nothing -{-# INLINE extraConsumerProps #-} +extraProps :: Map String String -> ConsumerProperties +extraProps m = ConsumerProperties m Nothing [] +{-# INLINE extraProps #-} -- | Any configuration options that are supported by /librdkafka/. -- The full list can be found -extraConsumerProp :: String -> String -> ConsumerProperties -extraConsumerProp k v = ConsumerProperties (M.singleton k v) Nothing Nothing Nothing -{-# INLINE extraConsumerProp #-} +extraProp :: String -> String -> ConsumerProperties +extraProp k v = ConsumerProperties (M.singleton k v) Nothing [] +{-# INLINE extraProp #-} -- | Sets debug features for the consumer. -- Usually is used with 'consumerLogLevel'. -consumerDebug :: [KafkaDebug] -> ConsumerProperties -consumerDebug [] = extraConsumerProps M.empty -consumerDebug d = +debugOptions :: [KafkaDebug] -> ConsumerProperties +debugOptions [] = extraProps M.empty +debugOptions d = let points = L.intercalate "," (kafkaDebugToString <$> d) - in extraConsumerProps $ M.fromList [("debug", points)] + in extraProps $ M.fromList [("debug", points)] -consumerQueuedMaxMessagesKBytes :: Int -> ConsumerProperties -consumerQueuedMaxMessagesKBytes kBytes = - extraConsumerProp "queued.max.messages.kbytes" (show kBytes) -{-# INLINE consumerQueuedMaxMessagesKBytes #-} +queuedMaxMessagesKBytes :: Int -> ConsumerProperties +queuedMaxMessagesKBytes kBytes = + extraProp "queued.max.messages.kbytes" (show kBytes) +{-# INLINE queuedMaxMessagesKBytes #-} diff --git a/src/Kafka/Consumer/Convert.hs b/src/Kafka/Consumer/Convert.hs --- a/src/Kafka/Consumer/Convert.hs +++ b/src/Kafka/Consumer/Convert.hs @@ -3,15 +3,16 @@ where import Control.Monad -import qualified Data.ByteString as BS +import qualified Data.ByteString as BS +import Data.Map.Strict (Map, fromListWith) +import qualified Data.Set as S import Foreign import Foreign.C.Error import Foreign.C.String import Kafka.Consumer.Types -import Kafka.Types import Kafka.Internal.RdKafka -import Kafka.Internal.RdKafkaEnum import Kafka.Internal.Shared +import Kafka.Types -- | Converts offsets sync policy to integer (the way Kafka understands it): -- @@ -23,8 +24,8 @@ offsetSyncToInt :: OffsetStoreSync -> Int offsetSyncToInt sync = case sync of - OffsetSyncDisable -> -1 - OffsetSyncImmediate -> 0 + OffsetSyncDisable -> -1 + OffsetSyncImmediate -> 0 OffsetSyncInterval ms -> ms {-# INLINE offsetSyncToInt #-} @@ -46,6 +47,10 @@ | otherwise = PartitionOffsetInvalid {-# INLINE int64ToOffset #-} +fromNativeTopicPartitionList'' :: RdKafkaTopicPartitionListTPtr -> IO [TopicPartition] +fromNativeTopicPartitionList'' ptr = + withForeignPtr ptr $ \fptr -> fromNativeTopicPartitionList' fptr + fromNativeTopicPartitionList' :: Ptr RdKafkaTopicPartitionListT -> IO [TopicPartition] fromNativeTopicPartitionList' ppl = peek ppl >>= fromNativeTopicPartitionList @@ -75,11 +80,21 @@ rdKafkaTopicPartitionListSetOffset pl tn tp to) ps return pl +toNativeTopicPartitionList' :: [(TopicName, PartitionId)] -> IO RdKafkaTopicPartitionListTPtr +toNativeTopicPartitionList' tps = do + let utps = S.toList . S.fromList $ tps + pl <- newRdKafkaTopicPartitionListT (length utps) + mapM_ (\(TopicName t, PartitionId p) -> rdKafkaTopicPartitionListAdd pl t p) utps + return pl + topicPartitionFromMessage :: ConsumerRecord k v -> TopicPartition topicPartitionFromMessage m = let (Offset moff) = crOffset m in TopicPartition (crTopic m) (crPartition m) (PartitionOffset moff) +toMap :: Ord k => [(k, v)] -> Map k [v] +toMap kvs = fromListWith (++) [(k, [v]) | (k, v) <- kvs] + fromMessagePtr :: RdKafkaMessageTPtr -> IO (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))) fromMessagePtr ptr = withForeignPtr ptr $ \realPtr -> @@ -87,30 +102,43 @@ else do s <- peek realPtr msg <- if err'RdKafkaMessageT s /= RdKafkaRespErrNoError - then return . Left . KafkaResponseError $ err'RdKafkaMessageT s - else Right <$> fromMessageStorable s + then return . Left . KafkaResponseError $ err'RdKafkaMessageT s + else Right <$> mkRecord s rdKafkaMessageDestroy realPtr return msg + where + mkRecord msg = do + topic <- readTopic msg + key <- readKey msg + payload <- readPayload msg + timestamp <- readTimestamp ptr + return ConsumerRecord + { crTopic = TopicName topic + , crPartition = PartitionId $ partition'RdKafkaMessageT msg + , crOffset = Offset $ offset'RdKafkaMessageT msg + , crTimestamp = timestamp + , crKey = key + , crValue = payload + } -fromMessageStorable :: RdKafkaMessageT -> IO (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)) -fromMessageStorable s = do - topic <- newForeignPtr_ (topic'RdKafkaMessageT s) >>= rdKafkaTopicName + readTopic msg = newForeignPtr_ (topic'RdKafkaMessageT msg) >>= rdKafkaTopicName + readPayload = readBS len'RdKafkaMessageT payload'RdKafkaMessageT + readKey = readBS keyLen'RdKafkaMessageT key'RdKafkaMessageT + readTimestamp msg = + alloca $ \p -> do + typeP <- newForeignPtr_ p + ts <- fromIntegral <$> rdKafkaMessageTimestamp msg typeP + tsType <- peek p + return $ case tsType of + RdKafkaTimestampCreateTime -> CreateTime (Millis ts) + RdKafkaTimestampLogAppendTime -> LogAppendTime (Millis ts) + RdKafkaTimestampNotAvailable -> NoTimestamp - payload <- if payload'RdKafkaMessageT s == nullPtr - then return Nothing - else Just <$> word8PtrToBS (len'RdKafkaMessageT s) (payload'RdKafkaMessageT s) + readBS flen fdata s = if fdata s == nullPtr + then return Nothing + else Just <$> word8PtrToBS (flen s) (fdata s) - key <- if key'RdKafkaMessageT s == nullPtr - then return Nothing - else Just <$> word8PtrToBS (keyLen'RdKafkaMessageT s) (key'RdKafkaMessageT s) - return $ ConsumerRecord - (TopicName topic) - (PartitionId $ partition'RdKafkaMessageT s) - (Offset $ offset'RdKafkaMessageT s) - key - payload - offsetCommitToBool :: OffsetCommit -> Bool -offsetCommitToBool OffsetCommit = False +offsetCommitToBool OffsetCommit = False offsetCommitToBool OffsetCommitAsync = True diff --git a/src/Kafka/Consumer/Types.hs b/src/Kafka/Consumer/Types.hs --- a/src/Kafka/Consumer/Types.hs +++ b/src/Kafka/Consumer/Types.hs @@ -4,26 +4,38 @@ where -import Data.Bifoldable -import Data.Bifunctor -import Data.Bitraversable -import Data.Int -import Data.Typeable -import Kafka.Internal.RdKafka -import Kafka.Types +import Data.Bifoldable +import Data.Bifunctor +import Data.Bitraversable +import Data.Int +import Data.Typeable +import Kafka.Types -data KafkaConsumer = KafkaConsumer { kcKafkaPtr :: !RdKafkaTPtr, kcKafkaConf :: !RdKafkaConfTPtr} deriving (Show) +data KafkaConsumer = KafkaConsumer { kcKafkaPtr :: !Kafka, kcKafkaConf :: !KafkaConf} deriving (Show) -newtype ReballanceCallback = ReballanceCallback (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) -newtype OffsetsCommitCallback = OffsetsCommitCallback (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) +instance HasKafka KafkaConsumer where + getKafka = kcKafkaPtr + {-# INLINE getKafka #-} +instance HasKafkaConf KafkaConsumer where + getKafkaConf = kcKafkaConf + {-# INLINE getKafkaConf #-} + newtype ConsumerGroupId = ConsumerGroupId String deriving (Show, Eq) newtype Offset = Offset Int64 deriving (Show, Eq, Read) -newtype PartitionId = PartitionId Int deriving (Show, Eq, Read, Ord) -newtype Millis = Millis Int deriving (Show, Eq, Ord, Num) -newtype ClientId = ClientId String deriving (Show, Eq, Ord) data OffsetReset = Earliest | Latest deriving (Show, Eq) +data SubscribedPartitions + = SubscribedPartitions [PartitionId] + | SubscribedPartitionsAll + deriving (Show, Eq) + +data Timestamp = + CreateTime !Millis + | LogAppendTime !Millis + | NoTimestamp + deriving (Show, Eq, Read) + -- | Offsets commit mode data OffsetCommit = OffsetCommit -- ^ Forces consumer to block until the broker offsets commit is done @@ -53,13 +65,14 @@ { crTopic :: !TopicName -- ^ Kafka topic this message was received from , crPartition :: !PartitionId -- ^ Kafka partition this message was received from , crOffset :: !Offset -- ^ Offset within the 'crPartition' Kafka partition + , crTimestamp :: !Timestamp -- ^ Message timestamp , crKey :: !k , crValue :: !v } deriving (Eq, Show, Read, Typeable) instance Bifunctor ConsumerRecord where - bimap f g (ConsumerRecord t p o k v) = ConsumerRecord t p o (f k) (g v) + bimap f g (ConsumerRecord t p o ts k v) = ConsumerRecord t p o ts (f k) (g v) {-# INLINE bimap #-} instance Functor (ConsumerRecord k) where @@ -134,3 +147,9 @@ | PartitionOffsetStored | PartitionOffsetInvalid deriving (Eq, Show) + +data CommitPartitionOffset = CommitPartitionOffset + { cpoTopicName :: !TopicName + , cpoPartitionId :: !PartitionId + , cpoOffset :: !Offset + } deriving (Eq, Show) diff --git a/src/Kafka/Internal/RdKafka.chs b/src/Kafka/Internal/RdKafka.chs --- a/src/Kafka/Internal/RdKafka.chs +++ b/src/Kafka/Internal/RdKafka.chs @@ -10,7 +10,6 @@ import Foreign.C.Error import Foreign.C.String import Foreign.C.Types -import Kafka.Internal.RdKafkaEnum import System.IO import System.Posix.IO import System.Posix.Types @@ -26,6 +25,11 @@ type Word8Ptr = Ptr Word8 type CCharBufPointer = Ptr CChar +{#enum rd_kafka_type_t as ^ {underscoreToCase} deriving (Show, Eq) #} +{#enum rd_kafka_conf_res_t as ^ {underscoreToCase} deriving (Show, Eq) #} +{#enum rd_kafka_resp_err_t as ^ {underscoreToCase} deriving (Show, Eq) #} +{#enum rd_kafka_timestamp_type_t as ^ {underscoreToCase} deriving (Show, Eq) #} + type RdKafkaMsgFlag = Int rdKafkaMsgFlagFree :: RdKafkaMsgFlag rdKafkaMsgFlagFree = 0x1 @@ -246,7 +250,7 @@ {`Int'} -> `RdKafkaTopicPartitionListTPtr' #} foreign import ccall unsafe "rdkafka.h &rd_kafka_topic_partition_list_destroy" - rdKafkaTopicPartitionListDestroy :: FunPtr (Ptr RdKafkaTopicPartitionListT -> IO ()) + rdKafkaTopicPartitionListDestroy :: FinalizerPtr RdKafkaTopicPartitionListT newRdKafkaTopicPartitionListT :: Int -> IO RdKafkaTopicPartitionListTPtr newRdKafkaTopicPartitionListT size = do @@ -275,7 +279,7 @@ ---- Rebalance Callback type RdRebalanceCallback' = Ptr RdKafkaT -> CInt -> Ptr RdKafkaTopicPartitionListT -> Ptr Word8 -> IO () -type RdRebalanceCallback = Ptr RdKafkaT -> RdKafkaRespErrT -> Ptr RdKafkaTopicPartitionListT -> Ptr Word8 -> IO () +type RdRebalanceCallback = Ptr RdKafkaT -> RdKafkaRespErrT -> Ptr RdKafkaTopicPartitionListT -> IO () foreign import ccall safe "wrapper" mkRebalanceCallback :: RdRebalanceCallback' -> IO (FunPtr RdRebalanceCallback') @@ -288,7 +292,7 @@ rdKafkaConfSetRebalanceCb :: RdKafkaConfTPtr -> RdRebalanceCallback -> IO () rdKafkaConfSetRebalanceCb conf cb = do - cb' <- mkRebalanceCallback (\k e p o -> cb k (cIntToEnum e) p o) + cb' <- mkRebalanceCallback (\k e p _ -> cb k (cIntToEnum e) p) withForeignPtr conf $ \c -> rdKafkaConfSetRebalanceCb' c cb' return () @@ -324,7 +328,7 @@ ---- Offset Commit Callback type OffsetCommitCallback' = Ptr RdKafkaT -> CInt -> Ptr RdKafkaTopicPartitionListT -> Word8Ptr -> IO () -type OffsetCommitCallback = Ptr RdKafkaT -> RdKafkaRespErrT -> Ptr RdKafkaTopicPartitionListT -> Word8Ptr -> IO () +type OffsetCommitCallback = Ptr RdKafkaT -> RdKafkaRespErrT -> Ptr RdKafkaTopicPartitionListT -> IO () foreign import ccall safe "wrapper" mkOffsetCommitCallback :: OffsetCommitCallback' -> IO (FunPtr OffsetCommitCallback') @@ -334,10 +338,26 @@ rdKafkaConfSetOffsetCommitCb :: RdKafkaConfTPtr -> OffsetCommitCallback -> IO () rdKafkaConfSetOffsetCommitCb conf cb = do - cb' <- mkOffsetCommitCallback (\k e p o -> cb k (cIntToEnum e) p o) + cb' <- mkOffsetCommitCallback (\k e p _ -> cb k (cIntToEnum e) p) withForeignPtr conf $ \c -> rdKafkaConfSetOffsetCommitCb' c cb' return () + +----- Error Callback +type ErrorCallback' = Ptr RdKafkaT -> CInt -> CString -> Word8Ptr -> IO () +type ErrorCallback = Ptr RdKafkaT -> RdKafkaRespErrT -> String -> IO () + +foreign import ccall safe "wrapper" + mkErrorCallback :: ErrorCallback' -> IO (FunPtr ErrorCallback') + +foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_error_cb" + rdKafkaConfSetErrorCb' :: Ptr RdKafkaConfT -> FunPtr ErrorCallback' -> IO () + +rdKafkaConfSetErrorCb :: RdKafkaConfTPtr -> ErrorCallback -> IO () +rdKafkaConfSetErrorCb conf cb = do + cb' <- mkErrorCallback (\k e r _ -> peekCAString r >>= cb k (cIntToEnum e)) + withForeignPtr conf $ \c -> rdKafkaConfSetErrorCb' c cb' + ---- Throttle Callback type ThrottleCallback = Ptr RdKafkaT -> CString -> Int -> Int -> Word8Ptr -> IO () @@ -353,18 +373,37 @@ withForeignPtr conf $ \c -> rdKafkaConfSetThrottleCb' c cb' return () +---- Log Callback +type LogCallback' = Ptr RdKafkaT -> CInt -> CString -> CString -> IO () +type LogCallback = Ptr RdKafkaT -> Int -> String -> String -> IO () + +foreign import ccall safe "wrapper" + mkLogCallback :: LogCallback' -> IO (FunPtr LogCallback') + +foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_log_cb" + rdKafkaConfSetLogCb' :: Ptr RdKafkaConfT -> FunPtr LogCallback' -> IO () + +rdKafkaConfSetLogCb :: RdKafkaConfTPtr -> LogCallback -> IO () +rdKafkaConfSetLogCb conf cb = do + cb' <- mkLogCallback $ \k l f b -> do + f' <- peekCAString f + b' <- peekCAString b + cb k (cIntConv l) f' b' + withForeignPtr conf $ \c -> rdKafkaConfSetLogCb' c cb' + ---- Stats Callback -type StatsCallback = Ptr RdKafkaT -> CString -> CSize -> Word8Ptr -> IO () +type StatsCallback' = Ptr RdKafkaT -> CString -> CSize -> Word8Ptr -> IO () +type StatsCallback = Ptr RdKafkaT -> String -> IO () foreign import ccall safe "wrapper" - mkStatsCallback :: StatsCallback -> IO (FunPtr StatsCallback) + mkStatsCallback :: StatsCallback' -> IO (FunPtr StatsCallback') foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_stats_cb" - rdKafkaConfSetStatsCb' :: Ptr RdKafkaConfT -> FunPtr StatsCallback -> IO () + rdKafkaConfSetStatsCb' :: Ptr RdKafkaConfT -> FunPtr StatsCallback' -> IO () rdKafkaConfSetStatsCb :: RdKafkaConfTPtr -> StatsCallback -> IO () rdKafkaConfSetStatsCb conf cb = do - cb' <- mkStatsCallback cb + cb' <- mkStatsCallback $ \k j jl _ -> peekCAStringLen (j, cIntConv jl) >>= cb k withForeignPtr conf $ \c -> rdKafkaConfSetStatsCb' c cb' return () @@ -466,7 +505,7 @@ {`RdKafkaTPtr'} -> `RdKafkaQueueTPtr' #} foreign import ccall unsafe "rdkafka.h &rd_kafka_queue_destroy" - rdKafkaQueueDestroy :: FunPtr (Ptr RdKafkaQueueT -> IO ()) + rdKafkaQueueDestroy :: FinalizerPtr RdKafkaQueueT newRdKafkaQueue :: RdKafkaTPtr -> IO RdKafkaQueueTPtr newRdKafkaQueue k = do @@ -484,10 +523,18 @@ {`RdKafkaTPtr'} -> `RdKafkaRespErrT' cIntToEnum #} -{#fun rd_kafka_subscription as ^ - {`RdKafkaTPtr', castPtr `Ptr (Ptr RdKafkaTopicPartitionListT)'} +{#fun rd_kafka_subscription as rdKafkaSubscription' + {`RdKafkaTPtr', alloca- `Ptr RdKafkaTopicPartitionListT' peekPtr*} -> `RdKafkaRespErrT' cIntToEnum #} +rdKafkaSubscription :: RdKafkaTPtr -> IO (Either RdKafkaRespErrT RdKafkaTopicPartitionListTPtr) +rdKafkaSubscription k = do + (err, sub) <- rdKafkaSubscription' k + case err of + RdKafkaRespErrNoError -> + Right <$> newForeignPtr rdKafkaTopicPartitionListDestroy sub + e -> return (Left e) + {#fun rd_kafka_consumer_poll as ^ {`RdKafkaTPtr', `Int'} -> `RdKafkaMessageTPtr' #} @@ -508,10 +555,18 @@ {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr'} -> `RdKafkaRespErrT' cIntToEnum #} -{#fun rd_kafka_assignment as ^ - {`RdKafkaTPtr', castPtr `Ptr (Ptr RdKafkaTopicPartitionListT)'} +{#fun rd_kafka_assignment as rdKafkaAssignment' + {`RdKafkaTPtr', alloca- `Ptr RdKafkaTopicPartitionListT' peekPtr* } -> `RdKafkaRespErrT' cIntToEnum #} +rdKafkaAssignment :: RdKafkaTPtr -> IO (Either RdKafkaRespErrT RdKafkaTopicPartitionListTPtr) +rdKafkaAssignment k = do + (err, ass) <- rdKafkaAssignment' k + case err of + RdKafkaRespErrNoError -> + Right <$> newForeignPtr rdKafkaTopicPartitionListDestroy ass + e -> return (Left e) + {#fun rd_kafka_commit as ^ {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr', boolToCInt `Bool'} -> `RdKafkaRespErrT' cIntToEnum #} @@ -559,16 +614,16 @@ {#set rd_kafka_group_member_info.member_assignment#} p (castPtr $ memberAssignment'RdKafkaGroupMemberInfoT x) {#set rd_kafka_group_member_info.member_assignment_size#} p (fromIntegral $ memberAssignmentSize'RdKafkaGroupMemberInfoT x) -{#pointer *rd_kafka_group_member_info as RdKafkaGroupMemberInfoTPtr foreign -> RdKafkaGroupMemberInfoT #} +{#pointer *rd_kafka_group_member_info as RdKafkaGroupMemberInfoTPtr -> RdKafkaGroupMemberInfoT #} data RdKafkaGroupInfoT = RdKafkaGroupInfoT - { broker'RdKafkaGroupInfoT :: Ptr RdKafkaMetadataBrokerT + { broker'RdKafkaGroupInfoT :: RdKafkaMetadataBrokerTPtr , group'RdKafkaGroupInfoT :: CString , err'RdKafkaGroupInfoT :: RdKafkaRespErrT , state'RdKafkaGroupInfoT :: CString , protocolType'RdKafkaGroupInfoT :: CString , protocol'RdKafkaGroupInfoT :: CString - , members'RdKafkaGroupInfoT :: Ptr RdKafkaGroupMemberInfoT + , members'RdKafkaGroupInfoT :: RdKafkaGroupMemberInfoTPtr , memberCnt'RdKafkaGroupInfoT :: Int } instance Storable RdKafkaGroupInfoT where @@ -611,38 +666,66 @@ {#pointer *rd_kafka_group_list as RdKafkaGroupListTPtr foreign -> RdKafkaGroupListT #} -{#fun rd_kafka_list_groups as ^ - {`RdKafkaTPtr', `String', castPtr `Ptr (Ptr RdKafkaGroupListT)', `Int'} +{#fun rd_kafka_list_groups as rdKafkaListGroups' + {`RdKafkaTPtr', `CString', alloca- `Ptr RdKafkaGroupListT' peek*, `Int'} -> `RdKafkaRespErrT' cIntToEnum #} -foreign import ccall unsafe "rdkafka.h &rd_kafka_list_groups" - rdKafkaGroupListDestroy :: FunPtr (Ptr RdKafkaGroupListT -> IO ()) +foreign import ccall "rdkafka.h &rd_kafka_group_list_destroy" + rdKafkaGroupListDestroy :: FinalizerPtr RdKafkaGroupListT --- listRdKafkaGroups :: RdKafkaTPtr -> String -> Int -> IO (Either RdKafkaRespErrT RdKafkaGroupListTPtr) --- listRdKafkaGroups k g t = alloca $ \lstDblPtr -> do --- err <- rdKafkaListGroups k g lstDblPtr t --- case err of --- RdKafkaRespErrNoError -> do --- lstPtr <- peek lstDblPtr --- lst <- peek lstPtr --- addForeignPtrFinalizer rdKafkaGroupListDestroy lstPtr --- return $ Right lstPtr --- e -> return $ Left e +rdKafkaListGroups :: RdKafkaTPtr -> Maybe String -> Int -> IO (Either RdKafkaRespErrT RdKafkaGroupListTPtr) +rdKafkaListGroups k g t = case g of + Nothing -> listGroups nullPtr + Just strGrp -> withCAString strGrp listGroups + where + listGroups grp = do + (err, res) <- rdKafkaListGroups' k grp t + case err of + RdKafkaRespErrNoError -> Right <$> newForeignPtr rdKafkaGroupListDestroy res + e -> return $ Left e ------------------------------------------------------------------------------------------------- -- rd_kafka_message foreign import ccall unsafe "rdkafka.h &rd_kafka_message_destroy" - rdKafkaMessageDestroyF :: FunPtr (Ptr RdKafkaMessageT -> IO ()) + rdKafkaMessageDestroyF :: FinalizerPtr RdKafkaMessageT foreign import ccall unsafe "rdkafka.h rd_kafka_message_destroy" rdKafkaMessageDestroy :: Ptr RdKafkaMessageT -> IO () +{#fun rd_kafka_query_watermark_offsets as rdKafkaQueryWatermarkOffsets' + {`RdKafkaTPtr', `String', cIntConv `CInt32T', + alloca- `Int64' peekInt64Conv*, alloca- `Int64' peekInt64Conv*, + cIntConv `Int' + } -> `RdKafkaRespErrT' cIntToEnum #} + + +rdKafkaQueryWatermarkOffsets :: RdKafkaTPtr -> String -> Int -> Int -> IO (Either RdKafkaRespErrT (Int64, Int64)) +rdKafkaQueryWatermarkOffsets kafka topic partition timeout = do + (err, l, h) <- rdKafkaQueryWatermarkOffsets' kafka topic (cIntConv partition) timeout + return $ case err of + RdKafkaRespErrNoError -> Right (cIntConv l, cIntConv h) + e -> Left e + +{#pointer *rd_kafka_timestamp_type_t as RdKafkaTimestampTypeTPtr foreign -> RdKafkaTimestampTypeT #} + +instance Storable RdKafkaTimestampTypeT where + sizeOf _ = {#sizeof rd_kafka_timestamp_type_t#} + alignment _ = {#alignof rd_kafka_timestamp_type_t#} + peek p = cIntToEnum <$> peek (castPtr p) + poke p x = poke (castPtr p) (enumToCInt x) + +{#fun rd_kafka_message_timestamp as ^ + {`RdKafkaMessageTPtr', `RdKafkaTimestampTypeTPtr'} -> `CInt64T' cIntConv #} + +{#fun rd_kafka_offsets_for_times as rdKafkaOffsetsForTimes + {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr', `Int'} -> `RdKafkaRespErrT' cIntToEnum #} + -- rd_kafka_conf {#fun rd_kafka_conf_new as ^ {} -> `RdKafkaConfTPtr' #} foreign import ccall unsafe "rdkafka.h &rd_kafka_conf_destroy" - rdKafkaConfDestroy :: FunPtr (Ptr RdKafkaConfT -> IO ()) + rdKafkaConfDestroy :: FinalizerPtr RdKafkaConfT {#fun rd_kafka_conf_dup as ^ {`RdKafkaConfTPtr'} -> `RdKafkaConfTPtr' #} @@ -674,7 +757,7 @@ {`RdKafkaTopicConfTPtr'} -> `RdKafkaTopicConfTPtr' #} foreign import ccall unsafe "rdkafka.h &rd_kafka_topic_conf_destroy" - rdKafkaTopicConfDestroy :: FunPtr (Ptr RdKafkaTopicConfT -> IO ()) + rdKafkaTopicConfDestroy :: FinalizerPtr RdKafkaTopicConfT {#fun rd_kafka_topic_conf_set as ^ {`RdKafkaTopicConfTPtr', `String', `String', id `CCharBufPointer', cIntConv `CSize'} @@ -728,6 +811,9 @@ {#fun rd_kafka_consume_stop as rdKafkaConsumeStopInternal {`RdKafkaTopicTPtr', cIntConv `CInt32T'} -> `Int' #} +{#fun rd_kafka_seek as rdKafkaSeek + {`RdKafkaTopicTPtr', `Int32', `Int64', `Int'} -> `RdKafkaRespErrT' cIntToEnum #} + {#fun rd_kafka_consume as ^ {`RdKafkaTopicTPtr', cIntConv `CInt32T', `Int'} -> `RdKafkaMessageTPtr' #} @@ -756,19 +842,25 @@ {#fun rd_kafka_produce_batch as ^ {`RdKafkaTopicTPtr', cIntConv `CInt32T', `Int', `RdKafkaMessageTPtr', `Int'} -> `Int' #} -castMetadata :: Ptr (Ptr RdKafkaMetadataT) -> Ptr (Ptr ()) -castMetadata ptr = castPtr ptr -- rd_kafka_metadata -{#fun rd_kafka_metadata as ^ +{#fun rd_kafka_metadata as rdKafkaMetadata' {`RdKafkaTPtr', boolToCInt `Bool', `RdKafkaTopicTPtr', - castMetadata `Ptr (Ptr RdKafkaMetadataT)', `Int'} + alloca- `Ptr RdKafkaMetadataT' peekPtr*, `Int'} -> `RdKafkaRespErrT' cIntToEnum #} -{# fun rd_kafka_metadata_destroy as ^ - {castPtr `Ptr RdKafkaMetadataT'} -> `()' #} +foreign import ccall unsafe "rdkafka.h &rd_kafka_metadata_destroy" + rdKafkaMetadataDestroy :: FinalizerPtr RdKafkaMetadataT +rdKafkaMetadata :: RdKafkaTPtr -> Bool -> Maybe RdKafkaTopicTPtr -> Int -> IO (Either RdKafkaRespErrT RdKafkaMetadataTPtr) +rdKafkaMetadata k allTopics mt timeout = do + tptr <- maybe (newForeignPtr_ nullPtr) pure mt + (err, res) <- rdKafkaMetadata' k allTopics tptr timeout + case err of + RdKafkaRespErrNoError -> Right <$> newForeignPtr rdKafkaMetadataDestroy res + e -> return (Left e) + {#fun rd_kafka_poll as ^ {`RdKafkaTPtr', `Int'} -> `Int' #} @@ -778,7 +870,6 @@ {#fun rd_kafka_dump as ^ {`CFilePtr', `RdKafkaTPtr'} -> `()' #} - -- rd_kafka_topic {#fun rd_kafka_topic_name as ^ {`RdKafkaTopicTPtr'} -> `String' #} @@ -794,36 +885,47 @@ withForeignPtr ptr rdKafkaTopicDestroy foreign import ccall unsafe "rdkafka.h &rd_kafka_topic_destroy" - rdKafkaTopicDestroy' :: FunPtr (Ptr RdKafkaTopicT -> IO ()) + rdKafkaTopicDestroy' :: FinalizerPtr RdKafkaTopicT -newUnmanagedRdKafkaTopicT :: RdKafkaTPtr -> String -> RdKafkaTopicConfTPtr -> IO (Either String RdKafkaTopicTPtr) +newUnmanagedRdKafkaTopicT :: RdKafkaTPtr -> String -> Maybe RdKafkaTopicConfTPtr -> IO (Either String RdKafkaTopicTPtr) newUnmanagedRdKafkaTopicT kafkaPtr topic topicConfPtr = do - duper <- rdKafkaTopicConfDup topicConfPtr + duper <- maybe (newForeignPtr_ nullPtr) rdKafkaTopicConfDup topicConfPtr ret <- rdKafkaTopicNew kafkaPtr topic duper withForeignPtr ret $ \realPtr -> if realPtr == nullPtr then kafkaErrnoString >>= return . Left else return $ Right ret -newRdKafkaTopicT :: RdKafkaTPtr -> String -> RdKafkaTopicConfTPtr -> IO (Either String RdKafkaTopicTPtr) +newRdKafkaTopicT :: RdKafkaTPtr -> String -> Maybe RdKafkaTopicConfTPtr -> IO (Either String RdKafkaTopicTPtr) newRdKafkaTopicT kafkaPtr topic topicConfPtr = do - duper <- rdKafkaTopicConfDup topicConfPtr - ret <- rdKafkaTopicNew kafkaPtr topic duper - withForeignPtr ret $ \realPtr -> - if realPtr == nullPtr then kafkaErrnoString >>= return . Left - else do - addForeignPtrFinalizer rdKafkaTopicDestroy' ret - return $ Right ret + res <- newUnmanagedRdKafkaTopicT kafkaPtr topic topicConfPtr + _ <- traverse (addForeignPtrFinalizer rdKafkaTopicDestroy') res + return res -- Marshall / Unmarshall enumToCInt :: Enum a => a -> CInt enumToCInt = fromIntegral . fromEnum +{-# INLINE enumToCInt #-} + cIntToEnum :: Enum a => CInt -> a cIntToEnum = toEnum . fromIntegral +{-# INLINE cIntToEnum #-} + cIntConv :: (Integral a, Num b) => a -> b cIntConv = fromIntegral +{-# INLINE cIntConv #-} + boolToCInt :: Bool -> CInt boolToCInt True = CInt 1 boolToCInt False = CInt 0 +{-# INLINE boolToCInt #-} + +peekInt64Conv :: (Storable a, Integral a) => Ptr a -> IO Int64 +peekInt64Conv = liftM cIntConv . peek +{-# INLINE peekInt64Conv #-} + +peekPtr :: Ptr a -> IO (Ptr b) +peekPtr = peek . castPtr +{-# INLINE peekPtr #-} -- Handle -> File descriptor diff --git a/src/Kafka/Internal/RdKafkaEnum.chs b/src/Kafka/Internal/RdKafkaEnum.chs deleted file mode 100644 --- a/src/Kafka/Internal/RdKafkaEnum.chs +++ /dev/null @@ -1,10 +0,0 @@ -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE EmptyDataDecls #-} - -module Kafka.Internal.RdKafkaEnum where - -#include "rdkafka.h" - -{#enum rd_kafka_type_t as ^ {underscoreToCase} deriving (Show, Eq) #} -{#enum rd_kafka_conf_res_t as ^ {underscoreToCase} deriving (Show, Eq) #} -{#enum rd_kafka_resp_err_t as ^ {underscoreToCase} deriving (Show, Eq) #} diff --git a/src/Kafka/Internal/Setup.hs b/src/Kafka/Internal/Setup.hs --- a/src/Kafka/Internal/Setup.hs +++ b/src/Kafka/Internal/Setup.hs @@ -1,22 +1,18 @@ module Kafka.Internal.Setup where -import Kafka.Types -import Kafka.Internal.RdKafka -import Kafka.Internal.RdKafkaEnum +import Kafka.Internal.RdKafka +import Kafka.Types -import Control.Exception -import Control.Monad -import Foreign -import Foreign.C.String +import Control.Exception +import Control.Monad +import Foreign +import Foreign.C.String -- -- Configuration -- newtype KafkaProps = KafkaProps [(String, String)] deriving (Show, Eq) newtype TopicProps = TopicProps [(String, String)] deriving (Show, Eq) - -newtype KafkaConf = KafkaConf RdKafkaConfTPtr deriving Show -newtype TopicConf = TopicConf RdKafkaTopicConfTPtr deriving Show newTopicConf :: IO TopicConf newTopicConf = TopicConf <$> newRdKafkaTopicConfT diff --git a/src/Kafka/Internal/Shared.hs b/src/Kafka/Internal/Shared.hs --- a/src/Kafka/Internal/Shared.hs +++ b/src/Kafka/Internal/Shared.hs @@ -2,11 +2,10 @@ where import Control.Exception -import qualified Data.ByteString as BS -import qualified Data.ByteString.Internal as BSI +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BSI import Foreign.C.Error import Kafka.Internal.RdKafka -import Kafka.Internal.RdKafkaEnum import Kafka.Types word8PtrToBS :: Int -> Word8Ptr -> IO BS.ByteString @@ -21,33 +20,33 @@ throwOnError action = do m <- action case m of - Just e -> throw $ KafkaError e + Just e -> throw $ KafkaError e Nothing -> return () hasError :: KafkaError -> Bool hasError err = case err of KafkaResponseError RdKafkaRespErrNoError -> False - _ -> True + _ -> True {-# INLINE hasError #-} +rdKafkaErrorToEither :: RdKafkaRespErrT -> Either KafkaError () +rdKafkaErrorToEither err = case err of + RdKafkaRespErrNoError -> Right () + _ -> Left (KafkaResponseError err) +{-# INLINE rdKafkaErrorToEither #-} + kafkaErrorToEither :: KafkaError -> Either KafkaError () kafkaErrorToEither err = case err of KafkaResponseError RdKafkaRespErrNoError -> Right () - _ -> Left err + _ -> Left err {-# INLINE kafkaErrorToEither #-} kafkaErrorToMaybe :: KafkaError -> Maybe KafkaError kafkaErrorToMaybe err = case err of KafkaResponseError RdKafkaRespErrNoError -> Nothing - _ -> Just err + _ -> Just err {-# INLINE kafkaErrorToMaybe #-} maybeToLeft :: Maybe a -> Either a () maybeToLeft = maybe (Right ()) Left {-# INLINE maybeToLeft #-} - -mapLeft :: (a -> a') -> Either a b -> Either a' b -mapLeft f e = case e of - Left a -> Left (f a) - Right b -> Right b -{-# INLINE mapLeft #-} diff --git a/src/Kafka/Metadata.hs b/src/Kafka/Metadata.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Metadata.hs @@ -0,0 +1,289 @@ +module Kafka.Metadata +( KafkaMetadata(..), BrokerMetadata(..), TopicMetadata(..), PartitionMetadata(..) +, WatermarkOffsets(..) +, GroupMemberId(..), GroupMemberInfo(..) +, GroupProtocolType(..), GroupProtocol(..), GroupState(..) +, GroupInfo(..) +, allTopicsMetadata, topicMetadata +, watermarkOffsets, watermarkOffsets' +, partitionWatermarkOffsets +, offsetsForTime, offsetsForTime', topicOffsetsForTime +, allConsumerGroupsInfo, consumerGroupInfo +) +where + +import Control.Arrow (left) +import Control.Exception (bracket) +import Control.Monad.IO.Class (MonadIO, liftIO) +import Data.Bifunctor +import Data.ByteString (ByteString, pack) +import Data.Monoid ((<>)) +import qualified Data.Set as S +import Foreign +import Foreign.C.String +import Kafka.Consumer.Convert +import Kafka.Consumer.Types +import Kafka.Internal.RdKafka +import Kafka.Internal.Shared +import Kafka.Types + +data KafkaMetadata = KafkaMetadata + { kmBrokers :: [BrokerMetadata] + , kmTopics :: [TopicMetadata] + , kmOrigBroker :: !BrokerId + } deriving (Show, Eq) + +data BrokerMetadata = BrokerMetadata + { bmBrokerId :: !BrokerId + , bmBrokerHost :: !String + , bmBrokerPort :: !Int + } deriving (Show, Eq) + +data PartitionMetadata = PartitionMetadata + { pmPartitionId :: !PartitionId + , pmError :: Maybe KafkaError + , pmLeader :: !BrokerId + , pmReplicas :: [PartitionId] + , pmInSyncReplicas :: [PartitionId] + } deriving (Show, Eq) + +data TopicMetadata = TopicMetadata + { tmTopicName :: !TopicName + , tmPartitions :: [PartitionMetadata] + , tmError :: Maybe KafkaError + } deriving (Show, Eq) + +data WatermarkOffsets = WatermarkOffsets + { woTopicName :: !TopicName + , woPartitionId :: !PartitionId + , woLowWatermark :: !Offset + , woHighWatermark :: !Offset + } deriving (Show, Eq) + +newtype GroupMemberId = GroupMemberId String deriving (Show, Eq, Read, Ord) +data GroupMemberInfo = GroupMemberInfo + { gmiMemberId :: !GroupMemberId + , gmiClientId :: !ClientId + , gmiClientHost :: !String + , gmiMetadata :: !ByteString + , gmiAssignment :: !ByteString + } deriving (Show, Eq) + +newtype GroupProtocolType = GroupProtocolType String deriving (Show, Eq, Read, Ord) +newtype GroupProtocol = GroupProtocol String deriving (Show, Eq, Read, Ord) +data GroupState + = GroupPreparingRebalance -- ^ Group is preparing to rebalance + | GroupEmpty -- ^ Group has no more members, but lingers until all offsets have expired + | GroupAwaitingSync -- ^ Group is awaiting state assignment from the leader + | GroupStable -- ^ Group is stable + | GroupDead -- ^ Group has no more members and its metadata is being removed + deriving (Show, Eq, Read, Ord) + +data GroupInfo = GroupInfo + { giGroup :: !ConsumerGroupId + , giError :: Maybe KafkaError + , giState :: !GroupState + , giProtocolType :: !GroupProtocolType + , giProtocol :: !GroupProtocol + , giMembers :: [GroupMemberInfo] + } deriving (Show, Eq) + +-- | Returns metadata for all topics in the cluster +allTopicsMetadata :: (MonadIO m, HasKafka k) => k -> Timeout -> m (Either KafkaError KafkaMetadata) +allTopicsMetadata k (Timeout timeout) = liftIO $ do + meta <- rdKafkaMetadata (getKafkaPtr k) True Nothing timeout + traverse fromKafkaMetadataPtr (left KafkaResponseError meta) + +-- | Returns metadata only for specified topic +topicMetadata :: (MonadIO m, HasKafka k) => k -> Timeout -> TopicName -> m (Either KafkaError KafkaMetadata) +topicMetadata k (Timeout timeout) (TopicName tn) = liftIO $ + bracket mkTopic clTopic $ \mbt -> case mbt of + Left err -> return (Left $ KafkaError err) + Right t -> do + meta <- rdKafkaMetadata (getKafkaPtr k) False (Just t) timeout + traverse fromKafkaMetadataPtr (left KafkaResponseError meta) + where + mkTopic = newUnmanagedRdKafkaTopicT (getKafkaPtr k) tn Nothing + clTopic = either (return . const ()) destroyUnmanagedRdKafkaTopic + +-- | Query broker for low (oldest/beginning) and high (newest/end) offsets for a given topic. +watermarkOffsets :: (MonadIO m, HasKafka k) => k -> Timeout -> TopicName -> m [Either KafkaError WatermarkOffsets] +watermarkOffsets k timeout t = do + meta <- topicMetadata k timeout t + case meta of + Left err -> return [Left err] + Right tm -> if null (kmTopics tm) + then return [] + else watermarkOffsets' k timeout (head $ kmTopics tm) + +-- | Query broker for low (oldest/beginning) and high (newest/end) offsets for a given topic. +watermarkOffsets' :: (MonadIO m, HasKafka k) => k -> Timeout -> TopicMetadata -> m [Either KafkaError WatermarkOffsets] +watermarkOffsets' k timeout tm = + let pids = pmPartitionId <$> tmPartitions tm + in liftIO $ traverse (partitionWatermarkOffsets k timeout (tmTopicName tm)) pids + +-- | Query broker for low (oldest/beginning) and high (newest/end) offsets for a specific partition +partitionWatermarkOffsets :: (MonadIO m, HasKafka k) => k -> Timeout -> TopicName -> PartitionId -> m (Either KafkaError WatermarkOffsets) +partitionWatermarkOffsets k (Timeout timeout) (TopicName t) (PartitionId p) = liftIO $ do + offs <- rdKafkaQueryWatermarkOffsets (getKafkaPtr k) t p timeout + return $ bimap KafkaResponseError toWatermark offs + where + toWatermark (l, h) = WatermarkOffsets (TopicName t) (PartitionId p) (Offset l) (Offset h) + +-- | Look up the offsets for the given topic by timestamp. +-- +-- The returned offset for each partition is the earliest offset whose +-- timestamp is greater than or equal to the given timestamp in the +-- corresponding partition. +topicOffsetsForTime :: (MonadIO m, HasKafka k) => k -> Timeout -> Millis -> TopicName -> m (Either KafkaError [TopicPartition]) +topicOffsetsForTime k timeout timestamp topic = do + meta <- topicMetadata k timeout topic + case meta of + Left err -> return $ Left err + Right meta' -> + let tps = [(tmTopicName t, pmPartitionId p)| t <- kmTopics meta', p <- tmPartitions t] + in offsetsForTime k timeout timestamp tps + +-- | Look up the offsets for the given metadata by timestamp. +-- +-- The returned offset for each partition is the earliest offset whose +-- timestamp is greater than or equal to the given timestamp in the +-- corresponding partition. +offsetsForTime' :: (MonadIO m, HasKafka k) => k -> Timeout -> Millis -> TopicMetadata -> m (Either KafkaError [TopicPartition]) +offsetsForTime' k timeout timestamp t = + let tps = [(tmTopicName t, pmPartitionId p) | p <- tmPartitions t] + in offsetsForTime k timeout timestamp tps + +-- | Look up the offsets for the given partitions by timestamp. +-- +-- The returned offset for each partition is the earliest offset whose +-- timestamp is greater than or equal to the given timestamp in the +-- corresponding partition. +offsetsForTime :: (MonadIO m, HasKafka k) => k -> Timeout -> Millis -> [(TopicName, PartitionId)] -> m (Either KafkaError [TopicPartition]) +offsetsForTime k (Timeout timeout) (Millis t) tps = liftIO $ do + ntps <- toNativeTopicPartitionList $ mkTopicPartition <$> uniqueTps + res <- rdKafkaOffsetsForTimes (getKafkaPtr k) ntps timeout + case res of + RdKafkaRespErrNoError -> Right <$> fromNativeTopicPartitionList'' ntps + err -> return $ Left (KafkaResponseError err) + where + uniqueTps = S.toList . S.fromList $ tps + -- rd_kafka_offsets_for_times reuses `offset` to specify timestamp :( + mkTopicPartition (tn, p) = TopicPartition tn p (PartitionOffset t) + +-- | List and describe all consumer groups in cluster. +allConsumerGroupsInfo :: (MonadIO m, HasKafka k) => k -> Timeout -> m (Either KafkaError [GroupInfo]) +allConsumerGroupsInfo k (Timeout t) = liftIO $ do + res <- rdKafkaListGroups (getKafkaPtr k) Nothing t + traverse fromGroupInfoListPtr (left KafkaResponseError res) + +-- | Describe a given consumer group. +consumerGroupInfo :: (MonadIO m, HasKafka k) => k -> Timeout -> ConsumerGroupId -> m (Either KafkaError [GroupInfo]) +consumerGroupInfo k (Timeout timeout) (ConsumerGroupId gn) = liftIO $ do + res <- rdKafkaListGroups (getKafkaPtr k) (Just gn) timeout + traverse fromGroupInfoListPtr (left KafkaResponseError res) + +------------------------------------------------------------------------------- +getKafkaPtr :: HasKafka k => k -> RdKafkaTPtr +getKafkaPtr k = let (Kafka k') = getKafka k in k' +{-# INLINE getKafkaPtr #-} + + +fromGroupInfoListPtr :: RdKafkaGroupListTPtr -> IO [GroupInfo] +fromGroupInfoListPtr ptr = + withForeignPtr ptr $ \realPtr -> do + gl <- peek realPtr + pgis <- peekArray (groupCnt'RdKafkaGroupListT gl) (groups'RdKafkaGroupListT gl) + traverse fromGroupInfoPtr pgis + +fromGroupInfoPtr :: RdKafkaGroupInfoT -> IO GroupInfo +fromGroupInfoPtr gi = do + --bmd <- peek (broker'RdKafkaGroupInfoT gi) -- >>= fromBrokerMetadataPtr + --xxx <- fromBrokerMetadataPtr bmd + cid <- peekCAString $ group'RdKafkaGroupInfoT gi + stt <- peekCAString $ state'RdKafkaGroupInfoT gi + prt <- peekCAString $ protocolType'RdKafkaGroupInfoT gi + pr <- peekCAString $ protocol'RdKafkaGroupInfoT gi + mbs <- peekArray (memberCnt'RdKafkaGroupInfoT gi) (members'RdKafkaGroupInfoT gi) + mbl <- mapM fromGroupMemberInfoPtr mbs + return GroupInfo + { --giBroker = bmd + giGroup = ConsumerGroupId cid + , giError = kafkaErrorToMaybe $ KafkaResponseError (err'RdKafkaGroupInfoT gi) + , giState = groupStateFromKafkaString stt + , giProtocolType = GroupProtocolType prt + , giProtocol = GroupProtocol pr + , giMembers = mbl + } + +fromGroupMemberInfoPtr :: RdKafkaGroupMemberInfoT -> IO GroupMemberInfo +fromGroupMemberInfoPtr mi = do + mid <- peekCAString $ memberId'RdKafkaGroupMemberInfoT mi + cid <- peekCAString $ clientId'RdKafkaGroupMemberInfoT mi + hst <- peekCAString $ clientHost'RdKafkaGroupMemberInfoT mi + mtd <- peekArray (memberMetadataSize'RdKafkaGroupMemberInfoT mi) (memberMetadata'RdKafkaGroupMemberInfoT mi) + ass <- peekArray (memberAssignmentSize'RdKafkaGroupMemberInfoT mi) (memberAssignment'RdKafkaGroupMemberInfoT mi) + return GroupMemberInfo + { gmiMemberId = GroupMemberId mid + , gmiClientId = ClientId cid + , gmiClientHost = hst + , gmiMetadata = pack mtd + , gmiAssignment = pack ass + } + +fromTopicMetadataPtr :: RdKafkaMetadataTopicT -> IO TopicMetadata +fromTopicMetadataPtr tm = do + tnm <- peekCAString (topic'RdKafkaMetadataTopicT tm) + pts <- peekArray (partitionCnt'RdKafkaMetadataTopicT tm) (partitions'RdKafkaMetadataTopicT tm) + pms <- mapM fromPartitionMetadataPtr pts + return TopicMetadata + { tmTopicName = TopicName tnm + , tmPartitions = pms + , tmError = kafkaErrorToMaybe $ KafkaResponseError (err'RdKafkaMetadataTopicT tm) + } + +fromPartitionMetadataPtr :: RdKafkaMetadataPartitionT -> IO PartitionMetadata +fromPartitionMetadataPtr pm = do + reps <- peekArray (replicaCnt'RdKafkaMetadataPartitionT pm) (replicas'RdKafkaMetadataPartitionT pm) + isrs <- peekArray (isrCnt'RdKafkaMetadataPartitionT pm) (isrs'RdKafkaMetadataPartitionT pm) + return PartitionMetadata + { pmPartitionId = PartitionId (id'RdKafkaMetadataPartitionT pm) + , pmError = kafkaErrorToMaybe $ KafkaResponseError (err'RdKafkaMetadataPartitionT pm) + , pmLeader = BrokerId (leader'RdKafkaMetadataPartitionT pm) + , pmReplicas = (PartitionId . fromIntegral) <$> reps + , pmInSyncReplicas = (PartitionId . fromIntegral) <$> isrs + } + + +fromBrokerMetadataPtr :: RdKafkaMetadataBrokerT -> IO BrokerMetadata +fromBrokerMetadataPtr bm = do + host <- peekCAString (host'RdKafkaMetadataBrokerT bm) + return BrokerMetadata + { bmBrokerId = BrokerId (id'RdKafkaMetadataBrokerT bm) + , bmBrokerHost = host + , bmBrokerPort = port'RdKafkaMetadataBrokerT bm + } + + +fromKafkaMetadataPtr :: RdKafkaMetadataTPtr -> IO KafkaMetadata +fromKafkaMetadataPtr ptr = + withForeignPtr ptr $ \realPtr -> do + km <- peek realPtr + pbms <- peekArray (brokerCnt'RdKafkaMetadataT km) (brokers'RdKafkaMetadataT km) + bms <- mapM fromBrokerMetadataPtr pbms + ptms <- peekArray (topicCnt'RdKafkaMetadataT km) (topics'RdKafkaMetadataT km) + tms <- mapM fromTopicMetadataPtr ptms + return KafkaMetadata + { kmBrokers = bms + , kmTopics = tms + , kmOrigBroker = BrokerId $ fromIntegral (origBrokerId'RdKafkaMetadataT km) + } + +groupStateFromKafkaString :: String -> GroupState +groupStateFromKafkaString s = case s of + "PreparingRebalance" -> GroupPreparingRebalance + "AwaitingSync" -> GroupAwaitingSync + "Stable" -> GroupStable + "Dead" -> GroupDead + "Empty" -> GroupEmpty + _ -> error $ "Unknown group state: " <> s diff --git a/src/Kafka/Producer.hs b/src/Kafka/Producer.hs --- a/src/Kafka/Producer.hs +++ b/src/Kafka/Producer.hs @@ -6,31 +6,28 @@ , produceMessage, produceMessageBatch , flushProducer , closeProducer -, RDE.RdKafkaRespErrT (..) +, RdKafkaRespErrT (..) ) where -import Control.Arrow ((&&&)) +import Control.Arrow ((&&&)) import Control.Exception import Control.Monad import Control.Monad.IO.Class -import qualified Data.Map as M -import qualified Data.ByteString as BS -import qualified Data.ByteString.Internal as BSI -import Foreign hiding (void) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BSI +import Data.Function (on) +import Data.List (groupBy, sortBy) +import qualified Data.Map as M +import Data.Ord (comparing) +import Foreign hiding (void) import Kafka.Internal.RdKafka -import Kafka.Internal.RdKafkaEnum import Kafka.Internal.Setup import Kafka.Producer.Convert -import Data.Function (on) -import Data.List (sortBy, groupBy) -import Data.Ord (comparing) -import qualified Kafka.Internal.RdKafkaEnum as RDE - -import Kafka.Types as X -import Kafka.Producer.Types as X import Kafka.Producer.ProducerProperties as X +import Kafka.Producer.Types as X +import Kafka.Types as X -- | Runs Kafka Producer. -- The callback provided is expected to call 'produceMessage' @@ -43,24 +40,28 @@ where mkProducer = newProducer props - clProducer (Left _) = return () + clProducer (Left _) = return () clProducer (Right prod) = closeProducer prod - runHandler (Left err) = return $ Left err + runHandler (Left err) = return $ Left err runHandler (Right prod) = f prod -- | Creates a new kafka producer -- A newly created producer must be closed with 'closeProducer' function. newProducer :: MonadIO m => ProducerProperties -> m (Either KafkaError KafkaProducer) -newProducer (ProducerProperties kp tp ll) = liftIO $ do - (KafkaConf kc) <- kafkaConf (KafkaProps $ M.toList kp) - (TopicConf tc) <- topicConf (TopicProps $ M.toList tp) - mbKafka <- newRdKafkaT RdKafkaProducer kc +newProducer (ProducerProperties kp tp ll cbs) = liftIO $ do + kc@(KafkaConf kc') <- kafkaConf (KafkaProps $ M.toList kp) + tc <- topicConf (TopicProps $ M.toList tp) + + -- set callbacks + forM_ cbs (\setCb -> setCb kc) + + mbKafka <- newRdKafkaT RdKafkaProducer kc' case mbKafka of Left err -> return . Left $ KafkaError err Right kafka -> do forM_ ll (rdKafkaSetLogLevel kafka . fromEnum) - return .Right $ KafkaProducer kafka kc tc + return .Right $ KafkaProducer (Kafka kafka) kc tc -- | Sends a single message. -- Since librdkafka is backed by a queue, this function can return before messages are sent. See @@ -69,10 +70,10 @@ => KafkaProducer -> ProducerRecord -> m (Maybe KafkaError) -produceMessage (KafkaProducer k _ tc) m = liftIO $ +produceMessage (KafkaProducer (Kafka k) _ (TopicConf tc)) m = liftIO $ bracket (mkTopic $ prTopic m) clTopic withTopic where - mkTopic (TopicName tn) = newUnmanagedRdKafkaTopicT k tn tc + mkTopic (TopicName tn) = newUnmanagedRdKafkaTopicT k tn (Just tc) clTopic = either (return . const ()) destroyUnmanagedRdKafkaTopic @@ -95,18 +96,18 @@ -> [ProducerRecord] -> m [(ProducerRecord, KafkaError)] -- ^ An empty list when the operation is successful, - -- otherwise a list of "failed" messages with corresponsing errors. -produceMessageBatch (KafkaProducer k _ tc) messages = liftIO $ + -- otherwise a list of "failed" messages with corresponsing errors. +produceMessageBatch (KafkaProducer (Kafka k) _ (TopicConf tc)) messages = liftIO $ concat <$> forM (mkBatches messages) sendBatch where mkSortKey = prTopic &&& prPartition mkBatches = groupBy ((==) `on` mkSortKey) . sortBy (comparing mkSortKey) - mkTopic (TopicName tn) = newUnmanagedRdKafkaTopicT k tn tc + mkTopic (TopicName tn) = newUnmanagedRdKafkaTopicT k tn (Just tc) clTopic = either (return . const ()) destroyUnmanagedRdKafkaTopic - sendBatch [] = return [] + sendBatch [] = return [] sendBatch batch = bracket (mkTopic $ prTopic (head batch)) clTopic (withTopic batch) withTopic ms (Left err) = return $ (, KafkaError err) <$> ms @@ -146,7 +147,7 @@ -- This function is also called automatically when the producer is closed -- with 'closeProducer' to ensure that all queued messages make it to Kafka. flushProducer :: MonadIO m => KafkaProducer -> m () -flushProducer kp@(KafkaProducer k _ _) = liftIO $ do +flushProducer kp@(KafkaProducer (Kafka k) _ _) = liftIO $ do pollEvents k 100 l <- outboundQueueLength k unless (l == 0) $ flushProducer kp diff --git a/src/Kafka/Producer/ProducerProperties.hs b/src/Kafka/Producer/ProducerProperties.hs --- a/src/Kafka/Producer/ProducerProperties.hs +++ b/src/Kafka/Producer/ProducerProperties.hs @@ -1,64 +1,72 @@ module Kafka.Producer.ProducerProperties +( module Kafka.Producer.ProducerProperties +, module Kafka.Callbacks +) where -import Control.Monad -import Data.Map (Map) -import Kafka.Types -import qualified Data.Map as M -import qualified Data.List as L +import Control.Monad +import qualified Data.List as L +import Data.Map (Map) +import qualified Data.Map as M +import Kafka.Callbacks +import Kafka.Types -- | Properties to create 'KafkaProducer'. data ProducerProperties = ProducerProperties { ppKafkaProps :: Map String String , ppTopicProps :: Map String String , ppLogLevel :: Maybe KafkaLogLevel - } deriving (Show) + , callbacks :: [KafkaConf -> IO ()] + } instance Monoid ProducerProperties where - mempty = ProducerProperties M.empty M.empty Nothing - mappend (ProducerProperties k1 t1 ll1) (ProducerProperties k2 t2 ll2) = - ProducerProperties (M.union k1 k2) (M.union t1 t2) (ll2 `mplus` ll1) + mempty = ProducerProperties M.empty M.empty Nothing [] + mappend (ProducerProperties k1 t1 ll1 cb1) (ProducerProperties k2 t2 ll2 cb2) = + ProducerProperties (M.union k1 k2) (M.union t1 t2) (ll2 `mplus` ll1) (cb2 `mplus` cb1) -producerBrokersList :: [BrokerAddress] -> ProducerProperties -producerBrokersList bs = +brokersList :: [BrokerAddress] -> ProducerProperties +brokersList bs = let bs' = L.intercalate "," ((\(BrokerAddress x) -> x) <$> bs) - in extraProducerProps $ M.fromList [("bootstrap.servers", bs')] + in extraProps $ M.fromList [("bootstrap.servers", bs')] +setCallback :: (KafkaConf -> IO ()) -> ProducerProperties +setCallback cb = ProducerProperties M.empty M.empty Nothing [cb] + -- | Sets the logging level. --- Usually is used with 'producerDebug' to configure which logs are needed. -producerLogLevel :: KafkaLogLevel -> ProducerProperties -producerLogLevel ll = ProducerProperties M.empty M.empty (Just ll) +-- Usually is used with 'debugOptions' to configure which logs are needed. +logLevel :: KafkaLogLevel -> ProducerProperties +logLevel ll = ProducerProperties M.empty M.empty (Just ll) [] -producerCompression :: KafkaCompressionCodec -> ProducerProperties -producerCompression c = - extraProducerProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) +compression :: KafkaCompressionCodec -> ProducerProperties +compression c = + extraProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) -producerTopicCompression :: KafkaCompressionCodec -> ProducerProperties -producerTopicCompression c = - extraProducerTopicProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) +topicCompression :: KafkaCompressionCodec -> ProducerProperties +topicCompression c = + extraTopicProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) -- | Any configuration options that are supported by /librdkafka/. -- The full list can be found -extraProducerProps :: Map String String -> ProducerProperties -extraProducerProps m = ProducerProperties m M.empty Nothing +extraProps :: Map String String -> ProducerProperties +extraProps m = ProducerProperties m M.empty Nothing [] -- | Suppresses producer disconnects logs. -- -- It might be useful to turn this off when interacting with brokers -- with an aggressive connection.max.idle.ms value. -producerSuppressDisconnectLogs :: ProducerProperties -producerSuppressDisconnectLogs = - extraProducerProps $ M.fromList [("log.connection.close", "false")] +suppressDisconnectLogs :: ProducerProperties +suppressDisconnectLogs = + extraProps $ M.fromList [("log.connection.close", "false")] -- | Any configuration options that are supported by /librdkafka/. -- The full list can be found -extraProducerTopicProps :: Map String String -> ProducerProperties -extraProducerTopicProps m = ProducerProperties M.empty m Nothing +extraTopicProps :: Map String String -> ProducerProperties +extraTopicProps m = ProducerProperties M.empty m Nothing [] -- | Sets debug features for the producer --- Usually is used with 'producerLogLevel'. -producerDebug :: [KafkaDebug] -> ProducerProperties -producerDebug [] = extraProducerProps M.empty -producerDebug d = +-- Usually is used with 'logLevel'. +debugOptions :: [KafkaDebug] -> ProducerProperties +debugOptions [] = extraProps M.empty +debugOptions d = let points = L.intercalate "," (kafkaDebugToString <$> d) - in extraProducerProps $ M.fromList [("debug", points)] + in extraProps $ M.fromList [("debug", points)] diff --git a/src/Kafka/Producer/Types.hs b/src/Kafka/Producer/Types.hs --- a/src/Kafka/Producer/Types.hs +++ b/src/Kafka/Producer/Types.hs @@ -6,14 +6,21 @@ import qualified Data.ByteString as BS import Data.Typeable import Kafka.Types -import Kafka.Internal.RdKafka -- | Main pointer to Kafka object, which contains our brokers data KafkaProducer = KafkaProducer - { kpKafkaPtr :: RdKafkaTPtr - , kpKafkaConf :: RdKafkaConfTPtr - , kpTopicConf :: RdKafkaTopicConfTPtr + { kpKafkaPtr :: !Kafka + , kpKafkaConf :: !KafkaConf + , kpTopicConf :: !TopicConf } deriving (Show) + +instance HasKafka KafkaProducer where + getKafka = kpKafkaPtr + {-# INLINE getKafka #-} + +instance HasKafkaConf KafkaProducer where + getKafkaConf = kpKafkaConf + {-# INLINE getKafkaConf #-} -- | Represents messages /to be enqueued/ onto a Kafka broker (i.e. used for a producer) data ProducerRecord = ProducerRecord diff --git a/src/Kafka/Types.hs b/src/Kafka/Types.hs --- a/src/Kafka/Types.hs +++ b/src/Kafka/Types.hs @@ -1,11 +1,39 @@ -{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} module Kafka.Types where import Control.Exception +import Data.Int import Data.Typeable -import Kafka.Internal.RdKafkaEnum +import Kafka.Internal.RdKafka +class HasKafka a where + getKafka :: a -> Kafka + +class HasKafkaConf a where + getKafkaConf :: a -> KafkaConf + +newtype BrokerId = + BrokerId Int + deriving (Show, Eq, Ord, Read) + +newtype Kafka = Kafka RdKafkaTPtr deriving Show +newtype KafkaConf = KafkaConf RdKafkaConfTPtr deriving Show +newtype TopicConf = TopicConf RdKafkaTopicConfTPtr deriving Show + +newtype PartitionId = PartitionId Int deriving (Show, Eq, Read, Ord) +newtype Millis = Millis Int64 deriving (Show, Read, Eq, Ord, Num) +newtype ClientId = ClientId String deriving (Show, Eq, Ord) + +instance HasKafkaConf KafkaConf where + getKafkaConf = id + {-# INLINE getKafkaConf #-} + +instance HasKafka Kafka where + getKafka = id + {-# INLINE getKafka #-} + -- | Topic name to be consumed -- -- Wildcard (regex) topics are supported by the /librdkafka/ assignor: @@ -39,14 +67,14 @@ toEnum 7 = KafkaLogDebug toEnum _ = undefined - fromEnum KafkaLogEmerg = 0 - fromEnum KafkaLogAlert = 1 - fromEnum KafkaLogCrit = 2 - fromEnum KafkaLogErr = 3 + fromEnum KafkaLogEmerg = 0 + fromEnum KafkaLogAlert = 1 + fromEnum KafkaLogCrit = 2 + fromEnum KafkaLogErr = 3 fromEnum KafkaLogWarning = 4 - fromEnum KafkaLogNotice = 5 - fromEnum KafkaLogInfo = 6 - fromEnum KafkaLogDebug = 7 + fromEnum KafkaLogNotice = 5 + fromEnum KafkaLogInfo = 6 + fromEnum KafkaLogDebug = 7 -- -- | Any Kafka errors @@ -79,18 +107,18 @@ kafkaDebugToString :: KafkaDebug -> String kafkaDebugToString d =case d of - DebugGeneric -> "generic" - DebugBroker -> "broker" - DebugTopic -> "topic" - DebugMetadata -> "metadata" - DebugQueue -> "queue" - DebugMsg -> "msg" - DebugProtocol -> "protocol" - DebugCgrp -> "cgrp" - DebugSecurity -> "security" - DebugFetch -> "fetch" - DebugFeature -> "feature" - DebugAll -> "all" + DebugGeneric -> "generic" + DebugBroker -> "broker" + DebugTopic -> "topic" + DebugMetadata -> "metadata" + DebugQueue -> "queue" + DebugMsg -> "msg" + DebugProtocol -> "protocol" + DebugCgrp -> "cgrp" + DebugSecurity -> "security" + DebugFetch -> "fetch" + DebugFeature -> "feature" + DebugAll -> "all" data KafkaCompressionCodec = NoCompression diff --git a/tests-it/Kafka/IntegrationSpec.hs b/tests-it/Kafka/IntegrationSpec.hs new file mode 100644 --- /dev/null +++ b/tests-it/Kafka/IntegrationSpec.hs @@ -0,0 +1,133 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Kafka.IntegrationSpec +where + +import Control.Monad (forM_) +import Control.Monad.Loops +import qualified Data.ByteString as BS +import Data.Either +import Data.Map + +import Kafka.Consumer as C +import Kafka.Metadata as M +import Kafka.Producer as P + +import Kafka.TestEnv +import Test.Hspec + +{-# ANN module ("HLint: ignore Redundant do" :: String) #-} + +spec :: Spec +spec = do + describe "Kafka.IntegrationSpec" $ do + specWithProducer "Run producer" $ do + it "sends messages to test topic" $ \prod -> do + res <- sendMessages (testMessages testTopic) prod + res `shouldBe` Right () + + specWithConsumer "Run consumer" $ do + it "should receive messages" $ \k -> do + res <- receiveMessages k + length <$> res `shouldBe` Right 2 + + let timestamps = crTimestamp <$> either (const []) id res + forM_ timestamps $ \ts -> + ts `shouldNotBe` NoTimestamp + + it "should get watermark offsets" $ \k -> do + res <- sequence <$> watermarkOffsets k (Timeout 1000) testTopic + res `shouldSatisfy` isRight + length <$> res `shouldBe` (Right 1) + + it "should return subscription" $ \k -> do + res <- subscription k + res `shouldSatisfy` isRight + length <$> res `shouldBe` Right 1 + + it "should return assignment" $ \k -> do + res <- assignment k + res `shouldSatisfy` isRight + res `shouldBe` Right (fromList [(testTopic, [PartitionId 0])]) + + it "should return all topics metadata" $ \k -> do + res <- allTopicsMetadata k (Timeout 1000) + res `shouldSatisfy` isRight + (length . kmBrokers) <$> res `shouldBe` Right 1 + (length . kmTopics) <$> res `shouldBe` Right 2 + + it "should return topic metadata" $ \k -> do + res <- topicMetadata k (Timeout 1000) testTopic + res `shouldSatisfy` isRight + (length . kmBrokers) <$> res `shouldBe` Right 1 + (length . kmTopics) <$> res `shouldBe` Right 1 + + it "should describe all consumer groups" $ \k -> do + res <- allConsumerGroupsInfo k (Timeout 1000) + fmap giGroup <$> res `shouldBe` Right [testGroupId] + + it "should describe a given consumer group" $ \k -> do + res <- consumerGroupInfo k (Timeout 1000) testGroupId + fmap giGroup <$> res `shouldBe` Right [testGroupId] + + it "should describe non-existent consumer group" $ \k -> do + res <- consumerGroupInfo k (Timeout 1000) (ConsumerGroupId "does-not-exist") + res `shouldBe` Right [] + + it "should read topic offsets for time" $ \k -> do + res <- topicOffsetsForTime k (Timeout 1000) (Millis 1904057189508) testTopic + res `shouldSatisfy` isRight + fmap tpOffset <$> res `shouldBe` Right [PartitionOffsetEnd] + + it "should seek and return no error" $ \k -> do + res <- seek k (Timeout 1000) [TopicPartition testTopic (PartitionId 0) (PartitionOffset 1)] + res `shouldBe` Nothing + msg <- pollMessage k (Timeout 1000) + crOffset <$> msg `shouldBe` Right (Offset 1) + + it "should seek to the beginning" $ \k -> do + res <- seek k (Timeout 1000) [TopicPartition testTopic (PartitionId 0) PartitionOffsetBeginning] + res `shouldBe` Nothing + msg <- pollMessage k (Timeout 1000) + crOffset <$> msg `shouldBe` Right (Offset 0) + + it "should seek to the end" $ \k -> do + res <- seek k (Timeout 1000) [TopicPartition testTopic (PartitionId 0) PartitionOffsetEnd] + res `shouldBe` Nothing + msg <- pollMessage k (Timeout 1000) + crOffset <$> msg `shouldBe` Left (KafkaResponseError RdKafkaRespErrPartitionEof) + + it "should respect out-of-bound offsets (invalid offset)" $ \k -> do + res <- seek k (Timeout 1000) [TopicPartition testTopic (PartitionId 0) PartitionOffsetInvalid] + res `shouldBe` Nothing + msg <- pollMessage k (Timeout 1000) + crOffset <$> msg `shouldBe` Right (Offset 0) + + it "should respect out-of-bound offsets (huge offset)" $ \k -> do + res <- seek k (Timeout 1000) [TopicPartition testTopic (PartitionId 0) (PartitionOffset 123456)] + res `shouldBe` Nothing + msg <- pollMessage k (Timeout 1000) + crOffset <$> msg `shouldBe` Right (Offset 0) + + +---------------------------------------------------------------------------------------------------------------- + +receiveMessages :: KafkaConsumer -> IO (Either KafkaError [ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)]) +receiveMessages kafka = + (Right . rights) <$> www + where + www = whileJust maybeMsg return + isOK msg = if msg /= Left (KafkaResponseError RdKafkaRespErrPartitionEof) then Just msg else Nothing + maybeMsg = isOK <$> get + get = pollMessage kafka (Timeout 1000) + +testMessages :: TopicName -> [ProducerRecord] +testMessages t = + [ ProducerRecord t UnassignedPartition Nothing (Just "test from producer") + , ProducerRecord t UnassignedPartition (Just "key") (Just "test from producer (with key)") + ] + +sendMessages :: [ProducerRecord] -> KafkaProducer -> IO (Either KafkaError ()) +sendMessages msgs prod = + Right <$> forM_ msgs (produceMessage prod) diff --git a/tests-it/Kafka/TestEnv.hs b/tests-it/Kafka/TestEnv.hs new file mode 100644 --- /dev/null +++ b/tests-it/Kafka/TestEnv.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Kafka.TestEnv +where + +import Control.Exception +import Control.Monad (void) +import Data.Monoid ((<>)) +import System.Environment +import System.IO.Unsafe + +import Kafka.Consumer as C +import Kafka.Metadata as M +import Kafka.Producer as P + +import Test.Hspec + +brokerAddress :: BrokerAddress +brokerAddress = unsafePerformIO $ + BrokerAddress <$> getEnv "KAFKA_TEST_BROKER" `catch` \(_ :: SomeException) -> (return "localhost:9092") +{-# NOINLINE brokerAddress #-} + +testTopic :: TopicName +testTopic = unsafePerformIO $ + TopicName <$> getEnv "KAFKA_TEST_TOPIC" `catch` \(_ :: SomeException) -> (return "kafka-client_tests") +{-# NOINLINE testTopic #-} + +testGroupId :: ConsumerGroupId +testGroupId = ConsumerGroupId "it_spec_03" + +consumerProps :: BrokerAddress -> ConsumerProperties +consumerProps broker = C.brokersList [broker] + <> groupId testGroupId + <> C.setCallback (logCallback (\l s1 s2 -> print $ show l <> ": " <> s1 <> ", " <> s2)) + <> C.setCallback (errorCallback (\e r -> print $ show e <> ": " <> r)) + <> noAutoCommit + +producerProps :: BrokerAddress -> ProducerProperties +producerProps broker = P.brokersList [broker] + <> P.setCallback (logCallback (\l s1 s2 -> print $ show l <> ": " <> s1 <> ", " <> s2)) + <> P.setCallback (errorCallback (\e r -> print $ show e <> ": " <> r)) + +testSubscription :: TopicName -> Subscription +testSubscription t = topics [t] + <> offsetReset Earliest + +mkProducer :: IO KafkaProducer +mkProducer = do + (Right p) <- newProducer (producerProps brokerAddress) + return p + +mkConsumer :: IO KafkaConsumer +mkConsumer = do + (Right c) <- newConsumer (consumerProps brokerAddress) (testSubscription testTopic) + return c + +specWithConsumer :: String -> SpecWith KafkaConsumer -> Spec +specWithConsumer s f = beforeAll mkConsumer $ afterAll (void . closeConsumer) $ describe s f + +specWithProducer :: String -> SpecWith KafkaProducer -> Spec +specWithProducer s f = beforeAll mkProducer $ afterAll (void . closeProducer) $ describe s f diff --git a/tests-it/Spec.hs b/tests-it/Spec.hs new file mode 100644 --- /dev/null +++ b/tests-it/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-} diff --git a/tests/Kafka/Consumer/ConsumerRecordMapSpec.hs b/tests/Kafka/Consumer/ConsumerRecordMapSpec.hs --- a/tests/Kafka/Consumer/ConsumerRecordMapSpec.hs +++ b/tests/Kafka/Consumer/ConsumerRecordMapSpec.hs @@ -4,8 +4,8 @@ ) where import Data.Bitraversable -import Kafka.Types import Kafka.Consumer.Types +import Kafka.Types import Test.Hspec testKey, testValue :: String @@ -17,6 +17,7 @@ { crTopic = TopicName "some-topic" , crPartition = PartitionId 0 , crOffset = Offset 5 + , crTimestamp = NoTimestamp , crKey = Just testKey , crValue = Just testValue } diff --git a/tests/Kafka/Consumer/ConsumerRecordTraverseSpec.hs b/tests/Kafka/Consumer/ConsumerRecordTraverseSpec.hs --- a/tests/Kafka/Consumer/ConsumerRecordTraverseSpec.hs +++ b/tests/Kafka/Consumer/ConsumerRecordTraverseSpec.hs @@ -3,10 +3,10 @@ ( spec ) where -import Data.Bitraversable import Data.Bifunctor -import Kafka.Types +import Data.Bitraversable import Kafka.Consumer.Types +import Kafka.Types import Test.Hspec testKey, testValue :: String @@ -18,6 +18,7 @@ { crTopic = TopicName "some-topic" , crPartition = PartitionId 0 , crOffset = Offset 5 + , crTimestamp = NoTimestamp , crKey = testKey , crValue = testValue } diff --git a/tests/Kafka/IntegrationSpec.hs b/tests/Kafka/IntegrationSpec.hs deleted file mode 100644 --- a/tests/Kafka/IntegrationSpec.hs +++ /dev/null @@ -1,80 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} - -module Kafka.IntegrationSpec -( spec -) where - -import Control.Exception -import Control.Monad (forM_) -import Control.Monad.Loops -import qualified Data.ByteString as BS -import Data.Either -import Data.Monoid ((<>)) -import System.Environment - -import Kafka - -import Test.Hspec - -brokerAddress :: IO BrokerAddress -brokerAddress = BrokerAddress <$> getEnv "KAFKA_TEST_BROKER" `catch` \(_ :: SomeException) -> (return "localhost:9092") - -testTopic :: IO TopicName -testTopic = TopicName <$> getEnv "KAFKA_TEST_TOPIC" `catch` \(_ :: SomeException) -> (return "kafka-client_tests") - -consumerProps :: BrokerAddress -> ConsumerProperties -consumerProps broker = consumerBrokersList [broker] - <> groupId (ConsumerGroupId "it_spec_02") - <> noAutoCommit - -producerProps :: BrokerAddress -> ProducerProperties -producerProps broker = producerBrokersList [broker] - -subscription :: TopicName -> Subscription -subscription t = topics [t] - <> offsetReset Earliest - -spec :: Spec -spec = describe "Kafka.IntegrationSpec" $ do - it "sends messages to test topic" $ do - broker <- brokerAddress - topic <- testTopic - let msgs = testMessages topic - res <- runProducer (producerProps broker) (sendMessages msgs) - res `shouldBe` Right () - - it "consumes messages from test topic" $ do - broker <- brokerAddress - topic <- testTopic - res <- runConsumer - (consumerProps broker) - (subscription topic) - receiveMessages - length <$> res `shouldBe` Right 2 - - it "Integration spec is finished" $ True `shouldBe` True - ----------------------------------------------------------------------------------------------------------------- - -receiveMessages :: KafkaConsumer -> IO (Either a [ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)]) -receiveMessages kafka = - (Right . rights) <$> www - where - www = whileJust maybeMsg return - isOK msg = if msg /= Left (KafkaResponseError RdKafkaRespErrPartitionEof) then Just msg else Nothing - maybeMsg = isOK <$> get - get = do - x <- pollMessage kafka (Timeout 1000) - print $ show x - return x - -testMessages :: TopicName -> [ProducerRecord] -testMessages t = - [ ProducerRecord t UnassignedPartition Nothing (Just "test from producer") - , ProducerRecord t UnassignedPartition (Just "key") (Just "test from producer (with key)") - ] - -sendMessages :: [ProducerRecord] -> KafkaProducer -> IO (Either KafkaError ()) -sendMessages msgs prod = - Right <$> forM_ msgs (produceMessage prod)