diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2016 Alexey Raga + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/example/ConsumerExample.hs b/example/ConsumerExample.hs new file mode 100644 --- /dev/null +++ b/example/ConsumerExample.hs @@ -0,0 +1,58 @@ +{-# LANGUAGE ScopedTypeVariables #-} +module ConsumerExample + +where + +import Control.Arrow ((&&&)) +import Data.Monoid ((<>)) +import Kafka.Consumer + +-- Global consumer properties +consumerProps :: ConsumerProperties +consumerProps = consumerBrokersList [BrokerAddress "localhost:9092"] + <> groupId (ConsumerGroupId "consumer_example_group") + <> noAutoCommit + <> reballanceCallback (ReballanceCallback printingRebalanceCallback) + <> offsetsCommitCallback (OffsetsCommitCallback printingOffsetCallback) + <> consumerLogLevel KafkaLogInfo + +-- Subscription to topics +consumerSub :: Subscription +consumerSub = topics [TopicName "kafka-client-example-topic"] + <> offsetReset Earliest + +runConsumerExample :: IO () +runConsumerExample = do + print $ cpLogLevel consumerProps + res <- runConsumer consumerProps consumerSub processMessages + print $ show res + +------------------------------------------------------------------- +processMessages :: KafkaConsumer -> IO (Either KafkaError ()) +processMessages kafka = do + mapM_ (\_ -> do + msg1 <- pollMessage kafka (Timeout 1000) + print $ "Message: " <> show msg1 + err <- commitAllOffsets OffsetCommit kafka + print $ "Offsets: " <> maybe "Committed." show err + ) [0 :: Integer .. 10] + return $ Right () + +printingRebalanceCallback :: KafkaConsumer -> KafkaError -> [TopicPartition] -> IO () +printingRebalanceCallback k e ps = do + putStrLn "Rebalance callback!" + print ("Rebalance Error: " ++ show e) + mapM_ (print . show . (tpTopicName &&& tpPartition &&& tpOffset)) ps + case e of + KafkaResponseError RdKafkaRespErrAssignPartitions -> + assign k ps >>= print . show + KafkaResponseError RdKafkaRespErrRevokePartitions -> + assign k [] >>= print . show + x -> print "UNKNOWN (and unlikely!)" >> print (show x) + + +printingOffsetCallback :: KafkaConsumer -> KafkaError -> [TopicPartition] -> IO () +printingOffsetCallback _ e ps = do + putStrLn "Offsets callback!" + print ("Offsets Error:" ++ show e) + mapM_ (print . show . (tpTopicName &&& tpPartition &&& tpOffset)) ps diff --git a/example/Main.hs b/example/Main.hs new file mode 100644 --- /dev/null +++ b/example/Main.hs @@ -0,0 +1,14 @@ +module Main where + +import ConsumerExample +import ProducerExample + +main :: IO () +main = do + putStrLn "Running producer example..." + runProducerExample + + putStrLn "Running consumer example..." + runConsumerExample + + putStrLn "Ok." diff --git a/example/ProducerExample.hs b/example/ProducerExample.hs new file mode 100644 --- /dev/null +++ b/example/ProducerExample.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE OverloadedStrings #-} + +module ProducerExample +where + +import Control.Monad (forM_) +import Data.Monoid +import Kafka.Producer + +-- Global producer properties +producerProps :: ProducerProperties +producerProps = producerBrokersList [BrokerAddress "localhost:9092"] + <> producerLogLevel KafkaLogDebug + +-- Topic to send messages to +targetTopic :: TopicName +targetTopic = TopicName "kafka-client-example-topic" + +-- Run an example +runProducerExample :: IO () +runProducerExample = do + res <- runProducer producerProps sendMessages + print $ show res + +sendMessages :: KafkaProducer -> IO (Either KafkaError String) +sendMessages prod = do + err1 <- produceMessage prod ProducerRecord + { prTopic = targetTopic + , prPartition = UnassignedPartition + , prKey = Nothing + , prValue = 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)" + } + forM_ err2 print + + return $ Right "All done, Sir." diff --git a/hw-kafka-client.cabal b/hw-kafka-client.cabal new file mode 100644 --- /dev/null +++ b/hw-kafka-client.cabal @@ -0,0 +1,96 @@ +name: hw-kafka-client +version: 1.0.0 +homepage: https://github.com/haskell-works/hw-kafka-client +bug-reports: https://github.com/haskell-works/hw-kafka-client/issues +license: MIT +license-file: LICENSE +author: Alexey Raga +maintainer: Alexey Raga +category: Database +build-type: Simple +cabal-version: >=1.10 +synopsis: Kafka bindings for Haskell +description: Apache Kafka bindings backed by the librdkafka C library. + . + Features include: + . + * Consumer groups: auto-rebalancing consumers + . + * Keyed and keyless messages producing/consuming + . + * Batch producing messages + +source-repository head + type: git + location: git://github.com/haskell-works/hw-kafka-client.git + +executable kafka-client-example + main-is: Main.hs + other-modules: ConsumerExample + , ProducerExample + hs-source-dirs: example + default-language: Haskell2010 + build-depends: base >=4.6 && < 5 + , bifunctors + , bytestring + , containers + , temporary + , transformers + , unix + , hw-kafka-client + +library + Build-tools: c2hs + build-depends: base >=4.6 && < 5 + , bifunctors + , bytestring + , containers + , temporary + , transformers + , unix + exposed-modules: + Kafka + Kafka.Types + Kafka.Consumer.ConsumerProperties + Kafka.Consumer.Subscription + Kafka.Consumer.Types + Kafka.Consumer + Kafka.Producer + Kafka.Producer.ProducerProperties + Kafka.Producer.Types + other-modules: + Kafka.Consumer.Convert + 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 + ghc-options: -Wall + include-dirs: /usr/local/include/librdkafka + , /usr/include/librdkafka + extra-lib-dirs: /usr/local/lib + extra-libraries: rdkafka + if os(darwin) + cpp-options: -D__attribute__(A)= -D_Nullable= -D_Nonnull= + +test-suite tests + type: exitcode-stdio-1.0 + default-language: Haskell2010 + hs-source-dirs: tests + other-modules: Kafka.IntegrationSpec + , 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 diff --git a/src/Kafka.hs b/src/Kafka.hs new file mode 100644 --- /dev/null +++ b/src/Kafka.hs @@ -0,0 +1,14 @@ +module Kafka +( module X +) where + +import Kafka.Consumer as X +import Kafka.Producer as X + + + + +-- -- | Sets library log level (noisiness) with respect to a kafka instance +-- setLogLevel :: Kafka -> KafkaLogLevel -> IO () +-- setLogLevel (Kafka kptr _) level = +-- rdKafkaSetLogLevel kptr (fromEnum level) diff --git a/src/Kafka/Consumer.hs b/src/Kafka/Consumer.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Consumer.hs @@ -0,0 +1,197 @@ +{-# LANGUAGE TupleSections #-} +module Kafka.Consumer +( module X +, runConsumer +, newConsumer +, assign +, pollMessage +, commitOffsetMessage +, commitAllOffsets +, closeConsumer + +-- ReExport Types +, TopicName (..) +, CIT.OffsetCommit (..) +, CIT.PartitionOffset (..) +, CIT.TopicPartition (..) +, CIT.ConsumerRecord (..) +, RDE.RdKafkaRespErrT (..) +) +where + +import Control.Exception +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 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 Kafka.Types as X +import Kafka.Consumer.Types as X +import Kafka.Consumer.Subscription as X +import Kafka.Consumer.ConsumerProperties as X + +-- | Runs high-level kafka consumer. +-- +-- A callback provided is expected to call 'pollMessage' when convenient. +runConsumer :: ConsumerProperties + -> Subscription + -> (KafkaConsumer -> IO (Either KafkaError a)) -- ^ A callback function to poll and handle messages + -> IO (Either KafkaError a) +runConsumer cp sub f = + bracket mkConsumer clConsumer runHandler + where + mkConsumer = newConsumer cp sub + + clConsumer (Left err) = return (Left err) + clConsumer (Right kc) = maybeToLeft <$> closeConsumer kc + + runHandler (Left err) = return (Left err) + runHandler (Right kc) = f kc + +-- | Creates a kafka consumer. +-- A new consumer MUST be closed with 'closeConsumer' function. +newConsumer :: MonadIO m + => ConsumerProperties + -> Subscription + -> m (Either KafkaError KafkaConsumer) +newConsumer cp (Subscription ts tp) = liftIO $ do + -- (Kafka kkk (KafkaConf kc)) <- newKafka (KafkaProps $ M.toList m) + (KafkaConf kc) <- newConsumerConf cp + tp' <- topicConf (TopicProps $ M.toList tp) + _ <- setDefaultTopicConf kc tp' + rdk <- mapLeft KafkaError <$> newRdKafkaT RdKafkaConsumer kc + case flip KafkaConsumer kc <$> rdk of + Left err -> return $ Left err + Right kafka -> do + forM_ (cpLogLevel cp) (setConsumerLogLevel kafka) + sub <- subscribe kafka ts + case sub of + Nothing -> return $ Right kafka + Just err -> closeConsumer kafka >> return (Left err) + +-- | Polls the next message from a subscription +pollMessage :: MonadIO m + => KafkaConsumer + -> Timeout -- ^ the timeout, in milliseconds (@10^3@ per second) + -> m (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))) -- ^ Left on error or timeout, right for success +pollMessage (KafkaConsumer k _) (Timeout ms) = + liftIO $ rdKafkaConsumerPoll k (fromIntegral ms) >>= fromMessagePtr + +-- | Commit message's offset on broker for the message's partition. +commitOffsetMessage :: MonadIO m + => OffsetCommit + -> KafkaConsumer + -> ConsumerRecord k v + -> m (Maybe KafkaError) +commitOffsetMessage o k m = + liftIO $ toNativeTopicPartitionList [topicPartitionFromMessage m] >>= commitOffsets o k + +-- | Commit offsets for all currently assigned partitions. +commitAllOffsets :: MonadIO m + => OffsetCommit + -> KafkaConsumer + -> m (Maybe KafkaError) +commitAllOffsets o k = + liftIO $ newForeignPtr_ nullPtr >>= commitOffsets o k + +-- | Assigns specified partitions to a current consumer. +-- Assigning an empty list means unassigning from all partitions that are currently assigned. +-- See 'setRebalanceCallback' for more details. +assign :: MonadIO m => KafkaConsumer -> [TopicPartition] -> m KafkaError +assign (KafkaConsumer k _) ps = + let pl = if null ps + then newForeignPtr_ nullPtr + else toNativeTopicPartitionList ps + in liftIO $ KafkaResponseError <$> (pl >>= rdKafkaAssign k) + + +-- | Closes the consumer and destroys it. +closeConsumer :: MonadIO m => KafkaConsumer -> m (Maybe KafkaError) +closeConsumer (KafkaConsumer k _) = + liftIO $ (kafkaErrorToMaybe . KafkaResponseError) <$> rdKafkaConsumerClose k + +----------------------------------------------------------------------------- +newConsumerConf :: ConsumerProperties -> IO KafkaConf +newConsumerConf (ConsumerProperties m rcb ccb _) = do + conf <- kafkaConf (KafkaProps $ M.toList m) + forM_ rcb (\(ReballanceCallback cb) -> setRebalanceCallback conf cb) + forM_ ccb (\(OffsetsCommitCallback cb) -> setOffsetCommitCallback conf cb) + 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: +-- any topic name in the topics list that is prefixed with @^@ will +-- 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 + 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) = + rdKafkaTopicConfDup tc >>= rdKafkaConfSetDefaultTopicConf kc + +commitOffsets :: OffsetCommit -> KafkaConsumer -> RdKafkaTopicPartitionListTPtr -> IO (Maybe KafkaError) +commitOffsets o (KafkaConsumer k _) pl = + (kafkaErrorToMaybe . KafkaResponseError) <$> rdKafkaCommit k pl (offsetCommitToBool o) + +setConsumerLogLevel :: KafkaConsumer -> KafkaLogLevel -> IO () +setConsumerLogLevel (KafkaConsumer k _) level = + liftIO $ rdKafkaSetLogLevel k (fromEnum level) diff --git a/src/Kafka/Consumer/ConsumerProperties.hs b/src/Kafka/Consumer/ConsumerProperties.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Consumer/ConsumerProperties.hs @@ -0,0 +1,85 @@ +module Kafka.Consumer.ConsumerProperties +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 + +data ConsumerProperties = ConsumerProperties + { cpProps :: Map String String + , cpRebalanceCallback :: Maybe ReballanceCallback + , cpOffsetsCallback :: Maybe OffsetsCommitCallback + , cpLogLevel :: Maybe KafkaLogLevel + } + +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) + +consumerBrokersList :: [BrokerAddress] -> ConsumerProperties +consumerBrokersList bs = + let bs' = L.intercalate "," ((\(BrokerAddress x) -> x) <$> bs) + in extraConsumerProps $ M.fromList [("bootstrap.servers", bs')] + +noAutoCommit :: ConsumerProperties +noAutoCommit = + extraConsumerProps $ M.fromList [("enable.auto.commit", "false")] + +groupId :: ConsumerGroupId -> ConsumerProperties +groupId (ConsumerGroupId cid) = + extraConsumerProps $ 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 + +-- | 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 + +consumerLogLevel :: KafkaLogLevel -> ConsumerProperties +consumerLogLevel ll = ConsumerProperties M.empty Nothing Nothing (Just ll) + +consumerCompression :: KafkaCompressionCodec -> ConsumerProperties +consumerCompression c = + extraConsumerProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) + +extraConsumerProps :: Map String String -> ConsumerProperties +extraConsumerProps m = ConsumerProperties m Nothing Nothing Nothing +{-# INLINE extraConsumerProps #-} + +-- | Sets debug features for the consumer +consumerDebug :: [KafkaDebug] -> ConsumerProperties +consumerDebug [] = extraConsumerProps M.empty +consumerDebug d = + let points = L.intercalate "," (kafkaDebugToString <$> d) + in extraConsumerProps $ M.fromList [("debug", points)] diff --git a/src/Kafka/Consumer/Convert.hs b/src/Kafka/Consumer/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Consumer/Convert.hs @@ -0,0 +1,116 @@ +module Kafka.Consumer.Convert + +where + +import Control.Monad +import qualified Data.ByteString as BS +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 + +-- | Converts offsets sync policy to integer (the way Kafka understands it): +-- +-- * @OffsetSyncDisable == -1@ +-- +-- * @OffsetSyncImmediate == 0@ +-- +-- * @OffsetSyncInterval ms == ms@ +offsetSyncToInt :: OffsetStoreSync -> Int +offsetSyncToInt sync = + case sync of + OffsetSyncDisable -> -1 + OffsetSyncImmediate -> 0 + OffsetSyncInterval ms -> ms +{-# INLINE offsetSyncToInt #-} + +offsetToInt64 :: PartitionOffset -> Int64 +offsetToInt64 o = case o of + PartitionOffsetBeginning -> -2 + PartitionOffsetEnd -> -1 + PartitionOffset off -> off + PartitionOffsetStored -> -1000 + PartitionOffsetInvalid -> -1001 +{-# INLINE offsetToInt64 #-} + +int64ToOffset :: Int64 -> PartitionOffset +int64ToOffset o + | o == -2 = PartitionOffsetBeginning + | o == -1 = PartitionOffsetEnd + | o == -1000 = PartitionOffsetStored + | o >= 0 = PartitionOffset o + | otherwise = PartitionOffsetInvalid +{-# INLINE int64ToOffset #-} + +fromNativeTopicPartitionList' :: Ptr RdKafkaTopicPartitionListT -> IO [TopicPartition] +fromNativeTopicPartitionList' ppl = peek ppl >>= fromNativeTopicPartitionList + +fromNativeTopicPartitionList :: RdKafkaTopicPartitionListT -> IO [TopicPartition] +fromNativeTopicPartitionList pl = + let count = cnt'RdKafkaTopicPartitionListT pl + elems = elems'RdKafkaTopicPartitionListT pl + in mapM (peekElemOff elems >=> toPart) [0..(fromIntegral count - 1)] + where + toPart :: RdKafkaTopicPartitionT -> IO TopicPartition + toPart p = do + topic <- peekCString $ topic'RdKafkaTopicPartitionT p + return TopicPartition { + tpTopicName = TopicName topic, + tpPartition = PartitionId $ partition'RdKafkaTopicPartitionT p, + tpOffset = int64ToOffset $ offset'RdKafkaTopicPartitionT p + } + +toNativeTopicPartitionList :: [TopicPartition] -> IO RdKafkaTopicPartitionListTPtr +toNativeTopicPartitionList ps = do + pl <- newRdKafkaTopicPartitionListT (length ps) + mapM_ (\p -> do + let TopicName tn = tpTopicName p + (PartitionId tp) = tpPartition p + to = offsetToInt64 $ tpOffset p + _ <- rdKafkaTopicPartitionListAdd pl tn tp + rdKafkaTopicPartitionListSetOffset pl tn tp to) ps + return pl + +topicPartitionFromMessage :: ConsumerRecord k v -> TopicPartition +topicPartitionFromMessage m = + let (Offset moff) = crOffset m + in TopicPartition (crTopic m) (crPartition m) (PartitionOffset moff) + +fromMessagePtr :: RdKafkaMessageTPtr -> IO (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))) +fromMessagePtr ptr = + withForeignPtr ptr $ \realPtr -> + if realPtr == nullPtr then (Left . kafkaRespErr) <$> getErrno + else do + s <- peek realPtr + msg <- if err'RdKafkaMessageT s /= RdKafkaRespErrNoError + then return . Left . KafkaResponseError $ err'RdKafkaMessageT s + else Right <$> fromMessageStorable s + rdKafkaMessageDestroy realPtr + return msg + +fromMessageStorable :: RdKafkaMessageT -> IO (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)) +fromMessageStorable s = do + topic <- newForeignPtr_ (topic'RdKafkaMessageT s) >>= rdKafkaTopicName + + payload <- if payload'RdKafkaMessageT s == nullPtr + then return Nothing + else Just <$> word8PtrToBS (len'RdKafkaMessageT s) (payload'RdKafkaMessageT 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 OffsetCommitAsync = True diff --git a/src/Kafka/Consumer/Subscription.hs b/src/Kafka/Consumer/Subscription.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Consumer/Subscription.hs @@ -0,0 +1,38 @@ +module Kafka.Consumer.Subscription +where + +-- +import Data.Map (Map) +import qualified Data.Map as M +import qualified Data.List as L +import Kafka.Types +import Kafka.Consumer.Types + +data Subscription = Subscription [TopicName] (Map String String) + +instance Monoid Subscription where + mempty = Subscription [] M.empty + mappend (Subscription ts1 m1) (Subscription ts2 m2) = + let ts' = L.nub $ L.union ts1 ts2 + ps' = M.union m1 m2 + in Subscription ts' ps' + +topics :: [TopicName] -> Subscription +topics ts = Subscription (L.nub ts) M.empty + +offsetReset :: OffsetReset -> Subscription +offsetReset o = + let o' = case o of + Earliest -> "earliest" + Latest -> "latest" + in Subscription [] (M.fromList [("auto.offset.reset", o')]) + +autoCommit :: Millis -> Subscription +autoCommit (Millis ms) = Subscription [] $ + M.fromList + [ ("enable.auto.commit", "true") + , ("auto.commit.interval.ms", show ms) + ] + +extraSubscriptionProps :: Map String String -> Subscription +extraSubscriptionProps = Subscription [] diff --git a/src/Kafka/Consumer/Types.hs b/src/Kafka/Consumer/Types.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Consumer/Types.hs @@ -0,0 +1,137 @@ +{-# LANGUAGE DeriveDataTypeable, GeneralizedNewtypeDeriving #-} +module Kafka.Consumer.Types + +where + +import Data.Bifunctor +import Data.Bifoldable +import Data.Bitraversable +import Data.Int +import Data.Typeable +import Kafka.Types +import Kafka.Internal.RdKafka + +data KafkaConsumer = KafkaConsumer { kcKafkaPtr :: !RdKafkaTPtr, kcKafkaConf :: !RdKafkaConfTPtr} deriving (Show) + +newtype ReballanceCallback = ReballanceCallback (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) +newtype OffsetsCommitCallback = OffsetsCommitCallback (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) + +newtype ConsumerGroupId = ConsumerGroupId String deriving (Show, Eq) +newtype Offset = Offset Int64 deriving (Show, Eq, Read) +newtype PartitionId = PartitionId Int deriving (Show, Eq, Read) +newtype Millis = Millis Int deriving (Show, Eq, Ord, Num) +newtype ClientId = ClientId String deriving (Show, Eq, Ord) +data OffsetReset = Earliest | Latest deriving (Show, Eq) + +-- | Offsets commit mode +data OffsetCommit = + OffsetCommit -- ^ Forces consumer to block until the broker offsets commit is done + | OffsetCommitAsync -- ^ Offsets will be committed in a non-blocking way + deriving (Show, Eq) + + +-- | Indicates how offsets are to be synced to disk +data OffsetStoreSync = + OffsetSyncDisable -- ^ Do not sync offsets (in Kafka: -1) + | OffsetSyncImmediate -- ^ Sync immediately after each offset commit (in Kafka: 0) + | OffsetSyncInterval Int -- ^ Sync after specified interval in millis + +-- | Indicates the method of storing the offsets +data OffsetStoreMethod = + OffsetStoreBroker -- ^ Offsets are stored in Kafka broker (preferred) + | OffsetStoreFile FilePath OffsetStoreSync -- ^ Offsets are stored in a file (and synced to disk according to the sync policy) + +-- | Kafka topic partition structure +data TopicPartition = TopicPartition + { tpTopicName :: TopicName + , tpPartition :: PartitionId + , tpOffset :: PartitionOffset } deriving (Show, Eq) + +-- | Represents /received/ messages from a Kafka broker (i.e. used in a consumer) +data ConsumerRecord k v = ConsumerRecord + { crTopic :: !TopicName + -- | Kafka partition this message was received from + , crPartition :: !PartitionId + -- | Offset within the 'crPartition' Kafka partition + , crOffset :: !Offset + , 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) + {-# INLINE bimap #-} + +instance Functor (ConsumerRecord k) where + fmap = second + {-# INLINE fmap #-} + +instance Foldable (ConsumerRecord k) where + foldMap f r = f (crValue r) + {-# INLINE foldMap #-} + +instance Traversable (ConsumerRecord k) where + traverse f r = (\v -> crMapValue (const v) r) <$> f (crValue r) + {-# INLINE traverse #-} + +instance Bifoldable ConsumerRecord where + bifoldMap f g r = f (crKey r) `mappend` g (crValue r) + {-# INLINE bifoldMap #-} + +instance Bitraversable ConsumerRecord where + bitraverse f g r = (\k v -> bimap (const k) (const v) r) <$> f (crKey r) <*> g (crValue r) + {-# INLINE bitraverse #-} + +crMapKey :: (k -> k') -> ConsumerRecord k v -> ConsumerRecord k' v +crMapKey = first +{-# INLINE crMapKey #-} + +crMapValue :: (v -> v') -> ConsumerRecord k v -> ConsumerRecord k v' +crMapValue = second +{-# INLINE crMapValue #-} + +crMapKV :: (k -> k') -> (v -> v') -> ConsumerRecord k v -> ConsumerRecord k' v' +crMapKV = bimap +{-# INLINE crMapKV #-} + +sequenceFirst :: (Bitraversable t, Applicative f) => t (f k) v -> f (t k v) +sequenceFirst = bitraverse id pure +{-# INLINE sequenceFirst #-} + +traverseFirst :: (Bitraversable t, Applicative f) + => (k -> f k') + -> t k v + -> f (t k' v) +traverseFirst f = bitraverse f pure +{-# INLINE traverseFirst #-} + +traverseFirstM :: (Bitraversable t, Applicative f, Monad m) + => (k -> m (f k')) + -> t k v + -> m (f (t k' v)) +traverseFirstM f r = bitraverse id pure <$> bitraverse f pure r +{-# INLINE traverseFirstM #-} + +traverseM :: (Traversable t, Applicative f, Monad m) + => (v -> m (f v')) + -> t v + -> m (f (t v')) +traverseM f r = sequenceA <$> traverse f r +{-# INLINE traverseM #-} + +bitraverseM :: (Bitraversable t, Applicative f, Monad m) + => (k -> m (f k')) + -> (v -> m (f v')) + -> t k v + -> m (f (t k' v')) +bitraverseM f g r = bisequenceA <$> bimapM f g r +{-# INLINE bitraverseM #-} + +data PartitionOffset = + PartitionOffsetBeginning + | PartitionOffsetEnd + | PartitionOffset Int64 + | PartitionOffsetStored + | PartitionOffsetInvalid + deriving (Eq, Show) diff --git a/src/Kafka/Internal/Dump.hs b/src/Kafka/Internal/Dump.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Internal/Dump.hs @@ -0,0 +1,50 @@ +module Kafka.Internal.Dump +where + +-- +import Kafka.Internal.RdKafka + +import Control.Monad +import Data.Map.Strict (Map) +import Foreign +import Foreign.C.String +import System.IO +import qualified Data.Map.Strict as Map + +-- | Prints out all supported Kafka conf properties to a handle +hPrintSupportedKafkaConf :: Handle -> IO () +hPrintSupportedKafkaConf h = handleToCFile h "w" >>= rdKafkaConfPropertiesShow + +-- | Prints out all data associated with a specific kafka object to a handle +hPrintKafka :: Handle -> RdKafkaTPtr -> IO () +hPrintKafka h k = handleToCFile h "w" >>= \f -> rdKafkaDump f k + +-- | Returns a map of the current kafka configuration +dumpConfFromKafka :: RdKafkaConfTPtr -> IO (Map String String) +dumpConfFromKafka = dumpKafkaConf + +-- | Returns a map of the current topic configuration +dumpConfFromKafkaTopic :: RdKafkaTopicConfTPtr -> IO (Map String String) +dumpConfFromKafkaTopic = dumpTopicConf + +dumpTopicConf :: RdKafkaTopicConfTPtr -> IO (Map String String) +dumpTopicConf kptr = parseDump (rdKafkaTopicConfDump kptr) + +dumpKafkaConf :: RdKafkaConfTPtr -> IO (Map String String) +dumpKafkaConf kptr = parseDump (rdKafkaConfDump kptr) + +parseDump :: (CSizePtr -> IO (Ptr CString)) -> IO (Map String String) +parseDump cstr = alloca $ \sizeptr -> do + strPtr <- cstr sizeptr + size <- peek sizeptr + + keysAndValues <- mapM (peekCString <=< peekElemOff strPtr) [0..(fromIntegral size - 1)] + + let ret = Map.fromList $ listToTuple keysAndValues + rdKafkaConfDumpFree strPtr size + return ret + +listToTuple :: [String] -> [(String, String)] +listToTuple [] = [] +listToTuple (k:v:t) = (k, v) : listToTuple t +listToTuple _ = error "list to tuple can only be called on even length lists" diff --git a/src/Kafka/Internal/RdKafka.chs b/src/Kafka/Internal/RdKafka.chs new file mode 100644 --- /dev/null +++ b/src/Kafka/Internal/RdKafka.chs @@ -0,0 +1,843 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +{-# LANGUAGE EmptyDataDecls #-} + +module Kafka.Internal.RdKafka where + +--import Control.Applicative +import Control.Monad +import Data.Word +import Foreign +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 + +#include "rdkafka.h" + +type CInt64T = {#type int64_t #} +type CInt32T = {#type int32_t #} + +{#pointer *FILE as CFilePtr -> CFile #} +{#pointer *size_t as CSizePtr -> CSize #} + +type Word8Ptr = Ptr Word8 +type CCharBufPointer = Ptr CChar + +type RdKafkaMsgFlag = Int +rdKafkaMsgFlagFree :: RdKafkaMsgFlag +rdKafkaMsgFlagFree = 0x1 +rdKafkaMsgFlagCopy :: RdKafkaMsgFlag +rdKafkaMsgFlagCopy = 0x2 + +-- Number of bytes allocated for an error buffer +nErrorBytes :: Int +nErrorBytes = 1024 * 8 + +-- Helper functions +{#fun pure rd_kafka_version as ^ + {} -> `Int' #} + +{#fun pure rd_kafka_version_str as ^ + {} -> `String' #} + +{#fun pure rd_kafka_err2str as ^ + {enumToCInt `RdKafkaRespErrT'} -> `String' #} + +{#fun pure rd_kafka_errno2err as ^ + {`Int'} -> `RdKafkaRespErrT' cIntToEnum #} + + +kafkaErrnoString :: IO (String) +kafkaErrnoString = do + (Errno num) <- getErrno + return $ rdKafkaErr2str $ rdKafkaErrno2err (fromIntegral num) + +-- Kafka Pointer Types +data RdKafkaConfT +{#pointer *rd_kafka_conf_t as RdKafkaConfTPtr foreign -> RdKafkaConfT #} + +data RdKafkaTopicConfT +{#pointer *rd_kafka_topic_conf_t as RdKafkaTopicConfTPtr foreign -> RdKafkaTopicConfT #} + +data RdKafkaT +{#pointer *rd_kafka_t as RdKafkaTPtr foreign -> RdKafkaT #} + +data RdKafkaTopicPartitionT = RdKafkaTopicPartitionT + { topic'RdKafkaTopicPartitionT :: CString + , partition'RdKafkaTopicPartitionT :: Int + , offset'RdKafkaTopicPartitionT :: Int64 + , metadata'RdKafkaTopicPartitionT :: Word8Ptr + , metadataSize'RdKafkaTopicPartitionT :: Int + , opaque'RdKafkaTopicPartitionT :: Word8Ptr + , err'RdKafkaTopicPartitionT :: RdKafkaRespErrT + } deriving (Show, Eq) + +instance Storable RdKafkaTopicPartitionT where + alignment _ = {#alignof rd_kafka_topic_partition_t#} + sizeOf _ = {#sizeof rd_kafka_topic_partition_t#} + peek p = RdKafkaTopicPartitionT + <$> liftM id ({#get rd_kafka_topic_partition_t->topic #} p) + <*> liftM fromIntegral ({#get rd_kafka_topic_partition_t->partition #} p) + <*> liftM fromIntegral ({#get rd_kafka_topic_partition_t->offset #} p) + <*> liftM castPtr ({#get rd_kafka_topic_partition_t->metadata #} p) + <*> liftM fromIntegral ({#get rd_kafka_topic_partition_t->metadata_size #} p) + <*> liftM castPtr ({#get rd_kafka_topic_partition_t->opaque #} p) + <*> liftM cIntToEnum ({#get rd_kafka_topic_partition_t->err #} p) + poke p x = do + {#set rd_kafka_topic_partition_t.topic#} p (id $ topic'RdKafkaTopicPartitionT x) + {#set rd_kafka_topic_partition_t.partition#} p (fromIntegral $ partition'RdKafkaTopicPartitionT x) + {#set rd_kafka_topic_partition_t.offset#} p (fromIntegral $ offset'RdKafkaTopicPartitionT x) + {#set rd_kafka_topic_partition_t.metadata#} p (castPtr $ metadata'RdKafkaTopicPartitionT x) + {#set rd_kafka_topic_partition_t.metadata_size#} p (fromIntegral $ metadataSize'RdKafkaTopicPartitionT x) + {#set rd_kafka_topic_partition_t.opaque#} p (castPtr $ opaque'RdKafkaTopicPartitionT x) + {#set rd_kafka_topic_partition_t.err#} p (enumToCInt $ err'RdKafkaTopicPartitionT x) + +{#pointer *rd_kafka_topic_partition_t as RdKafkaTopicPartitionTPtr foreign -> RdKafkaTopicPartitionT #} + +data RdKafkaTopicPartitionListT = RdKafkaTopicPartitionListT + { cnt'RdKafkaTopicPartitionListT :: Int + , size'RdKafkaTopicPartitionListT :: Int + , elems'RdKafkaTopicPartitionListT :: Ptr RdKafkaTopicPartitionT + } deriving (Show, Eq) + +{#pointer *rd_kafka_topic_partition_list_t as RdKafkaTopicPartitionListTPtr foreign -> RdKafkaTopicPartitionListT #} + +instance Storable RdKafkaTopicPartitionListT where + alignment _ = {#alignof rd_kafka_topic_partition_list_t#} + sizeOf _ = {#sizeof rd_kafka_topic_partition_list_t #} + peek p = RdKafkaTopicPartitionListT + <$> liftM fromIntegral ({#get rd_kafka_topic_partition_list_t->cnt #} p) + <*> liftM fromIntegral ({#get rd_kafka_topic_partition_list_t->size #} p) + <*> liftM castPtr ({#get rd_kafka_topic_partition_list_t->elems #} p) + poke p x = do + {#set rd_kafka_topic_partition_list_t.cnt#} p (fromIntegral $ cnt'RdKafkaTopicPartitionListT x) + {#set rd_kafka_topic_partition_list_t.size#} p (fromIntegral $ size'RdKafkaTopicPartitionListT x) + {#set rd_kafka_topic_partition_list_t.elems#} p (castPtr $ elems'RdKafkaTopicPartitionListT x) + +data RdKafkaTopicT +{#pointer *rd_kafka_topic_t as RdKafkaTopicTPtr foreign -> RdKafkaTopicT #} + +data RdKafkaMessageT = RdKafkaMessageT + { err'RdKafkaMessageT :: RdKafkaRespErrT + , topic'RdKafkaMessageT :: Ptr RdKafkaTopicT + , partition'RdKafkaMessageT :: Int + , len'RdKafkaMessageT :: Int + , keyLen'RdKafkaMessageT :: Int + , offset'RdKafkaMessageT :: Int64 + , payload'RdKafkaMessageT :: Word8Ptr + , key'RdKafkaMessageT :: Word8Ptr + } + deriving (Show, Eq) + +instance Storable RdKafkaMessageT where + alignment _ = {#alignof rd_kafka_message_t#} + sizeOf _ = {#sizeof rd_kafka_message_t#} + peek p = RdKafkaMessageT + <$> liftM cIntToEnum ({#get rd_kafka_message_t->err #} p) + <*> liftM castPtr ({#get rd_kafka_message_t->rkt #} p) + <*> liftM fromIntegral ({#get rd_kafka_message_t->partition #} p) + <*> liftM fromIntegral ({#get rd_kafka_message_t->len #} p) + <*> liftM fromIntegral ({#get rd_kafka_message_t->key_len #} p) + <*> liftM fromIntegral ({#get rd_kafka_message_t->offset #} p) + <*> liftM castPtr ({#get rd_kafka_message_t->payload #} p) + <*> liftM castPtr ({#get rd_kafka_message_t->key #} p) + poke p x = do + {#set rd_kafka_message_t.err#} p (enumToCInt $ err'RdKafkaMessageT x) + {#set rd_kafka_message_t.rkt#} p (castPtr $ topic'RdKafkaMessageT x) + {#set rd_kafka_message_t.partition#} p (fromIntegral $ partition'RdKafkaMessageT x) + {#set rd_kafka_message_t.len#} p (fromIntegral $ len'RdKafkaMessageT x) + {#set rd_kafka_message_t.key_len#} p (fromIntegral $ keyLen'RdKafkaMessageT x) + {#set rd_kafka_message_t.offset#} p (fromIntegral $ offset'RdKafkaMessageT x) + {#set rd_kafka_message_t.payload#} p (castPtr $ payload'RdKafkaMessageT x) + {#set rd_kafka_message_t.key#} p (castPtr $ key'RdKafkaMessageT x) + +{#pointer *rd_kafka_message_t as RdKafkaMessageTPtr foreign -> RdKafkaMessageT #} + +data RdKafkaMetadataBrokerT = RdKafkaMetadataBrokerT + { id'RdKafkaMetadataBrokerT :: Int + , host'RdKafkaMetadataBrokerT :: CString + , port'RdKafkaMetadataBrokerT :: Int + } deriving (Show, Eq) + +{#pointer *rd_kafka_metadata_broker_t as RdKafkaMetadataBrokerTPtr -> RdKafkaMetadataBrokerT #} + + +instance Storable RdKafkaMetadataBrokerT where + alignment _ = {#alignof rd_kafka_metadata_broker_t#} + sizeOf _ = {#sizeof rd_kafka_metadata_broker_t#} + peek p = RdKafkaMetadataBrokerT + <$> liftM fromIntegral ({#get rd_kafka_metadata_broker_t->id #} p) + <*> liftM id ({#get rd_kafka_metadata_broker_t->host #} p) + <*> liftM fromIntegral ({#get rd_kafka_metadata_broker_t->port #} p) + poke = undefined + +data RdKafkaMetadataPartitionT = RdKafkaMetadataPartitionT + { id'RdKafkaMetadataPartitionT :: Int + , err'RdKafkaMetadataPartitionT :: RdKafkaRespErrT + , leader'RdKafkaMetadataPartitionT :: Int + , replicaCnt'RdKafkaMetadataPartitionT :: Int + , replicas'RdKafkaMetadataPartitionT :: Ptr CInt32T + , isrCnt'RdKafkaMetadataPartitionT :: Int + , isrs'RdKafkaMetadataPartitionT :: Ptr CInt32T + } deriving (Show, Eq) + +instance Storable RdKafkaMetadataPartitionT where + alignment _ = {#alignof rd_kafka_metadata_partition_t#} + sizeOf _ = {#sizeof rd_kafka_metadata_partition_t#} + peek p = RdKafkaMetadataPartitionT + <$> liftM fromIntegral ({#get rd_kafka_metadata_partition_t->id#} p) + <*> liftM cIntToEnum ({#get rd_kafka_metadata_partition_t->err#} p) + <*> liftM fromIntegral ({#get rd_kafka_metadata_partition_t->leader#} p) + <*> liftM fromIntegral ({#get rd_kafka_metadata_partition_t->replica_cnt#} p) + <*> liftM castPtr ({#get rd_kafka_metadata_partition_t->replicas#} p) + <*> liftM fromIntegral ({#get rd_kafka_metadata_partition_t->isr_cnt#} p) + <*> liftM castPtr ({#get rd_kafka_metadata_partition_t->isrs#} p) + + poke = undefined + +{#pointer *rd_kafka_metadata_partition_t as RdKafkaMetadataPartitionTPtr -> RdKafkaMetadataPartitionT #} + +data RdKafkaMetadataTopicT = RdKafkaMetadataTopicT + { topic'RdKafkaMetadataTopicT :: CString + , partitionCnt'RdKafkaMetadataTopicT :: Int + , partitions'RdKafkaMetadataTopicT :: Ptr RdKafkaMetadataPartitionT + , err'RdKafkaMetadataTopicT :: RdKafkaRespErrT + } deriving (Show, Eq) + +instance Storable RdKafkaMetadataTopicT where + alignment _ = {#alignof rd_kafka_metadata_topic_t#} + sizeOf _ = {#sizeof rd_kafka_metadata_topic_t #} + peek p = RdKafkaMetadataTopicT + <$> liftM id ({#get rd_kafka_metadata_topic_t->topic #} p) + <*> liftM fromIntegral ({#get rd_kafka_metadata_topic_t->partition_cnt #} p) + <*> liftM castPtr ({#get rd_kafka_metadata_topic_t->partitions #} p) + <*> liftM cIntToEnum ({#get rd_kafka_metadata_topic_t->err #} p) + poke _ _ = undefined + +{#pointer *rd_kafka_metadata_topic_t as RdKafkaMetadataTopicTPtr -> RdKafkaMetadataTopicT #} + +data RdKafkaMetadataT = RdKafkaMetadataT + { brokerCnt'RdKafkaMetadataT :: Int + , brokers'RdKafkaMetadataT :: RdKafkaMetadataBrokerTPtr + , topicCnt'RdKafkaMetadataT :: Int + , topics'RdKafkaMetadataT :: RdKafkaMetadataTopicTPtr + , origBrokerId'RdKafkaMetadataT :: CInt32T + } deriving (Show, Eq) + +instance Storable RdKafkaMetadataT where + alignment _ = {#alignof rd_kafka_metadata_t#} + sizeOf _ = {#sizeof rd_kafka_metadata_t#} + peek p = RdKafkaMetadataT + <$> liftM fromIntegral ({#get rd_kafka_metadata_t->broker_cnt #} p) + <*> liftM castPtr ({#get rd_kafka_metadata_t->brokers #} p) + <*> liftM fromIntegral ({#get rd_kafka_metadata_t->topic_cnt #} p) + <*> liftM castPtr ({#get rd_kafka_metadata_t->topics #} p) + <*> liftM fromIntegral ({#get rd_kafka_metadata_t->orig_broker_id #} p) + poke _ _ = undefined + +{#pointer *rd_kafka_metadata_t as RdKafkaMetadataTPtr foreign -> RdKafkaMetadataT #} + +------------------------------------------------------------------------------------------------- +---- Partitions +{#fun rd_kafka_topic_partition_list_new as ^ + {`Int'} -> `RdKafkaTopicPartitionListTPtr' #} + +foreign import ccall unsafe "rdkafka.h &rd_kafka_topic_partition_list_destroy" + rdKafkaTopicPartitionListDestroy :: FunPtr (Ptr RdKafkaTopicPartitionListT -> IO ()) + +newRdKafkaTopicPartitionListT :: Int -> IO RdKafkaTopicPartitionListTPtr +newRdKafkaTopicPartitionListT size = do + ret <- rdKafkaTopicPartitionListNew size + addForeignPtrFinalizer rdKafkaTopicPartitionListDestroy ret + return ret + +{# fun rd_kafka_topic_partition_list_add as ^ + {`RdKafkaTopicPartitionListTPtr', `String', `Int'} -> `RdKafkaTopicPartitionTPtr' #} + +{# fun rd_kafka_topic_partition_list_add_range as ^ + {`RdKafkaTopicPartitionListTPtr', `String', `Int', `Int'} -> `()' #} + +{# fun rd_kafka_topic_partition_list_copy as ^ + {`RdKafkaTopicPartitionListTPtr'} -> `RdKafkaTopicPartitionListTPtr' #} + +copyRdKafkaTopicPartitionList :: RdKafkaTopicPartitionListTPtr -> IO RdKafkaTopicPartitionListTPtr +copyRdKafkaTopicPartitionList pl = do + cp <- rdKafkaTopicPartitionListCopy pl + addForeignPtrFinalizer rdKafkaTopicPartitionListDestroy cp + return cp + +{# fun rd_kafka_topic_partition_list_set_offset as ^ + {`RdKafkaTopicPartitionListTPtr', `String', `Int', `Int64'} + -> `RdKafkaRespErrT' cIntToEnum #} + +---- Rebalance Callback +type RdRebalanceCallback' = Ptr RdKafkaT -> CInt -> Ptr RdKafkaTopicPartitionListT -> Ptr Word8 -> IO () +type RdRebalanceCallback = Ptr RdKafkaT -> RdKafkaRespErrT -> Ptr RdKafkaTopicPartitionListT -> Ptr Word8 -> IO () + +foreign import ccall safe "wrapper" + mkRebalanceCallback :: RdRebalanceCallback' -> IO (FunPtr RdRebalanceCallback') + +foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_rebalance_cb" + rdKafkaConfSetRebalanceCb' :: + Ptr RdKafkaConfT + -> FunPtr RdRebalanceCallback' + -> IO () + +rdKafkaConfSetRebalanceCb :: RdKafkaConfTPtr -> RdRebalanceCallback -> IO () +rdKafkaConfSetRebalanceCb conf cb = do + cb' <- mkRebalanceCallback (\k e p o -> cb k (cIntToEnum e) p o) + withForeignPtr conf $ \c -> rdKafkaConfSetRebalanceCb' c cb' + return () + +---- Delivery Callback +type DeliveryCallback = Ptr RdKafkaT -> Ptr RdKafkaMessageT -> Word8Ptr -> IO () + +foreign import ccall safe "wrapper" + mkDeliveryCallback :: DeliveryCallback -> IO (FunPtr DeliveryCallback) + +foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_dr_msg_cb" + rdKafkaConfSetDrMsgCb' :: Ptr RdKafkaConfT -> FunPtr DeliveryCallback -> IO () + +rdKafkaConfSetDrMsgCb :: RdKafkaConfTPtr -> DeliveryCallback -> IO () +rdKafkaConfSetDrMsgCb conf cb = do + cb' <- mkDeliveryCallback cb + withForeignPtr conf $ \c -> rdKafkaConfSetDrMsgCb' c cb' + return () + +---- Consume Callback +type ConsumeCallback = Ptr RdKafkaMessageT -> Word8Ptr -> IO () + +foreign import ccall safe "wrapper" + mkConsumeCallback :: ConsumeCallback -> IO (FunPtr ConsumeCallback) + +foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_consume_cb" + rdKafkaConfSetConsumeCb' :: Ptr RdKafkaConfT -> FunPtr ConsumeCallback -> IO () + +rdKafkaConfSetConsumeCb :: RdKafkaConfTPtr -> ConsumeCallback -> IO () +rdKafkaConfSetConsumeCb conf cb = do + cb' <- mkConsumeCallback cb + withForeignPtr conf $ \c -> rdKafkaConfSetConsumeCb' c cb' + return () + +---- Offset Commit Callback +type OffsetCommitCallback' = Ptr RdKafkaT -> CInt -> Ptr RdKafkaTopicPartitionListT -> Word8Ptr -> IO () +type OffsetCommitCallback = Ptr RdKafkaT -> RdKafkaRespErrT -> Ptr RdKafkaTopicPartitionListT -> Word8Ptr -> IO () + +foreign import ccall safe "wrapper" + mkOffsetCommitCallback :: OffsetCommitCallback' -> IO (FunPtr OffsetCommitCallback') + +foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_offset_commit_cb" + rdKafkaConfSetOffsetCommitCb' :: Ptr RdKafkaConfT -> FunPtr OffsetCommitCallback' -> IO () + +rdKafkaConfSetOffsetCommitCb :: RdKafkaConfTPtr -> OffsetCommitCallback -> IO () +rdKafkaConfSetOffsetCommitCb conf cb = do + cb' <- mkOffsetCommitCallback (\k e p o -> cb k (cIntToEnum e) p o) + withForeignPtr conf $ \c -> rdKafkaConfSetOffsetCommitCb' c cb' + return () + +---- Throttle Callback +type ThrottleCallback = Ptr RdKafkaT -> CString -> Int -> Int -> Word8Ptr -> IO () + +foreign import ccall safe "wrapper" + mkThrottleCallback :: ThrottleCallback -> IO (FunPtr ThrottleCallback) + +foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_throttle_cb" + rdKafkaConfSetThrottleCb' :: Ptr RdKafkaConfT -> FunPtr ThrottleCallback -> IO () + +rdKafkaConfSetThrottleCb :: RdKafkaConfTPtr -> ThrottleCallback -> IO () +rdKafkaConfSetThrottleCb conf cb = do + cb' <- mkThrottleCallback cb + withForeignPtr conf $ \c -> rdKafkaConfSetThrottleCb' c cb' + return () + +---- Stats Callback +type StatsCallback = Ptr RdKafkaT -> CString -> CSize -> Word8Ptr -> IO () + +foreign import ccall safe "wrapper" + 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 :: RdKafkaConfTPtr -> StatsCallback -> IO () +rdKafkaConfSetStatsCb conf cb = do + cb' <- mkStatsCallback cb + withForeignPtr conf $ \c -> rdKafkaConfSetStatsCb' c cb' + return () + +---- Socket Callback +type SocketCallback = Int -> Int -> Int -> Word8Ptr -> IO () + +foreign import ccall safe "wrapper" + mkSocketCallback :: SocketCallback -> IO (FunPtr SocketCallback) + +foreign import ccall safe "rd_kafka.h rd_kafka_conf_set_socket_cb" + rdKafkaConfSetSocketCb' :: Ptr RdKafkaConfT -> FunPtr SocketCallback -> IO () + +rdKafkaConfSetSocketCb :: RdKafkaConfTPtr -> SocketCallback -> IO () +rdKafkaConfSetSocketCb conf cb = do + cb' <- mkSocketCallback cb + withForeignPtr conf $ \c -> rdKafkaConfSetSocketCb' c cb' + return () + +{#fun rd_kafka_conf_set_opaque as ^ + {`RdKafkaConfTPtr', castPtr `Word8Ptr'} -> `()' #} + +{#fun rd_kafka_opaque as ^ + {`RdKafkaTPtr'} -> `Word8Ptr' castPtr #} + +{#fun rd_kafka_conf_set_default_topic_conf as ^ + {`RdKafkaConfTPtr', `RdKafkaTopicConfTPtr'} -> `()' #} + +---- Partitioner Callback +type PartitionerCallback = + Ptr RdKafkaTopicTPtr + -> Word8Ptr -- keydata + -> Int -- keylen + -> Int -- partition_cnt + -> Word8Ptr -- topic_opaque + -> Word8Ptr -- msg_opaque + -> IO Int + +foreign import ccall safe "wrapper" + mkPartitionerCallback :: PartitionerCallback -> IO (FunPtr PartitionerCallback) + +foreign import ccall safe "rd_kafka.h rd_kafka_topic_conf_set_partitioner_cb" + rdKafkaTopicConfSetPartitionerCb' :: Ptr RdKafkaTopicConfT -> FunPtr PartitionerCallback -> IO () + +rdKafkaTopicConfSetPartitionerCb :: RdKafkaTopicConfTPtr -> PartitionerCallback -> IO () +rdKafkaTopicConfSetPartitionerCb conf cb = do + cb' <- mkPartitionerCallback cb + withForeignPtr conf $ \c -> rdKafkaTopicConfSetPartitionerCb' c cb' + return () + +---- Partition + +{#fun rd_kafka_topic_partition_available as ^ + {`RdKafkaTopicTPtr', cIntConv `CInt32T'} -> `Int' #} + +{#fun rd_kafka_msg_partitioner_random as ^ + { `RdKafkaTopicTPtr' + , castPtr `Word8Ptr' + , cIntConv `CSize' + , cIntConv `CInt32T' + , castPtr `Word8Ptr' + , castPtr `Word8Ptr'} + -> `CInt32T' cIntConv #} + +{#fun rd_kafka_msg_partitioner_consistent as ^ + { `RdKafkaTopicTPtr' + , castPtr `Word8Ptr' + , cIntConv `CSize' + , cIntConv `CInt32T' + , castPtr `Word8Ptr' + , castPtr `Word8Ptr'} + -> `CInt32T' cIntConv #} + +{#fun rd_kafka_msg_partitioner_consistent_random as ^ + { `RdKafkaTopicTPtr' + , castPtr `Word8Ptr' + , cIntConv `CSize' + , cIntConv `CInt32T' + , castPtr `Word8Ptr' + , castPtr `Word8Ptr'} + -> `CInt32T' cIntConv #} + +---- Poll / Yield + +{#fun rd_kafka_yield as ^ + {`RdKafkaTPtr'} -> `()' #} + +---- Pause / Resume +{#fun rd_kafka_pause_partitions as ^ + {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr'} -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_resume_partitions as ^ + {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr'} -> `RdKafkaRespErrT' cIntToEnum #} + +---- QUEUE +data RdKafkaQueueT +{#pointer *rd_kafka_queue_t as RdKafkaQueueTPtr foreign -> RdKafkaQueueT #} + +{#fun rd_kafka_queue_new as ^ + {`RdKafkaTPtr'} -> `RdKafkaQueueTPtr' #} + +foreign import ccall unsafe "rdkafka.h &rd_kafka_queue_destroy" + rdKafkaQueueDestroy :: FunPtr (Ptr RdKafkaQueueT -> IO ()) + +newRdKafkaQueue :: RdKafkaTPtr -> IO RdKafkaQueueTPtr +newRdKafkaQueue k = do + q <- rdKafkaQueueNew k + addForeignPtrFinalizer rdKafkaQueueDestroy q + return q +------------------------------------------------------------------------------------------------- +---- High-level KafkaConsumer + +{#fun rd_kafka_subscribe as ^ + {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_unsubscribe as ^ + {`RdKafkaTPtr'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_subscription as ^ + {`RdKafkaTPtr', castPtr `Ptr (Ptr RdKafkaTopicPartitionListT)'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_consumer_poll as ^ + {`RdKafkaTPtr', `Int'} -> `RdKafkaMessageTPtr' #} + +pollRdKafkaConsumer :: RdKafkaTPtr -> Int -> IO RdKafkaMessageTPtr +pollRdKafkaConsumer k t = do + m <- rdKafkaConsumerPoll k t + addForeignPtrFinalizer rdKafkaMessageDestroyF m + return m + +{#fun rd_kafka_consumer_close as ^ + {`RdKafkaTPtr'} -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_poll_set_consumer as ^ + {`RdKafkaTPtr'} -> `RdKafkaRespErrT' cIntToEnum #} + +-- rd_kafka_assign +{#fun rd_kafka_assign as ^ + {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_assignment as ^ + {`RdKafkaTPtr', castPtr `Ptr (Ptr RdKafkaTopicPartitionListT)'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_commit as ^ + {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr', boolToCInt `Bool'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_commit_message as ^ + {`RdKafkaTPtr', `RdKafkaMessageTPtr', boolToCInt `Bool'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_committed as ^ + {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr', `Int'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{#fun rd_kafka_position as ^ + {`RdKafkaTPtr', `RdKafkaTopicPartitionListTPtr'} + -> `RdKafkaRespErrT' cIntToEnum #} + +------------------------------------------------------------------------------------------------- +---- Groups +data RdKafkaGroupMemberInfoT = RdKafkaGroupMemberInfoT + { memberId'RdKafkaGroupMemberInfoT :: CString + , clientId'RdKafkaGroupMemberInfoT :: CString + , clientHost'RdKafkaGroupMemberInfoT :: CString + , memberMetadata'RdKafkaGroupMemberInfoT :: Word8Ptr + , memberMetadataSize'RdKafkaGroupMemberInfoT :: Int + , memberAssignment'RdKafkaGroupMemberInfoT :: Word8Ptr + , memberAssignmentSize'RdKafkaGroupMemberInfoT :: Int } + +instance Storable RdKafkaGroupMemberInfoT where + alignment _ = {#alignof rd_kafka_group_member_info#} + sizeOf _ = {#sizeof rd_kafka_group_member_info#} + peek p = RdKafkaGroupMemberInfoT + <$> liftM id ({#get rd_kafka_group_member_info->member_id #} p) + <*> liftM id ({#get rd_kafka_group_member_info->client_id #} p) + <*> liftM id ({#get rd_kafka_group_member_info->client_host #} p) + <*> liftM castPtr ({#get rd_kafka_group_member_info->member_metadata #} p) + <*> liftM fromIntegral ({#get rd_kafka_group_member_info->member_metadata_size #} p) + <*> liftM castPtr ({#get rd_kafka_group_member_info->member_assignment #} p) + <*> liftM fromIntegral ({#get rd_kafka_group_member_info->member_assignment_size #} p) + poke p x = do + {#set rd_kafka_group_member_info.member_id#} p (id $ memberId'RdKafkaGroupMemberInfoT x) + {#set rd_kafka_group_member_info.client_id#} p (id $ clientId'RdKafkaGroupMemberInfoT x) + {#set rd_kafka_group_member_info.client_host#} p (id $ clientHost'RdKafkaGroupMemberInfoT x) + {#set rd_kafka_group_member_info.member_metadata#} p (castPtr $ memberMetadata'RdKafkaGroupMemberInfoT x) + {#set rd_kafka_group_member_info.member_metadata_size#} p (fromIntegral $ memberMetadataSize'RdKafkaGroupMemberInfoT x) + {#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 #} + +data RdKafkaGroupInfoT = RdKafkaGroupInfoT + { broker'RdKafkaGroupInfoT :: Ptr RdKafkaMetadataBrokerT + , group'RdKafkaGroupInfoT :: CString + , err'RdKafkaGroupInfoT :: RdKafkaRespErrT + , state'RdKafkaGroupInfoT :: CString + , protocolType'RdKafkaGroupInfoT :: CString + , protocol'RdKafkaGroupInfoT :: CString + , members'RdKafkaGroupInfoT :: Ptr RdKafkaGroupMemberInfoT + , memberCnt'RdKafkaGroupInfoT :: Int } + +instance Storable RdKafkaGroupInfoT where + alignment _ = {#alignof rd_kafka_group_info #} + sizeOf _ = {#sizeof rd_kafka_group_info #} + peek p = RdKafkaGroupInfoT + <$> liftM castPtr ({#get rd_kafka_group_info->broker #} p) + <*> liftM id ({#get rd_kafka_group_info->group #} p) + <*> liftM cIntToEnum ({#get rd_kafka_group_info->err #} p) + <*> liftM id ({#get rd_kafka_group_info->state #} p) + <*> liftM id ({#get rd_kafka_group_info->protocol_type #} p) + <*> liftM id ({#get rd_kafka_group_info->protocol #} p) + <*> liftM castPtr ({#get rd_kafka_group_info->members #} p) + <*> liftM fromIntegral ({#get rd_kafka_group_info->member_cnt #} p) + poke p x = do + {#set rd_kafka_group_info.broker#} p (castPtr $ broker'RdKafkaGroupInfoT x) + {#set rd_kafka_group_info.group#} p (id $ group'RdKafkaGroupInfoT x) + {#set rd_kafka_group_info.err#} p (enumToCInt $ err'RdKafkaGroupInfoT x) + {#set rd_kafka_group_info.state#} p (id $ state'RdKafkaGroupInfoT x) + {#set rd_kafka_group_info.protocol_type#} p (id $ protocolType'RdKafkaGroupInfoT x) + {#set rd_kafka_group_info.protocol#} p (id $ protocol'RdKafkaGroupInfoT x) + {#set rd_kafka_group_info.members#} p (castPtr $ members'RdKafkaGroupInfoT x) + {#set rd_kafka_group_info.member_cnt#} p (fromIntegral $ memberCnt'RdKafkaGroupInfoT x) + +{#pointer *rd_kafka_group_info as RdKafkaGroupInfoTPtr foreign -> RdKafkaGroupInfoT #} + +data RdKafkaGroupListT = RdKafkaGroupListT + { groups'RdKafkaGroupListT :: Ptr RdKafkaGroupInfoT + , groupCnt'RdKafkaGroupListT :: Int } + +instance Storable RdKafkaGroupListT where + alignment _ = {#alignof rd_kafka_group_list #} + sizeOf _ = {#sizeof rd_kafka_group_list #} + peek p = RdKafkaGroupListT + <$> liftM castPtr ({#get rd_kafka_group_list->groups #} p) + <*> liftM fromIntegral ({#get rd_kafka_group_list->group_cnt #} p) + poke p x = do + {#set rd_kafka_group_list.groups#} p (castPtr $ groups'RdKafkaGroupListT x) + {#set rd_kafka_group_list.group_cnt#} p (fromIntegral $ groupCnt'RdKafkaGroupListT x) + +{#pointer *rd_kafka_group_list as RdKafkaGroupListTPtr foreign -> RdKafkaGroupListT #} + +{#fun rd_kafka_list_groups as ^ + {`RdKafkaTPtr', `String', castPtr `Ptr (Ptr RdKafkaGroupListT)', `Int'} + -> `RdKafkaRespErrT' cIntToEnum #} + +foreign import ccall unsafe "rdkafka.h &rd_kafka_list_groups" + rdKafkaGroupListDestroy :: FunPtr (Ptr RdKafkaGroupListT -> IO ()) + +-- 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 +------------------------------------------------------------------------------------------------- + +-- rd_kafka_message +foreign import ccall unsafe "rdkafka.h &rd_kafka_message_destroy" + rdKafkaMessageDestroyF :: FunPtr (Ptr RdKafkaMessageT -> IO ()) + +foreign import ccall unsafe "rdkafka.h rd_kafka_message_destroy" + rdKafkaMessageDestroy :: Ptr RdKafkaMessageT -> IO () + +-- 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 ()) + +{#fun rd_kafka_conf_dup as ^ + {`RdKafkaConfTPtr'} -> `RdKafkaConfTPtr' #} + +{#fun rd_kafka_conf_set as ^ + {`RdKafkaConfTPtr', `String', `String', id `CCharBufPointer', cIntConv `CSize'} + -> `RdKafkaConfResT' cIntToEnum #} + +newRdKafkaConfT :: IO RdKafkaConfTPtr +newRdKafkaConfT = do + ret <- rdKafkaConfNew + addForeignPtrFinalizer rdKafkaConfDestroy ret + return ret + +{#fun rd_kafka_conf_dump as ^ + {`RdKafkaConfTPtr', castPtr `CSizePtr'} -> `Ptr CString' id #} + +{#fun rd_kafka_conf_dump_free as ^ + {id `Ptr CString', cIntConv `CSize'} -> `()' #} + +{#fun rd_kafka_conf_properties_show as ^ + {`CFilePtr'} -> `()' #} + +-- rd_kafka_topic_conf +{#fun rd_kafka_topic_conf_new as ^ + {} -> `RdKafkaTopicConfTPtr' #} + +{#fun rd_kafka_topic_conf_dup as ^ + {`RdKafkaTopicConfTPtr'} -> `RdKafkaTopicConfTPtr' #} + +foreign import ccall unsafe "rdkafka.h &rd_kafka_topic_conf_destroy" + rdKafkaTopicConfDestroy :: FunPtr (Ptr RdKafkaTopicConfT -> IO ()) + +{#fun rd_kafka_topic_conf_set as ^ + {`RdKafkaTopicConfTPtr', `String', `String', id `CCharBufPointer', cIntConv `CSize'} + -> `RdKafkaConfResT' cIntToEnum #} + +newRdKafkaTopicConfT :: IO RdKafkaTopicConfTPtr +newRdKafkaTopicConfT = do + ret <- rdKafkaTopicConfNew + addForeignPtrFinalizer rdKafkaTopicConfDestroy ret + return ret + +{#fun rd_kafka_topic_conf_dump as ^ + {`RdKafkaTopicConfTPtr', castPtr `CSizePtr'} -> `Ptr CString' id #} + +-- rd_kafka +{#fun rd_kafka_new as ^ + {enumToCInt `RdKafkaTypeT', `RdKafkaConfTPtr', id `CCharBufPointer', cIntConv `CSize'} + -> `RdKafkaTPtr' #} + +foreign import ccall unsafe "rdkafka.h &rd_kafka_destroy" + rdKafkaDestroy :: FunPtr (Ptr RdKafkaT -> IO ()) + +newRdKafkaT :: RdKafkaTypeT -> RdKafkaConfTPtr -> IO (Either String RdKafkaTPtr) +newRdKafkaT kafkaType confPtr = + allocaBytes nErrorBytes $ \charPtr -> do + duper <- rdKafkaConfDup confPtr + ret <- rdKafkaNew kafkaType duper charPtr (fromIntegral nErrorBytes) + withForeignPtr ret $ \realPtr -> do + if realPtr == nullPtr then peekCString charPtr >>= return . Left + else do + addForeignPtrFinalizer rdKafkaDestroy ret + return $ Right ret + +{#fun rd_kafka_brokers_add as ^ + {`RdKafkaTPtr', `String'} -> `Int' #} + +{#fun rd_kafka_set_log_level as ^ + {`RdKafkaTPtr', `Int'} -> `()' #} + +-- rd_kafka consume + +{#fun rd_kafka_consume_start as rdKafkaConsumeStartInternal + {`RdKafkaTopicTPtr', cIntConv `CInt32T', cIntConv `CInt64T'} -> `Int' #} + +rdKafkaConsumeStart :: RdKafkaTopicTPtr -> Int -> Int64 -> IO (Maybe String) +rdKafkaConsumeStart topicPtr partition offset = do + i <- rdKafkaConsumeStartInternal topicPtr (fromIntegral partition) (fromIntegral offset) + case i of + -1 -> kafkaErrnoString >>= return . Just + _ -> return Nothing +{#fun rd_kafka_consume_stop as rdKafkaConsumeStopInternal + {`RdKafkaTopicTPtr', cIntConv `CInt32T'} -> `Int' #} + +{#fun rd_kafka_consume as ^ + {`RdKafkaTopicTPtr', cIntConv `CInt32T', `Int'} -> `RdKafkaMessageTPtr' #} + +{#fun rd_kafka_consume_batch as ^ + {`RdKafkaTopicTPtr', cIntConv `CInt32T', `Int', castPtr `Ptr (Ptr RdKafkaMessageT)', cIntConv `CSize'} + -> `CSize' cIntConv #} + +rdKafkaConsumeStop :: RdKafkaTopicTPtr -> Int -> IO (Maybe String) +rdKafkaConsumeStop topicPtr partition = do + i <- rdKafkaConsumeStopInternal topicPtr (fromIntegral partition) + case i of + -1 -> kafkaErrnoString >>= return . Just + _ -> return Nothing + +{#fun rd_kafka_offset_store as rdKafkaOffsetStore + {`RdKafkaTopicTPtr', cIntConv `CInt32T', cIntConv `CInt64T'} + -> `RdKafkaRespErrT' cIntToEnum #} + +-- rd_kafka produce + +{#fun rd_kafka_produce as ^ + {`RdKafkaTopicTPtr', cIntConv `CInt32T', `Int', castPtr `Word8Ptr', + cIntConv `CSize', castPtr `Word8Ptr', cIntConv `CSize', castPtr `Word8Ptr'} + -> `Int' #} + +{#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 ^ + {`RdKafkaTPtr', boolToCInt `Bool', `RdKafkaTopicTPtr', + castMetadata `Ptr (Ptr RdKafkaMetadataT)', `Int'} + -> `RdKafkaRespErrT' cIntToEnum #} + +{# fun rd_kafka_metadata_destroy as ^ + {castPtr `Ptr RdKafkaMetadataT'} -> `()' #} + +{#fun rd_kafka_poll as ^ + {`RdKafkaTPtr', `Int'} -> `Int' #} + +{#fun rd_kafka_outq_len as ^ + {`RdKafkaTPtr'} -> `Int' #} + +{#fun rd_kafka_dump as ^ + {`CFilePtr', `RdKafkaTPtr'} -> `()' #} + + +-- rd_kafka_topic +{#fun rd_kafka_topic_name as ^ + {`RdKafkaTopicTPtr'} -> `String' #} + +{#fun rd_kafka_topic_new as ^ + {`RdKafkaTPtr', `String', `RdKafkaTopicConfTPtr'} -> `RdKafkaTopicTPtr' #} + +{# fun rd_kafka_topic_destroy as ^ + {castPtr `Ptr RdKafkaTopicT'} -> `()' #} + +destroyUnmanagedRdKafkaTopic :: RdKafkaTopicTPtr -> IO () +destroyUnmanagedRdKafkaTopic ptr = + withForeignPtr ptr rdKafkaTopicDestroy + +foreign import ccall unsafe "rdkafka.h &rd_kafka_topic_destroy" + rdKafkaTopicDestroy' :: FunPtr (Ptr RdKafkaTopicT -> IO ()) + +newUnmanagedRdKafkaTopicT :: RdKafkaTPtr -> String -> RdKafkaTopicConfTPtr -> IO (Either String RdKafkaTopicTPtr) +newUnmanagedRdKafkaTopicT kafkaPtr topic topicConfPtr = do + duper <- 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 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 + +-- Marshall / Unmarshall +enumToCInt :: Enum a => a -> CInt +enumToCInt = fromIntegral . fromEnum +cIntToEnum :: Enum a => CInt -> a +cIntToEnum = toEnum . fromIntegral +cIntConv :: (Integral a, Num b) => a -> b +cIntConv = fromIntegral +boolToCInt :: Bool -> CInt +boolToCInt True = CInt 1 +boolToCInt False = CInt 0 + +-- Handle -> File descriptor + +foreign import ccall "" fdopen :: Fd -> CString -> IO (Ptr CFile) + +handleToCFile :: Handle -> String -> IO (CFilePtr) +handleToCFile h m = + do iomode <- newCString m + fd <- handleToFd h + fdopen fd iomode + +c_stdin :: IO CFilePtr +c_stdin = handleToCFile stdin "r" +c_stdout :: IO CFilePtr +c_stdout = handleToCFile stdout "w" +c_stderr :: IO CFilePtr +c_stderr = handleToCFile stderr "w" diff --git a/src/Kafka/Internal/RdKafkaEnum.chs b/src/Kafka/Internal/RdKafkaEnum.chs new file mode 100644 --- /dev/null +++ b/src/Kafka/Internal/RdKafkaEnum.chs @@ -0,0 +1,10 @@ +{-# 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 new file mode 100644 --- /dev/null +++ b/src/Kafka/Internal/Setup.hs @@ -0,0 +1,65 @@ +module Kafka.Internal.Setup where + +import Kafka.Types +import Kafka.Internal.RdKafka +import Kafka.Internal.RdKafkaEnum + +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 + +newKafkaConf :: IO KafkaConf +newKafkaConf = KafkaConf <$> newRdKafkaConfT + +kafkaConf :: KafkaProps -> IO KafkaConf +kafkaConf overrides = do + conf <- newKafkaConf + setAllKafkaConfValues conf overrides + return conf + +topicConf :: TopicProps -> IO TopicConf +topicConf overrides = do + conf <- newTopicConf + setAllTopicConfValues conf overrides + return conf + +checkConfSetValue :: RdKafkaConfResT -> CCharBufPointer -> IO () +checkConfSetValue err charPtr = case err of + RdKafkaConfOk -> return () + RdKafkaConfInvalid -> do + str <- peekCString charPtr + throw $ KafkaInvalidConfigurationValue str + RdKafkaConfUnknown -> do + str <- peekCString charPtr + throw $ KafkaUnknownConfigurationKey str + +setKafkaConfValue :: KafkaConf -> String -> String -> IO () +setKafkaConfValue (KafkaConf confPtr) key value = + allocaBytes nErrorBytes $ \charPtr -> do + err <- rdKafkaConfSet confPtr key value charPtr (fromIntegral nErrorBytes) + checkConfSetValue err charPtr + +setAllKafkaConfValues :: KafkaConf -> KafkaProps -> IO () +setAllKafkaConfValues conf (KafkaProps props) = forM_ props $ uncurry (setKafkaConfValue conf) + +setTopicConfValue :: TopicConf -> String -> String -> IO () +setTopicConfValue (TopicConf confPtr) key value = + allocaBytes nErrorBytes $ \charPtr -> do + err <- rdKafkaTopicConfSet confPtr key value charPtr (fromIntegral nErrorBytes) + checkConfSetValue err charPtr + +setAllTopicConfValues :: TopicConf -> TopicProps -> IO () +setAllTopicConfValues conf (TopicProps props) = forM_ props $ uncurry (setTopicConfValue conf) diff --git a/src/Kafka/Internal/Shared.hs b/src/Kafka/Internal/Shared.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Internal/Shared.hs @@ -0,0 +1,53 @@ +module Kafka.Internal.Shared +where + +import Control.Exception +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 +word8PtrToBS len ptr = BSI.create len $ \bsptr -> + BSI.memcpy bsptr ptr len + +kafkaRespErr :: Errno -> KafkaError +kafkaRespErr (Errno num) = KafkaResponseError $ rdKafkaErrno2err (fromIntegral num) +{-# INLINE kafkaRespErr #-} + +throwOnError :: IO (Maybe String) -> IO () +throwOnError action = do + m <- action + case m of + Just e -> throw $ KafkaError e + Nothing -> return () + +hasError :: KafkaError -> Bool +hasError err = case err of + KafkaResponseError RdKafkaRespErrNoError -> False + _ -> True +{-# INLINE hasError #-} + +kafkaErrorToEither :: KafkaError -> Either KafkaError () +kafkaErrorToEither err = case err of + KafkaResponseError RdKafkaRespErrNoError -> Right () + _ -> Left err +{-# INLINE kafkaErrorToEither #-} + +kafkaErrorToMaybe :: KafkaError -> Maybe KafkaError +kafkaErrorToMaybe err = case err of + KafkaResponseError RdKafkaRespErrNoError -> Nothing + _ -> 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/Producer.hs b/src/Kafka/Producer.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Producer.hs @@ -0,0 +1,152 @@ +module Kafka.Producer +( module X +, runProducer +, newProducer +, produceMessage +, flushProducer +, closeProducer +, RDE.RdKafkaRespErrT (..) +) +where + +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 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 + +runProducer :: ProducerProperties + -> (KafkaProducer -> IO (Either KafkaError a)) + -> IO (Either KafkaError a) +runProducer props f = + bracket mkProducer clProducer runHandler + where + mkProducer = newProducer props + + clProducer (Left _) = return () + clProducer (Right prod) = closeProducer prod + + runHandler (Left err) = return $ Left err + runHandler (Right prod) = f prod + +-- | Creates a new kafka producer +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 + case mbKafka of + Left err -> return . Left $ KafkaError err + Right kafka -> do + forM_ ll (rdKafkaSetLogLevel kafka . fromEnum) + return .Right $ KafkaProducer kafka kc tc + +-- | Produce a single message. +-- Since librdkafka is backed by a queue, this function can return before messages are sent. See +-- 'flushProducer' to wait for queue to empty. +produceMessage :: MonadIO m + => KafkaProducer + -> ProducerRecord + -> m (Maybe KafkaError) +produceMessage (KafkaProducer k _ tc) m = liftIO $ + bracket (mkTopic $ prTopic m) clTopic withTopic + where + mkTopic (TopicName tn) = newUnmanagedRdKafkaTopicT k tn tc + + clTopic (Left _) = return () + clTopic (Right t) = destroyUnmanagedRdKafkaTopic t + + withTopic (Left err) = return . Just . KafkaError $ err + withTopic (Right t) = + withBS (prValue m) $ \payloadPtr payloadLength -> + withBS (prKey m) $ \keyPtr keyLength -> + handleProduceErr =<< + rdKafkaProduce t (producePartitionCInt (prPartition m)) + copyMsgFlags payloadPtr (fromIntegral payloadLength) + keyPtr (fromIntegral keyLength) nullPtr +-- +-- let (topic, key, partition, payload) = keyAndPayload m +-- in withBS (Just payload) $ \payloadPtr payloadLength -> +-- withBS key $ \keyPtr keyLength -> +-- handleProduceErr =<< +-- rdKafkaProduce t (producePartitionCInt partition) +-- copyMsgFlags payloadPtr (fromIntegral payloadLength) +-- keyPtr (fromIntegral keyLength) nullPtr +-- +-- | Produce a batch of messages. Since librdkafka is backed by a queue, this function can return +-- before messages are sent. See 'flushProducer' to wait for the queue to be empty. +-- produceMessageBatch :: KafkaTopic -- ^ topic pointer +-- -> [ProducerRecord] -- ^ list of messages to enqueue. +-- -> IO [(ProducerRecord, KafkaError)] -- list of failed messages with their errors. This will be empty on success. +-- produceMessageBatch (KafkaTopic t _ _) pms = +-- concat <$> mapM sendBatch (partBatches pms) +-- where +-- partBatches msgs = (\x -> (pmPartition (head x), x)) <$> batches msgs +-- batches = groupBy ((==) `on` pmPartition) . sortBy (comparing pmPartition) +-- sendBatch (part, ms) = do +-- msgs <- forM ms toNativeMessage +-- let msgsCount = length msgs +-- withArray msgs $ \batchPtr -> do +-- batchPtrF <- newForeignPtr_ batchPtr +-- numRet <- rdKafkaProduceBatch t (producePartitionCInt part) copyMsgFlags batchPtrF msgsCount +-- if numRet == msgsCount then return [] +-- else do +-- errs <- mapM (return . err'RdKafkaMessageT <=< peekElemOff batchPtr) +-- [0..(fromIntegral $ msgsCount - 1)] +-- return [(m, KafkaResponseError e) | (m, e) <- zip pms errs, e /= RdKafkaRespErrNoError] +-- where +-- toNativeMessage msg = +-- let (key, partition, payload) = keyAndPayload msg +-- in withBS (Just payload) $ \payloadPtr payloadLength -> +-- withBS key $ \keyPtr keyLength -> +-- withForeignPtr t $ \ptrTopic -> +-- return RdKafkaMessageT +-- { err'RdKafkaMessageT = RdKafkaRespErrNoError +-- , topic'RdKafkaMessageT = ptrTopic +-- , partition'RdKafkaMessageT = producePartitionInt partition +-- , len'RdKafkaMessageT = payloadLength +-- , payload'RdKafkaMessageT = payloadPtr +-- , offset'RdKafkaMessageT = 0 +-- , keyLen'RdKafkaMessageT = keyLength +-- , key'RdKafkaMessageT = keyPtr +-- } + +closeProducer :: MonadIO m => KafkaProducer -> m () +closeProducer = flushProducer + +-- | Drains the outbound queue for a producer. This function is called automatically at the end of +-- 'runKafkaProducer' or 'runKafkaProducerConf' and usually doesn't need to be called directly. +flushProducer :: MonadIO m => KafkaProducer -> m () +flushProducer kp@(KafkaProducer k _ _) = liftIO $ do + pollEvents k 100 + l <- outboundQueueLength k + unless (l == 0) $ flushProducer kp + +------------------------------------------------------------------------------------ + +withBS :: Maybe BS.ByteString -> (Ptr a -> Int -> IO b) -> IO b +withBS Nothing f = f nullPtr 0 +withBS (Just bs) f = + let (d, o, l) = BSI.toForeignPtr bs + in withForeignPtr d $ \p -> f (p `plusPtr` o) l + +pollEvents :: RdKafkaTPtr -> Int -> IO () +pollEvents kPtr timeout = void (rdKafkaPoll kPtr timeout) + +outboundQueueLength :: RdKafkaTPtr -> IO Int +outboundQueueLength = rdKafkaOutqLen diff --git a/src/Kafka/Producer/Convert.hs b/src/Kafka/Producer/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Producer/Convert.hs @@ -0,0 +1,28 @@ +module Kafka.Producer.Convert +where + +import Foreign.C.Error +import Foreign.C.Types +import Kafka.Internal.RdKafka +import Kafka.Internal.Shared +import Kafka.Types +import Kafka.Producer.Types + +copyMsgFlags :: Int +copyMsgFlags = rdKafkaMsgFlagCopy +{-# INLINE copyMsgFlags #-} + +producePartitionInt :: ProducePartition -> Int +producePartitionInt UnassignedPartition = -1 +producePartitionInt (SpecifiedPartition n) = n +{-# INLINE producePartitionInt #-} + +producePartitionCInt :: ProducePartition -> CInt +producePartitionCInt = fromIntegral . producePartitionInt +{-# INLINE producePartitionCInt #-} + +handleProduceErr :: Int -> IO (Maybe KafkaError) +handleProduceErr (- 1) = (Just . kafkaRespErr) <$> getErrno +handleProduceErr 0 = return Nothing +handleProduceErr _ = return $ Just KafkaInvalidReturnValue +{-# INLINE handleProduceErr #-} diff --git a/src/Kafka/Producer/ProducerProperties.hs b/src/Kafka/Producer/ProducerProperties.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Producer/ProducerProperties.hs @@ -0,0 +1,49 @@ +module Kafka.Producer.ProducerProperties +where + +-- +import Control.Monad +import Data.Map (Map) +import Kafka.Types +import qualified Data.Map as M +import qualified Data.List as L + +data ProducerProperties = ProducerProperties + { ppKafkaProps :: Map String String + , ppTopicProps :: Map String String + , ppLogLevel :: Maybe KafkaLogLevel + } deriving (Show) + +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) + +producerBrokersList :: [BrokerAddress] -> ProducerProperties +producerBrokersList bs = + let bs' = L.intercalate "," ((\(BrokerAddress x) -> x) <$> bs) + in extraProducerProps $ M.fromList [("bootstrap.servers", bs')] + +producerLogLevel :: KafkaLogLevel -> ProducerProperties +producerLogLevel ll = ProducerProperties M.empty M.empty (Just ll) + +producerCompression :: KafkaCompressionCodec -> ProducerProperties +producerCompression c = + extraProducerProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) + +producerTopicCompression :: KafkaCompressionCodec -> ProducerProperties +producerTopicCompression c = + extraProducerTopicProps $ M.singleton "compression.codec" (kafkaCompressionCodecToString c) + +extraProducerProps :: Map String String -> ProducerProperties +extraProducerProps m = ProducerProperties m M.empty Nothing + +extraProducerTopicProps :: Map String String -> ProducerProperties +extraProducerTopicProps m = ProducerProperties M.empty m Nothing + +-- | Sets debug features for the producer +producerDebug :: [KafkaDebug] -> ProducerProperties +producerDebug [] = extraProducerProps M.empty +producerDebug d = + let points = L.intercalate "," (kafkaDebugToString <$> d) + in extraProducerProps $ M.fromList [("debug", points)] diff --git a/src/Kafka/Producer/Types.hs b/src/Kafka/Producer/Types.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Producer/Types.hs @@ -0,0 +1,29 @@ +{-# LANGUAGE DeriveDataTypeable #-} +module Kafka.Producer.Types + +where + +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 + } deriving (Show) + +-- | Represents messages /to be enqueued/ onto a Kafka broker (i.e. used for a producer) +data ProducerRecord = ProducerRecord + { prTopic :: !TopicName + , prPartition :: !ProducePartition + , prKey :: Maybe BS.ByteString + , prValue :: Maybe BS.ByteString + } deriving (Eq, Show, Typeable) + +data ProducePartition = + SpecifiedPartition {-# UNPACK #-} !Int -- the partition number of the topic + | UnassignedPartition + deriving (Show, Eq, Ord, Typeable) diff --git a/src/Kafka/Types.hs b/src/Kafka/Types.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Types.hs @@ -0,0 +1,107 @@ +{-# LANGUAGE DeriveDataTypeable #-} +module Kafka.Types +where + +import Control.Exception +import Data.Typeable +import Kafka.Internal.RdKafkaEnum + +-- | Topic name to be consumed +-- +-- Wildcard (regex) topics are supported by the librdkafka assignor: +-- any topic name in the topics list that is prefixed with @^@ will +-- be regex-matched to the full list of topics in the cluster and matching +-- topics will be added to the subscription list. +newtype TopicName = + TopicName String -- ^ a simple topic name or a regex if started with @^@ + deriving (Show, Eq, Read) + +-- | Kafka broker address string (e.g. @broker1:9092@) +newtype BrokerAddress = BrokerAddress String deriving (Show, Eq) + +-- | Timeout in milliseconds +newtype Timeout = Timeout Int deriving (Show, Eq, Read) + +-- | Log levels for the RdKafkaLibrary used in 'setKafkaLogLevel' +data KafkaLogLevel = + KafkaLogEmerg | KafkaLogAlert | KafkaLogCrit | KafkaLogErr | KafkaLogWarning | + KafkaLogNotice | KafkaLogInfo | KafkaLogDebug + deriving (Show, Eq) + +instance Enum KafkaLogLevel where + toEnum 0 = KafkaLogEmerg + toEnum 1 = KafkaLogAlert + toEnum 2 = KafkaLogCrit + toEnum 3 = KafkaLogErr + toEnum 4 = KafkaLogWarning + toEnum 5 = KafkaLogNotice + toEnum 6 = KafkaLogInfo + toEnum 7 = KafkaLogDebug + toEnum _ = undefined + + fromEnum KafkaLogEmerg = 0 + fromEnum KafkaLogAlert = 1 + fromEnum KafkaLogCrit = 2 + fromEnum KafkaLogErr = 3 + fromEnum KafkaLogWarning = 4 + fromEnum KafkaLogNotice = 5 + fromEnum KafkaLogInfo = 6 + fromEnum KafkaLogDebug = 7 + +-- +-- | Any Kafka errors +data KafkaError = + KafkaError String + | KafkaInvalidReturnValue + | KafkaBadSpecification String + | KafkaResponseError RdKafkaRespErrT + | KafkaInvalidConfigurationValue String + | KafkaUnknownConfigurationKey String + | KakfaBadConfiguration + deriving (Eq, Show, Typeable) + +instance Exception KafkaError + +data KafkaDebug = + DebugGeneric + | DebugBroker + | DebugTopic + | DebugMetadata + | DebugQueue + | DebugMsg + | DebugProtocol + | DebugCgrp + | DebugSecurity + | DebugFetch + | DebugFeature + | DebugAll + deriving (Eq, Show, Typeable) + +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" + +data KafkaCompressionCodec = + NoCompression + | Gzip + | Snappy + | Lz4 + deriving (Eq, Show, Typeable) + +kafkaCompressionCodecToString :: KafkaCompressionCodec -> String +kafkaCompressionCodecToString c = case c of + NoCompression -> "none" + Gzip -> "gzip" + Snappy -> "snappy" + Lz4 -> "lz4" diff --git a/tests/Kafka/Consumer/ConsumerRecordMapSpec.hs b/tests/Kafka/Consumer/ConsumerRecordMapSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Kafka/Consumer/ConsumerRecordMapSpec.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE OverloadedStrings #-} +module Kafka.Consumer.ConsumerRecordMapSpec +( spec +) where + +import Data.Bitraversable +import Kafka.Types +import Kafka.Consumer.Types +import Test.Hspec + +testKey, testValue :: String +testKey = "some-key" +testValue = "some-value" + +testRecord :: ConsumerRecord (Maybe String) (Maybe String) +testRecord = ConsumerRecord + { crTopic = TopicName "some-topic" + , crPartition = PartitionId 0 + , crOffset = Offset 5 + , crKey = Just testKey + , crValue = Just testValue + } + +spec :: Spec +spec = describe "Kafka.Consumer.ConsumerRecordSpec" $ do + it "should exract key" $ + sequenceFirst testRecord `shouldBe` Just (crMapKey (const testKey) testRecord) + + it "should extract value" $ + sequence testRecord `shouldBe` Just (crMapValue (const testValue) testRecord) + + it "should extract both key and value" $ + bisequence testRecord `shouldBe` Just (crMapKV (const testKey) (const testValue) testRecord) diff --git a/tests/Kafka/Consumer/ConsumerRecordTraverseSpec.hs b/tests/Kafka/Consumer/ConsumerRecordTraverseSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Kafka/Consumer/ConsumerRecordTraverseSpec.hs @@ -0,0 +1,86 @@ +{-# LANGUAGE OverloadedStrings #-} +module Kafka.Consumer.ConsumerRecordTraverseSpec +( spec +) where + +import Data.Bitraversable +import Data.Bifunctor +import Kafka.Types +import Kafka.Consumer.Types +import Test.Hspec + +testKey, testValue :: String +testKey = "some-key" +testValue = "some-value" + +testRecord :: ConsumerRecord String String +testRecord = ConsumerRecord + { crTopic = TopicName "some-topic" + , crPartition = PartitionId 0 + , crOffset = Offset 5 + , crKey = testKey + , crValue = testValue + } + +liftValue :: a -> Maybe a +liftValue = Just + +liftNothing :: a -> Maybe a +liftNothing _ = Nothing + +liftValueM :: a -> Maybe (Either String a) +liftValueM = pure . pure + +liftNothingM :: a -> Maybe (Either String a) +liftNothingM _ = Nothing + +testError :: Either String a +testError = Left "test error" + +liftErrorM :: a -> Maybe (Either String a) +liftErrorM _ = Just testError + +spec :: Spec +spec = describe "Kafka.Consumer.ConsumerRecordTraverseSpec" $ do + it "should sequence" $ do + sequence (liftValue <$> testRecord) `shouldBe` Just testRecord + sequenceA (liftValue <$> testRecord) `shouldBe` Just testRecord + sequence (liftNothing <$> testRecord) `shouldBe` Nothing + + it "should traverse" $ do + traverse liftValue testRecord `shouldBe` Just testRecord + traverse liftNothing testRecord `shouldBe` Nothing + + it "should bisequence" $ do + bisequence (bimap liftValue liftValue testRecord) `shouldBe` Just testRecord + bisequence (bimap liftNothing liftValue testRecord) `shouldBe` Nothing + bisequence (bimap liftValue liftNothing testRecord) `shouldBe` Nothing + bisequenceA (bimap liftValue liftValue testRecord) `shouldBe` Just testRecord + bisequenceA (bimap liftNothing liftValue testRecord) `shouldBe` Nothing + bisequenceA (bimap liftValue liftNothing testRecord) `shouldBe` Nothing + + it "should bitraverse" $ do + bitraverse liftValue liftValue testRecord `shouldBe` Just testRecord + bitraverse liftNothing liftValue testRecord `shouldBe` Nothing + bitraverse liftValue liftNothing testRecord `shouldBe` Nothing + + it "should traverse key (monadic)" $ + traverseFirstM liftValueM testRecord `shouldBe` Just (Right testRecord) + + it "should traverse value (monadic)" $ + traverseM liftValueM testRecord `shouldBe` Just (Right testRecord) + + it "should traverse KV (monadic)" $ + bitraverseM liftValueM liftValueM testRecord `shouldBe` Just (Right testRecord) + + it "should traverse and report error (monadic)" $ do + traverseFirstM liftErrorM testRecord `shouldBe` Just testError + traverseM liftErrorM testRecord `shouldBe` Just testError + bitraverseM liftValueM liftErrorM testRecord `shouldBe` Just testError + bitraverseM liftErrorM liftValueM testRecord `shouldBe` Just testError + + it "should traverse and report empty container (monadic)" $ do + traverseFirstM liftNothingM testRecord `shouldBe` Nothing + traverseM liftNothingM testRecord `shouldBe` Nothing + bitraverseM liftValueM liftNothingM testRecord `shouldBe` Nothing + bitraverseM liftNothingM liftValueM testRecord `shouldBe` Nothing diff --git a/tests/Kafka/IntegrationSpec.hs b/tests/Kafka/IntegrationSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Kafka/IntegrationSpec.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Kafka.IntegrationSpec +( spec +) where + +import Control.Exception +import Control.Monad (forM_) +import Control.Monad.Loops +import Data.Monoid ((<>)) +import Data.Either +import System.Environment +import qualified Data.ByteString as BS + +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 + +---------------------------------------------------------------------------------------------------------------- + +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) diff --git a/tests/Spec.hs b/tests/Spec.hs new file mode 100644 --- /dev/null +++ b/tests/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-}