everything is a factory (and the complete counts work)

johnny-schmidt · johnny-schmidt · commit f48b4864a991 · 2025-03-19T09:01:20.000-07:00
diff --git a/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/task/internal/LoadPipelineStepTask.kt b/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/task/internal/LoadPipelineStepTask.kt
@@ -55,7 +55,9 @@ class LoadPipelineStepTask<S : AutoCloseable, K1 : WithStream, T, K2 : WithStrea
     private val flushStrategy: PipelineFlushStrategy?,
     private val part: Int,
     private val numWorkers: Int,
-    private val streamCompletions: ConcurrentHashMap<DestinationStream.Descriptor, AtomicInteger>
+    private val taskIndex: Int,
+    private val streamCompletions:
+        ConcurrentHashMap<Pair<Int, DestinationStream.Descriptor>, AtomicInteger>
 ) : Task {
     override val terminalCondition: TerminalCondition = OnEndOfSync
 
@@ -150,7 +152,7 @@ class LoadPipelineStepTask<S : AutoCloseable, K1 : WithStream, T, K2 : WithStrea
                         // Only forward end-of-stream if ALL workers have seen end-of-stream.
                         if (
                             streamCompletions
-                                .getOrPut(input.stream) { AtomicInteger(0) }
+                                .getOrPut(Pair(taskIndex, input.stream)) { AtomicInteger(0) }
                                 .incrementAndGet() == numWorkers
                         ) {
                             outputQueue?.broadcast(PipelineEndOfStream(input.stream))
@@ -211,7 +213,10 @@ class LoadPipelineStepTaskFactory(
     @Value("\${airbyte.destination.core.record-batch-size-override:null}")
     val batchSizeOverride: Long? = null,
 ) {
-    private val streamCompletions = ConcurrentHashMap<DestinationStream.Descriptor, AtomicInteger>()
+    // A map of (TaskIndex, Stream) -> Count_of_closed streams to ensure eos is not forwared from
+    // task N to N+1 until all workers have seen eos.
+    private val streamCompletions =
+        ConcurrentHashMap<Pair<Int, DestinationStream.Descriptor>, AtomicInteger>()
 
     fun <S : AutoCloseable, K1 : WithStream, T, K2 : WithStream, U : Any> create(
         batchAccumulator: BatchAccumulator<S, K1, T, U>,
@@ -221,6 +226,7 @@ class LoadPipelineStepTaskFactory(
         flushStrategy: PipelineFlushStrategy?,
         part: Int,
         numWorkers: Int,
+        taskIndex: Int,
     ): LoadPipelineStepTask<S, K1, T, K2, U> {
         return LoadPipelineStepTask(
             batchAccumulator,
@@ -231,6 +237,7 @@ class LoadPipelineStepTaskFactory(
             flushStrategy,
             part,
             numWorkers,
+            taskIndex,
             streamCompletions
         )
     }
@@ -249,24 +256,26 @@ class LoadPipelineStepTaskFactory(
             outputQueue,
             batchSizeOverride?.let { RecordCountFlushStrategy(it) },
             part,
-            numWorkers
+            numWorkers,
+            taskIndex = 0
         )
     }
 
-    fun <S : AutoCloseable, K1 : WithStream, T, K2 : WithStream, U : Any> createFinalStep(
+    fun <S : AutoCloseable, K1 : WithStream, T, U : Any> createFinalStep(
         batchAccumulator: BatchAccumulator<S, K1, T, U>,
         inputQueue: PartitionedQueue<PipelineEvent<K1, T>>,
         part: Int,
         numWorkers: Int,
-    ): LoadPipelineStepTask<S, K1, T, K2, U> {
+    ): LoadPipelineStepTask<S, K1, T, K1, U> {
         return create(
             batchAccumulator,
             inputQueue.consume(part),
             null,
             null,
             null,
             part,
-            numWorkers
+            numWorkers,
+            taskIndex = -1
         )
     }
 
@@ -277,4 +286,25 @@ class LoadPipelineStepTaskFactory(
     ): LoadPipelineStepTask<S, StreamKey, DestinationRecordRaw, K2, U> {
         return createFirstStep(batchAccumulator, null, null, part, numWorkers)
     }
+
+    fun <S : AutoCloseable, K1 : WithStream, T, K2 : WithStream, U : Any> createIntermediateStep(
+        batchAccumulator: BatchAccumulator<S, K1, T, U>,
+        inputQueue: PartitionedQueue<PipelineEvent<K1, T>>,
+        outputPartitioner: OutputPartitioner<K1, T, K2, U>?,
+        outputQueue: PartitionedQueue<PipelineEvent<K2, U>>?,
+        part: Int,
+        numWorkers: Int,
+        taskIndex: Int,
+    ): LoadPipelineStepTask<S, K1, T, K2, U> {
+        return create(
+            batchAccumulator,
+            inputQueue.consume(part),
+            outputPartitioner,
+            outputQueue,
+            null,
+            part,
+            numWorkers,
+            taskIndex
+        )
+    }
 }
diff --git a/airbyte-cdk/bulk/toolkits/load-object-storage/src/main/kotlin/io/airbyte/cdk/load/pipline/object_storage/ObjectLoaderPartFormatterStep.kt b/airbyte-cdk/bulk/toolkits/load-object-storage/src/main/kotlin/io/airbyte/cdk/load/pipline/object_storage/ObjectLoaderPartFormatterStep.kt
@@ -4,53 +4,35 @@
 
 package io.airbyte.cdk.load.pipline.object_storage
 
-import io.airbyte.cdk.load.command.DestinationStream
-import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.PartitionedQueue
 import io.airbyte.cdk.load.message.PipelineEvent
-import io.airbyte.cdk.load.message.QueueWriter
-import io.airbyte.cdk.load.message.StreamKey
-import io.airbyte.cdk.load.pipeline.BatchUpdate
 import io.airbyte.cdk.load.pipeline.LoadPipelineStep
-import io.airbyte.cdk.load.pipeline.RecordCountFlushStrategy
 import io.airbyte.cdk.load.task.internal.LoadPipelineStepTask
+import io.airbyte.cdk.load.task.internal.LoadPipelineStepTaskFactory
 import io.airbyte.cdk.load.write.object_storage.ObjectLoader
 import io.micronaut.context.annotation.Requires
-import io.micronaut.context.annotation.Value
 import jakarta.inject.Named
 import jakarta.inject.Singleton
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.atomic.AtomicInteger
 
 @Singleton
 @Requires(bean = ObjectLoader::class)
 class ObjectLoaderPartFormatterStep(
     private val objectLoader: ObjectLoader,
-    private val recordToPartAccumulator: ObjectLoaderPartFormatter<*>,
-    @Named("recordQueue")
-    val inputQueue: PartitionedQueue<PipelineEvent<StreamKey, DestinationRecordRaw>>,
-    @Named("batchStateUpdateQueue") val batchQueue: QueueWriter<BatchUpdate>,
+    private val partFormatter: ObjectLoaderPartFormatter<*>,
     @Named("objectLoaderPartQueue")
     val outputQueue:
         PartitionedQueue<PipelineEvent<ObjectKey, ObjectLoaderPartFormatter.FormattedPart>>,
-    @Value("\${airbyte.destination.core.record-batch-size-override:null}")
-    val batchSizeOverride: Long? = null,
+    val taskFactory: LoadPipelineStepTaskFactory
 ) : LoadPipelineStep {
     override val numWorkers: Int = objectLoader.numPartWorkers
-    private val streamCompletionMap =
-        ConcurrentHashMap<DestinationStream.Descriptor, AtomicInteger>()
 
     override fun taskForPartition(partition: Int): LoadPipelineStepTask<*, *, *, *, *> {
-        return LoadPipelineStepTask(
-            recordToPartAccumulator,
-            inputQueue.consume(partition),
-            batchQueue,
+        return taskFactory.createFirstStep(
+            partFormatter,
             ObjectLoaderFormattedPartPartitioner(),
             outputQueue,
-            batchSizeOverride?.let { RecordCountFlushStrategy(it) },
             partition,
-            numWorkers,
-            streamCompletionMap
+            numWorkers
         )
     }
 }
diff --git a/airbyte-cdk/bulk/toolkits/load-object-storage/src/main/kotlin/io/airbyte/cdk/load/pipline/object_storage/ObjectLoaderPartLoaderStep.kt b/airbyte-cdk/bulk/toolkits/load-object-storage/src/main/kotlin/io/airbyte/cdk/load/pipline/object_storage/ObjectLoaderPartLoaderStep.kt
@@ -4,47 +4,39 @@
 
 package io.airbyte.cdk.load.pipline.object_storage
 
-import io.airbyte.cdk.load.command.DestinationStream
 import io.airbyte.cdk.load.message.PartitionedQueue
 import io.airbyte.cdk.load.message.PipelineEvent
-import io.airbyte.cdk.load.message.QueueWriter
-import io.airbyte.cdk.load.pipeline.BatchUpdate
 import io.airbyte.cdk.load.pipeline.LoadPipelineStep
 import io.airbyte.cdk.load.task.internal.LoadPipelineStepTask
+import io.airbyte.cdk.load.task.internal.LoadPipelineStepTaskFactory
 import io.airbyte.cdk.load.write.object_storage.ObjectLoader
 import io.micronaut.context.annotation.Requires
 import jakarta.inject.Named
 import jakarta.inject.Singleton
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.atomic.AtomicInteger
 
 @Singleton
 @Requires(bean = ObjectLoader::class)
 class ObjectLoaderPartLoaderStep(
     val loader: ObjectLoader,
-    val accumulator: ObjectLoaderPartLoader,
+    val partLoader: ObjectLoaderPartLoader,
     @Named("objectLoaderPartQueue")
     val inputQueue:
         PartitionedQueue<PipelineEvent<ObjectKey, ObjectLoaderPartFormatter.FormattedPart>>,
     @Named("objectLoaderLoadedPartQueue")
     val outputQueue: PartitionedQueue<PipelineEvent<ObjectKey, ObjectLoaderPartLoader.PartResult>>,
-    @Named("batchStateUpdateQueue") val batchQueue: QueueWriter<BatchUpdate>,
+    val taskFactory: LoadPipelineStepTaskFactory,
 ) : LoadPipelineStep {
     override val numWorkers: Int = loader.numUploadWorkers
-    private val streamCompletionMap =
-        ConcurrentHashMap<DestinationStream.Descriptor, AtomicInteger>()
 
     override fun taskForPartition(partition: Int): LoadPipelineStepTask<*, *, *, *, *> {
-        return LoadPipelineStepTask(
-            accumulator,
-            inputQueue.consume(partition),
-            batchQueue,
-            outputPartitioner = ObjectLoaderLoadedPartPartitioner(),
-            outputQueue = outputQueue,
-            flushStrategy = null,
+        return taskFactory.createIntermediateStep(
+            partLoader,
+            inputQueue,
+            ObjectLoaderLoadedPartPartitioner(),
+            outputQueue,
             partition,
             numWorkers,
-            streamCompletionMap
+            taskIndex = 2
         )
     }
 }
diff --git a/airbyte-cdk/bulk/toolkits/load-object-storage/src/main/kotlin/io/airbyte/cdk/load/pipline/object_storage/ObjectLoaderUploadCompleterStep.kt b/airbyte-cdk/bulk/toolkits/load-object-storage/src/main/kotlin/io/airbyte/cdk/load/pipline/object_storage/ObjectLoaderUploadCompleterStep.kt
@@ -4,19 +4,15 @@
 
 package io.airbyte.cdk.load.pipline.object_storage
 
-import io.airbyte.cdk.load.command.DestinationStream
 import io.airbyte.cdk.load.message.PartitionedQueue
 import io.airbyte.cdk.load.message.PipelineEvent
-import io.airbyte.cdk.load.message.QueueWriter
-import io.airbyte.cdk.load.pipeline.BatchUpdate
 import io.airbyte.cdk.load.pipeline.LoadPipelineStep
 import io.airbyte.cdk.load.task.internal.LoadPipelineStepTask
+import io.airbyte.cdk.load.task.internal.LoadPipelineStepTaskFactory
 import io.airbyte.cdk.load.write.object_storage.ObjectLoader
 import io.micronaut.context.annotation.Requires
 import jakarta.inject.Named
 import jakarta.inject.Singleton
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.atomic.AtomicInteger
 
 @Singleton
 @Requires(bean = ObjectLoader::class)
@@ -25,26 +21,16 @@ class ObjectLoaderUploadCompleterStep(
     val uploadCompleter: ObjectLoaderUploadCompleter,
     @Named("objectLoaderLoadedPartQueue")
     val inputQueue: PartitionedQueue<PipelineEvent<ObjectKey, ObjectLoaderPartLoader.PartResult>>,
-    @Named("batchStateUpdateQueue") val batchQueue: QueueWriter<BatchUpdate>,
+    val taskFactory: LoadPipelineStepTaskFactory
 ) : LoadPipelineStep {
     override val numWorkers: Int = objectLoader.numUploadCompleters
-    private val streamCompletionMap =
-        ConcurrentHashMap<DestinationStream.Descriptor, AtomicInteger>()
 
     override fun taskForPartition(partition: Int): LoadPipelineStepTask<*, *, *, *, *> {
-        return LoadPipelineStepTask(
+        return taskFactory.createFinalStep(
             uploadCompleter,
-            inputQueue.consume(partition),
-            batchQueue,
-            null,
-            null
-                as
-                PartitionedQueue<
-                    PipelineEvent<ObjectKey, ObjectLoaderUploadCompleter.UploadResult>>?,
-            null,
+            inputQueue,
             partition,
             numWorkers,
-            streamCompletionMap
         )
     }
 }
diff --git a/airbyte-cdk/bulk/toolkits/load-object-storage/src/test/kotlin/io/airbyte/cdk/load/pipeline/object_storage/ObjectLoaderPartPartitionerTest.kt b/airbyte-cdk/bulk/toolkits/load-object-storage/src/test/kotlin/io/airbyte/cdk/load/pipeline/object_storage/ObjectLoaderPartPartitionerTest.kt
@@ -17,14 +17,13 @@ import org.junit.jupiter.api.Test
 class ObjectLoaderPartPartitionerTest<K : WithStream, T> {
     @Test
     fun `partitioner always assigns 0`() {
-        val keys = (0 until 12).map(Int::toString)
         val partitioner = ObjectLoaderFormattedPartPartitioner<K, T>()
         val stream = DestinationStream.Descriptor("test", "stream")
         val numParts = 3
         var lastPartition: Int? = null
         (0 until 12).forEach {
             val partition = partitioner.getPart(ObjectKey(stream, it.toString()), numParts)
-            lastPartition?.let { last -> assertEquals(0, partition) }
+            lastPartition?.let { _ -> assertEquals(0, partition) }
             lastPartition = partition
         }
     }
diff --git a/airbyte-integrations/connectors/destination-s3/src/main/kotlin/S3V2ObjectLoader.kt b/airbyte-integrations/connectors/destination-s3/src/main/kotlin/S3V2ObjectLoader.kt
@@ -4,6 +4,7 @@
 
 package io.airbyte.integrations.destination.s3_v2
 
+import io.airbyte.cdk.load.pipeline.RoundRobinInputPartitioner
 import io.airbyte.cdk.load.write.object_storage.ObjectLoader
 
 /**
@@ -19,4 +20,5 @@ class S3V2ObjectLoader(config: S3V2Configuration<*>) : ObjectLoader {
     override val partSizeBytes: Long = config.partSizeBytes
 }
 
-// @Singleton class S3V2RoundRobinInputPartitioner : RoundRobinInputPartitioner()
+// @Singleton
+class S3V2RoundRobinInputPartitioner : RoundRobinInputPartitioner()
diff --git a/airbyte-integrations/connectors/destination-s3/src/test-integration/kotlin/io/airbyte/integrations/destination/s3_v2/S3V2PerformanceTest.kt b/airbyte-integrations/connectors/destination-s3/src/test-integration/kotlin/io/airbyte/integrations/destination/s3_v2/S3V2PerformanceTest.kt
@@ -8,7 +8,7 @@ import io.airbyte.cdk.load.write.BasicPerformanceTest
 import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
 
-// @Disabled("We don't want this to run in CI")
+@Disabled("We don't want this to run in CI")
 class S3V2JsonNoFrillsPerformanceTest :
     BasicPerformanceTest(
         configContents = S3V2TestUtils.getConfig(S3V2TestUtils.JSON_UNCOMPRESSED_CONFIG_PATH),
@@ -22,11 +22,6 @@ class S3V2JsonNoFrillsPerformanceTest :
     override fun testFileTransfer() {
         super.testFileTransfer()
     }
-
-    @Test
-    override fun testInsertRecords() {
-        super.testInsertRecords()
-    }
 }
 
 @Disabled("We don't want this to run in CI")
