Remove unnecessary poll for events

Author: AlexeyRaga

src/Kafka/Consumer.hs

@@ -19,9 +19,9 @@ module Kafka.Consumer
 where
 
 import           Control.Arrow                    (left, (&&&))
-import           Control.Concurrent               (forkIO, rtsSupportsBoundThreads)
+import           Control.Concurrent               (forkIO, rtsSupportsBoundThreads, withMVar)
 import           Control.Exception                (bracket)
-import           Control.Monad                    (forM_, void, when)
+import           Control.Monad                    (forM_, mapM_, void, when)
 import           Control.Monad.IO.Class           (MonadIO (liftIO))
 import           Control.Monad.Trans.Except       (ExceptT (ExceptT), runExceptT)
 import           Data.Bifunctor                   (bimap, first)
@@ -38,7 +38,7 @@ import           Kafka.Consumer.Convert           (fromMessagePtr, fromNativeTop
 import           Kafka.Consumer.Types             (KafkaConsumer (..))
 import           Kafka.Internal.CancellationToken as CToken
 import           Kafka.Internal.RdKafka           (RdKafkaRespErrT (..), RdKafkaTopicPartitionListTPtr, RdKafkaTypeT (..), newRdKafkaT, newRdKafkaTopicPartitionListT, newRdKafkaTopicT, rdKafkaAssignment, rdKafkaCommit, rdKafkaCommitted, rdKafkaConfSetDefaultTopicConf, rdKafkaConsumeBatchQueue, rdKafkaConsumeQueue, rdKafkaConsumerClose, rdKafkaConsumerPoll, rdKafkaOffsetsStore, rdKafkaPausePartitions, rdKafkaPollSetConsumer, rdKafkaPosition, rdKafkaQueueDestroy, rdKafkaQueueNew, rdKafkaResumePartitions, rdKafkaSeek, rdKafkaSetLogLevel, rdKafkaSubscribe, rdKafkaSubscription, rdKafkaTopicConfDup, rdKafkaTopicPartitionListAdd)
-import           Kafka.Internal.Setup             (Kafka (..), KafkaConf (..), KafkaProps (..), TopicConf (..), TopicProps (..), getRdKafka, kafkaConf, topicConf)
+import           Kafka.Internal.Setup             (Kafka (..), KafkaConf (..), KafkaProps (..), TopicConf (..), TopicProps (..), getKafkaConf, getRdKafka, kafkaConf, topicConf)
 import           Kafka.Internal.Shared            (kafkaErrorToMaybe, maybeToLeft, rdKafkaErrorToEither)
 
 import Kafka.Consumer.ConsumerProperties as X
@@ -69,17 +69,17 @@ newConsumer :: MonadIO m
             -> Subscription
             -> m (Either KafkaError KafkaConsumer)
 newConsumer props (Subscription ts tp) = liftIO $ do
-  let cp = case cpUserPolls props of
-            CallbackModeAsync -> setCallback (rebalanceCallback (\_ _ -> return ())) <> props
-            CallbackModeSync  -> props
-  kc@(KafkaConf kc' qref ct) <- newConsumerConf cp
+  let cp = case cpCallbackPollMode props of
+            CallbackPollModeAsync -> setCallback (rebalanceCallback (\_ _ -> return ())) <> props
+            CallbackPollModeSync  -> props
+  kc@(KafkaConf kc' qref _ ct) <- newConsumerConf cp
   tp' <- topicConf (TopicProps tp)
   _   <- setDefaultTopicConf kc tp'
   rdk <- newRdKafkaT RdKafkaConsumer kc'
   case rdk of
     Left err   -> return . Left $ KafkaError err
     Right rdk' -> do
-      when (cpUserPolls props == CallbackModeAsync) $ do
+      when (cpCallbackPollMode props == CallbackPollModeAsync) $ do
         msgq <- rdKafkaQueueNew rdk'
         writeIORef qref (Just msgq)
       let kafka = KafkaConsumer (Kafka rdk') kc
@@ -90,21 +90,19 @@ newConsumer props (Subscription ts tp) = liftIO $ do
           forM_ (cpLogLevel cp) (setConsumerLogLevel kafka)
           sub <- subscribe kafka ts
           case sub of
-            Nothing  -> (when (cpUserPolls props == CallbackModeAsync) $
+            Nothing  -> (when (cpCallbackPollMode props == CallbackPollModeAsync) $
               runConsumerLoop kafka ct (Just $ Timeout 100)) >> return (Right kafka)
             Just err -> closeConsumer kafka >> return (Left err)
 
 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 c@(KafkaConsumer _ (KafkaConf _ qr _)) (Timeout ms) = liftIO $ do
+pollMessage c@(KafkaConsumer _ (KafkaConf _ qr _ _)) (Timeout ms) = liftIO $ do
   mbq <- readIORef qr
   case mbq of
     Nothing -> rdKafkaConsumerPoll (getRdKafka c) ms >>= fromMessagePtr
-    Just q -> do
-      pollConsumerEvents c Nothing
-      rdKafkaConsumeQueue q (fromIntegral ms) >>= fromMessagePtr
+    Just q  -> rdKafkaConsumeQueue q (fromIntegral ms) >>= fromMessagePtr
 
 -- | Polls up to BatchSize messages.
 -- Unlike 'pollMessage' this function does not return usual "timeout" errors.
@@ -116,7 +114,7 @@ pollMessageBatch :: MonadIO m
                  -> Timeout
                  -> BatchSize
                  -> m [Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString))]
-pollMessageBatch c@(KafkaConsumer _ (KafkaConf _ qr _)) (Timeout ms) (BatchSize b) = liftIO $ do
+pollMessageBatch c@(KafkaConsumer _ (KafkaConf _ qr _ _)) (Timeout ms) (BatchSize b) = liftIO $ do
   pollConsumerEvents c Nothing
   mbq <- readIORef qr
   case mbq of
@@ -206,8 +204,7 @@ seek (KafkaConsumer (Kafka k) _) (Timeout timeout) tps = liftIO $
   where
     seekAll = runExceptT $ do
       tr <- traverse (ExceptT . topicPair) tps
-      void $ traverse (\(kt, p, o) -> ExceptT (rdSeek kt p o)) tr
-      return ()
+      mapM_ (\(kt, p, o) -> ExceptT (rdSeek kt p o)) tr
 
     rdSeek kt (PartitionId p) o =
       rdKafkaErrorToEither <$> rdKafkaSeek kt (fromIntegral p) (offsetToInt64 o) timeout
@@ -252,13 +249,14 @@ position (KafkaConsumer (Kafka k) _) tps = liftIO $ do
 -- when polling for events on each 'pollMessage' is not
 -- frequent enough.
 pollConsumerEvents :: KafkaConsumer -> Maybe Timeout -> IO ()
-pollConsumerEvents k timeout =
+pollConsumerEvents k timeout = do
   let (Timeout tm) = fromMaybe (Timeout 0) timeout
-  in void $ rdKafkaConsumerPoll (getRdKafka k) tm
+  let (KafkaConf _ _ rv _) = getKafkaConf k
+  withMVar rv . const $ void $ rdKafkaConsumerPoll (getRdKafka k) tm
 
 -- | Closes the consumer.
 closeConsumer :: MonadIO m => KafkaConsumer -> m (Maybe KafkaError)
-closeConsumer (KafkaConsumer (Kafka k) (KafkaConf _ qr ct)) = liftIO $ do
+closeConsumer (KafkaConsumer (Kafka k) (KafkaConf _ qr _ ct)) = liftIO $ do
   CToken.cancel ct
   mbq <- readIORef qr
   void $ traverse rdKafkaQueueDestroy mbq
@@ -285,7 +283,7 @@ subscribe (KafkaConsumer (Kafka k) _) ts = do
     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)

src/Kafka/Consumer/Callbacks.hs

@@ -12,30 +12,17 @@ import Data.Monoid            ((<>))
 import Foreign.ForeignPtr     (newForeignPtr_)
 import Foreign.Ptr            (nullPtr)
 import Kafka.Callbacks        as X
-import Kafka.Consumer.Convert (fromNativeTopicPartitionList', fromNativeTopicPartitionList'', toNativeTopicPartitionList)
+import Kafka.Consumer.Convert (fromNativeTopicPartitionList', fromNativeTopicPartitionList'')
 import Kafka.Consumer.Types   (KafkaConsumer (..), RebalanceEvent (..), TopicPartition (..))
 import Kafka.Internal.RdKafka
 import Kafka.Internal.Setup   (HasKafka (..), HasKafkaConf (..), Kafka (..), KafkaConf (..), getRdMsgQueue)
-import Kafka.Internal.Shared  (kafkaErrorToMaybe)
 import Kafka.Types            (KafkaError (..), PartitionId (..), TopicName (..))
 
 import qualified Data.Text as Text
 
 -- | Sets a callback that is called when rebalance is needed.
---
--- Callback implementations suppose to watch for 'KafkaResponseError' 'RdKafkaRespErrAssignPartitions' and
--- for 'KafkaResponseError' 'RdKafkaRespErrRevokePartitions'. Other error codes are not expected and would indicate
--- something really bad happening in a system, or bugs in @librdkafka@ itself.
---
--- A callback is expected to call 'assign' according to the error code it receives.
---
---     * When 'RdKafkaRespErrAssignPartitions' happens 'assign' should be called with all the partitions it was called with.
---       It is OK to alter partitions offsets before calling 'assign'.
---
---     * When 'RdKafkaRespErrRevokePartitions' happens 'assign' should be called with an empty list of partitions.
--- rebalanceCallback :: (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) -> KafkaConf -> IO ()
 rebalanceCallback :: (KafkaConsumer -> RebalanceEvent -> IO ()) -> KafkaConf -> IO ()
-rebalanceCallback callback kc@(KafkaConf conf _ _) = rdKafkaConfSetRebalanceCb conf realCb
+rebalanceCallback callback kc@(KafkaConf conf _ _ _) = rdKafkaConfSetRebalanceCb conf realCb
   where
     realCb k err pl = do
       k' <- newForeignPtr_ k
@@ -47,13 +34,11 @@ rebalanceCallback callback kc@(KafkaConf conf _ _) = rdKafkaConfSetRebalanceCb c
 -- The results of automatic or manual offset commits will be scheduled
 -- for this callback and is served by `pollMessage`.
 --
--- A callback is expected to call 'assign' according to the error code it receives.
---
 -- If no partitions had valid offsets to commit this callback will be called
 -- with `KafkaError` == `KafkaResponseError` `RdKafkaRespErrNoOffset` which is not to be considered
 -- an error.
 offsetCommitCallback :: (KafkaConsumer -> KafkaError -> [TopicPartition] -> IO ()) -> KafkaConf -> IO ()
-offsetCommitCallback callback kc@(KafkaConf conf _ _) = rdKafkaConfSetOffsetCommitCb conf realCb
+offsetCommitCallback callback kc@(KafkaConf conf _ _ _) = rdKafkaConfSetOffsetCommitCb conf realCb
   where
     realCb k err pl = do
       k' <- newForeignPtr_ k
@@ -75,46 +60,30 @@ setRebalanceCallback :: (KafkaConsumer -> RebalanceEvent -> IO ())
 setRebalanceCallback f k e pls = do
   ps <- fromNativeTopicPartitionList'' pls
   let assignment = (tpTopicName &&& tpPartition) <$> ps
+  let (Kafka kptr) = getKafka k
+
   case e of
     KafkaResponseError RdKafkaRespErrAssignPartitions -> do
-        mbq <- getRdMsgQueue $ getKafkaConf k
         f k (RebalanceBeforeAssign assignment)
+        void $ rdKafkaAssign kptr pls
+
+        mbq <- getRdMsgQueue $ getKafkaConf k
         case mbq of
           Nothing -> pure ()
           Just mq -> do
-            let (Kafka kptr) = getKafka k
-            -- Magnus Edenhill:
-            -- If you redirect after assign() it means some messages may be forwarded to the single consumer queue,
-            -- so either do it before assign() or do: assign(); pause(); redirect; resume()
-            void $ rdKafkaAssign kptr pls
+            {- Magnus Edenhill:
+                If you redirect after assign() it means some messages may be forwarded to the single consumer queue,
+                so either do it before assign() or do: assign(); pause(); redirect; resume()
+            -}
             void $ rdKafkaPausePartitions kptr pls
             forM_ ps (\tp -> redirectPartitionQueue (getKafka k) (tpTopicName tp) (tpPartition tp) mq)
             void $ rdKafkaResumePartitions kptr pls
+
         f k (RebalanceAssign assignment)
+
     KafkaResponseError RdKafkaRespErrRevokePartitions -> do
         f k (RebalanceBeforeRevoke assignment)
-        void $ assign k []
+        void $ newForeignPtr_ nullPtr >>= rdKafkaAssign kptr
+        -- void $ assign k []
         f k (RebalanceRevoke assignment)
     x -> error $ "Rebalance: UNKNOWN response: " <> show x
-
-
--- | Assigns specified partitions to a current consumer.
--- Assigning an empty list means unassigning from all partitions that are currently assigned.
-assign :: KafkaConsumer -> [TopicPartition] -> IO (Maybe KafkaError)
-assign (KafkaConsumer (Kafka k) _) ps =
-    let pl = if null ps
-                then newForeignPtr_ nullPtr
-                else toNativeTopicPartitionList ps
-        er = KafkaResponseError <$> (pl >>= rdKafkaAssign k)
-    in kafkaErrorToMaybe <$> er
-
--- -- | Assigns specified partitions to a current consumer.
--- -- Assigning an empty list means unassigning from all partitions that are currently assigned.
--- assign' :: KafkaConsumer -> RdKafkaTopicPartitionListTPtr -> IO (Maybe KafkaError)
--- assign' (KafkaConsumer (Kafka k) _) pls = do
-
-
-
---   where
---     asExcept f = ExceptT $ rdKafkaErrorToEither <$> f
-

src/Kafka/Consumer/ConsumerProperties.hs

@@ -2,7 +2,7 @@
 
 module Kafka.Consumer.ConsumerProperties
 ( ConsumerProperties(..)
-, CallbackMode(..)
+, CallbackPollMode(..)
 , brokersList
 , autoCommit
 , noAutoCommit
@@ -30,18 +30,18 @@ import           Data.Text            (Text)
 import qualified Data.Text            as Text
 import           Kafka.Consumer.Types (ConsumerGroupId (..))
 import           Kafka.Internal.Setup (KafkaConf (..))
-import           Kafka.Types          (BrokerAddress (..), ClientId (..), KafkaCompressionCodec (..), KafkaDebug (..), KafkaLogLevel (..), kafkaCompressionCodecToText, kafkaDebugToText, Millis(..))
+import           Kafka.Types          (BrokerAddress (..), ClientId (..), KafkaCompressionCodec (..), KafkaDebug (..), KafkaLogLevel (..), Millis (..), kafkaCompressionCodecToText, kafkaDebugToText)
 
 import Kafka.Consumer.Callbacks as X
 
-data CallbackMode = CallbackModeSync | CallbackModeAsync deriving (Show, Eq)
+data CallbackPollMode = CallbackPollModeSync | CallbackPollModeAsync deriving (Show, Eq)
 
 -- | Properties to create 'KafkaConsumer'.
 data ConsumerProperties = ConsumerProperties
-  { cpProps     :: Map Text Text
-  , cpLogLevel  :: Maybe KafkaLogLevel
-  , cpCallbacks :: [KafkaConf -> IO ()]
-  , cpUserPolls :: CallbackMode
+  { cpProps            :: Map Text Text
+  , cpLogLevel         :: Maybe KafkaLogLevel
+  , cpCallbacks        :: [KafkaConf -> IO ()]
+  , cpCallbackPollMode :: CallbackPollMode
   }
 
 instance Sem.Semigroup ConsumerProperties where
@@ -52,10 +52,10 @@ instance Sem.Semigroup ConsumerProperties where
 -- | /Right biased/ so we prefer newer properties over older ones.
 instance Monoid ConsumerProperties where
   mempty = ConsumerProperties
-    { cpProps          = M.empty
-    , cpLogLevel       = Nothing
-    , cpCallbacks      = []
-    , cpUserPolls      = CallbackModeAsync
+    { cpProps             = M.empty
+    , cpLogLevel          = Nothing
+    , cpCallbacks         = []
+    , cpCallbackPollMode  = CallbackPollModeAsync
     }
   {-# INLINE mempty #-}
   mappend = (Sem.<>)
@@ -140,14 +140,14 @@ queuedMaxMessagesKBytes kBytes =
 
 -- | Sets the callback poll mode.
 --
--- The default 'CallbackModeAsync' mode handles polling rebalance
+-- The default 'CallbackPollModeAsync' mode handles polling rebalance
 -- and keep alive events for you
 -- in a background thread.
 --
--- With 'CalalcacModeSync' the user will poll the consumer
+-- With 'CallbacPollModeSync' the user will poll the consumer
 -- frequently to handle new messages as well as rebalance and keep alive events.
--- 'CalalcacModeSync' lets you can simplify
+-- 'CallbacPollModeSync' lets you can simplify
 -- hw-kafka-client's footprint and have full control over when polling
 -- happens at the cost of having to manage this yourself.
-callbackPollMode :: CallbackMode -> ConsumerProperties
-callbackPollMode mode = mempty { cpUserPolls = mode }
+callbackPollMode :: CallbackPollMode -> ConsumerProperties
+callbackPollMode mode = mempty { cpCallbackPollMode = mode }

src/Kafka/Consumer/Subscription.hs

@@ -8,15 +8,14 @@ module Kafka.Consumer.Subscription
 )
 where
 
-import qualified Data.Text            as Text
-import           Data.Text            (Text)
 import           Data.Map             (Map)
 import qualified Data.Map             as M
 import           Data.Semigroup       as Sem
-import           Kafka.Consumer.Types (OffsetReset(..))
-import           Kafka.Types          (TopicName(..), Millis(..))
 import           Data.Set             (Set)
 import qualified Data.Set             as Set
+import           Data.Text            (Text)
+import           Kafka.Consumer.Types (OffsetReset (..))
+import           Kafka.Types          (TopicName (..))
 
 data Subscription = Subscription (Set TopicName) (Map Text Text)
 
@@ -41,7 +40,7 @@ offsetReset o =
   let o' = case o of
              Earliest -> "earliest"
              Latest   -> "latest"
-   in Subscription (Set.empty) (M.fromList [("auto.offset.reset", o')])
+   in Subscription Set.empty (M.fromList [("auto.offset.reset", o')])
 
 extraSubscriptionProps :: Map Text Text -> Subscription
-extraSubscriptionProps = Subscription (Set.empty)
+extraSubscriptionProps = Subscription Set.empty