diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -61,7 +61,7 @@ mapM_ (\_ -> do msg1 <- pollMessage kafka (Timeout 1000) putStrLn $ "Message: " <> show msg1 - err <- commitAllOffsets kafka OffsetCommit + err <- commitAllOffsets OffsetCommit kafka putStrLn $ "Offsets: " <> maybe "Committed." show err ) [0 .. 10] return $ Right () diff --git a/hw-kafka-client.cabal b/hw-kafka-client.cabal --- a/hw-kafka-client.cabal +++ b/hw-kafka-client.cabal @@ -1,5 +1,5 @@ name: hw-kafka-client -version: 2.0.4 +version: 2.1.0 homepage: https://github.com/haskell-works/hw-kafka-client bug-reports: https://github.com/haskell-works/hw-kafka-client/issues license: MIT @@ -78,6 +78,7 @@ Kafka.Callbacks Kafka.Consumer.Convert Kafka.Consumer.Callbacks + Kafka.Internal.CancellationToken Kafka.Internal.Shared Kafka.Internal.Setup Kafka.Producer.Convert diff --git a/src/Kafka/Callbacks.hs b/src/Kafka/Callbacks.hs --- a/src/Kafka/Callbacks.hs +++ b/src/Kafka/Callbacks.hs @@ -9,12 +9,12 @@ errorCallback :: HasKafkaConf k => (KafkaError -> String -> IO ()) -> k -> IO () errorCallback callback k = - let (KafkaConf c) = getKafkaConf k + let (KafkaConf c _) = getKafkaConf k realCb _ err = callback (KafkaResponseError err) in rdKafkaConfSetErrorCb c realCb logCallback :: HasKafkaConf k => (Int -> String -> String -> IO ()) -> k -> IO () logCallback callback k = - let (KafkaConf c) = getKafkaConf k + let (KafkaConf c _) = getKafkaConf k realCb _ = callback in rdKafkaConfSetLogCb c realCb diff --git a/src/Kafka/Consumer.hs b/src/Kafka/Consumer.hs --- a/src/Kafka/Consumer.hs +++ b/src/Kafka/Consumer.hs @@ -4,6 +4,7 @@ , runConsumer , newConsumer , assign, assignment, subscription +, pausePartitions, resumePartitions , committed, position, seek , pollMessage , commitOffsetMessage, commitAllOffsets, commitPartitionsOffsets @@ -21,14 +22,15 @@ import Control.Arrow import Control.Exception -import Control.Monad (forM_) +import Control.Monad (forM_) import Control.Monad.IO.Class import Control.Monad.Trans.Except import Data.Bifunctor -import qualified Data.ByteString as BS -import qualified Data.Map as M -import Foreign hiding (void) +import qualified Data.ByteString as BS +import qualified Data.Map as M +import Foreign hiding (void) import Kafka.Consumer.Convert +import Kafka.Internal.CancellationToken as CToken import Kafka.Internal.RdKafka import Kafka.Internal.Setup import Kafka.Internal.Shared @@ -41,8 +43,8 @@ import Kafka.Types as X -- | Runs high-level kafka consumer. --- -- A callback provided is expected to call 'pollMessage' when convenient. +{-# DEPRECATED runConsumer "Use 'newConsumer'/'closeConsumer' instead" #-} runConsumer :: ConsumerProperties -> Subscription -> (KafkaConsumer -> IO (Either KafkaError a)) -- ^ A callback function to poll and handle messages @@ -58,39 +60,33 @@ 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 - kc@(KafkaConf kc') <- newConsumerConf cp + kc@(KafkaConf kc' ct) <- newConsumerConf cp tp' <- topicConf (TopicProps $ M.toList tp) _ <- setDefaultTopicConf kc tp' rdk <- bimap KafkaError Kafka <$> newRdKafkaT RdKafkaConsumer kc' case flip KafkaConsumer kc <$> rdk of Left err -> return $ Left err Right kafka -> do - redErr <- redirectCallbacksPoll kafka - case redErr of + forM_ (cpLogLevel cp) (setConsumerLogLevel kafka) + sub <- subscribe kafka ts + case sub of + Nothing -> runEventLoop kafka ct (Just $ Timeout 100) >> return (Right kafka) Just err -> closeConsumer kafka >> return (Left err) - Nothing -> 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 -> m (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))) -- ^ Left on error or timeout, right for success -pollMessage (KafkaConsumer (Kafka k) _) (Timeout ms) = - liftIO $ rdKafkaConsumerPoll k (fromIntegral ms) >>= fromMessagePtr +pollMessage c@(KafkaConsumer (Kafka k) _) (Timeout ms) = + liftIO $ pollEvents c Nothing >> rdKafkaConsumerPoll k (fromIntegral ms) >>= fromMessagePtr + -- | Commit message's offset on broker for the message's partition. commitOffsetMessage :: MonadIO m => OffsetCommit @@ -147,6 +143,20 @@ subParts [PartitionId (-1)] = SubscribedPartitionsAll subParts ps = SubscribedPartitions ps +-- | Pauses specified partitions on the current consumer. +pausePartitions :: MonadIO m => KafkaConsumer -> [(TopicName, PartitionId)] -> m KafkaError +pausePartitions (KafkaConsumer (Kafka k) _) ps = liftIO $ do + pl <- newRdKafkaTopicPartitionListT (length ps) + mapM_ (\(TopicName topicName, PartitionId partitionId) -> rdKafkaTopicPartitionListAdd pl topicName partitionId) ps + KafkaResponseError <$> rdKafkaPausePartitions k pl + +-- | Resumes specified partitions on the current consumer. +resumePartitions :: MonadIO m => KafkaConsumer -> [(TopicName, PartitionId)] -> m KafkaError +resumePartitions (KafkaConsumer (Kafka k) _) ps = liftIO $ do + pl <- newRdKafkaTopicPartitionListT (length ps) + mapM_ (\(TopicName topicName, PartitionId partitionId) -> rdKafkaTopicPartitionListAdd pl topicName partitionId) ps + KafkaResponseError <$> rdKafkaResumePartitions k pl + seek :: MonadIO m => KafkaConsumer -> Timeout -> [TopicPartition] -> m (Maybe KafkaError) seek (KafkaConsumer (Kafka k) _) (Timeout timeout) tps = liftIO $ either Just (const Nothing) <$> seekAll @@ -187,8 +197,8 @@ -- | Closes the consumer. closeConsumer :: MonadIO m => KafkaConsumer -> m (Maybe KafkaError) -closeConsumer (KafkaConsumer (Kafka k) _) = - liftIO $ (kafkaErrorToMaybe . KafkaResponseError) <$> rdKafkaConsumerClose k +closeConsumer (KafkaConsumer (Kafka k) (KafkaConf _ ct)) = + liftIO $ CToken.cancel ct >> (kafkaErrorToMaybe . KafkaResponseError) <$> rdKafkaConsumerClose k ----------------------------------------------------------------------------- newConsumerConf :: ConsumerProperties -> IO KafkaConf @@ -211,7 +221,7 @@ return $ kafkaErrorToMaybe res setDefaultTopicConf :: KafkaConf -> TopicConf -> IO () -setDefaultTopicConf (KafkaConf kc) (TopicConf tc) = +setDefaultTopicConf (KafkaConf kc _) (TopicConf tc) = rdKafkaTopicConfDup tc >>= rdKafkaConfSetDefaultTopicConf kc commitOffsets :: OffsetCommit -> KafkaConsumer -> RdKafkaTopicPartitionListTPtr -> IO (Maybe KafkaError) @@ -221,7 +231,3 @@ setConsumerLogLevel :: KafkaConsumer -> KafkaLogLevel -> IO () setConsumerLogLevel (KafkaConsumer (Kafka k) _) level = liftIO $ rdKafkaSetLogLevel k (fromEnum level) - -redirectCallbacksPoll :: KafkaConsumer -> IO (Maybe KafkaError) -redirectCallbacksPoll (KafkaConsumer (Kafka k) _) = - (kafkaErrorToMaybe . KafkaResponseError) <$> rdKafkaPollSetConsumer k diff --git a/src/Kafka/Consumer/Callbacks.hs b/src/Kafka/Consumer/Callbacks.hs --- a/src/Kafka/Consumer/Callbacks.hs +++ b/src/Kafka/Consumer/Callbacks.hs @@ -25,12 +25,12 @@ -- -- * When 'RdKafkaRespErrRevokePartitions' happens 'assign' should be called with an empty list of partitions. rebalanceCallback :: (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) -> KafkaConf -> IO () -rebalanceCallback callback (KafkaConf conf) = rdKafkaConfSetRebalanceCb conf realCb +rebalanceCallback callback (KafkaConf conf ct) = rdKafkaConfSetRebalanceCb conf realCb where realCb k err pl = do k' <- newForeignPtr_ k pls <- fromNativeTopicPartitionList' pl - callback (KafkaConsumer (Kafka k') (KafkaConf conf)) (KafkaResponseError err) pls + callback (KafkaConsumer (Kafka k') (KafkaConf conf ct)) (KafkaResponseError err) pls -- | Sets a callback that is called when rebalance is needed. -- @@ -43,10 +43,10 @@ -- with `KafkaError` == `KafkaResponseError` `RdKafkaRespErrNoOffset` which is not to be considered -- an error. offsetCommitCallback :: (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) -> KafkaConf -> IO () -offsetCommitCallback callback (KafkaConf conf) = rdKafkaConfSetOffsetCommitCb conf realCb +offsetCommitCallback callback (KafkaConf conf ct) = rdKafkaConfSetOffsetCommitCb conf realCb where realCb k err pl = do k' <- newForeignPtr_ k pls <- fromNativeTopicPartitionList' pl - callback (KafkaConsumer (Kafka k') (KafkaConf conf)) (KafkaResponseError err) pls + callback (KafkaConsumer (Kafka k') (KafkaConf conf ct)) (KafkaResponseError err) pls diff --git a/src/Kafka/Consumer/Types.hs b/src/Kafka/Consumer/Types.hs --- a/src/Kafka/Consumer/Types.hs +++ b/src/Kafka/Consumer/Types.hs @@ -1,5 +1,4 @@ -{-# LANGUAGE DeriveDataTypeable #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE DeriveDataTypeable #-} module Kafka.Consumer.Types where @@ -11,7 +10,10 @@ import Data.Typeable import Kafka.Types -data KafkaConsumer = KafkaConsumer { kcKafkaPtr :: !Kafka, kcKafkaConf :: !KafkaConf} deriving (Show) +data KafkaConsumer = KafkaConsumer + { kcKafkaPtr :: !Kafka + , kcKafkaConf :: !KafkaConf + } instance HasKafka KafkaConsumer where getKafka = kcKafkaPtr diff --git a/src/Kafka/Internal/CancellationToken.hs b/src/Kafka/Internal/CancellationToken.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Internal/CancellationToken.hs @@ -0,0 +1,17 @@ +module Kafka.Internal.CancellationToken +where + +import Data.IORef + +data CancellationStatus = Cancelled | Running deriving (Show, Eq) +newtype CancellationToken = CancellationToken (IORef CancellationStatus) + +newCancellationToken :: IO CancellationToken +newCancellationToken = CancellationToken <$> newIORef Running + +status :: CancellationToken -> IO CancellationStatus +status (CancellationToken ref) = readIORef ref + +cancel :: CancellationToken -> IO () +cancel (CancellationToken ref) = atomicWriteIORef ref Cancelled + diff --git a/src/Kafka/Internal/Setup.hs b/src/Kafka/Internal/Setup.hs --- a/src/Kafka/Internal/Setup.hs +++ b/src/Kafka/Internal/Setup.hs @@ -7,6 +7,7 @@ import Control.Monad import Foreign import Foreign.C.String +import Kafka.Internal.CancellationToken -- -- Configuration @@ -18,7 +19,7 @@ newTopicConf = TopicConf <$> newRdKafkaTopicConfT newKafkaConf :: IO KafkaConf -newKafkaConf = KafkaConf <$> newRdKafkaConfT +newKafkaConf = KafkaConf <$> newRdKafkaConfT <*> newCancellationToken kafkaConf :: KafkaProps -> IO KafkaConf kafkaConf overrides = do @@ -43,7 +44,7 @@ throw $ KafkaUnknownConfigurationKey str setKafkaConfValue :: KafkaConf -> String -> String -> IO () -setKafkaConfValue (KafkaConf confPtr) key value = +setKafkaConfValue (KafkaConf confPtr _) key value = allocaBytes nErrorBytes $ \charPtr -> do err <- rdKafkaConfSet confPtr key value charPtr (fromIntegral nErrorBytes) checkConfSetValue err charPtr diff --git a/src/Kafka/Internal/Shared.hs b/src/Kafka/Internal/Shared.hs --- a/src/Kafka/Internal/Shared.hs +++ b/src/Kafka/Internal/Shared.hs @@ -1,12 +1,30 @@ module Kafka.Internal.Shared where +import Control.Concurrent (forkIO) import Control.Exception -import qualified Data.ByteString as BS -import qualified Data.ByteString.Internal as BSI +import Control.Monad (void) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BSI import Foreign.C.Error +import Kafka.Internal.CancellationToken as CToken import Kafka.Internal.RdKafka import Kafka.Types + +runEventLoop :: HasKafka a => a -> CancellationToken -> Maybe Timeout -> IO () +runEventLoop k ct timeout = void $ forkIO go + where + go = do + token <- CToken.status ct + case token of + Running -> pollEvents k timeout >> go + Cancelled -> return () + +pollEvents :: HasKafka a => a -> Maybe Timeout -> IO () +pollEvents a tm = + let timeout = maybe 0 (\(Timeout ms) -> ms) tm + (Kafka k) = getKafka a + in void (rdKafkaPoll k timeout) word8PtrToBS :: Int -> Word8Ptr -> IO BS.ByteString word8PtrToBS len ptr = BSI.create len $ \bsptr -> diff --git a/src/Kafka/Producer.hs b/src/Kafka/Producer.hs --- a/src/Kafka/Producer.hs +++ b/src/Kafka/Producer.hs @@ -10,19 +10,21 @@ ) where -import Control.Arrow ((&&&)) +import Control.Arrow ((&&&)) import Control.Exception import Control.Monad import Control.Monad.IO.Class -import qualified Data.ByteString as BS -import qualified Data.ByteString.Internal as BSI -import Data.Function (on) -import Data.List (groupBy, sortBy) -import qualified Data.Map as M -import Data.Ord (comparing) -import Foreign hiding (void) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BSI +import Data.Function (on) +import Data.List (groupBy, sortBy) +import qualified Data.Map as M +import Data.Ord (comparing) +import Foreign hiding (void) +import Kafka.Internal.CancellationToken as CToken import Kafka.Internal.RdKafka import Kafka.Internal.Setup +import Kafka.Internal.Shared import Kafka.Producer.Convert import Kafka.Producer.ProducerProperties as X @@ -32,6 +34,7 @@ -- | Runs Kafka Producer. -- The callback provided is expected to call 'produceMessage' -- or/and 'produceMessageBatch' to send messages to Kafka. +{-# DEPRECATED runProducer "Use 'newProducer'/'closeProducer' instead" #-} runProducer :: ProducerProperties -> (KafkaProducer -> IO (Either KafkaError a)) -> IO (Either KafkaError a) @@ -50,7 +53,7 @@ -- A newly created producer must be closed with 'closeProducer' function. newProducer :: MonadIO m => ProducerProperties -> m (Either KafkaError KafkaProducer) newProducer (ProducerProperties kp tp ll cbs) = liftIO $ do - kc@(KafkaConf kc') <- kafkaConf (KafkaProps $ M.toList kp) + kc@(KafkaConf kc' ct) <- kafkaConf (KafkaProps $ M.toList kp) tc <- topicConf (TopicProps $ M.toList tp) -- set callbacks @@ -61,7 +64,8 @@ Left err -> return . Left $ KafkaError err Right kafka -> do forM_ ll (rdKafkaSetLogLevel kafka . fromEnum) - return .Right $ KafkaProducer (Kafka kafka) kc tc + let prod = KafkaProducer (Kafka kafka) kc tc + runEventLoop prod ct (Just $ Timeout 100) >> return (Right prod) -- | Sends a single message. -- Since librdkafka is backed by a queue, this function can return before messages are sent. See @@ -141,15 +145,17 @@ -- | Closes the producer. -- Will wait until the outbound queue is drained before returning the control. closeProducer :: MonadIO m => KafkaProducer -> m () -closeProducer = flushProducer +closeProducer p = + let (KafkaConf _ ct) = kpKafkaConf p + in liftIO (CToken.cancel ct) >> flushProducer p -- | Drains the outbound queue for a producer. -- This function is also called automatically when the producer is closed -- with 'closeProducer' to ensure that all queued messages make it to Kafka. flushProducer :: MonadIO m => KafkaProducer -> m () -flushProducer kp@(KafkaProducer (Kafka k) _ _) = liftIO $ do - pollEvents k 100 - l <- outboundQueueLength k +flushProducer kp = liftIO $ do + pollEvents kp (Just $ Timeout 100) + l <- outboundQueueLength (kpKafkaPtr kp) unless (l == 0) $ flushProducer kp ------------------------------------------------------------------------------------ @@ -160,8 +166,5 @@ 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 +outboundQueueLength :: Kafka -> IO Int +outboundQueueLength (Kafka k) = rdKafkaOutqLen k diff --git a/src/Kafka/Producer/Types.hs b/src/Kafka/Producer/Types.hs --- a/src/Kafka/Producer/Types.hs +++ b/src/Kafka/Producer/Types.hs @@ -12,7 +12,7 @@ { kpKafkaPtr :: !Kafka , kpKafkaConf :: !KafkaConf , kpTopicConf :: !TopicConf - } deriving (Show) + } instance HasKafka KafkaProducer where getKafka = kpKafkaPtr diff --git a/src/Kafka/Types.hs b/src/Kafka/Types.hs --- a/src/Kafka/Types.hs +++ b/src/Kafka/Types.hs @@ -6,6 +6,7 @@ import Control.Exception import Data.Int import Data.Typeable +import Kafka.Internal.CancellationToken import Kafka.Internal.RdKafka class HasKafka a where @@ -19,7 +20,7 @@ deriving (Show, Eq, Ord, Read) newtype Kafka = Kafka RdKafkaTPtr deriving Show -newtype KafkaConf = KafkaConf RdKafkaConfTPtr deriving Show +data KafkaConf = KafkaConf RdKafkaConfTPtr CancellationToken newtype TopicConf = TopicConf RdKafkaTopicConfTPtr deriving Show newtype PartitionId = PartitionId Int deriving (Show, Eq, Read, Ord, Enum)