diff --git a/hw-kafka-client.cabal b/hw-kafka-client.cabal --- a/hw-kafka-client.cabal +++ b/hw-kafka-client.cabal @@ -1,11 +1,11 @@ --- This file has been generated from package.yaml by hpack version 0.20.0. +-- This file has been generated from package.yaml by hpack version 0.28.2. -- -- see: https://github.com/sol/hpack -- --- hash: ae89ef8878b399834f3550efb6ad9b04c27fad1d06d9d22fecfc872a4f0d23c4 +-- hash: c01d07e897b9cf5d0f6237227a00230ec841fe4a3a22766142dc015f4f703ba1 name: hw-kafka-client -version: 2.4.4 +version: 2.5.0 synopsis: Kafka bindings for Haskell description: Apache Kafka bindings backed by the librdkafka C library. . @@ -25,7 +25,6 @@ license-file: LICENSE build-type: Simple cabal-version: >= 1.10 - extra-source-files: README.md @@ -41,7 +40,7 @@ library hs-source-dirs: src - ghc-options: -Wall + ghc-options: -Wall -Wcompat -Wincomplete-record-updates -Wincomplete-uni-patterns -Wredundant-constraints extra-libraries: rdkafka build-depends: diff --git a/src/Kafka/Callbacks.hs b/src/Kafka/Callbacks.hs --- a/src/Kafka/Callbacks.hs +++ b/src/Kafka/Callbacks.hs @@ -1,6 +1,7 @@ module Kafka.Callbacks ( errorCallback , logCallback +, statsCallback ) where @@ -17,3 +18,8 @@ logCallback callback k = let realCb _ = callback in rdKafkaConfSetLogCb (getRdKafkaConf k) realCb + +statsCallback :: HasKafkaConf k => (String -> IO ()) -> k -> IO () +statsCallback callback k = + let realCb _ = callback + in rdKafkaConfSetStatsCb (getRdKafkaConf k) realCb diff --git a/src/Kafka/Consumer.hs b/src/Kafka/Consumer.hs --- a/src/Kafka/Consumer.hs +++ b/src/Kafka/Consumer.hs @@ -7,6 +7,7 @@ , pausePartitions, resumePartitions , committed, position, seek , pollMessage, pollConsumerEvents +, pollMessageBatch , commitOffsetMessage, commitAllOffsets, commitPartitionsOffsets , storeOffsets, storeOffsetMessage , closeConsumer @@ -90,12 +91,27 @@ -> Timeout -- ^ the timeout, in milliseconds -> m (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))) -- ^ Left on error or timeout, right for success pollMessage c@(KafkaConsumer _ (KafkaConf _ qr _)) (Timeout ms) = liftIO $ do - pollConsumerEvents c Nothing - mbq <- readIORef qr - case mbq of - Nothing -> return . Left $ KafkaBadSpecification "Messages queue is not configured, internal error, fatal." - Just q -> rdKafkaConsumeQueue q (fromIntegral ms) >>= fromMessagePtr + pollConsumerEvents c Nothing + mbq <- readIORef qr + case mbq of + Nothing -> return . Left $ KafkaBadSpecification "Messages queue is not configured, internal error, fatal." + Just q -> rdKafkaConsumeQueue q (fromIntegral ms) >>= fromMessagePtr +-- | Polls up to BatchSize messages. +-- Unlike 'pollMessage' this function does not return usual "timeout" errors. +-- An empty batch is returned when there are no messages available. +pollMessageBatch :: MonadIO m + => KafkaConsumer + -> Timeout + -> BatchSize + -> m [Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))] +pollMessageBatch c@(KafkaConsumer _ (KafkaConf _ qr _)) (Timeout ms) (BatchSize b) = liftIO $ do + pollConsumerEvents c Nothing + mbq <- readIORef qr + case mbq of + Nothing -> return [Left $ KafkaBadSpecification "Messages queue is not configured, internal error, fatal."] + Just q -> rdKafkaConsumeBatchQueue q ms b >>= traverse fromMessagePtr + -- | Commit message's offset on broker for the message's partition. commitOffsetMessage :: MonadIO m => OffsetCommit @@ -111,7 +127,7 @@ -> ConsumerRecord k v -> m (Maybe KafkaError) storeOffsetMessage k m = - liftIO $ toNativeTopicPartitionList [topicPartitionFromMessageForCommit m] >>= commitOffsetsStore k + liftIO $ toNativeTopicPartitionListNoDispose [topicPartitionFromMessageForCommit m] >>= commitOffsetsStore k -- | Stores offsets locally storeOffsets :: MonadIO m @@ -119,7 +135,7 @@ -> [TopicPartition] -> m (Maybe KafkaError) storeOffsets k ps = - liftIO $ toNativeTopicPartitionList ps >>= commitOffsetsStore k + liftIO $ toNativeTopicPartitionListNoDispose ps >>= commitOffsetsStore k -- | Commit offsets for all currently assigned partitions. commitAllOffsets :: MonadIO m 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 @@ -83,6 +83,17 @@ rdKafkaTopicPartitionListSetOffset pl tn tp to) ps return pl +toNativeTopicPartitionListNoDispose :: [TopicPartition] -> IO RdKafkaTopicPartitionListTPtr +toNativeTopicPartitionListNoDispose ps = do + pl <- rdKafkaTopicPartitionListNew (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 + toNativeTopicPartitionList' :: [(TopicName, PartitionId)] -> IO RdKafkaTopicPartitionListTPtr toNativeTopicPartitionList' tps = do let utps = S.toList . S.fromList $ tps @@ -100,8 +111,9 @@ -- the consumer reads from to process the next message. topicPartitionFromMessageForCommit :: ConsumerRecord k v -> TopicPartition topicPartitionFromMessageForCommit m = - let (TopicPartition t p (PartitionOffset moff)) = topicPartitionFromMessage m - in TopicPartition t p (PartitionOffset $ moff + 1) + case topicPartitionFromMessage m of + (TopicPartition t p (PartitionOffset moff)) -> TopicPartition t p (PartitionOffset $ moff + 1) + other -> other toMap :: Ord k => [(k, v)] -> Map k [v] toMap kvs = fromListWith (++) [(k, [v]) | (k, v) <- kvs] @@ -109,7 +121,7 @@ fromMessagePtr :: RdKafkaMessageTPtr -> IO (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))) fromMessagePtr ptr = withForeignPtr ptr $ \realPtr -> - if realPtr == nullPtr then (Left . kafkaRespErr) <$> getErrno + if realPtr == nullPtr then Left . kafkaRespErr <$> getErrno else do s <- peek realPtr msg <- if err'RdKafkaMessageT s /= RdKafkaRespErrNoError 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 @@ -537,7 +537,16 @@ addForeignPtrFinalizer rdKafkaQueueDestroyF ret return $ Just ret +{#fun rd_kafka_consume_batch_queue as rdKafkaConsumeBatchQueue' + {`RdKafkaQueueTPtr', `Int', castPtr `Ptr (Ptr RdKafkaMessageT)', cIntConv `CSize'} + -> `CSize' cIntConv #} +rdKafkaConsumeBatchQueue :: RdKafkaQueueTPtr -> Int -> Int -> IO [RdKafkaMessageTPtr] +rdKafkaConsumeBatchQueue qptr timeout batchSize = do + allocaArray batchSize $ \pArr -> do + rSize <- rdKafkaConsumeBatchQueue' qptr timeout pArr (fromIntegral batchSize) + peekArray (fromIntegral rSize) pArr >>= traverse newForeignPtr_ + ------------------------------------------------------------------------------------------------- ---- High-level KafkaConsumer @@ -739,6 +748,9 @@ 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 rdKafkaReadTimestamp' + {castPtr `Ptr RdKafkaMessageT', `RdKafkaTimestampTypeTPtr'} -> `CInt64T' cIntConv #} {#fun rd_kafka_message_timestamp as ^ {`RdKafkaMessageTPtr', `RdKafkaTimestampTypeTPtr'} -> `CInt64T' cIntConv #} diff --git a/src/Kafka/Types.hs b/src/Kafka/Types.hs --- a/src/Kafka/Types.hs +++ b/src/Kafka/Types.hs @@ -13,6 +13,7 @@ newtype PartitionId = PartitionId { unPartitionId :: Int} deriving (Show, Eq, Read, Ord, Enum) newtype Millis = Millis { unMillis :: Int64 } deriving (Show, Read, Eq, Ord, Num) newtype ClientId = ClientId { unClientId :: String} deriving (Show, Eq, Ord) +newtype BatchSize = BatchSize { unBatchSize :: Int } deriving (Show, Read, Eq, Ord, Num) -- | Topic name to be consumed -- diff --git a/tests-it/Kafka/IntegrationSpec.hs b/tests-it/Kafka/IntegrationSpec.hs --- a/tests-it/Kafka/IntegrationSpec.hs +++ b/tests-it/Kafka/IntegrationSpec.hs @@ -9,6 +9,7 @@ import qualified Data.ByteString as BS import Data.Either import Data.Map +import Data.Monoid ((<>)) import Kafka.Consumer as C import Kafka.Metadata as M @@ -92,6 +93,20 @@ length <$> storeRes `shouldBe` Right 2 comRes `shouldBe` Nothing + specWithKafka "Part 3 - Consume after committing stored offsets" consumerPropsNoStore $ do + it "5. sends 2 messages to test topic" $ \(_, prod) -> do + res <- sendMessages (testMessages testTopic) prod + res `shouldBe` Right () + + it "6. should receive 2 messages" $ \(k, _) -> do + res <- receiveMessages k + storeRes <- forM res $ mapM (storeOffsetMessage k) + comRes <- commitAllOffsets OffsetCommit k + + length <$> res `shouldBe` Right 2 + length <$> storeRes `shouldBe` Right 2 + comRes `shouldBe` Nothing + describe "Kafka.IntegrationSpec" $ do specWithProducer "Run producer" $ do it "sends messages to test topic" $ \prod -> do @@ -100,7 +115,7 @@ specWithConsumer "Run consumer" consumerProps $ do it "should get committed" $ \k -> do - res <- committed k (Timeout 10000) [(testTopic, PartitionId 0)] + res <- committed k (Timeout 1000) [(testTopic, PartitionId 0)] res `shouldSatisfy` isRight it "should get position" $ \k -> do @@ -194,14 +209,29 @@ msg <- pollMessage k (Timeout 1000) crOffset <$> msg `shouldBe` Right (Offset 0) + describe "Kafka.Consumer.BatchSpec" $ do + specWithConsumer "Batch consumer" (consumerProps <> groupId (ConsumerGroupId "batch-consumer")) $ do + it "should consume first batch" $ \k -> do + res <- pollMessageBatch k (Timeout 1000) (BatchSize 5) + length res `shouldBe` 5 + forM_ res (`shouldSatisfy` isRight) + it "should consume second batch with not enough messages" $ \k -> do + res <- pollMessageBatch k (Timeout 1000) (BatchSize 50) + let res' = Prelude.filter (/= Left (KafkaResponseError RdKafkaRespErrPartitionEof)) res + length res' `shouldSatisfy` (< 50) + forM_ res' (`shouldSatisfy` isRight) + + it "should consume empty batch when there are no messages" $ \k -> do + res <- pollMessageBatch k (Timeout 1000) (BatchSize 50) + length res `shouldBe` 0 ---------------------------------------------------------------------------------------------------------------- data ReadState = Skip | Read receiveMessages :: KafkaConsumer -> IO (Either KafkaError [ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)]) receiveMessages kafka = - (Right . rights) <$> allMessages + Right . rights <$> allMessages where allMessages = unfoldrM (\s -> do @@ -221,4 +251,4 @@ sendMessages :: [ProducerRecord] -> KafkaProducer -> IO (Either KafkaError ()) sendMessages msgs prod = - Right <$> forM_ msgs (produceMessage prod) + Right <$> (forM_ msgs (produceMessage prod) >> flushProducer prod) diff --git a/tests-it/Kafka/TestEnv.hs b/tests-it/Kafka/TestEnv.hs --- a/tests-it/Kafka/TestEnv.hs +++ b/tests-it/Kafka/TestEnv.hs @@ -9,8 +9,9 @@ import System.Environment import System.IO.Unsafe -import Kafka.Consumer as C -import Kafka.Producer as P +import Control.Concurrent +import Kafka.Consumer as C +import Kafka.Producer as P import Test.Hspec @@ -27,20 +28,20 @@ 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 +consumerProps :: ConsumerProperties +consumerProps = C.brokersList [brokerAddress] + <> groupId testGroupId + <> C.setCallback (logCallback (\l s1 s2 -> print $ show l <> ": " <> s1 <> ", " <> s2)) + <> C.setCallback (errorCallback (\e r -> print $ show e <> ": " <> r)) + <> noAutoCommit -consumerPropsNoStore :: BrokerAddress -> ConsumerProperties -consumerPropsNoStore broker = consumerProps broker <> noAutoOffsetStore +consumerPropsNoStore :: ConsumerProperties +consumerPropsNoStore = consumerProps <> noAutoOffsetStore -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)) +producerProps :: ProducerProperties +producerProps = P.brokersList [brokerAddress] + <> 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] @@ -48,16 +49,30 @@ mkProducer :: IO KafkaProducer mkProducer = do - (Right p) <- newProducer (producerProps brokerAddress) + (Right p) <- newProducer producerProps return p -mkConsumerWith :: (BrokerAddress -> ConsumerProperties) -> IO KafkaConsumer +mkConsumerWith :: ConsumerProperties -> IO KafkaConsumer mkConsumerWith props = do - (Right c) <- newConsumer (props brokerAddress) (testSubscription testTopic) - return c + waitVar <- newEmptyMVar + let props' = props <> C.setCallback (rebalanceCallback (\_ -> rebCallback waitVar)) + (Right c) <- newConsumer props' (testSubscription testTopic) + _ <- readMVar waitVar + return c + where + rebCallback var evt = case evt of + (RebalanceAssign _) -> putMVar var True + _ -> pure () -specWithConsumer :: String -> (BrokerAddress -> ConsumerProperties) -> SpecWith KafkaConsumer -> Spec + +specWithConsumer :: String -> ConsumerProperties -> SpecWith KafkaConsumer -> Spec specWithConsumer s p f = beforeAll (mkConsumerWith p) $ afterAll (void . closeConsumer) $ describe s f specWithProducer :: String -> SpecWith KafkaProducer -> Spec specWithProducer s f = beforeAll mkProducer $ afterAll (void . closeProducer) $ describe s f + +specWithKafka :: String -> ConsumerProperties -> SpecWith (KafkaConsumer, KafkaProducer) -> Spec +specWithKafka s p f = + beforeAll ((,) <$> mkConsumerWith p <*> mkProducer) + $ afterAll (\(consumer, producer) -> void $ closeProducer producer >> closeConsumer consumer) + $ describe s f