Log allocates and more data.

airbytehq · Oct 10, 2023 · 8605063 · 8605063
1 parent e39d6e2
commit 8605063
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 14 deletions.
diff --git a/...core/src/main/java/io/airbyte/cdk/integrations/destination_async/GlobalMemoryManager.java b/...core/src/main/java/io/airbyte/cdk/integrations/destination_async/GlobalMemoryManager.java
@@ -61,7 +61,7 @@ public long getCurrentMemoryBytes() {
    *
    * @return the size of the allocated block, in bytes
    */
-  public synchronized long requestMemory() {
+  public synchronized long requestMemory(final String name) {
     // todo(davin): what happens if the incoming record is larger than 30MB?
     if (currentMemoryBytes.get() >= maxMemoryBytes) {
       return 0L;
@@ -72,10 +72,11 @@ public synchronized long requestMemory() {
     final var toAllocateBytes = Math.min(freeMem, BLOCK_SIZE_BYTES);
     currentMemoryBytes.addAndGet(toAllocateBytes);
 
-    log.debug("Memory Requested: max: {}, allocated: {}, allocated in this request: {}",
+    log.info("Memory Requested: max: {}, allocated: {}, allocated in this request: {}, for {}",
         FileUtils.byteCountToDisplaySize(maxMemoryBytes),
         FileUtils.byteCountToDisplaySize(currentMemoryBytes.get()),
-        FileUtils.byteCountToDisplaySize(toAllocateBytes));
+        FileUtils.byteCountToDisplaySize(toAllocateBytes),
+        name);
     return toAllocateBytes;
   }
 

diff --git a/...re/src/main/java/io/airbyte/cdk/integrations/destination_async/buffers/BufferEnqueue.java b/...re/src/main/java/io/airbyte/cdk/integrations/destination_async/buffers/BufferEnqueue.java
@@ -48,14 +48,16 @@ public void addRecord(final PartialAirbyteMessage message, final Integer sizeInB
 
   private void handleRecord(final PartialAirbyteMessage message, final Integer sizeInBytes) {
     final StreamDescriptor streamDescriptor = extractStateFromRecord(message);
-    final var queue = buffers.computeIfAbsent(streamDescriptor, _k -> new StreamAwareQueue(memoryManager.requestMemory()));
+    final var streamName = streamDescriptor.getNamespace() == null ? streamDescriptor.getName() : streamDescriptor.getNamespace() + " - " + streamDescriptor.getName();
+    final var queue = buffers.computeIfAbsent(streamDescriptor, _k -> new StreamAwareQueue(memoryManager.requestMemory(streamName + " init")));
     final long stateId = stateManager.getStateIdAndIncrementCounter(streamDescriptor);
 
     var addedToQueue = queue.offer(message, sizeInBytes, stateId);
 
+
     int i = 0;
     while (!addedToQueue) {
-      final var newlyAllocatedMemory = memoryManager.requestMemory();
+      final var newlyAllocatedMemory = memoryManager.requestMemory(streamName + " record");
       if (newlyAllocatedMemory > 0) {
         queue.addMaxMemory(newlyAllocatedMemory);
       }

diff --git a/...re/src/main/java/io/airbyte/cdk/integrations/destination_async/buffers/BufferManager.java b/...re/src/main/java/io/airbyte/cdk/integrations/destination_async/buffers/BufferManager.java
@@ -101,8 +101,9 @@ private void printQueueInfo() {
     for (final var entry : buffers.entrySet()) {
       final var queue = entry.getValue();
       queueInfo.append(
-          String.format("  Queue name: %s, num records: %d, num bytes: %s",
-              entry.getKey().getName(), queue.size(), AirbyteFileUtils.byteCountToDisplaySize(queue.getCurrentMemoryUsage())))
+          String.format("  Queue name: %s, num records: %d, num bytes: %s, allocated bytes: %s",
+              entry.getKey().getName(), queue.size(), AirbyteFileUtils.byteCountToDisplaySize(queue.getCurrentMemoryUsage()),
+              AirbyteFileUtils.byteCountToDisplaySize(queue.getMaxMemoryUsage())))
           .append(System.lineSeparator());
     }
 

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
@@ -81,7 +81,7 @@ public class GlobalAsyncStateManager {
 
   public GlobalAsyncStateManager(final GlobalMemoryManager memoryManager) {
     this.memoryManager = memoryManager;
-    memoryAllocated = new AtomicLong(memoryManager.requestMemory());
+    memoryAllocated = new AtomicLong(memoryManager.requestMemory("init"));
     memoryUsed = new AtomicLong();
   }
 
@@ -286,7 +286,7 @@ private void closeState(final PartialAirbyteMessage message, final long sizeInBy
   private void allocateMemoryToState(final long sizeInBytes) {
     if (memoryAllocated.get() < memoryUsed.get() + sizeInBytes) {
       while (memoryAllocated.get() < memoryUsed.get() + sizeInBytes) {
-        memoryAllocated.addAndGet(memoryManager.requestMemory());
+        memoryAllocated.addAndGet(memoryManager.requestMemory("state"));
         try {
           LOGGER.debug("Insufficient memory to store state message. Allocated: {}, Used: {}, Size of State Msg: {}, Needed: {}",
               FileUtils.byteCountToDisplaySize(memoryAllocated.get()),

diff --git a/.../src/test/java/io/airbyte/cdk/integrations/destination_async/GlobalMemoryManagerTest.java b/.../src/test/java/io/airbyte/cdk/integrations/destination_async/GlobalMemoryManagerTest.java
@@ -16,14 +16,14 @@ public class GlobalMemoryManagerTest {
   void test() {
     final GlobalMemoryManager mgr = new GlobalMemoryManager(15 * BYTES_MB);
 
-    assertEquals(10 * BYTES_MB, mgr.requestMemory());
-    assertEquals(5 * BYTES_MB, mgr.requestMemory());
-    assertEquals(0, mgr.requestMemory());
+    assertEquals(10 * BYTES_MB, mgr.requestMemory(""));
+    assertEquals(5 * BYTES_MB, mgr.requestMemory(""));
+    assertEquals(0, mgr.requestMemory(""));
 
     mgr.free(10 * BYTES_MB);
-    assertEquals(10 * BYTES_MB, mgr.requestMemory());
+    assertEquals(10 * BYTES_MB, mgr.requestMemory(""));
     mgr.free(16 * BYTES_MB);
-    assertEquals(10 * BYTES_MB, mgr.requestMemory());
+    assertEquals(10 * BYTES_MB, mgr.requestMemory(""));
   }
 
 }