destination-async-framework: move the state emission logic into Globa…

…lAsyncStateManager (airbytehq#35240)
ollionorg · Feb 26, 2024 · b53d434 · b53d434
1 parent 40833d7
commit b53d434
Show file tree

Hide file tree

Showing 13 changed files with 312 additions and 157 deletions.
diff --git a/airbyte-cdk/java/airbyte-cdk/README.md b/airbyte-cdk/java/airbyte-cdk/README.md
@@ -166,6 +166,7 @@ MavenLocal debugging steps:
 
 | Version | Date       | Pull Request                                               | Subject                                                                                                                                                        |
 |:--------|:-----------|:-----------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| 0.20.9  | 2024-02-15 | [\#35240](https://github.com/airbytehq/airbyte/pull/35240) | Make state emission to platform inside state manager itself.                                                                                                   |
 | 0.20.8  | 2024-02-15 | [\#35285](https://github.com/airbytehq/airbyte/pull/35285) | Improve blobstore module structure.                                                                                                                            |
 | 0.20.7  | 2024-02-13 | [\#35236](https://github.com/airbytehq/airbyte/pull/35236) | output logs to files in addition to stdout when running tests                                                                                                  |
 | 0.20.6  | 2024-02-12 | [\#35036](https://github.com/airbytehq/airbyte/pull/35036) | Add trace utility to emit analytics messages.                                                                                                                  |

diff --git a/...te-cdk/core/src/main/java/io/airbyte/cdk/integrations/destination_async/FlushWorkers.java b/...te-cdk/core/src/main/java/io/airbyte/cdk/integrations/destination_async/FlushWorkers.java
@@ -8,11 +8,8 @@
 import io.airbyte.cdk.integrations.destination_async.buffers.StreamAwareQueue.MessageWithMeta;
 import io.airbyte.cdk.integrations.destination_async.state.FlushFailure;
 import io.airbyte.cdk.integrations.destination_async.state.GlobalAsyncStateManager;
-import io.airbyte.cdk.integrations.destination_async.state.PartialStateWithDestinationStats;
-import io.airbyte.commons.json.Jsons;
 import io.airbyte.protocol.models.v0.AirbyteMessage;
 import io.airbyte.protocol.models.v0.StreamDescriptor;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
@@ -67,8 +64,6 @@ public class FlushWorkers implements AutoCloseable {
   private final AtomicBoolean isClosing;
   private final GlobalAsyncStateManager stateManager;
 
-  private final Object LOCK = new Object();
-
   public FlushWorkers(final BufferDequeue bufferDequeue,
                       final DestinationFlushFunction flushFunction,
                       final Consumer<AirbyteMessage> outputRecordCollector,
@@ -172,7 +167,7 @@ private void flush(final StreamDescriptor desc, final UUID flushWorkerId) {
               AirbyteFileUtils.byteCountToDisplaySize(batch.getSizeInBytes()));
 
           flusher.flush(desc, batch.getData().stream().map(MessageWithMeta::message));
-          emitStateMessages(batch.flushStates(stateIdToCount));
+          batch.flushStates(stateIdToCount, outputRecordCollector);
         }
 
         log.info("Flush Worker ({}) -- Worker finished flushing. Current queue size: {}",
@@ -222,7 +217,7 @@ public void close() throws Exception {
     log.info("Closing flush workers -- all buffers flushed");
 
     // before shutting down the supervisor, flush all state.
-    emitStateMessages(stateManager.flushStates());
+    stateManager.flushStates(outputRecordCollector);
     supervisorThread.shutdown();
     while (!supervisorThread.awaitTermination(5L, TimeUnit.MINUTES)) {
       log.info("Waiting for flush worker supervisor to shut down");
@@ -239,17 +234,6 @@ public void close() throws Exception {
     debugLoop.shutdownNow();
   }
 
-  private void emitStateMessages(final List<PartialStateWithDestinationStats> partials) {
-    synchronized (LOCK) {
-      for (final PartialStateWithDestinationStats partial : partials) {
-        final AirbyteMessage message = Jsons.deserialize(partial.stateMessage().getSerialized(), AirbyteMessage.class);
-        message.getState().setDestinationStats(partial.stats());
-        log.info("State with arrival number {} emitted from thread {}", partial.stateArrivalNumber(), Thread.currentThread().getName());
-        outputRecordCollector.accept(message);
-      }
-    }
-  }
-
   private static String humanReadableFlushWorkerId(final UUID flushWorkerId) {
     return flushWorkerId.toString().substring(0, 5);
   }

diff --git a/...n/java/io/airbyte/cdk/integrations/destination_async/buffers/MemoryAwareMessageBatch.java b/...n/java/io/airbyte/cdk/integrations/destination_async/buffers/MemoryAwareMessageBatch.java
@@ -7,9 +7,10 @@
 import io.airbyte.cdk.integrations.destination_async.GlobalMemoryManager;
 import io.airbyte.cdk.integrations.destination_async.buffers.StreamAwareQueue.MessageWithMeta;
 import io.airbyte.cdk.integrations.destination_async.state.GlobalAsyncStateManager;
-import io.airbyte.cdk.integrations.destination_async.state.PartialStateWithDestinationStats;
+import io.airbyte.protocol.models.v0.AirbyteMessage;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -57,16 +58,13 @@ public void close() throws Exception {
   }
 
   /**
-   * For the batch, marks all the states that have now been flushed. Also returns states that can be
-   * flushed. This method is descriptrive, it assumes that whatever consumes the state messages emits
-   * them, internally it purges the states it returns. message that it can.
+   * For the batch, marks all the states that have now been flushed. Also writes the states that can
+   * be flushed back to platform via stateManager.
    * <p>
-   *
-   * @return list of states that can be flushed
    */
-  public List<PartialStateWithDestinationStats> flushStates(final Map<Long, Long> stateIdToCount) {
+  public void flushStates(final Map<Long, Long> stateIdToCount, final Consumer<AirbyteMessage> outputRecordCollector) {
     stateIdToCount.forEach(stateManager::decrement);
-    return stateManager.flushStates();
+    stateManager.flushStates(outputRecordCollector);
   }
 
 }
diff --git a/...ain/java/io/airbyte/cdk/integrations/destination_async/state/GlobalAsyncStateManager.java b/...ain/java/io/airbyte/cdk/integrations/destination_async/state/GlobalAsyncStateManager.java
@@ -10,13 +10,13 @@
 import com.google.common.base.Strings;
 import io.airbyte.cdk.integrations.destination_async.GlobalMemoryManager;
 import io.airbyte.cdk.integrations.destination_async.partial_messages.PartialAirbyteMessage;
+import io.airbyte.commons.json.Jsons;
 import io.airbyte.protocol.models.v0.AirbyteMessage;
 import io.airbyte.protocol.models.v0.AirbyteStateMessage;
 import io.airbyte.protocol.models.v0.AirbyteStateStats;
 import io.airbyte.protocol.models.v0.StreamDescriptor;
-import java.util.ArrayList;
+import java.time.Instant;
 import java.util.Collection;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -25,6 +25,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang3.tuple.ImmutablePair;
@@ -153,16 +154,12 @@ public void decrement(final long stateId, final long count) {
   }
 
   /**
-   * Returns state messages with no more inflight records i.e. counter = 0 across all streams.
+   * Flushes state messages with no more inflight records i.e. counter = 0 across all streams.
    * Intended to be called by {@link io.airbyte.cdk.integrations.destination_async.FlushWorkers} after
    * a worker has finished flushing its record batch.
    * <p>
-   * The return list of states should be emitted back to the platform.
-   *
-   * @return list of state messages with no more inflight records.
    */
-  public List<PartialStateWithDestinationStats> flushStates() {
-    final List<PartialStateWithDestinationStats> output = new ArrayList<>();
+  public void flushStates(final Consumer<AirbyteMessage> outputRecordCollector) {
     Long bytesFlushed = 0L;
     synchronized (LOCK) {
       for (final Map.Entry<StreamDescriptor, LinkedBlockingDeque<Long>> entry : descToStateIdQ.entrySet()) {
@@ -195,8 +192,13 @@ public List<PartialStateWithDestinationStats> flushStates() {
           if (allRecordsCommitted) {
             final StateMessageWithArrivalNumber stateMessage = oldestState.getLeft();
             final double flushedRecordsAssociatedWithState = stateIdToCounterForPopulatingDestinationStats.get(oldestStateId).doubleValue();
-            output.add(new PartialStateWithDestinationStats(stateMessage.partialAirbyteStateMessage(),
-                new AirbyteStateStats().withRecordCount(flushedRecordsAssociatedWithState), stateMessage.arrivalNumber()));
+
+            log.info("State with arrival number {} emitted from thread {} at {}", stateMessage.arrivalNumber(), Thread.currentThread().getName(),
+                Instant.now().toString());
+            final AirbyteMessage message = Jsons.deserialize(stateMessage.partialAirbyteStateMessage.getSerialized(), AirbyteMessage.class);
+            message.getState().setDestinationStats(new AirbyteStateStats().withRecordCount(flushedRecordsAssociatedWithState));
+            outputRecordCollector.accept(message);
+
             bytesFlushed += oldestState.getRight();
 
             // cleanup
@@ -212,7 +214,6 @@ public List<PartialStateWithDestinationStats> flushStates() {
     }
 
     freeBytes(bytesFlushed);
-    return output;
   }
 
   private Long getStateIdAndIncrement(final StreamDescriptor streamDescriptor, final long increment) {

diff --git a/airbyte-cdk/java/airbyte-cdk/core/src/main/resources/version.properties b/airbyte-cdk/java/airbyte-cdk/core/src/main/resources/version.properties
@@ -1 +1 @@
-version=0.20.8
+version=0.20.9
diff --git a/...rc/test/java/io/airbyte/cdk/integrations/destination_async/buffers/BufferDequeueTest.java b/...rc/test/java/io/airbyte/cdk/integrations/destination_async/buffers/BufferDequeueTest.java
@@ -22,7 +22,6 @@ public class BufferDequeueTest {
 
   private static final int RECORD_SIZE_20_BYTES = 20;
   private static final String DEFAULT_NAMESPACE = "foo_namespace";
-  public static final String RECORD_20_BYTES = "abc";
   private static final String STREAM_NAME = "stream1";
   private static final StreamDescriptor STREAM_DESC = new StreamDescriptor().withName(STREAM_NAME);
   private static final PartialAirbyteMessage RECORD_MSG_20_BYTES = new PartialAirbyteMessage()