diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -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 diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -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" +``` diff --git a/nri-kafka.cabal b/nri-kafka.cabal new file mode 100644 --- /dev/null +++ b/nri-kafka.cabal @@ -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 . +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 diff --git a/src/Kafka.hs b/src/Kafka.hs new file mode 100644 --- /dev/null +++ b/src/Kafka.hs @@ -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 diff --git a/src/Kafka/Internal.hs b/src/Kafka/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Internal.hs @@ -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 diff --git a/src/Kafka/Settings.hs b/src/Kafka/Settings.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Settings.hs @@ -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) diff --git a/src/Kafka/Settings/Internal.hs b/src/Kafka/Settings/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Settings/Internal.hs @@ -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 diff --git a/src/Kafka/Test.hs b/src/Kafka/Test.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Test.hs @@ -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 diff --git a/src/Kafka/Worker.hs b/src/Kafka/Worker.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Worker.hs @@ -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 diff --git a/src/Kafka/Worker/Analytics.hs b/src/Kafka/Worker/Analytics.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Worker/Analytics.hs @@ -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) diff --git a/src/Kafka/Worker/Fetcher.hs b/src/Kafka/Worker/Fetcher.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Worker/Fetcher.hs @@ -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 diff --git a/src/Kafka/Worker/Internal.hs b/src/Kafka/Worker/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Worker/Internal.hs @@ -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 () diff --git a/src/Kafka/Worker/Partition.hs b/src/Kafka/Worker/Partition.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Worker/Partition.hs @@ -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 diff --git a/src/Kafka/Worker/Settings.hs b/src/Kafka/Worker/Settings.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Worker/Settings.hs @@ -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 + ) + ) diff --git a/src/Kafka/Worker/Stopping.hs b/src/Kafka/Worker/Stopping.hs new file mode 100644 --- /dev/null +++ b/src/Kafka/Worker/Stopping.hs @@ -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 + ) diff --git a/test/Helpers.hs b/test/Helpers.hs new file mode 100644 --- /dev/null +++ b/test/Helpers.hs @@ -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 diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -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 + ] diff --git a/test/Spec/Kafka/Worker/Integration.hs b/test/Spec/Kafka/Worker/Integration.hs new file mode 100644 --- /dev/null +++ b/test/Spec/Kafka/Worker/Integration.hs @@ -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 diff --git a/test/Spec/Kafka/Worker/Partition.hs b/test/Spec/Kafka/Worker/Partition.hs new file mode 100644 --- /dev/null +++ b/test/Spec/Kafka/Worker/Partition.hs @@ -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 + ] + ]