Bulk load CDK: log diagnostics when unflushed states at end of sync (#68668)

Author: edgao

airbyte-cdk/bulk/changelog.md

@@ -1,3 +1,6 @@
+## Version 0.1.58
+
+load cdk: log unflushed state diagnostic info
 ## Version 0.1.57
 
 **Extract CDK**

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

@@ -67,7 +67,8 @@ class PipelineRunner(
         if (store.hasStates()) {
             val stateException =
                 IllegalStateException("Sync completed, but unflushed states were detected.")
-            log.info { "Destination Pipeline Completed — Exceptionally: $stateException" }
+            store.logStateInfo()
+            log.error { "Destination Pipeline Completed — Exceptionally: $stateException" }
             throw stateException
         }
 

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/StateHistogramStore.kt

@@ -32,6 +32,14 @@ class StateHistogramStore {
         return expectedCount == flushedCount
     }
 
+    // mirrors isComplete. Purely for debugging purposes.
+    fun whyIsStateIncomplete(key: StateKey): String {
+        val expectedCount = expected.get(key)
+        val partitionFlushCounts = key.partitionKeys.map { flushed.get(it) ?: 0 }
+        val flushedCount = partitionFlushCounts.sum()
+        return "expectedCount $expectedCount does not equal flushedCount $flushedCount (by partition: $partitionFlushCounts)"
+    }
+
     fun remove(key: StateKey): Long? {
         key.partitionKeys.forEach { flushed.remove(it) }
         return expected.remove(key)

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/StateStore.kt

@@ -47,4 +47,23 @@ class StateStore(
     }
 
     fun hasStates(): Boolean = states.isNotEmpty()
+
+    /** Log some information about why state messages cannot be flushed. */
+    // implementation-wise, this is largely just mirroring getNextComplete(),
+    // except it logs things instead of actually doing anything.
+    fun logStateInfo() {
+        val message = StringBuilder("State diagnostic information:")
+        val key = states.firstKey()
+        val state = states.get(key)
+        message.append("\nFirst state key: $key (full state message: $state)")
+        if (key.id != stateSequence.get()) {
+            message.append("\nKey ID did not match state sequence ($stateSequence.get()")
+        } else if (!histogramStore.isComplete(key)) {
+            message.append(
+                "\nhistogram store says key is incomplete: ${histogramStore.whyIsStateIncomplete(key)}"
+            )
+        }
+        log.info { message }
+        return
+    }
 }

airbyte-cdk/bulk/version.properties

@@ -1 +1 @@
-version=0.1.57
+version=0.1.58