Remove batch producer & add headers to the record type

Author: Levashov, Mykyta

example/ProducerExample.hs

@@ -31,6 +31,7 @@ mkMessage k v = ProducerRecord
                   , prPartition = UnassignedPartition
                   , prKey = k
                   , prValue = v
+                  , prHeaders = mempty
                   }
 
 -- Run an example
@@ -62,13 +63,7 @@ sendMessages prod = do
   msg3 <- getLine
   err3 <- produceMessage prod (mkMessage (Just "key3") (Just $ pack msg3))
 
-  err4 <- produceMessageWithHeaders prod (headersFromList [("fancy", "header")]) (mkMessage (Just "key4") (Just $ pack msg3))
-
-  -- errs <- produceMessageBatch prod
-  --           [ mkMessage (Just "b-1") (Just "batch-1")
-  --           , mkMessage (Just "b-2") (Just "batch-2")
-  --           , mkMessage Nothing      (Just "batch-3")
-  --           ]
+  err4 <- produceMessage prod ((mkMessage (Just "key4") (Just $ pack msg3)) { prHeaders = headersFromList [("fancy", "header")]})
 
   -- forM_ errs (print . snd)
 

nix/sources.json

@@ -1,26 +1,26 @@
 {
-  "niv": {
-    "branch": "master",
-    "description": "Easy dependency management for Nix projects",
-    "homepage": "https://github.com/nmattia/niv",
-    "owner": "nmattia",
-    "repo": "niv",
-    "rev": "e0ca65c81a2d7a4d82a189f1e23a48d59ad42070",
-    "sha256": "1pq9nh1d8nn3xvbdny8fafzw87mj7gsmp6pxkdl65w2g18rmcmzx",
-    "type": "tarball",
-    "url": "https://github.com/nmattia/niv/archive/e0ca65c81a2d7a4d82a189f1e23a48d59ad42070.tar.gz",
-    "url_template": "https://github.com/<owner>/<repo>/archive/<rev>.tar.gz"
-  },
-  "nixpkgs": {
-    "branch": "nixos-21.05",
-    "description": "Nix Packages collection",
-    "homepage": "",
-    "owner": "NixOS",
-    "repo": "nixpkgs",
-    "rev": "d5aadbefd650cb0a05ba9c788a26327afce2396c",
-    "sha256": "1h9nvs7nknczpz15c41p3irbykc5qzwr8wlcpzxndc25mnqi257w",
-    "type": "tarball",
-    "url": "https://github.com/NixOS/nixpkgs/archive/d5aadbefd650cb0a05ba9c788a26327afce2396c.tar.gz",
-    "url_template": "https://github.com/<owner>/<repo>/archive/<rev>.tar.gz"
-  }
+    "niv": {
+        "branch": "master",
+        "description": "Easy dependency management for Nix projects",
+        "homepage": "https://github.com/nmattia/niv",
+        "owner": "nmattia",
+        "repo": "niv",
+        "rev": "65a61b147f307d24bfd0a5cd56ce7d7b7cc61d2e",
+        "sha256": "17mirpsx5wyw262fpsd6n6m47jcgw8k2bwcp1iwdnrlzy4dhcgqh",
+        "type": "tarball",
+        "url": "https://github.com/nmattia/niv/archive/65a61b147f307d24bfd0a5cd56ce7d7b7cc61d2e.tar.gz",
+        "url_template": "https://github.com/<owner>/<repo>/archive/<rev>.tar.gz"
+    },
+    "nixpkgs": {
+        "branch": "nixos-21.05",
+        "description": "Nix Packages collection",
+        "homepage": "",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ce7a1190a0fa4ba3465b5f5471b08567060ca14c",
+        "sha256": "1zr1s9gp0h5g4arlba1bpb9yqfaaby5195ydm6a2psaxhm748li9",
+        "type": "tarball",
+        "url": "https://github.com/NixOS/nixpkgs/archive/ce7a1190a0fa4ba3465b5f5471b08567060ca14c.tar.gz",
+        "url_template": "https://github.com/<owner>/<repo>/archive/<rev>.tar.gz"
+    }
 }

src/Kafka/Consumer/Convert.hs

@@ -159,16 +159,16 @@ fromMessagePtr ptr =
         s <- peek realPtr
         msg <- if err'RdKafkaMessageT s /= RdKafkaRespErrNoError
                 then return . Left . KafkaResponseError $ err'RdKafkaMessageT s
-                else Right <$> mkRecord s
+                else Right <$> mkRecord s realPtr
         rdKafkaMessageDestroy realPtr
         return msg
     where
-        mkRecord msg = do
+        mkRecord msg rptr = do
             topic     <- readTopic msg
             key       <- readKey msg
             payload   <- readPayload msg
             timestamp <- readTimestamp ptr
-            headers   <- fromRight mempty <$> readHeaders ptr
+            headers   <- fromRight mempty <$> readHeaders rptr
             return ConsumerRecord
                 { crTopic     = TopicName topic
                 , crPartition = PartitionId $ partition'RdKafkaMessageT msg

src/Kafka/Internal/RdKafka.chs

@@ -993,7 +993,7 @@ data RdKafkaHeadersT
     {`RdKafkaHeadersTPtr', cIntConv `CSize', castPtr `Ptr CString', castPtr `Ptr Word8Ptr', `CSizePtr'} -> `RdKafkaRespErrT' cIntToEnum #}
 
 {#fun rd_kafka_message_headers as ^
-    {`RdKafkaMessageTPtr', alloca- `RdKafkaHeadersTPtr' peekPtr*} -> `RdKafkaRespErrT' cIntToEnum #}
+    {castPtr `Ptr RdKafkaMessageT', alloca- `RdKafkaHeadersTPtr' peekPtr*} -> `RdKafkaRespErrT' cIntToEnum #}
 
 --- Produceva api
 

src/Kafka/Internal/Shared.hs

@@ -106,13 +106,13 @@ readTimestamp msg =
                RdKafkaTimestampNotAvailable  -> NoTimestamp
 
 
-readHeaders :: RdKafkaMessageTPtr -> IO (Either RdKafkaRespErrT Headers)
+readHeaders :: Ptr RdKafkaMessageT -> IO (Either RdKafkaRespErrT Headers)
 readHeaders msg = do
     (err, headersPtr) <- rdKafkaMessageHeaders msg
     case err of
         RdKafkaRespErrNoent -> return $ Right mempty
         RdKafkaRespErrNoError -> fmap headersFromList <$> extractHeaders headersPtr
-        e -> return $ Left e
+        e -> return . Left $ e
     where extractHeaders ptHeaders =
             alloca $ \nptr ->
                 alloca $ \vptr ->
@@ -126,7 +126,7 @@ readHeaders msg = do
                                     hn <- BS.packCString cstr
                                     hv <- word8PtrToBS (fromIntegral csize) wptr
                                     go ((hn, hv) : acc) (idx + 1)
-                                _ -> error "Unexpected error code"
+                                _ -> error "Unexpected error code while extracting headers"
                         in go [] 0
 
 readBS :: (t -> Int) -> (t -> Ptr Word8) -> t -> IO (Maybe BS.ByteString)

src/Kafka/Producer.hs

@@ -58,35 +58,29 @@ module Kafka.Producer
 , module X
 , runProducer
 , newProducer
-, produceMessage, produceMessageBatch, produceMessageWithHeaders
-, produceMessage', produceMessageWithHeaders'
+, produceMessage
+, produceMessage'
 , flushProducer
 , closeProducer
 , RdKafkaRespErrT (..)
 )
 where
 
-import           Control.Arrow            ((&&&))
 import           Control.Exception        (bracket)
-import           Control.Monad            (forM, forM_, (<=<))
+import           Control.Monad            (forM_)
 import           Control.Monad.IO.Class   (MonadIO (liftIO))
 import qualified Data.ByteString          as BS
 import qualified Data.ByteString.Internal as BSI
-import           Data.Function            (on)
-import           Data.List                (groupBy, sortBy)
-import           Data.Ord                 (comparing)
 import qualified Data.Text                as Text
 import           Foreign.C.String         (withCString)
-import           Foreign.ForeignPtr       (newForeignPtr_, withForeignPtr)
-import           Foreign.Marshal.Array    (withArrayLen)
+import           Foreign.ForeignPtr       (withForeignPtr)
 import           Foreign.Marshal.Utils    (withMany)
 import           Foreign.Ptr              (Ptr, nullPtr, plusPtr)
-import           Foreign.Storable         (Storable (..))
 import           Foreign.StablePtr        (newStablePtr, castStablePtrToPtr)
-import           Kafka.Internal.RdKafka   (RdKafkaMessageT (..), RdKafkaRespErrT (..), RdKafkaTypeT (..), RdKafkaVuT(..), destroyUnmanagedRdKafkaTopic, newRdKafkaT, newUnmanagedRdKafkaTopicT, rdKafkaErrorCode, rdKafkaErrorDestroy, rdKafkaOutqLen, rdKafkaProduceBatch, rdKafkaMessageProduceVa, rdKafkaSetLogLevel)
-import           Kafka.Internal.Setup     (Kafka (..), KafkaConf (..), KafkaProps (..), TopicConf (..), TopicProps (..), kafkaConf, topicConf, Callback(..))
+import           Kafka.Internal.RdKafka   (RdKafkaRespErrT (..), RdKafkaTypeT (..), RdKafkaVuT(..), newRdKafkaT, rdKafkaErrorCode, rdKafkaErrorDestroy, rdKafkaOutqLen, rdKafkaMessageProduceVa, rdKafkaSetLogLevel)
+import           Kafka.Internal.Setup     (Kafka (..), KafkaConf (..), KafkaProps (..), TopicProps (..), kafkaConf, topicConf, Callback(..))
 import           Kafka.Internal.Shared    (pollEvents)
-import           Kafka.Producer.Convert   (copyMsgFlags, handleProduceErrT, producePartitionCInt, producePartitionInt)
+import           Kafka.Producer.Convert   (copyMsgFlags, handleProduceErrT, producePartitionCInt)
 import           Kafka.Producer.Types     (KafkaProducer (..))
 
 import Kafka.Producer.ProducerProperties as X
@@ -95,7 +89,7 @@ import Kafka.Types                       as X
 
 -- | Runs Kafka Producer.
 -- The callback provided is expected to call 'produceMessage'
--- or/and 'produceMessageBatch' to send messages to Kafka.
+-- to send messages to Kafka.
 {-# DEPRECATED runProducer "Use 'newProducer'/'closeProducer' instead" #-}
 runProducer :: ProducerProperties
             -> (KafkaProducer -> IO (Either KafkaError a))
@@ -146,42 +140,16 @@ produceMessage kp m = produceMessage' kp m (pure . mempty) >>= adjustRes
       Right () -> pure Nothing
       Left (ImmediateError err) -> pure (Just err)
 
--- | Sends a single message with a registered callback and headers.
-produceMessageWithHeaders :: MonadIO m
-                => KafkaProducer
-                -> Headers
-                -> ProducerRecord
-                -> m (Maybe KafkaError)
-produceMessageWithHeaders kp headers msg = produceMessageWithHeaders' kp headers msg (pure . mempty) >>= adjustRes
-  where
-    adjustRes = \case
-      Right () -> pure Nothing
-      Left (ImmediateError err) -> pure (Just err)
-
 -- | Sends a single message with a registered callback.
 --
 --   The callback can be a long running process, as it is forked by the thread
 --   that handles the delivery reports.
---
 produceMessage' :: MonadIO m
                 => KafkaProducer
                 -> ProducerRecord
                 -> (DeliveryReport -> IO ())
                 -> m (Either ImmediateError ())
-produceMessage' kp = produceMessageWithHeaders' kp mempty
-
--- | Sends a single message with a registered callback and headers.
---
---   The callback can be a long running process, as it is forked by the thread
---   that handles the delivery reports.
---
-produceMessageWithHeaders' :: MonadIO m
-                => KafkaProducer
-                -> Headers
-                -> ProducerRecord
-                -> (DeliveryReport -> IO ())
-                -> m (Either ImmediateError ())
-produceMessageWithHeaders' kp@(KafkaProducer (Kafka k) _ _) headers msg cb = liftIO $
+produceMessage' kp@(KafkaProducer (Kafka k) _ _) msg cb = liftIO $
   fireCallbacks >> produceIt
   where
     fireCallbacks =
@@ -190,7 +158,7 @@ produceMessageWithHeaders' kp@(KafkaProducer (Kafka k) _ _) headers msg cb = lif
     produceIt =
       withBS (prValue msg) $ \payloadPtr payloadLength ->
         withBS (prKey msg) $ \keyPtr keyLength ->
-          withHeaders headers $ \hdrs ->
+          withHeaders (prHeaders msg) $ \hdrs ->
             withCString (Text.unpack . unTopicName . prTopic $ msg) $ \topicName -> do
               callbackPtr <- newStablePtr cb
               let opts = [
@@ -208,59 +176,6 @@ produceMessageWithHeaders' kp@(KafkaProducer (Kafka k) _ _) headers msg cb = lif
                 Just err -> Left . ImmediateError $ err
                 Nothing -> Right ()
 
--- | Sends a batch of messages.
--- Returns a list of messages which it was unable to send with corresponding errors.
--- Since librdkafka is backed by a queue, this function can return before messages are sent. See
--- 'flushProducer' to wait for queue to empty.
-produceMessageBatch :: MonadIO m
-                    => KafkaProducer
-                    -> [ProducerRecord]
-                    -> m [(ProducerRecord, KafkaError)]
-                    -- ^ An empty list when the operation is successful,
-                    -- otherwise a list of "failed" messages with corresponsing errors.
-produceMessageBatch kp@(KafkaProducer (Kafka k) _ (TopicConf tc)) messages = liftIO $ do
-  pollEvents kp (Just $ Timeout 0) -- fire callbacks if any exist (handle delivery reports)
-  concat <$> forM (mkBatches messages) sendBatch
-  where
-    mkSortKey = prTopic &&& prPartition
-    mkBatches = groupBy ((==) `on` mkSortKey) . sortBy (comparing mkSortKey)
-
-    mkTopic (TopicName tn) = newUnmanagedRdKafkaTopicT k (Text.unpack tn) (Just tc)
-
-    clTopic = either (return . const ()) destroyUnmanagedRdKafkaTopic
-
-    sendBatch []    = return []
-    sendBatch batch = bracket (mkTopic $ prTopic (head batch)) clTopic (withTopic batch)
-
-    withTopic ms (Left err) = return $ (, KafkaError (Text.pack err)) <$> ms
-    withTopic ms (Right t) = do
-      let (partInt, partCInt) = (producePartitionInt &&& producePartitionCInt) $ prPartition (head ms)
-      withForeignPtr t $ \topicPtr -> do
-        nativeMs <- forM ms (toNativeMessage topicPtr partInt)
-        withArrayLen nativeMs $ \len batchPtr -> do
-          batchPtrF <- newForeignPtr_ batchPtr
-          numRet    <- rdKafkaProduceBatch t partCInt copyMsgFlags batchPtrF len
-          if numRet == len then return []
-          else do
-            errs <- mapM (return . err'RdKafkaMessageT <=< peekElemOff batchPtr)
-                         [0..(fromIntegral $ len - 1)]
-            return [(m, KafkaResponseError e) | (m, e) <- zip messages errs, e /= RdKafkaRespErrNoError]
-
-    toNativeMessage t p m =
-      withBS (prValue m) $ \payloadPtr payloadLength ->
-        withBS (prKey m) $ \keyPtr keyLength ->
-          return RdKafkaMessageT
-            { err'RdKafkaMessageT       = RdKafkaRespErrNoError
-            , topic'RdKafkaMessageT     = t
-            , partition'RdKafkaMessageT = p
-            , len'RdKafkaMessageT       = payloadLength
-            , payload'RdKafkaMessageT   = payloadPtr
-            , offset'RdKafkaMessageT    = 0
-            , keyLen'RdKafkaMessageT    = keyLength
-            , key'RdKafkaMessageT       = keyPtr
-            , opaque'RdKafkaMessageT    = nullPtr
-            }
-
 -- | Closes the producer.
 -- Will wait until the outbound queue is drained before returning the control.
 closeProducer :: MonadIO m => KafkaProducer -> m ()

src/Kafka/Producer/Callbacks.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE LambdaCase #-}
 module Kafka.Producer.Callbacks
 ( deliveryCallback
 , module X
@@ -15,10 +16,11 @@ import           Foreign.StablePtr      (castPtrToStablePtr, deRefStablePtr, fre
 import           Kafka.Callbacks        as X
 import           Kafka.Consumer.Types   (Offset(..))
 import           Kafka.Internal.RdKafka (RdKafkaMessageT(..), RdKafkaRespErrT(..), rdKafkaConfSetDrMsgCb)
-import           Kafka.Internal.Setup   (KafkaConf(..), getRdKafkaConf, Callback(..))
-import           Kafka.Internal.Shared  (kafkaRespErr, readTopic, readKey, readPayload)
+import           Kafka.Internal.Setup   (getRdKafkaConf, Callback(..))
+import           Kafka.Internal.Shared  (kafkaRespErr, readTopic, readKey, readPayload, readHeaders)
 import           Kafka.Producer.Types   (ProducerRecord(..), DeliveryReport(..), ProducePartition(..))
 import           Kafka.Types            (KafkaError(..), TopicName(..))
+import Data.Either (fromRight)
 
 -- | Sets the callback for delivery reports.
 --
@@ -36,10 +38,12 @@ deliveryCallback callback = Callback $ \kc -> rdKafkaConfSetDrMsgCb (getRdKafkaC
         then getErrno >>= (callback . NoMessageError . kafkaRespErr)
         else do
           s <- peek mptr
+          prodRec <- mkProdRec mptr
           let cbPtr = opaque'RdKafkaMessageT s
-          if err'RdKafkaMessageT s /= RdKafkaRespErrNoError
-            then mkErrorReport s   >>= callbacks cbPtr
-            else mkSuccessReport s >>= callbacks cbPtr
+          callbacks cbPtr $ 
+            if err'RdKafkaMessageT s /= RdKafkaRespErrNoError
+              then mkErrorReport s prodRec  
+              else mkSuccessReport s prodRec
 
     callbacks cbPtr rep = do
       callback rep
@@ -51,24 +55,23 @@ deliveryCallback callback = Callback $ \kc -> rdKafkaConfSetDrMsgCb (getRdKafkaC
         -- blocking here would block librdkafka from continuing its execution
         void . forkIO $ msgCb rep
 
-mkErrorReport :: RdKafkaMessageT -> IO DeliveryReport
-mkErrorReport msg = do
-  prodRec <- mkProdRec msg
-  pure $ DeliveryFailure prodRec (KafkaResponseError (err'RdKafkaMessageT msg))
+mkErrorReport :: RdKafkaMessageT -> ProducerRecord -> DeliveryReport
+mkErrorReport msg prodRec = DeliveryFailure prodRec (KafkaResponseError (err'RdKafkaMessageT msg))
 
-mkSuccessReport :: RdKafkaMessageT -> IO DeliveryReport
-mkSuccessReport msg = do
-  prodRec <- mkProdRec msg
-  pure $ DeliverySuccess prodRec (Offset $ offset'RdKafkaMessageT msg)
+mkSuccessReport :: RdKafkaMessageT -> ProducerRecord -> DeliveryReport
+mkSuccessReport msg prodRec = DeliverySuccess prodRec (Offset $ offset'RdKafkaMessageT msg)
 
-mkProdRec :: RdKafkaMessageT -> IO ProducerRecord
-mkProdRec msg = do
-  topic     <- readTopic msg
-  key       <- readKey msg
-  payload   <- readPayload msg
-  pure ProducerRecord
-    { prTopic = TopicName topic
-    , prPartition = SpecifiedPartition (partition'RdKafkaMessageT msg)
-    , prKey = key
-    , prValue = payload
-    }
+mkProdRec :: Ptr RdKafkaMessageT -> IO ProducerRecord
+mkProdRec pmsg = do
+  msg         <- peek pmsg  
+  topic       <- readTopic msg
+  key         <- readKey msg
+  payload     <- readPayload msg
+  flip fmap (fromRight mempty <$> readHeaders pmsg) $ \headers -> 
+    ProducerRecord
+      { prTopic = TopicName topic
+      , prPartition = SpecifiedPartition (partition'RdKafkaMessageT msg)
+      , prKey = key
+      , prValue = payload
+      , prHeaders = headers
+      }

src/Kafka/Producer/Types.hs

@@ -21,7 +21,7 @@ import Data.Typeable        (Typeable)
 import GHC.Generics         (Generic)
 import Kafka.Consumer.Types (Offset (..))
 import Kafka.Internal.Setup (HasKafka (..), HasKafkaConf (..), HasTopicConf (..), Kafka (..), KafkaConf (..), TopicConf (..))
-import Kafka.Types          (KafkaError (..), TopicName (..))
+import Kafka.Types          (KafkaError (..), TopicName (..), Headers)
 
 -- | The main type for Kafka message production, used e.g. to send messages.
 --
@@ -50,6 +50,7 @@ data ProducerRecord = ProducerRecord
   , prPartition :: !ProducePartition
   , prKey       :: Maybe ByteString
   , prValue     :: Maybe ByteString
+  , prHeaders   :: !Headers
   } deriving (Eq, Show, Typeable, Generic)
 
 -- | 

src/Kafka/Types.hs

@@ -21,7 +21,7 @@ module Kafka.Types
 , KafkaDebug(..)
 , KafkaCompressionCodec(..)
 , TopicType(..)
-, Headers(unHeaders), headersFromList, headersToList
+, Headers, headersFromList, headersToList
 , topicType
 , kafkaDebugToText
 , kafkaCompressionCodecToText