run the kcl logs at INFO level to inspect lease handover

build.sbt

@@ -154,12 +154,13 @@ lazy val metadataEditor = playProject("metadata-editor", 9007)
 lazy val thrall = playProject("thrall", 9002)
   .dependsOn(commonLib % "compile;test->test")
   .settings(
-    pipelineStages := Seq(digest, gzip),
-    libraryDependencies ++= Seq(
+      pipelineStages := Seq(digest, gzip),
+      resolvers += "Sonatype snapshots" at "https://central.sonatype.com/repository/maven-snapshots",
+      libraryDependencies ++= Seq(
       "org.codehaus.groovy" % "groovy-json" % "3.0.7",
       // TODO upgrading kcl to v3? check if you can remove avro override below
-      "software.amazon.kinesis" % "amazon-kinesis-client" % "2.6.1",
-      "com.gu" %% "kcl-pekko-stream" % "0.1.0",
+      "io.github.andrew-nowak.kinesis" % "amazon-kinesis-client" % "3.0.3-SNAPSHOT",
+      "com.gu" %% "kcl-pekko-stream" % "0.1.1-PREVIEW.anrelease-leadership-before-shutdown.2025-06-05T1055.e9880914" exclude ("software.amazon.kinesis", "amazon-kinesis-client"),
       "org.testcontainers" % "elasticsearch" % "1.19.2" % Test,
       "com.google.protobuf" % "protobuf-java" % "3.19.6"
     ),
@@ -237,6 +238,7 @@ def playProject(projectName: String, port: Int, path: Option[String] = None): Pr
       Linux / packageSummary := description.value,
       packageDescription := description.value,
 
+      Debian / killTimeout := 60,
       bashScriptEnvConfigLocation := Some("/etc/environment"),
       Debian / makeEtcDefault := None,
       Debian / packageBin := {

common-lib/src/main/resources/logback.xml

@@ -37,6 +37,7 @@
     <logger name="play" level="INFO" />
     <logger name="application" level="INFO" />
     <logger name="request" level="INFO" />
+    <logger name="software.amazon.kinesis" level="INFO" />
 
     <root level="INFO">
       <appender-ref ref="LOGFILE"/>

thrall/app/ThrallComponents.scala

@@ -12,6 +12,8 @@ import controllers.{AssetsComponents, HealthCheck, ReaperController, ThrallContr
 import lib._
 import lib.elasticsearch._
 import lib.kinesis.{KinesisConfig, ThrallEventConsumer}
+import org.apache.pekko.actor.CoordinatedShutdown
+import org.apache.pekko.stream.KillSwitch
 import play.api.ApplicationLoader.Context
 import router.Routes
 
@@ -71,7 +73,14 @@ class ThrallComponents(context: Context) extends GridComponents(context, new Thr
     actorSystem
   )
 
-  val streamRunning: Future[Done] = thrallStreamProcessor.run()
+  private val streamRunning: Future[Done] = thrallStreamProcessor.run()
+  streamRunning.onComplete(_ => logger.info("Thrall streams completed"))
+
+  coordinatedShutdown.addTask(CoordinatedShutdown.PhaseServiceRequestsDone, "shutdown-thrall-streams")(() => {
+    logger.info("thrall stream is closing; saying byebye")
+    thrallStreamProcessor.killSwitch.shutdown()
+    streamRunning
+  })
 
   val s3 = S3Ops.buildS3Client(config)
   val syncChecker = new SyncChecker(

thrall/app/lib/ThrallStreamProcessor.scala

@@ -1,17 +1,16 @@
 package lib
 
-import java.time.Instant
-import org.apache.pekko.actor.ActorSystem
-import org.apache.pekko.stream.scaladsl.{GraphDSL, MergePreferred, MergePrioritized, Source}
-import org.apache.pekko.stream.{Materializer, SourceShape}
-import org.apache.pekko.{Done, NotUsed}
 import com.gu.kinesis.KinesisRecord
 import com.gu.mediaservice.lib.DateTimeUtils
-import com.gu.mediaservice.lib.aws.UpdateMessage
 import com.gu.mediaservice.lib.logging._
-import com.gu.mediaservice.model.{ExternalThrallMessage, InternalThrallMessage, ThrallMessage}
-import lib.kinesis.{MessageTranslator, ThrallEventConsumer}
+import com.gu.mediaservice.model.{ExternalThrallMessage, ThrallMessage}
+import lib.kinesis.ThrallEventConsumer
+import org.apache.pekko.Done
+import org.apache.pekko.actor.ActorSystem
+import org.apache.pekko.stream.scaladsl.{GraphDSL, MergePreferred, Source}
+import org.apache.pekko.stream.{FlowShape, KillSwitches, Materializer}
 
+import java.time.Instant
 import scala.concurrent.{ExecutionContextExecutor, Future}
 import scala.util.{Failure, Success}
 
@@ -55,65 +54,72 @@ class ThrallStreamProcessor(
   actorSystem: ActorSystem
  ) extends GridLogging {
 
+  import GraphDSL.Implicits._
+
+  private def waitForBoth(f1: Future[Done], f2: Future[Done]): Future[Done] = {
+    f1.zip(f2).map(_ => Done)
+  }
+
+  val killSwitch = KillSwitches.shared("thrall-kill-switch")
+
   implicit val mat: Materializer = Materializer.matFromSystem(actorSystem)
   implicit val dispatcher: ExecutionContextExecutor = actorSystem.getDispatcher
 
-  val mergedKinesisSource: Source[TaggedRecord[ThrallMessage], NotUsed] = Source.fromGraph(GraphDSL.create() { implicit graphBuilder =>
-    import GraphDSL.Implicits._
+  val uiRecordSource = uiSource.via(killSwitch.flow).map(kinesisRecord =>
+    TaggedRecord(kinesisRecord.data.toArray, kinesisRecord.approximateArrivalTimestamp, UiPriority, kinesisRecord.markProcessed))
+
+  val automationRecordSource = automationSource.via(killSwitch.flow).map(kinesisRecord =>
+    TaggedRecord(kinesisRecord.data.toArray, kinesisRecord.approximateArrivalTimestamp, AutomationPriority, kinesisRecord.markProcessed))
 
-    val uiRecordSource = uiSource.map(kinesisRecord =>
-      TaggedRecord(kinesisRecord.data.toArray, kinesisRecord.approximateArrivalTimestamp, UiPriority, kinesisRecord.markProcessed))
+  val migrationMessagesSource = migrationSource.via(killSwitch.flow).map { case MigrationRecord(internalThrallMessage, time) =>
+    TaggedRecord(internalThrallMessage, time, MigrationPriority, () => {})
+  }
 
-    val automationRecordSource = automationSource.map(kinesisRecord =>
-      TaggedRecord(kinesisRecord.data.toArray, kinesisRecord.approximateArrivalTimestamp, AutomationPriority, kinesisRecord.markProcessed))
+  val uiAndAutomationRecordsMerge: Source[TaggedRecord[Array[Byte]], Future[Done]] = {
+    val mergeGraph = GraphDSL.createGraph(automationRecordSource) { implicit b => r =>
+      val merge = b.add(MergePreferred[TaggedRecord[Array[Byte]]](1, eagerComplete = false))
+      r ~> merge.in(0)
+      FlowShape(merge.in(1), merge.out)
+    }
+    uiRecordSource.viaMat(mergeGraph)(waitForBoth)
+  }
+
+  val uiAndAutomationMessagesSource: Source[TaggedRecord[ExternalThrallMessage], Future[Done]] =
+    uiAndAutomationRecordsMerge
+      .map { taggedRecord =>
+        val parsedRecord = ThrallEventConsumer
+          .parseRecord(taggedRecord.payload, taggedRecord.arrivalTimestamp)
+          .map(
+            message => taggedRecord.copy(payload = message)
+          )
+        // If we failed to parse the record (Left), we'll drop it below because we can't process it.
+        // However we still need to mark the record as processed, otherwise the kinesis stream can't progress
+        // and checkpoint will be stuck at this message forevermore.
+        parsedRecord.left.foreach(_ => taggedRecord.markProcessed())
+        parsedRecord
+      }.collect {
+      // drop unparseable records
+        case Right(taggedRecord) => taggedRecord
+      }
 
-    val migrationMessagesSource = migrationSource.map { case MigrationRecord(internalThrallMessage, time) =>
-      TaggedRecord(internalThrallMessage, time, MigrationPriority, () => {})
+  val allSourcesMerged: Source[TaggedRecord[ThrallMessage], Future[Done]] = {
+    val mergeGraph = GraphDSL.createGraph(migrationMessagesSource) { implicit b => r =>
+      val merge = b.add(MergePreferred[TaggedRecord[ThrallMessage]](1, eagerComplete = false))
+      r ~> merge.in(0)
+      FlowShape(merge.in(1), merge.out)
     }
+    uiAndAutomationMessagesSource.viaMat(mergeGraph)(waitForBoth)
+  }
 
-    // merge together ui and automation kinesis records
-    val uiAndAutomationRecordsMerge = graphBuilder.add(MergePreferred[TaggedRecord[Array[Byte]]](1))
-    uiRecordSource ~> uiAndAutomationRecordsMerge.preferred
-    automationRecordSource  ~> uiAndAutomationRecordsMerge.in(0)
-
-    // parse the kinesis records into thrall update messages (dropping those that fail)
-    val uiAndAutomationMessagesSource: PortOps[TaggedRecord[ExternalThrallMessage]] =
-      uiAndAutomationRecordsMerge.out
-        .map { taggedRecord =>
-          val parsedRecord = ThrallEventConsumer
-            .parseRecord(taggedRecord.payload, taggedRecord.arrivalTimestamp)
-            .map(
-              message => taggedRecord.copy(payload = message)
-            )
-          // If we failed to parse the record (Left), we'll drop it below because we can't process it.
-          // However we still need to mark the record as processed, otherwise the kinesis stream can't progress
-          // and checkpoint will be stuck at this message forevermore.
-          parsedRecord.left.foreach(_ => taggedRecord.markProcessed())
-          parsedRecord
-        }
-        // drop unparseable records
-        .collect {
-          case Right(taggedRecord) => taggedRecord
-        }
-
-    // merge in the re-ingestion source (preferring ui/automation)
-    val mergePreferred = graphBuilder.add(MergePreferred[TaggedRecord[ThrallMessage]](1))
-    uiAndAutomationMessagesSource ~> mergePreferred.preferred
-    migrationMessagesSource ~> mergePreferred.in(0)
-
-    SourceShape(mergePreferred.out)
-  })
-
-  def createStream(): Source[(TaggedRecord[ThrallMessage], Stopwatch, ThrallMessage), NotUsed] = {
-    mergedKinesisSource.mapAsync(1) { result =>
+  def createStream(): Source[(TaggedRecord[ThrallMessage], Stopwatch, ThrallMessage), Future[Done]] = {
+    allSourcesMerged.mapAsync(1) { result =>
       val stopwatch = Stopwatch.start
       consumer.processMessage(result.payload)
         .recover { case _ => () }
         .map(_ => (result, stopwatch, result.payload))
       }
-
-
   }
+
   def run(): Future[Done] = {
     val stream = this.createStream().runForeach {
       case (taggedRecord, stopwatch, _) =>
@@ -124,7 +130,7 @@ class ThrallStreamProcessor(
 
     stream.onComplete {
       case Failure(exception) => logger.error("Thrall stream completed with failure", exception)
-      case Success(_) => logger.info("Thrall stream completed with done, probably shutting down")
+      case Success(Done) => logger.info("Thrall stream completed with done, probably shutting down")
     }
 
     stream

thrall/app/lib/kinesis/KinesisConfig.scala

@@ -3,7 +3,9 @@ package lib.kinesis
 import com.gu.kinesis.ConsumerConfig
 import com.gu.mediaservice.lib.logging.GridLogging
 import lib.KinesisReceiverConfig
+import software.amazon.awssdk.services.dynamodb.model.BillingMode
 import software.amazon.kinesis.common.{InitialPositionInStream, InitialPositionInStreamExtended}
+import software.amazon.kinesis.leases.LeaseManagementConfig
 import software.amazon.kinesis.metrics.MetricsConfig
 import software.amazon.kinesis.retrieval.RetrievalConfig
 import software.amazon.kinesis.retrieval.polling.PollingConfig
@@ -28,6 +30,19 @@ object KinesisConfig extends GridLogging {
     val metricsConfig = new MetricsConfig(config.cloudwatchClient, config.streamName)
       .metricsLevel(config.metricsLevel)
 
+    val handoffConfig = LeaseManagementConfig.GracefulLeaseHandoffConfig.builder().isGracefulLeaseHandoffEnabled(false).build()
+
+    val leaseManagementConfig = new LeaseManagementConfig(
+      config.streamName,
+      config.streamName,
+      config.dynamoClient,
+      config.kinesisClient,
+      workerId,
+    ).billingMode(BillingMode.PAY_PER_REQUEST)
+      .gracefulLeaseHandoffConfig(handoffConfig)
+      .dynamoLockBasedLeaderLeaseDurationInSeconds(10)
+      .dynamoLockBasedLeaderHeartbeatPeriodInSeconds(2)
+
     val clientConfig = ConsumerConfig(
       streamName = config.streamName,
       appName = config.streamName,
@@ -37,7 +52,7 @@ object KinesisConfig extends GridLogging {
       cloudwatchClient = config.cloudwatchClient,
       initialPositionInStreamExtended = initialPosition,
       coordinatorConfig = None,
-      leaseManagementConfig = None,
+      leaseManagementConfig = Some(leaseManagementConfig),
       metricsConfig = Some(metricsConfig),
       retrievalConfig = Some(retrievalConfig)
     )