Convert.hs

module Kafka.Consumer.Convert
  ( offsetSyncToInt,
    offsetToInt64,
    int64ToOffset,
    fromNativeTopicPartitionList'',
    nativeTopicPartitionErrors,
    fromNativeTopicPartitionList',
    fromNativeTopicPartitionList,
    toNativeTopicPartitionList,
    toNativeTopicPartitionListNoDispose,
    toNativeTopicPartitionList',
    topicPartitionFromMessage,
    topicPartitionFromMessageForCommit,
    toMap,
    fromMessagePtr,
    fromMessagePtrAttributed,
    copyMessagePtr,
    copyMessagePtrAttributed,
    withPolledMessage,
    offsetCommitToBool,
  )
where

import Control.Exception (bracket, bracketOnError)
import Control.Monad (when, (>=>))
import Data.Bifunctor (first)
import Data.ByteString qualified as BS
import Data.Int (Int64)
import Data.Map.Strict (Map, fromListWith)
import Data.Set qualified as S
import Data.Text qualified as Text
import Foreign.C.Error (getErrno)
import Foreign.ForeignPtr (finalizeForeignPtr, withForeignPtr)
import Foreign.Ptr (Ptr, nullPtr)
import Foreign.Storable (Storable (..))
import Kafka.Consumer.Types (ConsumerMessageError (..), ConsumerRecord (..), Offset (..), OffsetCommit (..), OffsetStoreSync (..), PartitionOffset (..), TopicPartition (..))
import Kafka.Internal.RdKafka
  ( RdKafkaMessageT (..),
    RdKafkaMessageTPtr,
    RdKafkaRespErrT (..),
    RdKafkaTopicPartitionListT (..),
    RdKafkaTopicPartitionListTPtr,
    RdKafkaTopicPartitionT (..),
    newRdKafkaTopicPartitionListT,
    peekCText,
    rdKafkaMessageDestroy,
    rdKafkaTopicPartitionListAdd,
    rdKafkaTopicPartitionListNew,
    rdKafkaTopicPartitionListSetOffset,
  )
import Kafka.Internal.Shared (kafkaRespErr, readHeaders, readKey, readPayload, readTimestamp, readTopic)
import Kafka.Types (KafkaError (..), PartitionId (..), TopicName (..))

-- | Converts offsets sync policy to integer (the way Kafka understands it):
--
--     * @OffsetSyncDisable == -1@
--
--     * @OffsetSyncImmediate == 0@
--
--     * @OffsetSyncInterval ms == ms@
offsetSyncToInt :: OffsetStoreSync -> Int
offsetSyncToInt sync =
  case sync of
    OffsetSyncDisable -> -1
    OffsetSyncImmediate -> 0
    OffsetSyncInterval ms -> ms
{-# INLINE offsetSyncToInt #-}

offsetToInt64 :: PartitionOffset -> Int64
offsetToInt64 o = case o of
  PartitionOffsetBeginning -> -2
  PartitionOffsetEnd -> -1
  PartitionOffset off -> off
  PartitionOffsetStored -> -1000
  PartitionOffsetInvalid -> -1001
{-# INLINE offsetToInt64 #-}

int64ToOffset :: Int64 -> PartitionOffset
int64ToOffset o
  | o == -2 = PartitionOffsetBeginning
  | o == -1 = PartitionOffsetEnd
  | o == -1000 = PartitionOffsetStored
  | o >= 0 = PartitionOffset o
  | otherwise = PartitionOffsetInvalid
{-# INLINE int64ToOffset #-}

fromNativeTopicPartitionList'' :: RdKafkaTopicPartitionListTPtr -> IO [TopicPartition]
fromNativeTopicPartitionList'' ptr =
  withForeignPtr ptr $ \fptr -> fromNativeTopicPartitionList' fptr

fromNativeTopicPartitionList' :: Ptr RdKafkaTopicPartitionListT -> IO [TopicPartition]
fromNativeTopicPartitionList' ppl =
  if ppl == nullPtr
    then return []
    else peek ppl >>= fromNativeTopicPartitionList

fromNativeTopicPartitionList :: RdKafkaTopicPartitionListT -> IO [TopicPartition]
fromNativeTopicPartitionList pl =
  let count = cnt'RdKafkaTopicPartitionListT pl
      elems = elems'RdKafkaTopicPartitionListT pl
   in mapM (peekElemOff elems >=> toPart) [0 .. (count - 1)]
  where
    toPart :: RdKafkaTopicPartitionT -> IO TopicPartition
    toPart p = do
      topic <- peekCText $ topic'RdKafkaTopicPartitionT p
      return
        TopicPartition
          { tpTopicName = TopicName topic,
            tpPartition = PartitionId $ partition'RdKafkaTopicPartitionT p,
            tpOffset = int64ToOffset $ offset'RdKafkaTopicPartitionT p
          }

-- | Read member errors in native/input order while retaining the owning list.
-- This deliberately leaves the generic public TopicPartition conversion alone.
nativeTopicPartitionErrors :: RdKafkaTopicPartitionListTPtr -> IO [RdKafkaRespErrT]
nativeTopicPartitionErrors list = withForeignPtr list $ \ptr -> do
  value <- peek ptr
  mapM
    (fmap err'RdKafkaTopicPartitionT . peekElemOff (elems'RdKafkaTopicPartitionListT value))
    [0 .. cnt'RdKafkaTopicPartitionListT value - 1]

toNativeTopicPartitionList :: [TopicPartition] -> IO RdKafkaTopicPartitionListTPtr
toNativeTopicPartitionList ps = bracketOnError (newRdKafkaTopicPartitionListT (length ps)) finalizeForeignPtr $ \pl -> do
  mapM_
    ( \p -> do
        let TopicName tn = tpTopicName p
            (PartitionId tp) = tpPartition p
            to = offsetToInt64 $ tpOffset p
            tnS = Text.unpack tn
        _ <- rdKafkaTopicPartitionListAdd pl tnS tp
        rdKafkaTopicPartitionListSetOffset pl tnS 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
            tnS = Text.unpack tn
        _ <- rdKafkaTopicPartitionListAdd pl tnS tp
        rdKafkaTopicPartitionListSetOffset pl tnS tp to
    )
    ps
  return pl

toNativeTopicPartitionList' :: [(TopicName, PartitionId)] -> IO RdKafkaTopicPartitionListTPtr
toNativeTopicPartitionList' tps =
  let utps = S.toList . S.fromList $ tps
   in bracketOnError (newRdKafkaTopicPartitionListT (length utps)) finalizeForeignPtr $ \pl -> do
        mapM_ (\(TopicName t, PartitionId p) -> rdKafkaTopicPartitionListAdd pl (Text.unpack t) p) utps
        return pl

topicPartitionFromMessage :: ConsumerRecord k v -> TopicPartition
topicPartitionFromMessage m =
  let (Offset moff) = crOffset m
   in TopicPartition (crTopic m) (crPartition m) (PartitionOffset moff)

-- | Creates a topic partition message for use with the offset commit message.
-- We increment the offset by 1 here because when we commit, the offset is the position
-- the consumer reads from to process the next message.
topicPartitionFromMessageForCommit :: ConsumerRecord k v -> TopicPartition
topicPartitionFromMessageForCommit m =
  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]

-- | Acquire and copy one message within the same ownership scope. Masking
-- acquisition closes the window between native return and installing cleanup.
fromMessagePtr :: IO RdKafkaMessageTPtr -> IO (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)))
fromMessagePtr acquire = first cmeError <$> fromMessagePtrAttributed acquire

fromMessagePtrAttributed :: IO RdKafkaMessageTPtr -> IO (Either ConsumerMessageError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)))
fromMessagePtrAttributed acquire = withPolledMessage acquire copyMessagePtrAttributed

withPolledMessage :: IO RdKafkaMessageTPtr -> (RdKafkaMessageTPtr -> IO a) -> IO a
withPolledMessage acquire = bracket acquire releaseMessage
  where
    releaseMessage pointer = withForeignPtr pointer $ \raw ->
      when (raw /= nullPtr) (rdKafkaMessageDestroy raw)

-- | Copy a borrowed batch member. The caller owns and releases the complete
-- batch, including members not reached if a conversion is interrupted.
copyMessagePtr :: RdKafkaMessageTPtr -> IO (Either KafkaError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)))
copyMessagePtr pointer = first cmeError <$> copyMessagePtrAttributed pointer

copyMessagePtrAttributed :: RdKafkaMessageTPtr -> IO (Either ConsumerMessageError (ConsumerRecord (Maybe BS.ByteString) (Maybe BS.ByteString)))
copyMessagePtrAttributed ptr =
  withForeignPtr ptr $ \realPtr ->
    if realPtr == nullPtr
      then (\err -> Left (ConsumerMessageError err Nothing Nothing Nothing)) . kafkaRespErr <$> getErrno
      else do
        message <- peek realPtr
        if err'RdKafkaMessageT message /= RdKafkaRespErrNoError
          then Left <$> attributed (err'RdKafkaMessageT message) message
          else mkRecord message realPtr
  where
    attributed code msg = do
      topic <-
        if topic'RdKafkaMessageT msg == nullPtr
          then pure Nothing
          else Just . TopicName <$> readTopic msg
      let partition = partition'RdKafkaMessageT msg
          offset = offset'RdKafkaMessageT msg
      pure
        ConsumerMessageError
          { cmeError = KafkaResponseError code,
            cmeTopic = topic,
            cmePartition = if partition >= 0 then Just (PartitionId partition) else Nothing,
            cmeOffset = if offset >= 0 then Just (Offset offset) else Nothing
          }
    mkRecord msg rptr = do
      topic <- readTopic msg
      key <- readKey msg
      payload <- readPayload msg
      timestamp <- readTimestamp ptr
      readHeaders rptr >>= \case
        Left code -> Left <$> attributed code msg
        Right headers ->
          pure $
            Right
              ConsumerRecord
                { crTopic = TopicName topic,
                  crPartition = PartitionId $ partition'RdKafkaMessageT msg,
                  crOffset = Offset $ offset'RdKafkaMessageT msg,
                  crTimestamp = timestamp,
                  crHeaders = headers,
                  crKey = key,
                  crValue = payload
                }

offsetCommitToBool :: OffsetCommit -> Bool
offsetCommitToBool OffsetCommit = False
offsetCommitToBool OffsetCommitAsync = True