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