Destination S3V2: Restore Thread-Safety to Unique Key Counter

airbytehq · Jan 8, 2025 · af7d487 · af7d487
1 parent b8f4274
commit af7d487
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 8 deletions.
diff --git a/...n/kotlin/io/airbyte/cdk/load/state/object_storage/ObjectStorageDestinationStateManager.kt b/...n/kotlin/io/airbyte/cdk/load/state/object_storage/ObjectStorageDestinationStateManager.kt
@@ -21,6 +21,7 @@ import io.micronaut.context.annotation.Secondary
 import jakarta.inject.Singleton
 import java.nio.file.Paths
 import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicLong
 import kotlinx.coroutines.flow.mapNotNull
 import kotlinx.coroutines.flow.toList
 
@@ -30,7 +31,8 @@ class ObjectStorageDestinationState(
     @JsonProperty("generations_by_state")
     var generationMap: ConcurrentHashMap<State, ConcurrentHashMap<Long, MutableMap<String, Long>>> =
         ConcurrentHashMap(),
-    @JsonProperty("count_by_key") var countByKey: MutableMap<String, Long> = mutableMapOf()
+    @JsonProperty("count_by_key")
+    var countByKey: ConcurrentHashMap<String, AtomicLong> = ConcurrentHashMap()
 ) : DestinationState {
     enum class State {
         STAGED,
@@ -121,9 +123,10 @@ class ObjectStorageDestinationState(
 
     /** Used to guarantee the uniqueness of a key */
     suspend fun ensureUnique(key: String): String {
-        val ordinal = countByKey.merge(key, 0L) { old, new -> maxOf(old + 1, new) } ?: 0L
-        return if (ordinal > 0L) {
-            "$key-$ordinal"
+        val counter = countByKey.getOrPut(key) { AtomicLong(0L) }
+        val count = counter.incrementAndGet()
+        return if (count > 1L) {
+            "$key-$count"
         } else {
             key
         }
@@ -207,7 +210,7 @@ class ObjectStorageFallbackPersister(
                             ObjectStorageDestinationState.State.FINALIZED to ConcurrentHashMap(it)
                         )
                     ),
-                    countByKey
+                    ConcurrentHashMap(countByKey.mapValues { (_, v) -> AtomicLong(v) })
                 )
             }
     }

diff --git a/airbyte-integrations/connectors/destination-s3/metadata.yaml b/airbyte-integrations/connectors/destination-s3/metadata.yaml
@@ -2,7 +2,7 @@ data:
   connectorSubtype: file
   connectorType: destination
   definitionId: 4816b78f-1489-44c1-9060-4b19d5fa9362
-  dockerImageTag: 1.5.0-rc.4
+  dockerImageTag: 1.5.0-rc.5
   dockerRepository: airbyte/destination-s3
   githubIssueLabel: destination-s3
   icon: s3.svg

diff --git a/airbyte-integrations/connectors/destination-s3/src/main/kotlin/S3V2Checker.kt b/airbyte-integrations/connectors/destination-s3/src/main/kotlin/S3V2Checker.kt
@@ -50,8 +50,7 @@ class S3V2Checker<T : OutputStream>(private val timeProvider: TimeProvider) :
                 log.info { "Successfully wrote test file: $results" }
             } finally {
                 s3Object?.also { s3Client.delete(it) }
-                val results = s3Client.list(path).toList()
-                log.info { "Successfully removed test tile: $results" }
+                log.info { "Successfully removed test file" }
             }
         }
     }

diff --git a/docs/integrations/destinations/s3.md b/docs/integrations/destinations/s3.md
@@ -544,6 +544,7 @@ To see connector limitations, or troubleshoot your S3 connector, see more [in ou
 
 | Version    | Date       | Pull Request                                               | Subject                                                                                                                                              |
 |:-----------|:-----------|:-----------------------------------------------------------|:-----------------------------------------------------------------------------------------------------------------------------------------------------|
+| 1.5.0-rc.4 | 2025-01-06 | [50954](https://github.com/airbytehq/airbyte/pull/50954)   | Bug fix: transient failure due to bug in filename clash prevention                                                                                   |
 | 1.5.0-rc.4 | 2025-01-06 | [50954](https://github.com/airbytehq/airbyte/pull/50954)   | Bug fix: StreamLoader::close dispatched multiple times per stream                                                                                    |
 | 1.5.0-rc.3 | 2025-01-06 | [50949](https://github.com/airbytehq/airbyte/pull/50949)   | Bug fix: parquet types/values nested in union of objects do not convert properly                                                                     |
 | 1.5.0-rc.2 | 2025-01-02 | [50857](https://github.com/airbytehq/airbyte/pull/50857)   | Migrate to Bulk Load CDK: cost reduction, perf increase, bug fix for filename clashes                                                                |