Merge branch 'master' into daryna/source-mixpanel/set-up-concurrency

Author: darynaishchenko

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/message/MemoryReservingPartitionedQueue.kt

@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.message
+
+import io.airbyte.cdk.load.state.ReservationManager
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.runBlocking
+import org.jetbrains.annotations.VisibleForTesting
+
+class ResourceReservingPartitionedQueue<T>(
+    val reservationManager: ReservationManager,
+    val ratioOfTotalMemoryToReserve: Double,
+    val numConsumers: Int,
+    val numProducers: Int,
+    val expectedResourceUsagePerUnit: Long
+) : PartitionedQueue<T> {
+
+    private val requestedResourceAmount =
+        (ratioOfTotalMemoryToReserve * reservationManager.totalCapacityBytes).toLong()
+    private val reservation = runBlocking {
+        reservationManager.reserveOrThrow(requestedResourceAmount, this)
+    }
+    private val minNumUnits: Int = numProducers + numConsumers * 2
+    private val maxMessageSize = reservation.bytesReserved / minNumUnits
+
+    val clampedMessageSize = expectedResourceUsagePerUnit.coerceAtMost(maxMessageSize)
+    private val maxNumUnits = (reservation.bytesReserved / clampedMessageSize).toInt()
+
+    private val totalQueueCapacity: Int = (maxNumUnits - (numProducers + numConsumers))
+
+    // Our earlier calculations should ensure this is always at least 1, but
+    // we'll clamp it to be safe.
+    @VisibleForTesting
+    val queuePartitionCapacity: Int = (totalQueueCapacity / numConsumers).coerceAtLeast(1)
+
+    private val underlying =
+        StrictPartitionedQueue<T>(
+            (0 until numConsumers)
+                .map { ChannelMessageQueue<T>(Channel(queuePartitionCapacity)) }
+                .toTypedArray()
+        )
+
+    override val partitions: Int = numConsumers
+
+    override fun consume(partition: Int): Flow<T> = underlying.consume(partition)
+
+    override suspend fun close() {
+        underlying.close()
+        reservation.release()
+    }
+
+    override suspend fun broadcast(value: T) = underlying.broadcast(value)
+
+    override suspend fun publish(value: T, partition: Int) = underlying.publish(value, partition)
+}

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/pipeline/ByPrimaryKeyInputPartitioner.kt

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.pipeline
+
+import io.airbyte.cdk.load.command.Dedupe
+import io.airbyte.cdk.load.message.DestinationRecordRaw
+import kotlin.random.Random
+
+class ByPrimaryKeyInputPartitioner : InputPartitioner {
+    private val random = Random(System.currentTimeMillis())
+
+    override fun getPartition(record: DestinationRecordRaw, numParts: Int): Int {
+        if (numParts == 1) {
+            return 0
+        }
+
+        if (record.stream.importType !is Dedupe) {
+            return random.nextInt(numParts)
+        }
+
+        val primaryKey = (record.stream.importType).primaryKey
+        val jsonData = record.asRawJson()
+
+        val primaryKeyValues =
+            primaryKey.map { keys ->
+                keys.map { key -> if (jsonData.has(key)) jsonData.get(key) else null }
+            }
+        return Math.floorMod(primaryKeyValues.hashCode(), numParts)
+    }
+}

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/pipeline/OutputPartitioner.kt

@@ -13,5 +13,5 @@ import io.airbyte.cdk.load.message.WithStream
  */
 interface OutputPartitioner<K1 : WithStream, T, K2 : WithStream, U> {
     fun getOutputKey(inputKey: K1, output: U): K2
-    fun getPart(outputKey: K2, numParts: Int): Int
+    fun getPart(outputKey: K2, inputPart: Int, numParts: Int): Int
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/pipeline/RandomInputPartitioner.kt

@@ -0,0 +1,16 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.pipeline
+
+import io.airbyte.cdk.load.message.DestinationRecordRaw
+import kotlin.random.Random
+
+class RandomInputPartitioner : InputPartitioner {
+    private val prng = Random(System.currentTimeMillis())
+
+    override fun getPartition(record: DestinationRecordRaw, numParts: Int): Int {
+        return prng.nextInt(numParts)
+    }
+}

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/spec/DestinationSpecificationInternal.kt

@@ -4,8 +4,10 @@
 
 package io.airbyte.cdk.load.spec
 
+import com.fasterxml.jackson.databind.node.ObjectNode
 import io.airbyte.cdk.spec.IdentitySpecificationExtender
 import io.airbyte.cdk.spec.SpecificationExtender
+import io.airbyte.cdk.util.Jsons
 import io.airbyte.protocol.models.v0.ConnectorSpecification
 import io.airbyte.protocol.models.v0.DestinationSyncMode
 import io.micronaut.context.annotation.Replaces
@@ -18,13 +20,47 @@ import jakarta.inject.Singleton
 class DestinationSpecificationExtender(private val spec: DestinationSpecificationExtension) :
     SpecificationExtender {
     override fun invoke(specification: ConnectorSpecification): ConnectorSpecification {
+        if (spec.groups.isNotEmpty()) {
+            val schema = specification.connectionSpecification as ObjectNode
+            schema.set<ObjectNode>(
+                "groups",
+                Jsons.arrayNode().apply {
+                    spec.groups.forEach { group ->
+                        add(
+                            Jsons.objectNode().apply {
+                                put("id", group.id)
+                                put("title", group.title)
+                            }
+                        )
+                    }
+                }
+            )
+        }
+
         return specification
             .withSupportedDestinationSyncModes(spec.supportedSyncModes)
             .withSupportsIncremental(spec.supportsIncremental)
     }
 }
 
 interface DestinationSpecificationExtension {
+    /**
+     * A connector's spec can specify "groups", which the UI will use to put related spec options in
+     * the same place. To do this, you should:
+     * * add a [Group] to the [groups] list (e.g. `Group(id = "foo", title = "Foo")`
+     * * inject `{"group": "foo"}` to the generated JSONSchema for the relevant spec options
+     * (`@JsonSchemaInject(json = """{"group": "foo"}""") val theOption: String`
+     * * note that this should be the id of the group, not the title.
+     */
+    data class Group(
+        /** A computer-friendly ID for the group */
+        val id: String,
+        /** A human-readable name for the group */
+        val title: String
+    )
+
     val supportedSyncModes: List<DestinationSyncMode>
     val supportsIncremental: Boolean
+    val groups: List<Group>
+        get() = emptyList()
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/task/internal/LoadPipelineStepTask.kt

@@ -263,7 +263,7 @@ class LoadPipelineStepTask<S : AutoCloseable, K1 : WithStream, T, K2 : WithStrea
         outputQueue?.let {
             val outputKey = outputPartitioner!!.getOutputKey(inputKey, output)
             val message = PipelineMessage(checkpointCounts.toMap(), outputKey, output)
-            val outputPart = outputPartitioner.getPart(outputKey, it.partitions)
+            val outputPart = outputPartitioner.getPart(outputKey, part, it.partitions)
             it.publish(message, outputPart)
         }
 

airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/message/ResourceReservingPartitionedQueueTest.kt

@@ -0,0 +1,83 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.message
+
+import io.airbyte.cdk.load.state.ReservationManager
+import io.airbyte.cdk.load.state.Reserved
+import io.mockk.coEvery
+import io.mockk.coVerify
+import io.mockk.mockk
+import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.Test
+
+class ResourceReservingPartitionedQueueTest {
+
+    @Test
+    fun `part queue respects memory available`() = runTest {
+        val reservationManager = mockk<ReservationManager>(relaxed = true)
+        coEvery { reservationManager.totalCapacityBytes } returns 1000
+        val reservation = mockk<Reserved<ResourceReservingPartitionedQueue<Unit>>>(relaxed = true)
+        coEvery { reservation.bytesReserved } returns 500
+        coEvery {
+            reservationManager.reserveOrThrow<ResourceReservingPartitionedQueue<Unit>>(any(), any())
+        } returns reservation
+        ResourceReservingPartitionedQueue<Unit>(
+            reservationManager,
+            0.5,
+            1, // not relevant
+            1, // not relevant
+            1, // not relevant
+        )
+        coVerify {
+            reservationManager.reserveOrThrow<ResourceReservingPartitionedQueue<Unit>>(500, any())
+        }
+    }
+
+    @Test
+    fun `part queue clamps part size if too many workers`() {
+        val reservationManager = ReservationManager(1000)
+        val queue =
+            ResourceReservingPartitionedQueue<Unit>(
+                reservationManager,
+                0.8,
+                3,
+                5,
+                100,
+            )
+        val clampedSize = queue.clampedMessageSize
+        Assertions.assertEquals(800 / 11, clampedSize)
+    }
+
+    @Test
+    fun `part queue does not clamp part size if not too many workers`() {
+        val reservationManager = ReservationManager(1000)
+        val queue =
+            ResourceReservingPartitionedQueue<Unit>(
+                reservationManager,
+                0.8,
+                1,
+                5,
+                100,
+            )
+        val clampedSize = queue.clampedMessageSize
+        Assertions.assertEquals(100, clampedSize)
+    }
+
+    @Test
+    fun `queue capacity is derived from clamped size and available memory`() {
+        val reservationManager = ReservationManager(1000)
+        val queue =
+            ResourceReservingPartitionedQueue<Unit>(
+                reservationManager,
+                0.75,
+                1,
+                3,
+                100,
+            )
+        Assertions.assertEquals(100L, queue.clampedMessageSize)
+        Assertions.assertEquals(3, queue.queuePartitionCapacity)
+    }
+}

airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/task/internal/LoadPipelineStepTaskUTest.kt

@@ -429,7 +429,11 @@ class LoadPipelineStepTaskUTest {
                             return TestKey(output, inputKey.stream)
                         }
 
-                        override fun getPart(outputKey: TestKey, numParts: Int): Int {
+                        override fun getPart(
+                            outputKey: TestKey,
+                            inputPart: Int,
+                            numParts: Int
+                        ): Int {
                             if (outputKey.output) return 1
                             return 0
                         }

airbyte-cdk/bulk/toolkits/load-db/src/main/kotlin/io/airbyte/cdk/load/pipeline/db/BulkLoadCompletedUploadPartitioner.kt

@@ -26,7 +26,7 @@ class BulkLoadCompletedUploadPartitioner<T : RemoteObject<*>> :
         return StreamKey(inputKey.stream)
     }
 
-    override fun getPart(outputKey: StreamKey, numParts: Int): Int {
+    override fun getPart(outputKey: StreamKey, inputPart: Int, numParts: Int): Int {
         return Math.floorMod(outputKey.stream.hashCode(), numParts)
     }
 }

airbyte-cdk/bulk/toolkits/load-db/src/main/kotlin/io/airbyte/cdk/load/pipeline/db/InsertLoaderInputPartitioner.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.pipeline.db
+
+import io.airbyte.cdk.load.pipeline.ByPrimaryKeyInputPartitioner
+import io.airbyte.cdk.load.pipeline.ByStreamInputPartitioner
+import io.airbyte.cdk.load.pipeline.InputPartitioner
+import io.airbyte.cdk.load.pipeline.RandomInputPartitioner
+import io.airbyte.cdk.load.write.db.InsertLoader
+import io.micronaut.context.annotation.Factory
+import io.micronaut.context.annotation.Requires
+import jakarta.inject.Singleton
+
+@Factory
+class InsertLoaderInputPartitionerFactory {
+    @Singleton
+    @Requires(bean = InsertLoader::class)
+    fun insertLoaderInputPartitioner(insertLoader: InsertLoader<*>): InputPartitioner {
+        return when (insertLoader.partitioningStrategy) {
+            InsertLoader.PartitioningStrategy.ByStream -> ByStreamInputPartitioner()
+            InsertLoader.PartitioningStrategy.ByPrimaryKey -> ByPrimaryKeyInputPartitioner()
+            InsertLoader.PartitioningStrategy.Random -> RandomInputPartitioner()
+        }
+    }
+}