Create table concurrently with subscribing to stream of events (#77)

For KCL apps, we want the KCL to get initialized as early as possible, so the worker can claim shard leases before they get stolen by other workers. Before this PR, the loader initialized the destination table first, and then subscribed to the stream afterwards. Initializing the destination table can be fairly slow, especially because we do things like syncing to the external catalog, and possibly cleaning up aborted commits. After this PR, the loader subscribes to the stream concurrently with initializing the destination table. This lets the KCL claim leases before they get stolen.
snowplow-incubator · Aug 8, 2024 · b67cb86 · b67cb86
1 parent 2fe599a
commit b67cb86
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 46 deletions.
diff --git a/modules/core/src/main/scala/com.snowplowanalytics.snowplow.lakes/processing/Processing.scala b/modules/core/src/main/scala/com.snowplowanalytics.snowplow.lakes/processing/Processing.scala
@@ -13,7 +13,7 @@ package com.snowplowanalytics.snowplow.lakes.processing
 import cats.implicits._
 import cats.data.NonEmptyList
 import cats.{Applicative, Foldable, Functor}
-import cats.effect.{Async, Sync}
+import cats.effect.{Async, Deferred, Sync}
 import cats.effect.kernel.{Ref, Unique}
 import fs2.{Chunk, Pipe, Stream}
 import org.typelevel.log4cats.Logger
@@ -48,10 +48,19 @@ object Processing {
   private implicit def logger[F[_]: Sync]: Logger[F] = Slf4jLogger.getLogger[F]
 
   def stream[F[_]: Async](env: Environment[F]): Stream[F, Nothing] =
-    Stream.eval(env.lakeWriter.createTable).flatMap { _ =>
+    Stream.eval(Deferred[F, Unit]).flatMap { deferredTableExists =>
+      val runInBackground =
+        // Create the table in a background stream, so it does not block subscribing to the stream.
+        // Needed for Kinesis, where we want to subscribe to the stream as early as possible, so that other workers don't steal our shard leases
+        Stream.eval(env.lakeWriter.createTable *> deferredTableExists.complete(()))
+
       implicit val lookup: RegistryLookup[F]           = Http4sRegistryLookup(env.httpClient)
       val eventProcessingConfig: EventProcessingConfig = EventProcessingConfig(env.windowing)
-      env.source.stream(eventProcessingConfig, eventProcessor(env))
+
+      env.source
+        .stream(eventProcessingConfig, eventProcessor(env, deferredTableExists.get))
+        .concurrently(runInBackground)
+
     }
 
   /** Model used between stages of the processing pipeline */
@@ -74,9 +83,11 @@ object Processing {
   )
 
   private def eventProcessor[F[_]: Async: RegistryLookup](
-    env: Environment[F]
+    env: Environment[F],
+    deferredTableExists: F[Unit]
   ): EventProcessor[F] = { in =>
     val resources = for {
+      _ <- Stream.eval(deferredTableExists)
       windowState <- Stream.eval(WindowState.build[F])
       stateRef <- Stream.eval(Ref[F].of(windowState))
       _ <- manageDataFrame(env, windowState.viewName)

diff --git a/modules/core/src/test/scala/com.snowplowanalytics.snowplow.lakes/MockEnvironment.scala b/modules/core/src/test/scala/com.snowplowanalytics.snowplow.lakes/MockEnvironment.scala
@@ -33,11 +33,15 @@ object MockEnvironment {
   /** All tests can use the same window duration */
   val WindowDuration = 42.seconds
 
+  /** All tests can use the same time taken to create table */
+  val TimeTakenToCreateTable = 10.seconds
+
   sealed trait Action
   object Action {
+    case object SubscribedToStream extends Action
+    case object CreatedTable extends Action
     case class Checkpointed(tokens: List[Unique.Token]) extends Action
     case class SentToBad(count: Int) extends Action
-    case object CreatedTable extends Action
     case class InitializedLocalDataFrame(viewName: String) extends Action
     case class AppendedRowsToDataFrame(viewName: String, numEvents: Int) extends Action
     case class RemovedDataFrameFromDisk(viewName: String) extends Action
@@ -89,7 +93,7 @@ object MockEnvironment {
 
   private def testLakeWriter(state: Ref[IO, Vector[Action]]): LakeWriter.WithHandledErrors[IO] = new LakeWriter.WithHandledErrors[IO] {
     def createTable: IO[Unit] =
-      state.update(_ :+ CreatedTable)
+      IO.sleep(TimeTakenToCreateTable) *> state.update(_ :+ CreatedTable)
 
     def initializeLocalDataFrame(viewName: String): IO[Unit] =
       state.update(_ :+ InitializedLocalDataFrame(viewName))
@@ -111,17 +115,18 @@ object MockEnvironment {
   private def testSourceAndAck(windows: List[List[TokenedEvents]], state: Ref[IO, Vector[Action]]): SourceAndAck[IO] =
     new SourceAndAck[IO] {
       def stream(config: EventProcessingConfig, processor: EventProcessor[IO]): Stream[IO, Nothing] =
-        Stream.emits(windows).flatMap { batches =>
-          Stream
-            .emits(batches)
-            .onFinalize(IO.sleep(WindowDuration))
-            .through(processor)
-            .chunks
-            .evalMap { chunk =>
-              state.update(_ :+ Checkpointed(chunk.toList))
-            }
-            .drain
-        }
+        Stream.eval(state.update(_ :+ SubscribedToStream)).drain ++
+          Stream.emits(windows).flatMap { batches =>
+            Stream
+              .emits(batches)
+              .onFinalize(IO.sleep(WindowDuration))
+              .through(processor)
+              .chunks
+              .evalMap { chunk =>
+                state.update(_ :+ Checkpointed(chunk.toList))
+              }
+              .drain
+          }
 
       def isHealthy(maxAllowedProcessingLatency: FiniteDuration): IO[SourceAndAck.HealthStatus] =
         IO.pure(SourceAndAck.Healthy)

diff --git a/.../core/src/test/scala/com.snowplowanalytics.snowplow.lakes/processing/ProcessingSpec.scala b/.../core/src/test/scala/com.snowplowanalytics.snowplow.lakes/processing/ProcessingSpec.scala
@@ -46,16 +46,17 @@ class ProcessingSpec extends Specification with CatsEffect {
       state <- control.state.get
     } yield state should beEqualTo(
       Vector(
+        Action.SubscribedToStream,
         Action.CreatedTable,
-        Action.InitializedLocalDataFrame("v19700101000000"),
+        Action.InitializedLocalDataFrame("v19700101000010"),
         Action.AddedReceivedCountMetric(2),
         Action.AddedReceivedCountMetric(2),
-        Action.AppendedRowsToDataFrame("v19700101000000", 4),
-        Action.CommittedToTheLake("v19700101000000"),
+        Action.AppendedRowsToDataFrame("v19700101000010", 4),
+        Action.CommittedToTheLake("v19700101000010"),
         Action.AddedCommittedCountMetric(4),
         Action.SetProcessingLatencyMetric(MockEnvironment.WindowDuration),
         Action.Checkpointed(inputs.map(_.ack)),
-        Action.RemovedDataFrameFromDisk("v19700101000000")
+        Action.RemovedDataFrameFromDisk("v19700101000010")
       )
     )
 
@@ -70,8 +71,9 @@ class ProcessingSpec extends Specification with CatsEffect {
       state <- control.state.get
     } yield state should beEqualTo(
       Vector(
+        Action.SubscribedToStream,
         Action.CreatedTable,
-        Action.InitializedLocalDataFrame("v19700101000000"),
+        Action.InitializedLocalDataFrame("v19700101000010"),
         Action.AddedReceivedCountMetric(2),
         Action.AddedBadCountMetric(2),
         Action.SentToBad(2),
@@ -82,7 +84,7 @@ class ProcessingSpec extends Specification with CatsEffect {
         Action.AddedBadCountMetric(2),
         Action.SentToBad(2),
         Action.Checkpointed(inputs.map(_.ack)),
-        Action.RemovedDataFrameFromDisk("v19700101000000")
+        Action.RemovedDataFrameFromDisk("v19700101000010")
       )
     )
 
@@ -99,40 +101,41 @@ class ProcessingSpec extends Specification with CatsEffect {
       state <- control.state.get
     } yield state should beEqualTo(
       Vector(
+        Action.SubscribedToStream,
         Action.CreatedTable,
 
         /* window 1 */
-        Action.InitializedLocalDataFrame("v19700101000000"),
+        Action.InitializedLocalDataFrame("v19700101000010"),
         Action.AddedReceivedCountMetric(2),
-        Action.AppendedRowsToDataFrame("v19700101000000", 2),
-        Action.CommittedToTheLake("v19700101000000"),
+        Action.AppendedRowsToDataFrame("v19700101000010", 2),
+        Action.CommittedToTheLake("v19700101000010"),
         Action.AddedCommittedCountMetric(2),
         Action.SetProcessingLatencyMetric(MockEnvironment.WindowDuration),
         Action.Checkpointed(window1.map(_.ack)),
-        Action.RemovedDataFrameFromDisk("v19700101000000"),
+        Action.RemovedDataFrameFromDisk("v19700101000010"),
 
         /* window 2 */
-        Action.InitializedLocalDataFrame("v19700101000042"),
+        Action.InitializedLocalDataFrame("v19700101000052"),
         Action.AddedReceivedCountMetric(2),
         Action.AddedReceivedCountMetric(2),
         Action.AddedReceivedCountMetric(2),
-        Action.AppendedRowsToDataFrame("v19700101000042", 6),
-        Action.CommittedToTheLake("v19700101000042"),
+        Action.AppendedRowsToDataFrame("v19700101000052", 6),
+        Action.CommittedToTheLake("v19700101000052"),
         Action.AddedCommittedCountMetric(6),
         Action.SetProcessingLatencyMetric(MockEnvironment.WindowDuration),
         Action.Checkpointed(window2.map(_.ack)),
-        Action.RemovedDataFrameFromDisk("v19700101000042"),
+        Action.RemovedDataFrameFromDisk("v19700101000052"),
 
         /* window 3 */
-        Action.InitializedLocalDataFrame("v19700101000124"),
+        Action.InitializedLocalDataFrame("v19700101000134"),
         Action.AddedReceivedCountMetric(2),
         Action.AddedReceivedCountMetric(2),
-        Action.AppendedRowsToDataFrame("v19700101000124", 4),
-        Action.CommittedToTheLake("v19700101000124"),
+        Action.AppendedRowsToDataFrame("v19700101000134", 4),
+        Action.CommittedToTheLake("v19700101000134"),
         Action.AddedCommittedCountMetric(4),
         Action.SetProcessingLatencyMetric(MockEnvironment.WindowDuration),
         Action.Checkpointed(window3.map(_.ack)),
-        Action.RemovedDataFrameFromDisk("v19700101000124")
+        Action.RemovedDataFrameFromDisk("v19700101000134")
       )
     )
 
@@ -148,19 +151,20 @@ class ProcessingSpec extends Specification with CatsEffect {
       state <- control.state.get
     } yield state should beEqualTo(
       Vector(
+        Action.SubscribedToStream,
         Action.CreatedTable,
-        Action.InitializedLocalDataFrame("v19700101000000"),
+        Action.InitializedLocalDataFrame("v19700101000010"),
         Action.AddedReceivedCountMetric(2),
-        Action.AppendedRowsToDataFrame("v19700101000000", 2),
+        Action.AppendedRowsToDataFrame("v19700101000010", 2),
         Action.AddedReceivedCountMetric(2),
-        Action.AppendedRowsToDataFrame("v19700101000000", 2),
+        Action.AppendedRowsToDataFrame("v19700101000010", 2),
         Action.AddedReceivedCountMetric(2),
-        Action.AppendedRowsToDataFrame("v19700101000000", 2),
-        Action.CommittedToTheLake("v19700101000000"),
+        Action.AppendedRowsToDataFrame("v19700101000010", 2),
+        Action.CommittedToTheLake("v19700101000010"),
         Action.AddedCommittedCountMetric(6),
         Action.SetProcessingLatencyMetric(MockEnvironment.WindowDuration),
         Action.Checkpointed(inputs.map(_.ack)),
-        Action.RemovedDataFrameFromDisk("v19700101000000")
+        Action.RemovedDataFrameFromDisk("v19700101000010")
       )
     )
 
@@ -178,8 +182,9 @@ class ProcessingSpec extends Specification with CatsEffect {
       state <- control.state.get
     } yield state should beEqualTo(
       Vector(
+        Action.SubscribedToStream,
         Action.CreatedTable,
-        Action.InitializedLocalDataFrame("v19700101000000"),
+        Action.InitializedLocalDataFrame("v19700101000010"),
         Action.AddedReceivedCountMetric(2),
         Action.AddedBadCountMetric(2),
         Action.SentToBad(2),
@@ -196,12 +201,12 @@ class ProcessingSpec extends Specification with CatsEffect {
         Action.AddedBadCountMetric(2),
         Action.SentToBad(2),
         Action.AddedReceivedCountMetric(2),
-        Action.AppendedRowsToDataFrame("v19700101000000", 8),
-        Action.CommittedToTheLake("v19700101000000"),
+        Action.AppendedRowsToDataFrame("v19700101000010", 8),
+        Action.CommittedToTheLake("v19700101000010"),
         Action.AddedCommittedCountMetric(8),
         Action.SetProcessingLatencyMetric(MockEnvironment.WindowDuration),
         Action.Checkpointed((bads1 ::: goods1 ::: bads2 ::: goods2).map(_.ack)),
-        Action.RemovedDataFrameFromDisk("v19700101000000")
+        Action.RemovedDataFrameFromDisk("v19700101000010")
       )
     )
 
@@ -210,7 +215,7 @@ class ProcessingSpec extends Specification with CatsEffect {
 
   def e6 = {
     val messageTime = Instant.parse("2023-10-24T10:00:00.000Z")
-    val processTime = Instant.parse("2023-10-24T10:00:42.123Z")
+    val processTime = Instant.parse("2023-10-24T10:00:42.123Z").minusMillis(MockEnvironment.TimeTakenToCreateTable.toMillis)
 
     val io = for {
       inputs <- generateEvents.take(2).compile.toList.map {
@@ -224,6 +229,7 @@ class ProcessingSpec extends Specification with CatsEffect {
       state <- control.state.get
     } yield state should beEqualTo(
       Vector(
+        Action.SubscribedToStream,
         Action.CreatedTable,
         Action.InitializedLocalDataFrame("v20231024100042"),
         Action.SetLatencyMetric(42123.millis),