packages feed

nri-kafka (empty) → 0.1.0.0

raw patch · 20 files changed

+2636/−0 lines, 20 filesdep +aesondep +asyncdep +base

Dependencies added: aeson, async, base, bytestring, conduit, containers, hw-kafka-client, nri-env-parser, nri-observability, nri-prelude, safe-exceptions, stm, text, time, unix, uuid

Files

+ CHANGELOG.md view
@@ -0,0 +1,4 @@+# 0.1.0.0++- First release, but we've battle-tested it against significant load for months now!+  Hope you enjoy
+ LICENSE view
@@ -0,0 +1,29 @@+BSD 3-Clause License++Copyright (c) 2021, NoRedInk+All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++* Redistributions of source code must retain the above copyright notice, this+  list of conditions and the following disclaimer.++* Redistributions in binary form must reproduce the above copyright notice,+  this list of conditions and the following disclaimer in the documentation+  and/or other materials provided with the distribution.++* Neither the name of the copyright holder nor the names of its+  contributors may be used to endorse or promote products derived from+  this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,51 @@+# Kafka integration++_Reviewed last on 2021-05-28_++This library exposes an Elm-like API to Kafka. It exports two main modules:++- `Kafka`, for writing to Kafka.+- `Kafka.Worker`, For building long-running worker apps that process Haskell+  messages.++At NoRedInk, we use this to power our high-throughput quiz-engine service. If+you work at NoRedInk: look there for a simple example app.++Otherwise: here's the gist of it:++```+import qualified Environment -- from nri-env-parser+import qualified Kafka.Worker++-- your long running app+main :: IO ()+main =+  settings <- Environment.decode Kafka.Worker.decoder+  Kafka.Worker.process+    Kafka.Worker.Description+      settings+      "this worker's group id"+      (Kafka.Worker.subscription "my.topic" processMessage,)++data MyKafkaMessageType =+  ReticulateSplines Int+  AddHiddenAgenda Text+  CalculateLlamaExpectorationTrajectory Llamas+  deriving (generic)++instance Aeson.ToJSON Envelope+instance Aeson.FromJSON Envelope++-- the meat and potatoes: handles all MyKafkaMessageTypes+processMessage ::+  Kafka.Worker.Envelope MyKafkaMessageType ->+  Task Text ()+processMessage record myMessage =+  -- process your message in here+  -- because of our usage of `Task` you probably want to pass in any handlers+  case myMessage of+    AddHiddenAgenda agenda ->+    	Debug.todo "Add the agenda"+    _ ->+    	Debug.todo "and also handle the other cases"+```
+ nri-kafka.cabal view
@@ -0,0 +1,141 @@+cabal-version: 1.12++-- This file has been generated from package.yaml by hpack version 0.34.4.+--+-- see: https://github.com/sol/hpack++name:           nri-kafka+version:        0.1.0.0+synopsis:       Functions for working with Kafka+description:    Please see the README at <https://github.com/NoRedInk/haskell-libraries/tree/trunk/nri-kafka#readme>.+category:       Web+homepage:       https://github.com/NoRedInk/haskell-libraries/tree/trunk/nri-kafka#readme+bug-reports:    https://github.com/NoRedInk/haskell-libraries/issues+author:         NoRedInk+maintainer:     haskell-open-source@noredink.com+copyright:      2021 NoRedInk Corp.+license:        BSD3+license-file:   LICENSE+build-type:     Simple+extra-source-files:+    README.md+    LICENSE+    CHANGELOG.md++source-repository head+  type: git+  location: https://github.com/NoRedInk/haskell-libraries+  subdir: nri-kafka++library+  exposed-modules:+      Kafka+      Kafka.Worker+      Kafka.Test+  other-modules:+      Kafka.Internal+      Kafka.Settings+      Kafka.Settings.Internal+      Kafka.Worker.Analytics+      Kafka.Worker.Fetcher+      Kafka.Worker.Internal+      Kafka.Worker.Partition+      Kafka.Worker.Settings+      Kafka.Worker.Stopping+      Paths_nri_kafka+  hs-source-dirs:+      src+  default-extensions:+      DataKinds+      DeriveGeneric+      ExtendedDefaultRules+      FlexibleContexts+      FlexibleInstances+      GeneralizedNewtypeDeriving+      MultiParamTypeClasses+      NamedFieldPuns+      NoImplicitPrelude+      NumericUnderscores+      OverloadedStrings+      PartialTypeSignatures+      ScopedTypeVariables+      Strict+      TypeOperators+  ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wpartial-fields -Wredundant-constraints -Wincomplete-uni-patterns -fno-warn-type-defaults -fplugin=NriPrelude.Plugin+  build-depends:+      aeson >=1.4.6.0 && <1.6+    , async >=2.2.2 && <2.3+    , base >=4.12.0.0 && <4.16+    , bytestring >=0.10.8.2 && <0.12+    , conduit >=1.3.0 && <1.4+    , containers >=0.6.0.1 && <0.7+    , hw-kafka-client >=4.0.3 && <5.0+    , nri-env-parser >=0.1.0.0 && <0.2+    , nri-observability >=0.1.1.1 && <0.2+    , nri-prelude >=0.1.0.0 && <0.7+    , safe-exceptions >=0.1.7.0 && <1.3+    , stm >=2.4 && <2.6+    , text >=1.2.3.1 && <1.3+    , time >=1.8.0.2 && <1.13+    , unix >=2.7.2.2 && <2.8.0.0+    , uuid >=1.3.0 && <1.4+  default-language: Haskell2010++test-suite tests+  type: exitcode-stdio-1.0+  main-is: Main.hs+  other-modules:+      Kafka+      Kafka.Internal+      Kafka.Settings+      Kafka.Settings.Internal+      Kafka.Test+      Kafka.Worker+      Kafka.Worker.Analytics+      Kafka.Worker.Fetcher+      Kafka.Worker.Internal+      Kafka.Worker.Partition+      Kafka.Worker.Settings+      Kafka.Worker.Stopping+      Helpers+      Spec.Kafka.Worker.Integration+      Spec.Kafka.Worker.Partition+      Paths_nri_kafka+  hs-source-dirs:+      src+      test+  default-extensions:+      DataKinds+      DeriveGeneric+      ExtendedDefaultRules+      FlexibleContexts+      FlexibleInstances+      GeneralizedNewtypeDeriving+      MultiParamTypeClasses+      NamedFieldPuns+      NoImplicitPrelude+      NumericUnderscores+      OverloadedStrings+      PartialTypeSignatures+      ScopedTypeVariables+      Strict+      TypeOperators+  ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wpartial-fields -Wredundant-constraints -Wincomplete-uni-patterns -fno-warn-type-defaults -fplugin=NriPrelude.Plugin -threaded -rtsopts "-with-rtsopts=-N -T" -fno-warn-type-defaults+  build-depends:+      aeson >=1.4.6.0 && <1.6+    , async >=2.2.2 && <2.3+    , base >=4.12.0.0 && <4.16+    , bytestring >=0.10.8.2 && <0.12+    , conduit >=1.3.0 && <1.4+    , containers >=0.6.0.1 && <0.7+    , hw-kafka-client >=4.0.3 && <5.0+    , nri-env-parser >=0.1.0.0 && <0.2+    , nri-observability >=0.1.1.1 && <0.2+    , nri-prelude >=0.1.0.0 && <0.7+    , safe-exceptions >=0.1.7.0 && <1.3+    , stm >=2.4 && <2.6+    , text >=1.2.3.1 && <1.3+    , time >=1.8.0.2 && <1.13+    , unix >=2.7.2.2 && <2.8.0.0+    , uuid >=1.3.0 && <1.4+  default-language: Haskell2010
+ src/Kafka.hs view
@@ -0,0 +1,276 @@+{-# LANGUAGE RankNTypes #-}+{-# OPTIONS_GHC -Wno-redundant-constraints #-}++-- | Kafka is a module for _writing_ to Kafka+--+-- See Kafka.Worker for the basic building blocks of a CLI app that will poll &+-- process kafka messages+module Kafka+  ( -- * Setup+    Internal.Handler,+    Settings.Settings,+    Settings.decoder,+    handler,++    -- * Creating messages+    Internal.Msg,+    emptyMsg,+    addPayload,+    addKey,++    -- * Sending messags+    Internal.sendAsync,+    Internal.sendSync,++    -- * Reading messages+    topic,+    payload,+    key,+  )+where++import qualified Conduit+import qualified Control.Concurrent+import qualified Control.Concurrent.Async as Async+import qualified Control.Concurrent.STM as STM+import qualified Control.Concurrent.STM.TMVar as TMVar+import qualified Control.Exception.Safe as Exception+import Control.Monad.IO.Class (liftIO)+import qualified Data.Aeson as Aeson+import qualified Data.ByteString.Lazy as ByteString.Lazy+import qualified Data.Text.Encoding+import qualified Dict+import qualified Kafka.Internal as Internal+import qualified Kafka.Producer as Producer+import qualified Kafka.Settings as Settings+import qualified Platform+import qualified Prelude++data Details = Details+  { detailsBrokers :: List Text,+    detailsMsg :: Internal.Msg+  }+  deriving (Generic, Show)++instance Aeson.ToJSON Details++instance Platform.TracingSpanDetails Details++newtype DeliveryReportDetails = DeliveryReportDetails+  { deliveryReportProducerRecord :: Text+  }+  deriving (Generic, Show)++instance Aeson.ToJSON DeliveryReportDetails++instance Platform.TracingSpanDetails DeliveryReportDetails++-- | Creates a Kafka-writable message for a topic.+--+-- > msg =+-- >   emptyMsg "groceries"+-- >     |> addPayload "broccoli"+-- >     |> addKey "vegetables"+emptyMsg :: Text -> Internal.Msg+emptyMsg topic' =+  Internal.Msg+    { Internal.topic = Internal.Topic topic',+      Internal.payload = Nothing,+      Internal.key = Nothing+    }++-- Add a payload to a message.+--+-- Message payloads aren't mandatory in Kafka, so using this function really is+-- optional. A counter is an example of an application that doesn't require+-- message payloads. Just knowing an increment event took place would be enough+-- for it to work.+--+-- We ask for JSON decodability to ensure the Kafka worker can later read the message+addPayload :: (Aeson.FromJSON a, Aeson.ToJSON a) => a -> Internal.Msg -> Internal.Msg+addPayload contents msg =+  msg {Internal.payload = (Just (Internal.Encodable contents))}++-- Add a key to a message.+--+-- Kafka divides messages in a topic in different partitions. Kafka workers can+-- collaborate on a topic by each processing messages from a couple of the+-- topic's partitions. Within a partition messages will never overtake each+-- other.+--+-- By default each message is assigned to a random partition. Setting a key on+-- the message gives you more control over this process. Messages with the same+-- key are guaranteed to end up in the same partition.+--+-- Example: if each message is related to a single user and you need to ensure+-- messagse for a user don't overtake each other, you can set the key to be the+-- user's id.+addKey :: Text -> Internal.Msg -> Internal.Msg+addKey key' msg = msg {Internal.key = Just (Internal.Key key')}++record :: Internal.Msg -> Task e Producer.ProducerRecord+record msg = do+  requestId <- Platform.requestId+  Task.succeed+    Producer.ProducerRecord+      { Producer.prTopic =+          Internal.topic msg+            |> Internal.unTopic+            |> Producer.TopicName,+        Producer.prPartition = Producer.UnassignedPartition,+        Producer.prKey =+          Maybe.map+            (Data.Text.Encoding.encodeUtf8 << Internal.unKey)+            (Internal.key msg),+        Producer.prValue =+          Maybe.map+            ( \payload' ->+                Internal.MsgWithMetaData+                  { Internal.metaData =+                      Internal.MetaData+                        { Internal.requestId+                        },+                    Internal.value = payload'+                  }+                  |> Aeson.encode+                  |> ByteString.Lazy.toStrict+            )+            (Internal.payload msg)+      }++-- | The topic of a message. This function might sometimes be useful in tests.+topic :: Internal.Msg -> Text+topic msg = Internal.unTopic (Internal.topic msg)++-- | The payload of a message. This function might sometimes be useful in tests.+payload :: (Aeson.FromJSON a) => Internal.Msg -> Maybe a+payload msg =+  Internal.payload msg+    |> Maybe.andThen (Aeson.decode << Aeson.encode)++-- | The key of a message. This function might sometimes be useful in tests.+key :: Internal.Msg -> Maybe Text+key msg = Maybe.map Internal.unKey (Internal.key msg)++-- | Function for creating a Kafka handler.+--+-- See 'Kafka.Settings' for potential customizations.+handler :: Settings.Settings -> Conduit.Acquire Internal.Handler+handler settings = do+  producer <- Conduit.mkAcquire (mkProducer settings) Producer.closeProducer+  _ <- Conduit.mkAcquire (startPollEventLoop producer) (\terminator -> STM.atomically (TMVar.putTMVar terminator Terminate))+  liftIO (mkHandler settings producer)++data Terminate = Terminate++-- | By default events only get polled right before sending a record to kafka.+-- This means that the deliveryCallback only gets fired on the next call to produceMessage'.+-- We want to be informed about delivery status as soon as possible though.+startPollEventLoop :: Producer.KafkaProducer -> Prelude.IO (TMVar.TMVar b)+startPollEventLoop producer = do+  terminator <- STM.atomically TMVar.newEmptyTMVar+  _ <-+    Async.race_+      (pollEvents producer)+      (STM.atomically <| TMVar.readTMVar terminator)+      |> Async.async+  Prelude.pure terminator++-- | We use a little trick here to poll events, by sending an empty message batch.+-- This will call the internal pollEvent function in hw-kafka-client.+pollEvents :: Producer.KafkaProducer -> Prelude.IO ()+pollEvents producer = do+  Producer.produceMessageBatch producer []+    |> map (\_ -> ())+  Control.Concurrent.threadDelay 100_000 {- 100ms -}+  pollEvents producer++-- |+mkHandler :: Settings.Settings -> Producer.KafkaProducer -> Prelude.IO Internal.Handler+mkHandler settings producer = do+  doAnything <- Platform.doAnythingHandler+  Prelude.pure+    Internal.Handler+      { Internal.sendAsync = \onDeliveryCallback msg' ->+          Platform.tracingSpan "Async send Kafka messages" <| do+            let details = Details (List.map Producer.unBrokerAddress (Settings.brokerAddresses settings)) msg'+            Platform.setTracingSpanDetails details+            sendHelperAsync producer doAnything onDeliveryCallback msg'+              |> Task.mapError Internal.errorToText,+        Internal.sendSync = \msg' ->+          Platform.tracingSpan "Sync send Kafka messages" <| do+            let details = Details (List.map Producer.unBrokerAddress (Settings.brokerAddresses settings)) msg'+            Platform.setTracingSpanDetails details+            terminator <- doSTM doAnything TMVar.newEmptyTMVar+            let onDeliveryCallback = doSTM doAnything (TMVar.putTMVar terminator Terminate)+            sendHelperAsync producer doAnything onDeliveryCallback msg'+              |> Task.mapError Internal.errorToText+            Terminate <- doSTM doAnything (TMVar.readTMVar terminator)+            Task.succeed ()+      }++doSTM :: Platform.DoAnythingHandler -> STM.STM a -> Task e a+doSTM doAnything stm =+  STM.atomically stm+    |> map Ok+    |> Platform.doAnything doAnything++mkProducer :: Settings.Settings -> Prelude.IO Producer.KafkaProducer+mkProducer Settings.Settings {Settings.brokerAddresses, Settings.deliveryTimeout, Settings.logLevel, Settings.batchNumMessages} = do+  let properties =+        Producer.brokersList brokerAddresses+          ++ Producer.sendTimeout deliveryTimeout+          ++ Producer.logLevel logLevel+          ++ Producer.compression Producer.Snappy+          ++ Producer.extraProps+            ( Dict.fromList+                [ ( "batch.num.messages",+                    batchNumMessages+                      |> Settings.unBatchNumMessages+                      |> Text.fromInt+                  ),+                  -- Enable idemptent producers+                  -- See https://www.cloudkarafka.com/blog/apache-kafka-idempotent-producer-avoiding-message-duplication.html for reference+                  ("enable.idempotence", "true"),+                  ("acks", "all")+                ]+            )+  eitherProducer <- Producer.newProducer properties+  case eitherProducer of+    Prelude.Left err ->+      -- We create the handler as part of starting the application. Throwing+      -- means that if there's a problem with the settings the application will+      -- fail immediately upon start. It won't result in runtime errors during+      -- operation.+      Exception.throwIO err+    Prelude.Right producer ->+      Prelude.pure producer++sendHelperAsync ::+  Producer.KafkaProducer ->+  Platform.DoAnythingHandler ->+  Task Never () ->+  Internal.Msg ->+  Task Internal.Error ()+sendHelperAsync producer doAnything onDeliveryCallback msg' = do+  record' <- record msg'+  Exception.handleAny+    (\exception -> Prelude.pure (Err (Internal.Uncaught exception)))+    ( do+        maybeFailedMessages <-+          Producer.produceMessage'+            producer+            record'+            ( \deliveryReport -> do+                log <- Platform.silentHandler+                Task.perform log+                  <| case deliveryReport of+                    Producer.DeliverySuccess _producerRecord _offset -> onDeliveryCallback+                    _ -> Task.succeed ()+            )+        Prelude.pure <| case maybeFailedMessages of+          Prelude.Right _ -> Ok ()+          Prelude.Left (Producer.ImmediateError failure) ->+            Err (Internal.SendingFailed (record', failure))+    )+    |> Platform.doAnything doAnything
+ src/Kafka/Internal.hs view
@@ -0,0 +1,83 @@+{-# LANGUAGE GADTs #-}++module Kafka.Internal where++import qualified Control.Exception.Safe as Exception+import qualified Data.Aeson as Aeson+import qualified Kafka.Producer as Producer+import qualified Prelude++-- | A handler for writing to Kafka+data Handler = Handler+  { -- | sends messages asynchronously with to Kafka+    --+    -- This is the recommended approach for high throughput. The C++ library+    -- behind hte scenes, librdkafka, will batch messages together.+    sendAsync :: Task Never () -> Msg -> Task Text (),+    -- | sends messages synchronously with to Kafka+    --+    -- This can have a large negative impact on throughput. Use sparingly!+    sendSync :: Msg -> Task Text ()+  }++-- | A message that can be written to Kafka+data Msg = Msg+  { topic :: Topic,+    key :: Maybe Key,+    payload :: Maybe Encodable+  }+  deriving (Generic, Show)++instance Aeson.ToJSON Msg++data Encodable where+  Encodable :: (Aeson.FromJSON a, Aeson.ToJSON a) => a -> Encodable++instance Aeson.ToJSON Encodable where+  toJSON (Encodable x) = Aeson.toJSON x+  toEncoding (Encodable x) = Aeson.toEncoding x++instance Aeson.FromJSON Encodable where+  parseJSON x = do+    val <- Aeson.parseJSON x+    Prelude.pure (Encodable (val :: Aeson.Value))++instance Show Encodable where+  show (Encodable x) = Prelude.show (Aeson.toJSON x)++-- | Errors.+-- If you experience an 'Uncaught' exception, please wrap it here type here!+data Error+  = SendingFailed (Producer.ProducerRecord, Producer.KafkaError)+  | Uncaught Exception.SomeException+  deriving (Show)++errorToText :: Error -> Text+errorToText err = Text.fromList (Prelude.show err)++-- | A kafka topic+newtype Topic = Topic {unTopic :: Text} deriving (Aeson.ToJSON, Show)++-- | A kafka key+newtype Key = Key {unKey :: Text} deriving (Show, Aeson.ToJSON, Eq, Ord)++data MsgWithMetaData = MsgWithMetaData+  { metaData :: MetaData,+    value :: Encodable+  }+  deriving (Generic)++instance Aeson.ToJSON MsgWithMetaData++instance Aeson.FromJSON MsgWithMetaData++newtype MetaData = MetaData+  { requestId :: Text+  }+  deriving (Generic)++instance Aeson.ToJSON MetaData++instance Aeson.FromJSON MetaData++newtype Offset = Offset Int
+ src/Kafka/Settings.hs view
@@ -0,0 +1,66 @@+-- | Kafka.Settings+module Kafka.Settings+  ( Settings (..),+    decoder,+    BatchNumMessages,+    unBatchNumMessages,+    exampleBatchNumMessages,+  )+where++import qualified Environment+import qualified Kafka.Producer+import qualified Kafka.Settings.Internal as Internal++-- | Settings required to write to Kafka+data Settings = Settings+  { -- | broker addresses. See hw-kafka's documentation for more info+    brokerAddresses :: [Kafka.Producer.BrokerAddress],+    -- | client log level. See hw-kafka's documentation for more info+    logLevel :: Internal.KafkaLogLevel,+    -- | Message delivery timeout. See hw-kafka's documentation for more info+    deliveryTimeout :: Kafka.Producer.Timeout,+    -- | Number of messages to batch together before sending to Kafka.+    batchNumMessages :: BatchNumMessages+  }++-- | Number of messages to batch together before sending to Kafka.+newtype BatchNumMessages = BatchNumMessages {unBatchNumMessages :: Int}++-- |  example BatchNumMessages to use in tests+exampleBatchNumMessages :: BatchNumMessages+exampleBatchNumMessages = BatchNumMessages 1++-- | decodes Settings from environmental variables+-- KAFKA_BROKER_ADDRESSES=localhost:9092 # comma delimeted list+-- KAFKA_LOG_LEVEL=Debug+-- KAFKA_DELIVERY_TIMEOUT=120000+-- KAFKA_BATCH_SIZE=10000+decoder :: Environment.Decoder Settings+decoder =+  map4+    Settings+    Internal.decoderBrokerAddresses+    Internal.decoderKafkaLogLevel+    decoderDeliveryTimeout+    decoderBatchNumMessages++decoderDeliveryTimeout :: Environment.Decoder Kafka.Producer.Timeout+decoderDeliveryTimeout =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_DELIVERY_TIMEOUT",+        Environment.description = "Delivery timout for producer. Aka 'delivery.timeout.ms'",+        Environment.defaultValue = "120000"+      }+    (map Kafka.Producer.Timeout Environment.int)++decoderBatchNumMessages :: Environment.Decoder BatchNumMessages+decoderBatchNumMessages =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_BATCH_SIZE",+        Environment.description = "Kafka Producer 'batch.num.messages'",+        Environment.defaultValue = "10000"+      }+    (map BatchNumMessages Environment.int)
+ src/Kafka/Settings/Internal.hs view
@@ -0,0 +1,42 @@+module Kafka.Settings.Internal+  ( Kafka.Types.KafkaLogLevel (..),+    decoderBrokerAddresses,+    decoderKafkaLogLevel,+  )+where++import qualified Environment+import qualified Kafka.Types++decoderBrokerAddresses :: Environment.Decoder [Kafka.Types.BrokerAddress]+decoderBrokerAddresses =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_BROKER_ADDRESSES",+        Environment.description = "A comma-separated list of broker addresses",+        Environment.defaultValue = "localhost:9092"+      }+    (map (List.map Kafka.Types.BrokerAddress << Text.split ",") Environment.text)++decoderKafkaLogLevel :: Environment.Decoder Kafka.Types.KafkaLogLevel+decoderKafkaLogLevel =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_LOG_LEVEL",+        Environment.description = "Kafka log level",+        Environment.defaultValue = "Debug"+      }+    (map kafkaLogLevelFromText Environment.text)++kafkaLogLevelFromText :: Text -> Kafka.Types.KafkaLogLevel+kafkaLogLevelFromText text =+  case text of+    "Emerg" -> Kafka.Types.KafkaLogEmerg+    "Alert" -> Kafka.Types.KafkaLogAlert+    "Crit" -> Kafka.Types.KafkaLogCrit+    "Err" -> Kafka.Types.KafkaLogErr+    "Warning" -> Kafka.Types.KafkaLogWarning+    "Notice" -> Kafka.Types.KafkaLogNotice+    "Info" -> Kafka.Types.KafkaLogInfo+    "Debug" -> Kafka.Types.KafkaLogDebug+    _ -> Kafka.Types.KafkaLogDebug
+ src/Kafka/Test.hs view
@@ -0,0 +1,35 @@+-- | Helpers for testing code that sends messages to Kafka.+module Kafka.Test+  ( stub,+  )+where++import qualified Data.IORef+import qualified Expect+import qualified GHC.Stack as Stack+import qualified Kafka+import qualified Kafka.Internal as Internal+import qualified Platform++-- | Can be used to test your Kafka writer.+-- yields a mock Kafka handler, and returns an expectation wrapping a list of+-- messages that would have been written if the handler was real+stub ::+  Stack.HasCallStack =>+  (Internal.Handler -> Expect.Expectation) ->+  Expect.Expectation' (List Kafka.Msg)+stub stubbed = do+  logRef <- Expect.fromIO (Data.IORef.newIORef [])+  doAnything <- Expect.fromIO Platform.doAnythingHandler+  let sendStub = \msg' -> do+        Data.IORef.modifyIORef' logRef (\prev -> msg' : prev)+          |> map Ok+          |> Platform.doAnything doAnything+  let mockHandler =+        Internal.Handler+          { Internal.sendAsync = \_ -> sendStub,+            Internal.sendSync = sendStub+          }+  Expect.around (\f -> f mockHandler) (Stack.withFrozenCallStack stubbed)+  Expect.fromIO (Data.IORef.readIORef logRef)+    |> map List.reverse
+ src/Kafka/Worker.hs view
@@ -0,0 +1,34 @@+{-# LANGUAGE GADTs #-}++-- The Kafka worker has the following concurrent workflows:+-- 1. The main thread which handles+--   - gracefully quitting+--   - a single-thread `pollingLoop` that reads new messages from kafka+-- 2. The Kafka `rebalanceCallback` which handles rebalancing partitions, in turn+--    turning on and off worker threads.+-- 3. A single-thread `pauseAndAnalyticsLoop` that tells the kafka library to pause & resume+--    sending us messages from specific partitions (based on a number of factors)+-- 4. Multiple worker-threads (`Partition.processMsgLoop`), one per+--    (topic,partition)++-- | Kafka.Worker is a module for processing a Kafka log.+-- It can be used to build a CLI that will consume and process a user-defined message type+module Kafka.Worker+  ( Internal.process,++    -- * Settings+    Settings.Settings,+    Settings.decoder,++    -- * Subscriptions+    Internal.TopicSubscription,+    Internal.subscription,+    Internal.subscriptionManageOwnOffsets,+    Internal.PartitionOffset (..),+    Partition.SeekCmd (..),+  )+where++import qualified Kafka.Worker.Internal as Internal+import qualified Kafka.Worker.Partition as Partition+import qualified Kafka.Worker.Settings as Settings
+ src/Kafka/Worker/Analytics.hs view
@@ -0,0 +1,54 @@+module Kafka.Worker.Analytics+  ( Analytics,+    AssignedPartitions (AssignedPartitions),+    PausedPartitions (PausedPartitions),+    TimeOfLastRebalance (TimeOfLastRebalance),+    init,+    read,+    updatePaused,+    updateTimeOfLastRebalance,+  )+where++import qualified Control.Concurrent.STM as STM+import qualified Control.Concurrent.STM.TVar as TVar+import qualified Prelude++newtype PausedPartitions = PausedPartitions Int++newtype AssignedPartitions = AssignedPartitions Int++newtype TimeOfLastRebalance = TimeOfLastRebalance Float++data Analytics = Analytics+  { pausedPartitions :: TVar.TVar PausedPartitions,+    timeOfLastRebalance :: TVar.TVar TimeOfLastRebalance,+    assignedPartitions :: Prelude.IO AssignedPartitions+  }++init :: Prelude.IO Int -> Prelude.IO Analytics+init assignedPartitions' = do+  pausedPartitions <- TVar.newTVarIO (PausedPartitions 0)+  timeOfLastRebalance <- TVar.newTVarIO (TimeOfLastRebalance 0)+  Prelude.pure+    ( Analytics+        { pausedPartitions,+          timeOfLastRebalance,+          assignedPartitions = map AssignedPartitions assignedPartitions'+        }+    )++read :: Analytics -> Prelude.IO (PausedPartitions, AssignedPartitions, TimeOfLastRebalance)+read Analytics {pausedPartitions, assignedPartitions, timeOfLastRebalance} = do+  analyticsPausedPartitions <- TVar.readTVarIO pausedPartitions+  analyticsTimeOfLastRebalance <- TVar.readTVarIO timeOfLastRebalance+  analyticsAssignedPartitions <- assignedPartitions+  Prelude.pure (analyticsPausedPartitions, analyticsAssignedPartitions, analyticsTimeOfLastRebalance)++updatePaused :: Int -> Analytics -> Prelude.IO ()+updatePaused numPaused Analytics {pausedPartitions} =+  STM.atomically <| TVar.writeTVar pausedPartitions (PausedPartitions numPaused)++updateTimeOfLastRebalance :: Float -> Analytics -> Prelude.IO ()+updateTimeOfLastRebalance now Analytics {timeOfLastRebalance} = do+  STM.atomically <| TVar.writeTVar timeOfLastRebalance (TimeOfLastRebalance now)
+ src/Kafka/Worker/Fetcher.hs view
@@ -0,0 +1,204 @@+module Kafka.Worker.Fetcher (pollingLoop) where++import qualified Control.Concurrent+import qualified Control.Exception.Safe as Exception+import qualified Data.ByteString as ByteString+import qualified Dict+import qualified GHC.Clock+import qualified Kafka.Consumer as Consumer+import qualified Kafka.Worker.Analytics as Analytics+import qualified Kafka.Worker.Partition as Partition+import qualified Kafka.Worker.Settings as Settings+import qualified Prelude++type EnqueueRecord = (ConsumerRecord -> Prelude.IO Partition.SeekCmd)++-- | pollingLoop+-- our long-running event loop that+-- - polls for new messages+-- - for each message, spawns a thread for its partition if it doesn't yet exist+-- - appends the message to an in-memory queue that's being worked on by a partition-specific thread+pollingLoop ::+  Settings.Settings ->+  EnqueueRecord ->+  Analytics.Analytics ->+  Consumer.KafkaConsumer ->+  Prelude.IO ()+pollingLoop settings enqueueRecord analytics consumer = do+  now <- nextPollingTimestamp+  pollingLoop' settings enqueueRecord analytics consumer (pollTimeIsOld now)++newtype LastPollingTimestamp = LastPollingTimestamp Float++newtype NextPollingTimestamp = NextPollingTimestamp Float++pollTimeIsOld :: NextPollingTimestamp -> LastPollingTimestamp+pollTimeIsOld (NextPollingTimestamp time) = LastPollingTimestamp time++nextPollingTimestamp :: Prelude.IO NextPollingTimestamp+nextPollingTimestamp = do+  map NextPollingTimestamp GHC.Clock.getMonotonicTime++type ConsumerRecord = Consumer.ConsumerRecord (Maybe ByteString.ByteString) (Maybe ByteString.ByteString)++pollingLoop' ::+  Settings.Settings ->+  EnqueueRecord ->+  Analytics.Analytics ->+  Consumer.KafkaConsumer ->+  LastPollingTimestamp ->+  Prelude.IO ()+pollingLoop'+  settings@Settings.Settings+    { Settings.pollingTimeout,+      Settings.pollBatchSize,+      Settings.maxMsgsPerSecondPerPartition,+      Settings.maxPollIntervalMs+    }+  enqueueRecord+  analytics+  consumer+  lastPollTimestamp = do+    -- we block here if we're actively revoking+    -- Check whether we need to shut down while long-polling for new messages.+    eitherMsgs <- Consumer.pollMessageBatch consumer pollingTimeout pollBatchSize+    msgs <- Prelude.traverse handleKafkaError eitherMsgs+    assignment <-+      Consumer.assignment consumer+        |> andThen handleKafkaError+    appendResults <-+      msgs+        -- We occasionally get a message here for a partition that based on+        -- internal state we believed to be revoked. We feel uneasy just+        -- dropping those messages, for what if our internal state is wrong? We+        -- might be dropping messages we really should be processing.+        -- So instead we ask librdkafka to tell us what our current assignment+        -- is. If we receive messages for partitions outside of that+        -- assignment, then we can confidently drop them.+        |> List.filter (msgIsForAssignedPartition assignment)+        -- Enqueue messages in per-partition queues.+        |> Prelude.traverse enqueueRecord+    List.map2 (,) (List.map getPartitionKey msgs) appendResults+      |> groupDictAndMap identity+      |> Dict.toList+      |> List.filterMap toSeekPartition+      |> seek consumer+    now <- nextPollingTimestamp+    throttle maxMsgsPerSecondPerPartition maxPollIntervalMs (List.length appendResults) analytics now lastPollTimestamp+    pollingLoop' settings enqueueRecord analytics consumer (pollTimeIsOld now)++getPartitionKey :: Consumer.ConsumerRecord k v -> (Consumer.TopicName, Consumer.PartitionId)+getPartitionKey record =+  ( Consumer.crTopic record,+    Consumer.crPartition record+  )++toSeekPartition ::+  ( (Consumer.TopicName, Consumer.PartitionId),+    List Partition.SeekCmd+  ) ->+  Maybe Consumer.TopicPartition+toSeekPartition ((topicName, partitionId), appendResults) =+  -- Among they last batch of fetched messages might have been multiple messages+  -- for this partition, which we subsequently tried to enqueue. It's possible+  -- the first message might have had an offset smaller than the one that we+  -- were looking for, but that somewhere in the middle of the series we caught+  -- up. That's why we consider only the last message appended to the patition+  -- in this batch. If that one was succesfull then there's nothing for us to+  -- do. If it had an unexpected offset then we should seek.+  case last appendResults of+    Nothing -> Nothing+    Just Partition.NoSeek -> Nothing+    Just (Partition.SeekToOffset offset) ->+      Just+        Consumer.TopicPartition+          { Consumer.tpTopicName = topicName,+            Consumer.tpPartition = partitionId,+            Consumer.tpOffset = Consumer.PartitionOffset offset+          }++last :: List a -> Maybe a+last list = List.head (List.reverse list)++seek :: Consumer.KafkaConsumer -> List Consumer.TopicPartition -> Prelude.IO ()+seek consumer partitions = do+  let goSeek = Consumer.seek consumer (Consumer.Timeout 5000 {- 5 seconds -}) partitions+  maybeSeekError <- goSeek+  case maybeSeekError of+    Nothing -> Prelude.pure ()+    Just _ -> do+      -- Retry once, after a delay, because we're seeing reports+      -- that attempting to `seek` after just having called+      -- `assign` (which hw-kafka-client does for us before running+      -- this callback) might result in seek failing. See:+      -- https://github.com/confluentinc/confluent-kafka-dotnet/issues/1303+      Control.Concurrent.threadDelay 5_000_000 {- 5 seconds -}+      maybeSeekError2 <- goSeek+      case maybeSeekError2 of+        Nothing -> Prelude.pure ()+        Just seekError2 -> Exception.throwIO seekError2++msgIsForAssignedPartition ::+  Dict.Dict Consumer.TopicName [Consumer.PartitionId] ->+  ConsumerRecord ->+  Bool+msgIsForAssignedPartition assignment msg =+  case Dict.get (Consumer.crTopic msg) assignment of+    Nothing -> False+    Just partitionIds ->+      List.member (Consumer.crPartition msg) partitionIds++handleKafkaError :: Prelude.Either Consumer.KafkaError a -> Prelude.IO a+handleKafkaError eitherMsg = do+  case eitherMsg of+    Prelude.Left err ->+      -- Kill the worker process if polling for messages results in an error.+      -- Every individual message of a batch can contain an error. It's unclear+      -- what the implications of that are. Specifically: could such errors+      -- result in holes in the stream of messages for a partition? That would+      -- be bad, it could result in some messages being ignored.+      --+      -- Crashing the worker seems a "safe" option at least. If such crashes+      -- are very rare this solution might be good enough. If it happens+      -- regularly we should figure out a better solution.+      Exception.throwIO err+    Prelude.Right record ->+      Prelude.pure record++-- | Call on the poll thread after fetching a new batch of messages. If we're+-- ahead of our quotum this function will sleep for a bit, delaying the fetch of+-- the next batch.+throttle ::+  Settings.MaxMsgsPerSecondPerPartition ->+  Settings.MaxPollIntervalMs ->+  Int ->+  Analytics.Analytics ->+  NextPollingTimestamp ->+  LastPollingTimestamp ->+  Prelude.IO ()+throttle Settings.DontThrottle _ _ _ _ _ = Prelude.pure ()+throttle (Settings.ThrottleAt maxMsgsPerSecondPerPartition) maxPollIntervalMs newPolledMessages analytics (NextPollingTimestamp now) (LastPollingTimestamp lastPollTimestamp) = do+  (_, Analytics.AssignedPartitions numPartitions, _) <- Analytics.read analytics+  let timeDiff = Prelude.floor (now - lastPollTimestamp)+  let quotumPerSecond = maxMsgsPerSecondPerPartition * numPartitions+  let quotum = timeDiff * quotumPerSecond+  let overQuotum = newPolledMessages - quotum+  let secondsToSleep =+        Prelude.fromIntegral overQuotum / Prelude.fromIntegral quotumPerSecond+  let microSecondsToSleep =+        Prelude.floor (secondsToSleep * 1e6)+          |> min (Prelude.fromIntegral <| Settings.unMaxPollIntervalMs maxPollIntervalMs - 100) -- -100ms so that it has time to loop.+  if microSecondsToSleep > 0+    then Control.Concurrent.threadDelay microSecondsToSleep+    else Prelude.pure ()++groupDictAndMap :: Ord b => (a -> (b, c)) -> List a -> Dict.Dict b (List c)+groupDictAndMap f =+  List.foldr+    ( \x ->+        Dict.update (Tuple.first (f x)) <| \val ->+          case val of+            Nothing -> Just [Tuple.second (f x)]+            Just y -> Just (Tuple.second (f x) : y)+    )+    Dict.empty
+ src/Kafka/Worker/Internal.hs view
@@ -0,0 +1,588 @@+{-# LANGUAGE GADTs #-}++module Kafka.Worker.Internal where++import qualified Conduit+import qualified Control.Concurrent+import qualified Control.Concurrent.Async as Async+import qualified Control.Concurrent.STM as STM+import qualified Control.Concurrent.STM.TVar as TVar+import qualified Control.Exception.Safe as Exception+import qualified Data.Aeson as Aeson+import qualified Data.UUID+import qualified Data.UUID.V4+import qualified Dict+import qualified GHC.Clock+import qualified Kafka.Consumer as Consumer+import qualified Kafka.Internal as Kafka+import qualified Kafka.Metadata+import qualified Kafka.Worker.Analytics as Analytics+import qualified Kafka.Worker.Fetcher as Fetcher+import qualified Kafka.Worker.Partition as Partition+import qualified Kafka.Worker.Settings as Settings+import qualified Kafka.Worker.Stopping as Stopping+import qualified Observability+import qualified Set+import qualified System.Environment+import qualified System.Exit+import qualified System.Posix.Process+import qualified System.Posix.Signals as Signals+import qualified Prelude++-- | Alias for a TopicName and PartitionId, something every message will have+type PartitionKey = (Consumer.TopicName, Consumer.PartitionId)++type AllPartitions = TVar.TVar (Dict.Dict PartitionKey Partition.Partition)++data Rebalance = Assign | Revoking | Revoked deriving (Show)++type RebalanceInfo = TVar.TVar (Dict.Dict PartitionKey (Rebalance, Float))++data State = State+  { partitions :: AllPartitions,+    -- | When we receive a shutdown signal this variable will change. Threads+    -- should stop processing new messages to allow the process to shut down.+    stopping :: Stopping.Stopping,+    analytics :: Analytics.Analytics,+    rebalanceInfo :: RebalanceInfo+  }++-- | The topics this worker should subscribe too. At the moment this library+-- only supports subscribing to a single topic.+data TopicSubscription = TopicSubscription+  { topic :: Kafka.Topic,+    onMessage :: Partition.MessageCallback,+    offsetSource :: OffsetSource+  }++-- | Params needed to write / read offsets to another data store+data PartitionOffset = PartitionOffset+  { -- | The partition of a topic.+    partitionId :: Int,+    -- | The partition's offset.+    offset :: Int+  }++-- | Create a subscription for a topic.+--+-- > main :: IO ()+-- > main = do+-- >   settings <- Environment.decode decoder+-- >   let subscription =+-- >         subscription+-- >           "the-topic"+-- >           (\msg -> Debug.todo "Process your message here!")+-- >   process settings subscription+subscription ::+  (Aeson.FromJSON msg, Aeson.ToJSON msg) =>+  Text ->+  (msg -> Task Text ()) ->+  TopicSubscription+subscription topic callback =+  TopicSubscription+    { topic = Kafka.Topic topic,+      onMessage =+        Partition.MessageCallback+          ( \_ msg -> do+              callback msg+              Task.succeed Partition.NoSeek+          ),+      offsetSource = InKafka+    }++-- | Create a subscription for a topic and manage offsets for that topic+-- yourself.+--+-- You'll need to tell Kafka where it can read starting offsets. When passed+-- a message you can also tell Kafka to seek to a different offset.+--+-- > main :: IO ()+-- > main = do+-- >   settings <- Environment.decode decoder+-- >   let subscription =+-- >         subscriptionManageOwnOffsets+-- >           "the-topic"+-- >           (\partitions ->+-- >              sql+-- >                "SELECT partition, offset FROM offsets WHERE partition = %"+-- >                [partitions] )+-- >           (\msg -> Debug.todo "Process your message here!")+-- >   process settings subscription+subscriptionManageOwnOffsets ::+  (Aeson.FromJSON msg, Aeson.ToJSON msg) =>+  Text ->+  ([Int] -> Task Text (List PartitionOffset)) ->+  (PartitionOffset -> msg -> Task Text Partition.SeekCmd) ->+  TopicSubscription+subscriptionManageOwnOffsets topic fetchOffsets callback =+  TopicSubscription+    { topic = Kafka.Topic topic,+      onMessage =+        Partition.MessageCallback+          ( \record msg -> do+              let offsetParams =+                    PartitionOffset+                      { partitionId =+                          Consumer.crPartition record+                            |> partitionIdToInt,+                        offset = Consumer.unOffset (Consumer.crOffset record)+                      }+              callback offsetParams msg+          ),+      offsetSource =+        Elsewhere+          ( \partitionKeys -> do+              let partitionIds =+                    partitionKeys+                      |> List.map (partitionIdToInt << Tuple.second)+              offsets <- fetchOffsets partitionIds+              offsets+                |> List.map toPartitionKey+                |> Task.succeed+          )+    }+  where+    toPartitionKey :: PartitionOffset -> (PartitionKey, Int)+    toPartitionKey (PartitionOffset {partitionId, offset}) =+      ( ( Consumer.TopicName topic,+          Consumer.PartitionId (Prelude.fromIntegral partitionId)+        ),+        offset+      )++    partitionIdToInt :: Consumer.PartitionId -> Int+    partitionIdToInt (Consumer.PartitionId int) = Prelude.fromIntegral int++-- | This determines how a worker that was just assigned a partition should+-- decide at which message offset to continue processing.+data OffsetSource where+  -- | Use Kafka's own offset storage mechanism.+  InKafka :: OffsetSource+  -- | Store offsets somewhere else, in which case you need to provide a+  -- function that the worker can use to load initial offsets. Storing offsets+  -- outside Kafka can be used to implement exactly-once-delivery schemes.+  -- Using this requires the message itself to commit the offset.+  Elsewhere ::+    ([PartitionKey] -> Task Text [(PartitionKey, Int)]) ->+    OffsetSource++-- | Starts the kafka worker handling messages.+process :: Settings.Settings -> Text -> TopicSubscription -> Prelude.IO ()+process settings groupIdText topicSubscriptions = do+  processWithoutShutdownEnsurance settings (Consumer.ConsumerGroupId groupIdText) topicSubscriptions+  -- Start an ensurance policy to make sure we exit in 5 seconds. We've seen+  -- cases where our graceful shutdown seems to hang, resulting in a worker+  -- that's not doing anything. We should try to fix those failures, but for the+  -- ones that remain this is our fallback.+  --+  -- Running it using `Async.async` makes it so we won't wait for this thread to+  -- complete. If the regular shutdown completes before this thread is done we+  -- will exit early.+  _ <-+    Async.async <| do+      Control.Concurrent.threadDelay 5_000_000 {- 5 seconds -}+      Prelude.putStrLn "Something is holding up shutdown. Going to die ungracefully now."+      System.Posix.Process.exitImmediately (System.Exit.ExitFailure 1)+  Prelude.pure ()++-- | Like `process`, but doesn't exit the current process by itself. This risks+-- leaving zombie processes when used in production but is safer in tests, where+-- the worker shares the OS process with other test code and the test runner.+processWithoutShutdownEnsurance :: Settings.Settings -> Consumer.ConsumerGroupId -> TopicSubscription -> Prelude.IO ()+processWithoutShutdownEnsurance settings groupId topicSubscriptions = do+  let TopicSubscription {onMessage, topic, offsetSource} = topicSubscriptions+  state <- initState+  onQuitSignal (Stopping.stopTakingRequests (stopping state))+  Conduit.withAcquire (Observability.handler (Settings.observability settings)) <| \observabilityHandler -> do+    Exception.bracketWithError+      (createConsumer settings groupId observabilityHandler offsetSource onMessage topic state)+      (cleanUp observabilityHandler (rebalanceInfo state) (stopping state))+      (runThreads settings state)++initState :: Prelude.IO State+initState = do+  stopping <- Stopping.init+  partitions <- TVar.newTVarIO Dict.empty+  analytics <- Analytics.init (map Dict.size (TVar.readTVarIO partitions))+  rebalanceInfo <- TVar.newTVarIO Dict.empty+  Prelude.pure+    State+      { partitions,+        stopping,+        analytics,+        rebalanceInfo+      }++-- | Goes to Kafka and registers a consumer on the Topic+-- Kafka whould give the consumer some number of partitions to be responsible for+createConsumer ::+  Settings.Settings ->+  Consumer.ConsumerGroupId ->+  Observability.Handler ->+  OffsetSource ->+  Partition.MessageCallback ->+  Kafka.Topic ->+  State ->+  Prelude.IO Consumer.KafkaConsumer+createConsumer+  Settings.Settings+    { Settings.brokerAddresses,+      Settings.logLevel,+      Settings.maxPollIntervalMs,+      Settings.onProcessMessageSkip+    }+  groupId+  observability+  offsetSource+  callback+  topic+  state = do+    let rebalance =+          rebalanceCallback+            onProcessMessageSkip+            observability+            callback+            offsetSource+            state+    let properties =+          Consumer.brokersList+            brokerAddresses+            ++ Consumer.groupId groupId+            ++ Consumer.noAutoCommit+            ++ Consumer.logLevel logLevel+            ++ Consumer.setCallback (Consumer.rebalanceCallback rebalance)+            ++ Consumer.compression Consumer.Snappy+            ++ Consumer.extraProps+              ( Dict.fromList+                  [("max.poll.interval.ms", Text.fromInt (Settings.unMaxPollIntervalMs maxPollIntervalMs))]+              )+    let subscription' =+          Consumer.topics [Consumer.TopicName (Kafka.unTopic topic)]+            ++ Consumer.offsetReset Consumer.Earliest+    eitherConsumer <- Consumer.newConsumer properties subscription'+    case eitherConsumer of+      Prelude.Left err ->+        -- We create the worker as part of starting the application. Throwing+        -- means that if there's a problem with the settings the application will+        -- fail immediately upon start. It won't result in runtime errors during+        -- operation.+        Exception.throwIO err+      Prelude.Right consumer ->+        Prelude.pure consumer++-- | Triggered when a rebalance happens, due to workers going offline or coming+-- online. This typically happens when we're scaling up or down, or when a+-- worker dies because of an unexpected exception.+--+-- The following provides reasonable documentation for these events. It's Java+-- documentation, but it should apply to what's happening here fairly well (the+-- `hw-kafka-client` library we use is a wrapper over librdkafka).+-- https://docs.confluent.io/2.0.0/clients/librdkafka/classRdKafka_1_1RebalanceCb.html#a490a91c52724382a72380af621958741+rebalanceCallback ::+  Settings.SkipOrNot ->+  Observability.Handler ->+  Partition.MessageCallback ->+  OffsetSource ->+  State ->+  Consumer.KafkaConsumer ->+  Consumer.RebalanceEvent ->+  Prelude.IO ()+rebalanceCallback skipOrNot observability callback offsetSource state consumer rebalanceEvent = do+  now <- GHC.Clock.getMonotonicTime+  Analytics.updateTimeOfLastRebalance now (analytics state)+  case rebalanceEvent of+    Consumer.RebalanceBeforeAssign newPartitions -> do+      keysWithOffsets <-+        case offsetSource of+          InKafka ->+            Prelude.pure <| List.map (\partitionKey -> (partitionKey, Partition.ToKafka)) newPartitions+          Elsewhere fetch -> do+            log <- Platform.silentHandler+            fetchResult <- Task.attempt log (fetch newPartitions)+            case fetchResult of+              Err err -> Exception.throwString (Text.toList err)+              Ok fetched -> do+                let storedOffsets =+                      List.map (Tuple.mapSecond Partition.Elsewhere) fetched+                        |> Dict.fromList+                let storedKeys =+                      fetched+                        |> List.map Tuple.first+                        |> Set.fromList+                let missingPartitions = Set.diff (Set.fromList newPartitions) storedKeys+                -- NOTE: we fallback to the end of the queue so we can reset+                -- offsets in staging.+                -- in production, we should never hit this code. Perhaps we+                -- should explicitly guard against it?+                waterMarkInfo :: Prelude.Either Consumer.KafkaError (List Kafka.Metadata.WatermarkOffsets) <-+                  missingPartitions+                    |> Set.toList+                    |> Prelude.traverse+                      ( \(topicName, partitionId) ->+                          Kafka.Metadata.partitionWatermarkOffsets+                            consumer+                            (Consumer.Timeout 5000 {- 5 seconds -})+                            topicName+                            partitionId+                      )+                    |> map Prelude.sequence+                case waterMarkInfo of+                  Prelude.Left err -> Exception.throwIO err+                  Prelude.Right waterMarks ->+                    let fallbackOffsets =+                          waterMarks+                            |> List.map+                              ( \Kafka.Metadata.WatermarkOffsets {Kafka.Metadata.woTopicName, Kafka.Metadata.woPartitionId, Kafka.Metadata.woHighWatermark} ->+                                  ((woTopicName, woPartitionId), Partition.Elsewhere (Consumer.unOffset woHighWatermark))+                              )+                            |> Dict.fromList+                     in Dict.union storedOffsets fallbackOffsets+                          |> Dict.toList+                          |> Prelude.pure+      -- We are being assigned new partitions. Lets prepare threads for the new+      -- messages we're about to receive. Lib-rdkafka hasn't yet started to send+      -- messages, so the queues exist but should be idle+      keysWithOffsets+        |> Prelude.traverse+          ( \(partitionKey, offset) -> do+              initPartition+                skipOrNot+                offset+                observability+                consumer+                callback+                state+                partitionKey+              STM.atomically+                <| TVar.modifyTVar' (rebalanceInfo state) (Dict.insert partitionKey (Assign, now))+          )+        |> map (\_ -> ())+    Consumer.RebalanceAssign _ -> Prelude.pure ()+    Consumer.RebalanceBeforeRevoke revokedPartitions -> do+      -- This callback is intended to allow us to finish committing any+      -- ongoing work before rebalancing. This avoids processing messages twice.+      --+      -- When the callback returns, the rebalancing occurs.+      --+      -- We will tell workers to wrap up current work, and wait for them to+      -- complete before rebalancing.+      --+      -- First: mark partitions as Stopping. Worker threads will stop.+      _ <-+        revokedPartitions+          |> Prelude.traverse+            ( \partitionKey -> do+                STM.atomically <| do+                  partitions <- TVar.readTVar (partitions state)+                  case Dict.get partitionKey partitions of+                    Nothing -> Prelude.pure ()+                    Just partition -> Partition.revoke partition+                STM.atomically <| TVar.modifyTVar' (rebalanceInfo state) (Dict.insert partitionKey (Revoking, now))+            )+      -- Second: Wait for the workers to stop working.+      -- (They will only remove themselves from the partitions when done processing+      -- any ongoing message, so checking non-existence should suffice.)+      _ <-+        revokedPartitions+          |> Prelude.traverse+            ( \partitionKey -> do+                STM.atomically <| do+                  partitions <- TVar.readTVar (partitions state)+                  case Dict.get partitionKey partitions of+                    Nothing -> Prelude.pure ()+                    Just _ -> do+                      -- we cannot block on work happening in the FETCHER+                      -- but here, we're blocking work happening in the WORKER+                      -- which is fine!+                      STM.retry+                STM.atomically <| TVar.modifyTVar' (rebalanceInfo state) (Dict.insert partitionKey (Revoked, now))+            )+      -- Now workers have stopped and it's safe to rebalance.+      -- Returning from this callback starts the rebalance.+      Prelude.pure ()+    Consumer.RebalanceRevoke _revokedPartitions -> do+      Prelude.pure ()++-- | Disconnects our Consumer / yields back partitions on quit / node shutdown+cleanUp :: Observability.Handler -> RebalanceInfo -> Stopping.Stopping -> Maybe Exception.SomeException -> Consumer.KafkaConsumer -> Prelude.IO ()+cleanUp observabilityHandler rebalanceInfo stopping maybeException consumer = do+  Prelude.putStrLn "Cleaning up"+  _ <- Consumer.closeConsumer consumer+  -- ensure workers shut down+  Stopping.stopTakingRequests stopping+  requestId <- map Data.UUID.toText Data.UUID.V4.nextRandom+  -- at some point, k8s should report system crashes. In the mean time, we'll do it.+  Platform.rootTracingSpanIO+    requestId+    (Observability.report observabilityHandler requestId)+    "Kafka consumer shutting down"+    <| \log -> do+      case maybeException of+        Nothing -> Prelude.pure ()+        Just exception -> do+          rebalanceInfo' <- TVar.readTVarIO rebalanceInfo+          Log.error+            "Kafka consumer crashed"+            [ Log.context "triage" ("The consumer should automatically restart. If we see lots of these in a short period of time we should try to figure out what's wrong" :: Text),+              Log.context "error" (Debug.toString exception),+              Log.context "rebalance info" (Debug.toString rebalanceInfo')+            ]+            |> Task.perform log++  writeCrashLogOnError maybeException+  Prelude.putStrLn "Bye!"++-- | Handle crash logging+writeCrashLogOnError :: Maybe Exception.SomeException -> Prelude.IO ()+writeCrashLogOnError maybeException = do+  -- Not using the nri-env-parser lib for this configuration option because it would+  -- require us to run code to make it available. If that code failed it+  -- wouldn't end up in the crash log! Using only `base` functionality allows us+  -- to put the crashlog reporting in the very root of the application.+  crashLogPath <- System.Environment.lookupEnv "CRASHLOG_PATH"+  let crashLog =+        case maybeException of+          Nothing -> "System exited in response to signal"+          Just exception -> Exception.displayException exception+  case crashLogPath of+    Nothing -> Prelude.pure ()+    Just "" -> Prelude.pure ()+    Just path -> Prelude.writeFile path crashLog++-- Adds a partition to our partitions dict. These partitions will be idle until+-- lib-rdkafka actually starts sending us new messages+-- (after the Consumer.rebalanceassign event occurs)+initPartition ::+  Settings.SkipOrNot ->+  Partition.CommitOffsets ->+  Observability.Handler ->+  Consumer.KafkaConsumer ->+  Partition.MessageCallback ->+  State ->+  PartitionKey ->+  Prelude.IO ()+initPartition skipOrNot commitOffset observabilityHandler consumer callback state key = do+  -- # Start worker thread for handling messages in partition.+  Partition.spawnWorkerThread+    skipOrNot+    commitOffset+    observabilityHandler+    (analytics state)+    (stopping state)+    consumer+    callback+    -- startup function+    ( Partition.OnStartup+        ( \partition ->+            STM.atomically <| do+              queues <- TVar.readTVar (partitions state)+              case Dict.get key queues of+                Just _ ->+                  -- We never expect to be asked to create a partition that already exists.+                  -- We expect that revoke has completely removed the partition before+                  -- re-assign.+                  --+                  -- If it happens anyway, it seems safer to crash (and restart) than to+                  -- try continue into the unknown.+                  STM.throwSTM (AskedToInitPartitionThatAlreadyExists key)+                Nothing -> do+                  TVar.writeTVar (partitions state) (Dict.insert key partition queues)+        )+    )+    -- cleanup function+    ( Partition.OnCleanup+        ( do+            -- Remove the partition from the dict to clean up memory+            STM.atomically <| TVar.modifyTVar' (partitions state) (Dict.remove key)+            Prelude.putStrLn ("Stop processing messages for partition: " ++ Prelude.show key)+        )+    )++runThreads ::+  Settings.Settings ->+  State ->+  Consumer.KafkaConsumer ->+  Prelude.IO ()+runThreads settings state consumer = do+  Stopping.runUnlessStopping+    (stopping state)+    ()+    ( Async.race+        (pauseAndAnalyticsLoop (Settings.maxMsgsPerPartitionBufferedLocally settings) consumer state Set.empty)+        (Fetcher.pollingLoop settings (enqueueRecord (partitions state)) (analytics state) consumer)+        |> map (\_ -> ())+    )++data RuntimeExceptions+  = ReceivedMsgNotInAssignedPartitions (Consumer.TopicName, Consumer.PartitionId)+  | AskedToInitPartitionThatAlreadyExists (Consumer.TopicName, Consumer.PartitionId)+  deriving (Show)++instance Exception.Exception RuntimeExceptions++enqueueRecord ::+  AllPartitions ->+  Partition.ConsumerRecord ->+  Prelude.IO Partition.SeekCmd+enqueueRecord partitions record =+  STM.atomically <| do+    let key = (Consumer.crTopic record, Consumer.crPartition record)+    partitions' <- TVar.readTVar partitions+    let maybePartition = Dict.get key partitions'+    case maybePartition of+      Nothing -> STM.throwSTM (ReceivedMsgNotInAssignedPartitions key)+      Just partition -> Partition.append record partition++-- | Intermittently updates+-- - paused partitions to reflect desired state.+-- - analytics, so that the worker node can report up-to-date data to honeycomb+pauseAndAnalyticsLoop ::+  Settings.MaxMsgsPerPartitionBufferedLocally ->+  Consumer.KafkaConsumer ->+  State ->+  Set.Set PartitionKey ->+  Prelude.IO ()+pauseAndAnalyticsLoop maxBufferSize consumer state pausedPartitions = do+  desiredPausedPartitions <- pausedPartitionKeys maxBufferSize (partitions state)+  Analytics.updatePaused (Set.size desiredPausedPartitions) (analytics state)+  let newlyPaused = Set.diff desiredPausedPartitions pausedPartitions+  _ <- Consumer.pausePartitions consumer (Set.toList newlyPaused)+  let newlyResumed = Set.diff pausedPartitions desiredPausedPartitions+  _ <- Consumer.resumePartitions consumer (Set.toList newlyResumed)+  Control.Concurrent.threadDelay 1_000_000 {- 1 second -}+  pauseAndAnalyticsLoop maxBufferSize consumer state desiredPausedPartitions++pausedPartitionKeys :: Settings.MaxMsgsPerPartitionBufferedLocally -> AllPartitions -> Prelude.IO (Set.Set PartitionKey)+pausedPartitionKeys (Settings.MaxMsgsPerPartitionBufferedLocally maxBufferSize) partitions = do+  partitions' <- TVar.readTVarIO partitions+  partitions'+    |> Dict.toList+    |> Prelude.traverse+      ( \(key, partition) -> do+          maybeLen <- Partition.length partition+          Prelude.pure+            <| case maybeLen of+              Nothing -> Nothing+              Just length ->+                if length > maxBufferSize+                  then Just key+                  else Nothing+      )+    |> map (List.filterMap identity >> Set.fromList)++quitSignals :: [Signals.Signal]+quitSignals =+  [ Signals.sigINT, -- ctrl-c+    Signals.sigQUIT, -- ctrl-\ ???+    Signals.sigTERM+  ]++onQuitSignal :: Prelude.IO () -> Prelude.IO ()+onQuitSignal release = do+  let handleQuit signal =+        Signals.installHandler+          signal+          (Signals.Catch release)+          Nothing+  _ <- Prelude.traverse handleQuit quitSignals+  Prelude.pure ()
+ src/Kafka/Worker/Partition.hs view
@@ -0,0 +1,503 @@+{-# LANGUAGE GADTs #-}++module Kafka.Worker.Partition+  ( spawnWorkerThread,+    append,+    length,+    revoke,+    ConsumerRecord,+    Partition,+    MessageCallback (..),+    SeekCmd (..),+    CommitOffsets (..),+    -- just exported for tests+    microSecondsDelayForAttempt,+    OnStartup (OnStartup),+    OnCleanup (OnCleanup),+  )+where++import qualified Control.Concurrent+import qualified Control.Concurrent.Async as Async+import qualified Control.Concurrent.STM as STM+import qualified Control.Concurrent.STM.TVar as TVar+import qualified Control.Exception.Safe as Exception+import qualified Data.Aeson as Aeson+import qualified Data.ByteString as ByteString+import qualified Data.Sequence as Seq+import qualified Data.Text.Encoding+import qualified Data.Time.Clock as Clock+import qualified Data.Time.Clock.POSIX as Clock.POSIX+import qualified Data.UUID+import qualified Data.UUID.V4+import qualified GHC.Clock+import qualified Kafka.Consumer as Consumer+import qualified Kafka.Internal as Internal+import qualified Kafka.Worker.Analytics as Analytics+import qualified Kafka.Worker.Settings as Settings+import qualified Kafka.Worker.Stopping as Stopping+import qualified Log.Kafka+import qualified Observability+import qualified Platform+import qualified Prelude++data WorkerError e+  = WorkerCallbackFailed e+  | WorkerCallbackThrew Exception.SomeException+  | MsgDecodingFailed Text+  | SeekFailed Consumer.KafkaError+  deriving (Show)++data State = State+  { analytics :: Analytics.Analytics,+    stopping :: Stopping.Stopping,+    partition :: Partition+  }++type ConsumerRecord = Consumer.ConsumerRecord (Maybe ByteString.ByteString) (Maybe ByteString.ByteString)++newtype ProcessAttemptsCount = ProcessAttemptsCount Int+  deriving (Num)++newtype Partition = Partition (TVar.TVar Backlog)++-- | SeekCmd is the expected response of the MessageCallback.+-- Return NoSeek if processing succeeded and the offset was correct.+-- Return SeekToOffset with the expected offset if the offset of the message+-- processed was wrong (this only makes sense when Kafka isn't managing the offset)+-- The MessageCallback will throw if proccessing throws.+data SeekCmd+  = NoSeek+  | SeekToOffset Int++-- for each partition, we keep a local backlog of messages we've read from Kafka+-- that are not yet processed+data Backlog+  = AwaitingSeekTo Int+  | -- | Use a `Sequence` type to store records. It's a bit like a list, but unlike+    -- lists it has O(1) appends to both ends, and an O(1) length function as well.+    -- We use this a lot considering we're appending new messages on one end, and+    -- processes messages from the other.+    Assigned (Seq.Seq (ProcessAttemptsCount, ConsumerRecord))+  | -- we set the Backlog to this to tell the partition's thread to stop+    Stopping++-- using these types as a proxy for named parameters+-- this is a function that runs with the partition once it's started+newtype OnStartup = OnStartup (Partition -> Prelude.IO ())++-- using this types as a proxy for named parameters+-- this is a function that runs with the partition when it's done to cleanup+newtype OnCleanup = OnCleanup (Prelude.IO ())++-- | MessageCallback is used by your worker to process messages peeled off the queue.+data MessageCallback where+  MessageCallback ::+    (Show e, Aeson.ToJSON msg, Aeson.FromJSON msg) =>+    (Consumer.ConsumerRecord () () -> msg -> Task e SeekCmd) ->+    MessageCallback++data CommitOffsets+  = ToKafka+  | -- | Commit offsets elsewhere. Takes the offset of the last committed+    -- message so we can skip old messages.+    Elsewhere Int++-- | A thread that processes messages for a particular partition. Cleans itself+-- up if it ever runs out.+--+-- When processing a message fails we have a couple of non-appealing options:+--+-- - We could throw an error. That would kill the worker process, preventing it+--   from working on unrelated partitions that might be fine.+-- - We can skip the message. This might be okay for some domains, but we cannot+--   generally assume ignorning messages is fine.+-- - We can park the message in a separate topic for processing later (a dead+--   letter queue) This might be fine for some domains, but we cannot generally+--   assume handling messages out of order is fine.+-- - We can retry until it works. If we're lucky this will get the message+--   unstuck, but some logic errors (such as JSON decoding errors) won't go+--   away by themselves and will require new code to be deployed. Until a retry+--   succeeds we'll be blocked from handling any messages on the same partition.+--+-- The retrying solution is unsatisfying, but at least is simple to implement,+-- and will not attempt to fix matters in a way that might also make them worse.+--+-- In domains where we skipping messages or resubmitting them out of order is+-- okay we already have the option of passing in a callback function that+-- catches and handles errors itself and always succeeds.+--+-- In case message order needs to be maintainted, there might be strategies we+-- can devise where being blocked on one key would allow processing on other+-- keys within the same partition to continue. It's not clear however how often+-- a logic error will affect only some keys in a partition and not others, and+-- so whether it's worth it to expend the effort and complexity to implement+-- such a scheme. As we run this code we'll gather data that can help us decide.+spawnWorkerThread ::+  Settings.SkipOrNot ->+  CommitOffsets ->+  Observability.Handler ->+  Analytics.Analytics ->+  Stopping.Stopping ->+  Consumer.KafkaConsumer ->+  MessageCallback ->+  OnStartup ->+  OnCleanup ->+  Prelude.IO ()+spawnWorkerThread skipOrNot commitOffsets observabilityHandler analytics stopping consumer callback (OnStartup onStartup) (OnCleanup onCleanup) = do+  -- Synchronously create the queue that will come to contain messages for the+  -- partition. This way we'll be able to start receiving messages for this+  -- partition as soon as this function returns, even if the processing thread+  -- we start below still needs boot.+  partition <-+    map Partition <| TVar.newTVarIO+      <| case commitOffsets of+        ToKafka -> Assigned Seq.empty+        Elsewhere offset -> AwaitingSeekTo offset+  onStartup partition+  Exception.finally+    (processMsgLoop skipOrNot commitOffsets observabilityHandler State {analytics, stopping, partition} consumer callback)+    onCleanup+    |> Async.async+    -- If the async process spawned here throws an exception, rethrow it+    -- in the main thread so we bring down the worker.+    --+    -- We take care to prevent exceptions from the user-provided callback+    -- to bring down the processing thread. Should an exception slip+    -- through somehow, or should the code in this module produce one,+    -- that could result in a partition thread being killed without us+    -- knowing, and without it being restarted. This would result in no+    -- messages for this partition being processed.+    --+    -- Linking the partition processing threads to the main one will+    -- ensure that if one thread goes down, they all go down. We'll get a+    -- loud crash, which isn't nice, but at least we'll know something bad+    -- happened.+    |> andThen Async.link++processMsgLoop ::+  Settings.SkipOrNot ->+  CommitOffsets ->+  Observability.Handler ->+  State ->+  Consumer.KafkaConsumer ->+  MessageCallback ->+  Prelude.IO ()+processMsgLoop skipOrNot commitOffsets observabilityHandler state consumer callback@(MessageCallback runCallback) = do+  -- # Get the next message from the queue.+  peekResponse <- peekRecord state+  case peekResponse of+    StopThread ->+      Prelude.pure ()+    (NextMsg processAttempts record) -> do+      case processAttempts of+        (ProcessAttemptsCount 0) -> Prelude.pure ()+        (ProcessAttemptsCount attempts) ->+          -- Wait a bit if this is a retry, to prevent putting a lot of retry+          -- stress on downstream systems or generating huge numbers of error+          -- messages.+          microSecondsDelayForAttempt attempts+            |> Prelude.fromIntegral+            |> Control.Concurrent.threadDelay++      doAnything <- Platform.doAnythingHandler+      let commit processResult =+            case processResult of+              SeekToOffset offset ->+                awaitingSeekTo (partition state) offset+                  |> map Ok+                  |> Platform.doAnything doAnything+              NoSeek -> do+                -- Still around? That means things must have gone well. Let's mark+                -- this message as succesfully processed.+                case commitOffsets of+                  ToKafka ->+                    commitRecord doAnything consumer record+                  Elsewhere _ ->+                    -- The user of the Kafka module is responsible for+                    -- comitting the offsets. To help them do so we pass+                    -- them the record in the callback function.+                    --+                    -- It's important the module user can determine when to+                    -- commit, for example to allow them to commit as part+                    -- of a larger database transaction as part of an+                    -- exactly-once-delivery scheme.+                    Prelude.pure ()++                -- finally, let's remove it from the queue+                dequeueRecord (partition state) record+                  |> map Ok+                  |> Platform.doAnything doAnything++      -- # Process message.+      (RequestId requestId, details) <- getTracingDetails (analytics state) processAttempts record+      Platform.rootTracingSpanIO+        requestId+        (Observability.report observabilityHandler requestId)+        "Assigned Kafka message"+        ( \log -> do+            -- Setting the tracing details first. If anything goes wrong below+            -- at least we'll have nice context in logs!+            Platform.setTracingSpanDetailsIO log details+            handleFailures log <| do+              msg <- decodeMessage record+              runCallback record {Consumer.crKey = (), Consumer.crValue = ()} msg+                |> Task.mapError WorkerCallbackFailed+                |> Task.onError+                  ( \err -> do+                      case skipOrNot of+                        Settings.Skip -> Task.succeed NoSeek+                        Settings.DoNotSkip -> Task.fail err+                  )+                |> Task.andThen commit+        )++      -- # Loop for the next message+      processMsgLoop+        skipOrNot+        commitOffsets+        observabilityHandler+        state+        consumer+        callback++microSecondsDelayForAttempt :: Int -> Int+microSecondsDelayForAttempt attempts =+  min+    3_600_000_000 {- 1 hour in microseconds -}+    ((10 Prelude.^ attempts) * 1000_000 {- 1 second in microseconds -})++handleFailures ::+  Show e =>+  Platform.LogHandler ->+  Task (WorkerError e) a ->+  Prelude.IO ()+handleFailures logHandler task = do+  result <-+    -- Catch any synchronous exceptions the callback might have thrown, to+    -- prevent them from propagating further and killing the entire worker+    -- process.+    Exception.handleAny+      (Prelude.pure << Err << WorkerCallbackThrew)+      (Task.attempt logHandler task)+  case result of+    Ok _ -> Prelude.pure ()+    Err err -> do+      Log.error+        "Assigned Kafka message failed"+        [ Log.context "triage" ("We'll automatically attempt to retry processing of the message. Until a retry succeeds no messages for the same topic partition as the failing message can be processed." :: Text),+          Log.context "error" (Debug.toString err)+        ]+        |> Task.perform logHandler++newtype RequestId = RequestId Text++getTracingDetails ::+  Analytics.Analytics ->+  ProcessAttemptsCount ->+  ConsumerRecord ->+  Prelude.IO (RequestId, Log.Kafka.Details)+getTracingDetails analytics (ProcessAttemptsCount processAttempt) record = do+  let (createTime, logAppendTime) =+        case Consumer.crTimestamp record of+          Consumer.CreateTime millis -> (Just (millisToSecs millis), Nothing)+          Consumer.LogAppendTime millis -> (Nothing, Just (millisToSecs millis))+          Consumer.NoTimestamp -> (Nothing, Nothing)+  let eitherMsg =+        Consumer.crValue record+          -- We'll accept the absence of a message if the worker expects a message+          -- of type `()`. The default JSON encoding for `()` is "[]".+          |> Maybe.withDefault "[]"+          |> Aeson.eitherDecodeStrict+  let (requestId, contents) = case eitherMsg of+        Prelude.Right Internal.MsgWithMetaData {Internal.metaData, Internal.value} ->+          (Just (Internal.requestId metaData), Log.Kafka.mkContents value)+        Prelude.Left _ ->+          ( Nothing,+            Consumer.crValue record+              |> Maybe.andThen (Data.Text.Encoding.decodeUtf8' >> Prelude.either (\_ -> Nothing) Just)+              |> Log.Kafka.mkContents+          )+  ( Analytics.PausedPartitions pausedPartitions,+    Analytics.AssignedPartitions assignedPartitions,+    Analytics.TimeOfLastRebalance timeOfLastRebalance+    ) <-+    Analytics.read analytics+  now <- GHC.Clock.getMonotonicTime+  let timeSinceLastRebalance = now - timeOfLastRebalance+  requestIdForReturn <-+    case requestId of+      Nothing ->+        -- if the message doens't contain a request id, create a new one+        map Data.UUID.toText Data.UUID.V4.nextRandom+      Just requestId' -> Prelude.pure requestId'+      |> map RequestId+  Prelude.pure+    ( requestIdForReturn,+      Log.Kafka.emptyDetails+        { Log.Kafka.topic = Just (Consumer.unTopicName (Consumer.crTopic record)),+          Log.Kafka.partitionId = Just (Prelude.fromIntegral (Consumer.unPartitionId (Consumer.crPartition record))),+          Log.Kafka.key =+            Consumer.crKey record+              |> Maybe.andThen+                ( \keyBytes ->+                    case Data.Text.Encoding.decodeUtf8' keyBytes of+                      Prelude.Left _ -> Nothing+                      Prelude.Right keyText -> Just keyText+                ),+          Log.Kafka.contents = Just contents,+          Log.Kafka.processAttempt = Just processAttempt,+          Log.Kafka.createTime,+          Log.Kafka.assignedPartitions = Just assignedPartitions,+          Log.Kafka.pausedPartitions = Just pausedPartitions,+          Log.Kafka.timeSinceLastRebalance = Just timeSinceLastRebalance,+          Log.Kafka.logAppendTime,+          Log.Kafka.requestId+        }+    )++millisToSecs :: Consumer.Millis -> Clock.UTCTime+millisToSecs (Consumer.Millis millis) = fromPosix (millis // 1000)++decodeMessage :: (Aeson.FromJSON msg) => ConsumerRecord -> Task (WorkerError e) msg+decodeMessage record = do+  let eitherMsg =+        Consumer.crValue record+          -- We'll accept the absence of a message if the worker expects a message+          -- of type `()`. The default JSON encoding for `()` is "[]".+          |> Maybe.withDefault "[]"+          |> Aeson.eitherDecodeStrict+  case eitherMsg of+    Prelude.Left err ->+      Task.fail (MsgDecodingFailed (Text.fromList err))+    Prelude.Right msgWithMetaData ->+      case Internal.value msgWithMetaData of+        (Internal.Encodable value) ->+          case Aeson.fromJSON (Aeson.toJSON value) of+            Aeson.Error err ->+              Task.fail (MsgDecodingFailed (Text.fromList err))+            Aeson.Success msg ->+              Task.succeed msg++commitRecord ::+  Platform.DoAnythingHandler ->+  Consumer.KafkaConsumer ->+  ConsumerRecord ->+  Task e ()+commitRecord doAnything consumer record = do+  commitResult <-+    Consumer.commitOffsetMessage Consumer.OffsetCommit consumer record+      |> map Ok+      |> Platform.doAnything doAnything+  case commitResult of+    Just err ->+      Log.error+        "Failed to commit offset to Kafka after succesfully processing message."+        [ Log.context "err" (Debug.toString err),+          Log.context "context" ("We failed to commit progress on the message, which means there is a risk of us processing it again. If the message is not idempotent this will be a problem. If we see a lot of these errors it might mean no commits are happening at all, in which cases our queues are not making forward progress." :: Text)+        ]+    Nothing -> Task.succeed ()++data PollResponse+  = NextMsg ProcessAttemptsCount ConsumerRecord+  | StopThread++-- | Read the next message for a particular partition, but keep it on the local+-- queue. We should only remove the message if we finish processing it.+peekRecord ::+  State ->+  Prelude.IO PollResponse+peekRecord state =+  Stopping.runUnlessStopping+    (stopping state)+    StopThread+    ( STM.atomically+        <| do+          let (Partition partition') = partition state+          backlog' <- TVar.readTVar partition'+          case backlog' of+            AwaitingSeekTo _ ->+              STM.retry+            Stopping -> do+              Prelude.pure StopThread+            Assigned Seq.Empty ->+              STM.retry+            Assigned ((processAttemptsCount, first) Seq.:<| rest) -> do+              -- Bump the retry count so that the next time we read this message, we+              -- know we've read it before.+              TVar.writeTVar+                partition'+                (Assigned ((processAttemptsCount + 1, first) Seq.:<| rest))+              Prelude.pure (NextMsg processAttemptsCount first)+    )++awaitingSeekTo :: Partition -> Int -> Prelude.IO ()+awaitingSeekTo (Partition partition) offset =+  STM.atomically (TVar.writeTVar partition (AwaitingSeekTo offset))++-- | removes the record from the Backlog if it's still assigned+-- if not assigned, doesn't matter, the current thread will die in the next+-- loop+dequeueRecord ::+  Partition ->+  ConsumerRecord ->+  Prelude.IO ()+dequeueRecord (Partition partition) record =+  STM.atomically <| do+    backlog' <- TVar.readTVar partition+    case backlog' of+      AwaitingSeekTo _ ->+        Prelude.pure ()+      Stopping ->+        Prelude.pure ()+      Assigned Seq.Empty ->+        Prelude.pure ()+      Assigned ((_, first) Seq.:<| rest) -> do+        if Consumer.crOffset first == Consumer.crOffset record+          then TVar.writeTVar partition (Assigned rest)+          else -- why would this ever be the case??? should we log here?+            Prelude.pure ()++append :: ConsumerRecord -> Partition -> STM.STM SeekCmd+append item (Partition partition) =+  TVar.stateTVar+    partition+    ( \queue' ->+        case queue' of+          AwaitingSeekTo offset ->+            if offset == Consumer.unOffset (Consumer.crOffset item)+              then+                ( NoSeek,+                  Assigned (Seq.singleton (ProcessAttemptsCount 0, item))+                )+              else+                ( SeekToOffset offset,+                  AwaitingSeekTo offset+                )+          Stopping -> (NoSeek, Stopping)+          Assigned queue ->+            ( NoSeek,+              Assigned (queue Seq.:|> (ProcessAttemptsCount 0, item))+            )+    )++length :: Partition -> Prelude.IO (Maybe Int)+length (Partition partition) = do+  backlog <- TVar.readTVarIO partition+  Prelude.pure+    <| case backlog of+      AwaitingSeekTo _ -> Nothing+      Stopping -> Nothing+      Assigned queue ->+        Just (Prelude.fromIntegral (Seq.length queue))++revoke :: Partition -> STM.STM ()+revoke (Partition partition) = TVar.writeTVar partition Stopping++-- | Create a time from a posix timestamp, a number of seconds since the Linux+-- epoch. This provides us a way to create constant timetamps for tests.+fromPosix :: Int -> Clock.UTCTime+fromPosix secondsSinceEpoch =+  secondsSinceEpoch+    |> Prelude.fromIntegral+    |> Clock.POSIX.posixSecondsToUTCTime
+ src/Kafka/Worker/Settings.hs view
@@ -0,0 +1,151 @@+module Kafka.Worker.Settings+  ( Settings (..),+    decoder,+    MaxMsgsPerSecondPerPartition (..),+    MaxMsgsPerPartitionBufferedLocally (..),+    MaxPollIntervalMs (..),+    SkipOrNot (..),+  )+where++import qualified Environment+import qualified Kafka.Consumer as Consumer+import qualified Kafka.Settings.Internal as Internal+import qualified Observability+import qualified Prelude++-- | Settings required to process kafka messages+data Settings = Settings+  { -- | broker addresses. See hw-kafka's documentation for more info+    brokerAddresses :: [Consumer.BrokerAddress],+    -- | Worker will poll Kafka for new messages. This is the timeout+    pollingTimeout :: Consumer.Timeout,+    -- | Used for throttling. Turn this down to give Kafka a speed limit.+    maxMsgsPerSecondPerPartition :: MaxMsgsPerSecondPerPartition,+    logLevel :: Internal.KafkaLogLevel,+    observability :: Observability.Settings,+    -- | Provides backpressure from message-workers to the queue-reader worker.+    -- Ensures that the thread responsible for pulling messages off of kafka+    -- doesn't race ahead / steal resources from the threads executing messages.+    maxMsgsPerPartitionBufferedLocally :: MaxMsgsPerPartitionBufferedLocally,+    pollBatchSize :: Consumer.BatchSize,+    -- | Time between polling+    maxPollIntervalMs :: MaxPollIntervalMs,+    -- | This option provides us the possibility to skip messages on failure.+    -- Useful for testing Kafka worker. DoNotSkip is a reasonable default!+    onProcessMessageSkip :: SkipOrNot+  }++-- | This option provides us the possibility to skip messages on failure.+-- Useful for testing Kafka worker. DoNotSkip is a reasonable default!+data SkipOrNot = Skip | DoNotSkip++-- | Used for throttling. Turn this down to give Kafka a speed limit.+data MaxMsgsPerSecondPerPartition = ThrottleAt Int | DontThrottle++-- | Provides backpressure from message-workers to the queue-reader worker.+-- Ensures that the thread responsible for pulling messages off of kafka+-- doesn't race ahead / steal resources from the threads executing messages.+newtype MaxMsgsPerPartitionBufferedLocally = MaxMsgsPerPartitionBufferedLocally {unMaxMsgsPerPartitionBufferedLocally :: Int}++-- | Time between polling+newtype MaxPollIntervalMs = MaxPollIntervalMs {unMaxPollIntervalMs :: Int}++-- | decodes Settings from environmental variables+-- Also consumes Observability env variables (see nri-observability)+-- KAFKA_BROKER_ADDRESSES=localhost:9092 # comma delimeted list+-- KAFKA_LOG_LEVEL=Debug+-- KAFKA_POLLING_TIMEOUT=1000+-- KAFKA_MAX_MESSAGES_PER_SECOND_PER_PARTITION=0 (disabled)+-- KAFKA_MAX_POLL_INTERVAL_MS=300000+-- KAFKA_MAX_MSGS_PER_PARTITION_BUFFERED_LOCALLY=100+-- KAFKA_POLL_BATCH_SIZE=100+-- KAFKA_SKIP_ON_PROCESS_MESSAGE_FAILURE=0+-- KAFKA_GROUP_ID=0+decoder :: Environment.Decoder Settings+decoder =+  Prelude.pure Settings+    |> andMap Internal.decoderBrokerAddresses+    |> andMap decoderPollingTimeout+    |> andMap decoderMaxMessagesPerSecondPerPartition+    |> andMap Internal.decoderKafkaLogLevel+    |> andMap Observability.decoder+    |> andMap decoderMaxMsgsPerPartitionBufferedLocally+    |> andMap decoderPollBatchSize+    |> andMap decoderMaxPollIntervalMs+    |> andMap decoderOnProcessMessageFailure++decoderPollingTimeout :: Environment.Decoder Consumer.Timeout+decoderPollingTimeout =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_POLLING_TIMEOUT",+        Environment.description = "Polling timout for consumers",+        Environment.defaultValue = "1000"+      }+    (map Consumer.Timeout Environment.int)++decoderMaxMessagesPerSecondPerPartition :: Environment.Decoder MaxMsgsPerSecondPerPartition+decoderMaxMessagesPerSecondPerPartition =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_MAX_MESSAGES_PER_SECOND_PER_PARTITION",+        Environment.description = "This is how we throttle workers. Sets the maximum amount of messages this worker should process per second per partition. 0 is disabled.",+        Environment.defaultValue = "0"+      }+    ( map+        ( \maxPerSecond ->+            ( if maxPerSecond == 0+                then DontThrottle+                else ThrottleAt maxPerSecond+            )+        )+        Environment.int+    )++decoderMaxPollIntervalMs :: Environment.Decoder MaxPollIntervalMs+decoderMaxPollIntervalMs =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_MAX_POLL_INTERVAL_MS",+        Environment.description = "This is used to set max.poll.interval.ms",+        Environment.defaultValue = "300000"+      }+    (map MaxPollIntervalMs Environment.int)++decoderMaxMsgsPerPartitionBufferedLocally :: Environment.Decoder MaxMsgsPerPartitionBufferedLocally+decoderMaxMsgsPerPartitionBufferedLocally =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_MAX_MSGS_PER_PARTITION_BUFFERED_LOCALLY",+        Environment.description = "Pausing reading from kafka when we have this many messages queued up but not yet processed",+        Environment.defaultValue = "100"+      }+    (map MaxMsgsPerPartitionBufferedLocally Environment.int)++decoderPollBatchSize :: Environment.Decoder Consumer.BatchSize+decoderPollBatchSize =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_POLL_BATCH_SIZE",+        Environment.description = "The amount of messages we request in a single poll request to Kafka",+        Environment.defaultValue = "100"+      }+    (map Consumer.BatchSize Environment.int)++decoderOnProcessMessageFailure :: Environment.Decoder SkipOrNot+decoderOnProcessMessageFailure =+  Environment.variable+    Environment.Variable+      { Environment.name = "KAFKA_SKIP_ON_PROCESS_MESSAGE_FAILURE",+        Environment.description = "Whether to skip message that are failing processing. 1 means on, 0 means off.",+        Environment.defaultValue = "0"+      }+    ( Environment.custom+        Environment.int+        ( \int ->+            if int >= 1+              then Ok Skip+              else Ok DoNotSkip+        )+    )
+ src/Kafka/Worker/Stopping.hs view
@@ -0,0 +1,34 @@+module Kafka.Worker.Stopping+  ( init,+    stopTakingRequests,+    runUnlessStopping,+    Stopping,+  )+where++import qualified Control.Concurrent.Async as Async+import qualified Control.Concurrent.MVar as MVar+import qualified Prelude++newtype Stopping = Stopping (MVar.MVar ())++init :: Prelude.IO Stopping+init = MVar.newEmptyMVar |> map Stopping++stopTakingRequests :: Stopping -> Prelude.IO ()+stopTakingRequests (Stopping stopping) = do+  Prelude.putStrLn "Gracefully shutting down..."+  MVar.tryPutMVar stopping ()+    |> map (\_ -> ())++runUnlessStopping :: Stopping -> a -> Prelude.IO a -> Prelude.IO a+runUnlessStopping (Stopping stopping) stoppingVal action =+  Async.race+    (MVar.readMVar stopping)+    action+    |> map+      ( \either ->+          case either of+            Prelude.Left () -> stoppingVal+            Prelude.Right r -> r+      )
+ test/Helpers.hs view
@@ -0,0 +1,217 @@+-- | Tests, exposed so we can run them against GHCID+module Helpers+  ( spawnWorker,+    stopWorker,+    test,+    sendSync,+  )+where++import qualified Control.Concurrent.Async as Async+import qualified Control.Concurrent.MVar as MVar+import qualified Control.Concurrent.STM as STM+import qualified Control.Exception.Safe as Exception+import qualified Data.Aeson as Aeson+import qualified Data.ByteString.Lazy+import qualified Data.UUID+import qualified Data.UUID.V4+import qualified Dict+import qualified Environment+import qualified Expect+import qualified GHC.Stack as Stack+import qualified Kafka.Consumer as Consumer+import qualified Kafka.Internal as Internal+import qualified Kafka.Producer as Producer+import qualified Kafka.Settings as Settings+import qualified Kafka.Worker as Worker+import qualified Kafka.Worker.Internal+import qualified Kafka.Worker.Settings as Worker.Settings+import qualified Platform+import qualified Test+import qualified Prelude++-- | A reference to a Kafka worker. In practice, each worker will be running+-- in its own container. For testing purposes, we sometimes launch multiple+-- workers in a single test thread.+newtype Worker = Worker (Async.Async ())++data TestHandler = TestHandler+  { producer :: Producer.KafkaProducer,+    terminator :: MVar.MVar (),+    doAnything :: Platform.DoAnythingHandler+  }++returnWhenTerminating :: TestHandler -> Prelude.IO ()+returnWhenTerminating TestHandler {terminator} = MVar.readMVar terminator++terminate :: TestHandler -> Prelude.IO ()+terminate TestHandler {terminator} = MVar.putMVar terminator ()++-- | Spawns a worker, guarded by the terminator 🦾+spawnWorker ::+  (Aeson.ToJSON msg, Aeson.FromJSON msg) =>+  TestHandler ->+  Internal.Topic ->+  (msg -> STM.STM ()) ->+  Expect.Expectation' Worker+spawnWorker handler' topic callback =+  Expect.fromIO <| do+    settings <-+      case Environment.decodeDefaults Worker.Settings.decoder of+        Ok settings' -> Prelude.pure settings'+        Err err -> Prelude.fail (Text.toList err)+    async <-+      Kafka.Worker.Internal.processWithoutShutdownEnsurance+        settings+        (Consumer.ConsumerGroupId "group")+        ( Worker.subscription+            (Internal.unTopic topic)+            ( \msg -> do+                callback msg+                  |> STM.atomically+                  |> map Ok+                  |> Platform.doAnything (doAnything handler')+            )+        )+        |> Async.race_ (returnWhenTerminating handler')+        |> Async.async+    Async.link async+    Prelude.pure (Worker async)++-- | Stops a single worker+stopWorker :: Worker -> Expect.Expectation+stopWorker (Worker async) =+  Async.cancel async+    |> Expect.fromIO++-- | creates a test handler+testHandler :: Settings.Settings -> Prelude.IO TestHandler+testHandler Settings.Settings {Settings.brokerAddresses, Settings.deliveryTimeout, Settings.logLevel, Settings.batchNumMessages} = do+  doAnything <- Platform.doAnythingHandler+  let properties =+        Producer.brokersList brokerAddresses+          ++ Producer.sendTimeout deliveryTimeout+          ++ Producer.logLevel logLevel+          ++ Producer.compression Producer.Snappy+          ++ Producer.extraProps+            ( Dict.fromList+                [ ( "batch.num.messages",+                    batchNumMessages+                      |> Settings.unBatchNumMessages+                      |> Text.fromInt+                  ),+                  -- Enable idemptent producers+                  -- See https://www.cloudkarafka.com/blog/apache-kafka-idempotent-producer-avoiding-message-duplication.html for reference+                  ("enable.idempotence", "true"),+                  ("acks", "all")+                ]+            )+  eitherProducer <- Producer.newProducer properties+  case eitherProducer of+    Prelude.Left err ->+      -- We create the handler as part of starting the application. Throwing+      -- means that if there's a problem with the settings the application will+      -- fail immediately upon start. It won't result in runtime errors during+      -- operation.+      Exception.throwIO err+    Prelude.Right producer -> do+      terminator <- MVar.newEmptyMVar+      Prelude.pure TestHandler {producer, doAnything, terminator}++-- | puts a message synchronously onto a topic-partition+sendSync :: Aeson.ToJSON a => TestHandler -> Internal.Topic -> Int -> a -> Expect.Expectation+sendSync handler topicName partitionId msg' =+  Platform.tracingSpan+    "Sync send Kafka messages"+    ( sendHelperSync+        (producer handler)+        (doAnything handler)+        topicName+        partitionId+        (Aeson.toJSON msg')+    )+    |> Expect.succeeds++sendHelperSync ::+  Producer.KafkaProducer ->+  Platform.DoAnythingHandler ->+  Internal.Topic ->+  Int ->+  Aeson.Value ->+  Task Text ()+sendHelperSync producer doAnything topicName partitionId msg' =+  Exception.handleAny+    (\exception -> Prelude.pure (Err (Debug.toString exception)))+    ( do+        res <- Producer.produceMessage producer (record topicName partitionId msg')+        case res of+          Nothing -> Prelude.pure ()+          Just err -> Exception.throwIO err+        -- by flushing the producer immediately after producing a message,+        -- we make this function synchronous. Without flush it's by default asynchronous.+        Producer.flushProducer producer+        Prelude.pure (Ok ())+    )+    |> Platform.doAnything doAnything++record :: Internal.Topic -> Int -> Aeson.Value -> Producer.ProducerRecord+record topicName partitionId val =+  Producer.ProducerRecord+    { Producer.prTopic = Producer.TopicName (Internal.unTopic topicName),+      Producer.prPartition = Producer.SpecifiedPartition (Prelude.fromIntegral partitionId),+      Producer.prKey = Nothing,+      Producer.prValue =+        Internal.MsgWithMetaData+          { Internal.metaData =+              Internal.MetaData+                { Internal.requestId = "test-request"+                },+            Internal.value = Internal.Encodable val+          }+          |> Aeson.encode+          |> Data.ByteString.Lazy.toStrict+          |> Just+    }++-- | test helper, that yields a new @Kafka.Topic@ and @TestHandler@+test ::+  Stack.HasCallStack =>+  Text ->+  ((Internal.Topic, TestHandler) -> Expect.Expectation) ->+  Test.Test+test description body =+  Stack.withFrozenCallStack Test.test description <| \_ -> do+    doAnything <- Expect.fromIO Platform.doAnythingHandler+    Expect.around+      ( \task' ->+          Platform.bracketWithError+            ( -- create handler+              Platform.doAnything doAnything+                <| case Environment.decodeDefaults Settings.decoder of+                  Ok settings ->+                    map+                      Ok+                      ( testHandler+                          settings+                            { Settings.batchNumMessages = Settings.exampleBatchNumMessages+                            }+                      )+                  Err err -> Debug.todo ("Failed to to decode worker settings" ++ err)+            )+            ( -- terminate all workers hanging around after the test is over+              \_maybeErr handler' ->+                terminate handler'+                  |> map Ok+                  |> Platform.doAnything doAnything+            )+            ( -- yield the test+              \handler' -> do+                uuid <-+                  map Data.UUID.toText Data.UUID.V4.nextRandom+                    |> map Ok+                    |> Platform.doAnything doAnything+                let topic = Internal.Topic uuid+                task' (topic, handler')+            )+      )+      body
+ test/Main.hs view
@@ -0,0 +1,17 @@+module Main (main) where++import qualified Spec.Kafka.Worker.Integration+import qualified Spec.Kafka.Worker.Partition+import qualified Test+import qualified Prelude++main :: Prelude.IO ()+main = Test.run tests++tests :: Test.Test+tests =+  Test.describe+    "lib/kafka"+    [ Spec.Kafka.Worker.Integration.tests,+      Spec.Kafka.Worker.Partition.tests+    ]
+ test/Spec/Kafka/Worker/Integration.hs view
@@ -0,0 +1,81 @@+module Spec.Kafka.Worker.Integration (tests) where++import qualified Control.Concurrent.STM as STM+import qualified Dict+import qualified Expect+import qualified Helpers+import qualified Set+import qualified Test+import qualified Prelude++tests :: Test.Test+tests =+  Test.describe+    "Worker"+    [ Test.describe+        "Integration"+        [ Helpers.test "We receive what we send" <| \(topic, handler) -> do+            Helpers.sendSync handler topic 1 1+            msgsTVar <- atomically (STM.newTVar Set.empty)+            _ <- Helpers.spawnWorker handler topic (\msg -> STM.modifyTVar' msgsTVar (Set.insert msg))+            Helpers.sendSync handler topic 2 2+            msgs' <- waitFor msgsTVar (\items -> Set.size items == 2)+            msgs' |> Expect.equal (Set.fromList [1, 2]),+          Helpers.test "two workers process all messages once" <| \(topic, handler) -> do+            msgsTVar <- atomically (STM.newTVar [])+            _ <- Helpers.spawnWorker handler topic (\msg -> STM.modifyTVar' msgsTVar (\msgs -> msg : msgs))+            _ <- Helpers.spawnWorker handler topic (\msg -> STM.modifyTVar' msgsTVar (\msgs -> msg : msgs))+            Helpers.sendSync handler topic 1 (1, 1)+            Helpers.sendSync handler topic 1 (1, 2)+            Helpers.sendSync handler topic 2 (2, 3)+            msgs' <- waitFor msgsTVar (\items -> List.length items == 3)+            msgs' |> groupDictAndMap identity+              |> Expect.equal+                ( Dict.fromList+                    [ (1, [2, 1]),+                      (2, [3])+                    ]+                ),+          Helpers.test "second worker takes over after first worker gets stopped" <| \(topic, handler) -> do+            msgsTVar <- atomically (STM.newTVar [])+            worker1 <- Helpers.spawnWorker handler topic (\msg -> STM.modifyTVar' msgsTVar (\msgs -> msg : msgs))+            _ <- Helpers.spawnWorker handler topic (\msg -> STM.modifyTVar' msgsTVar (\msgs -> msg : msgs))+            Helpers.sendSync handler topic 1 (1, 1)+            Helpers.sendSync handler topic 2 (2, 1)+            _ <- waitFor msgsTVar (\items -> List.length items == 2)+            Helpers.stopWorker worker1+            Helpers.sendSync handler topic 1 (1, 2)+            Helpers.sendSync handler topic 2 (2, 2)+            msgsAfterStoppingWorker <- waitFor msgsTVar (\items -> List.length items == 4)+            msgsAfterStoppingWorker+              |> groupDictAndMap identity+              |> Expect.equal+                ( Dict.fromList+                    [ (1, [2, 1]),+                      (2, [2, 1])+                    ]+                )+        ]+    ]++atomically :: STM.STM a -> Expect.Expectation' a+atomically = STM.atomically >> Expect.fromIO++waitFor :: STM.TVar a -> (a -> Bool) -> Expect.Expectation' a+waitFor tVar pred =+  atomically <| do+    val <- STM.readTVar tVar+    if pred val+      then Prelude.pure val+      else STM.retry++groupDictAndMap :: Ord b => (a -> (b, c)) -> List a -> Dict.Dict b (List c)+groupDictAndMap f =+  List.foldr+    ( \x ->+        Dict.update (Tuple.first (f x)) <| \val ->+          case val of+            Nothing -> Just [Tuple.second (f x)]+            Just y -> Just (Tuple.second (f x) : y)+    )+    Dict.empty
+ test/Spec/Kafka/Worker/Partition.hs view
@@ -0,0 +1,26 @@+module Spec.Kafka.Worker.Partition (tests) where++import qualified Expect+import qualified Kafka.Worker.Partition as Partition+import qualified Test++tests :: Test.Test+tests =+  Test.describe+    "Worker"+    [ Test.describe+        "microSecondsDelayForAttempt"+        [ Test.test "1 attempt" <| \() ->+            Partition.microSecondsDelayForAttempt 1+              |> Expect.equal 10_000_000,+          Test.test "2 attempts" <| \() ->+            Partition.microSecondsDelayForAttempt 2+              |> Expect.equal 100_000_000,+          Test.test "3 attempts" <| \() ->+            Partition.microSecondsDelayForAttempt 3+              |> Expect.equal 1000_000_000,+          Test.test "4 attempts" <| \() ->+            Partition.microSecondsDelayForAttempt 4+              |> Expect.equal 3_600_000_000+        ]+    ]