Flush remaining data to S3 during shutdown (#6424)

flush remaining data to S3 during shutdown

Signed-off-by: Subrahmanyam-Gollapalli <subrahmanyam.gollapalli@freshworks.com>

Author: Subrahmanyam-Gollapalli

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3Sink.java

@@ -175,4 +175,23 @@ private void doInitializeInternal() {
     public void doOutput(final Collection<Record<Event>> records) {
         s3SinkService.output(records);
     }
+
+    /**
+     * Shutdown the S3 sink by flushing all remaining data to S3 before shutting down.
+     * This ensures no data is lost during shutdown.
+     */
+    @Override
+    public void shutdown() {
+        LOG.info("S3 sink shutdown initiated. Flushing all remaining data to S3.");
+        try {
+            // Force flush all remaining groups to S3
+            s3SinkService.flushAllRemainingGroups();
+        } catch (Exception e) {
+            LOG.error("Error occurred while flushing remaining data during shutdown", e);
+        } finally {
+            // Call parent shutdown to stop retry thread
+            super.shutdown();
+            LOG.info("S3 sink shutdown completed.");
+        }
+    }
 }

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkService.java

@@ -248,4 +248,44 @@ private void checkAggregateThresholdsAndFlushIfNeeded(final List<CompletableFutu
             }
         }
     }
+
+    /**
+     * Force flush all remaining S3 groups to S3. This is called during shutdown
+     * to ensure all buffered data is uploaded before the sink shuts down.
+     */
+    void flushAllRemainingGroups() {
+        if (s3GroupManager.hasNoGroups()) {
+            LOG.debug("No remaining groups to flush during shutdown");
+            return;
+        }
+
+        LOG.info("Shutdown initiated. Flushing {} remaining S3 groups to S3.", s3GroupManager.getNumberOfGroups());
+        reentrantLock.lock();
+        try {
+            final List<CompletableFuture<?>> completableFutures = new ArrayList<>();
+
+            // Force flush all remaining groups
+            for (final S3Group s3Group : s3GroupManager.getS3GroupEntries()) {
+                LOG.info("Force flushing group with key {} containing {} events and {} bytes during shutdown",
+                        s3Group.getBuffer().getKey(), s3Group.getBuffer().getEventCount(), s3Group.getBuffer().getSize());
+                flushToS3IfNeeded(completableFutures, s3Group, true);
+                numberOfObjectsForceFlushed.increment();
+            }
+
+            // Wait for all uploads to complete
+            if (!completableFutures.isEmpty()) {
+                try {
+                    LOG.info("Waiting for {} uploads to complete during shutdown", completableFutures.size());
+                    CompletableFuture.allOf(completableFutures.toArray(new CompletableFuture[0]))
+                            .thenRun(() -> LOG.info("All {} uploads completed during shutdown", completableFutures.size()))
+                            .join();
+                    LOG.info("All remaining groups have been flushed to S3 during shutdown");
+                } catch (final Exception e) {
+                    LOG.error("Exception occurred while waiting for uploads to complete during shutdown", e);
+                }
+            }
+        } finally {
+            reentrantLock.unlock();
+        }
+    }
 }

data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceTest.java

@@ -665,6 +665,88 @@ void output_will_flush_the_largest_group_until_below_aggregate_threshold_when_ag
         verify(s3ObjectsForceFlushedCounter, times(2)).increment();
     }
 
+    @Test
+    void flushAllRemainingGroups_when_no_groups_does_not_flush() {
+        when(s3GroupManager.hasNoGroups()).thenReturn(true);
+
+        final S3SinkService s3SinkService = createObjectUnderTest();
+        s3SinkService.flushAllRemainingGroups();
+
+        verify(s3GroupManager).hasNoGroups();
+        verify(s3GroupManager, never()).getS3GroupEntries();
+        verify(s3ObjectsForceFlushedCounter, never()).increment();
+    }
+
+    @Test
+    void flushAllRemainingGroups_when_groups_exist_flushes_all() throws IOException {
+        when(s3GroupManager.hasNoGroups()).thenReturn(false);
+        when(s3GroupManager.getNumberOfGroups()).thenReturn(2);
+
+        final S3Group s3Group1 = mock(S3Group.class);
+        final S3Group s3Group2 = mock(S3Group.class);
+        final Buffer buffer1 = mock(Buffer.class);
+        final Buffer buffer2 = mock(Buffer.class);
+        final OutputStream outputStream1 = mock(OutputStream.class);
+        final OutputStream outputStream2 = mock(OutputStream.class);
+
+        when(s3Group1.getBuffer()).thenReturn(buffer1);
+        when(s3Group1.getOutputCodec()).thenReturn(codec);
+        when(buffer1.getOutputStream()).thenReturn(outputStream1);
+        when(buffer1.getKey()).thenReturn(UUID.randomUUID().toString());
+        when(buffer1.getEventCount()).thenReturn(3);
+        when(buffer1.getSize()).thenReturn(100L);
+        when(buffer1.getDuration()).thenReturn(Duration.ZERO);
+
+        when(s3Group2.getBuffer()).thenReturn(buffer2);
+        when(s3Group2.getOutputCodec()).thenReturn(codec);
+        when(buffer2.getOutputStream()).thenReturn(outputStream2);
+        when(buffer2.getKey()).thenReturn(UUID.randomUUID().toString());
+        when(buffer2.getEventCount()).thenReturn(5);
+        when(buffer2.getSize()).thenReturn(200L);
+        when(buffer2.getDuration()).thenReturn(Duration.ZERO);
+
+        final CompletableFuture<Void> completedFuture = CompletableFuture.completedFuture(null);
+        when(buffer1.flushToS3(any(Consumer.class), any(Consumer.class))).thenReturn(Optional.of(completedFuture));
+        when(buffer2.flushToS3(any(Consumer.class), any(Consumer.class))).thenReturn(Optional.of(completedFuture));
+
+        when(s3GroupManager.getS3GroupEntries()).thenReturn(List.of(s3Group1, s3Group2));
+        doNothing().when(codec).complete(any(OutputStream.class));
+
+        final S3SinkService s3SinkService = createObjectUnderTest();
+
+        try (final MockedStatic<CompletableFuture> completableFutureMockedStatic = mockStatic(CompletableFuture.class)) {
+            final CompletableFuture<Void> mockCompletableFuture = mock(CompletableFuture.class);
+            when(mockCompletableFuture.thenRun(any(Runnable.class))).thenReturn(mockCompletableFuture);
+            when(mockCompletableFuture.join()).thenReturn(null);
+            completableFutureMockedStatic.when(() -> CompletableFuture.allOf(any())).thenReturn(mockCompletableFuture);
+            s3SinkService.flushAllRemainingGroups();
+        }
+
+        verify(s3GroupManager).hasNoGroups();
+        verify(s3GroupManager).getNumberOfGroups();
+        verify(s3GroupManager).getS3GroupEntries();
+        verify(s3GroupManager).removeGroup(s3Group1);
+        verify(s3GroupManager).removeGroup(s3Group2);
+        verify(codec, times(2)).complete(any(OutputStream.class));
+        verify(buffer1).flushToS3(any(Consumer.class), any(Consumer.class));
+        verify(buffer2).flushToS3(any(Consumer.class), any(Consumer.class));
+        verify(s3ObjectsForceFlushedCounter, times(2)).increment();
+    }
+
+    @Test
+    void flushAllRemainingGroups_when_empty_groups_list_does_not_increment_force_flush_counter() {
+        when(s3GroupManager.hasNoGroups()).thenReturn(false);
+        when(s3GroupManager.getNumberOfGroups()).thenReturn(0);
+        when(s3GroupManager.getS3GroupEntries()).thenReturn(Collections.emptyList());
+
+        final S3SinkService s3SinkService = createObjectUnderTest();
+        s3SinkService.flushAllRemainingGroups();
+
+        verify(s3GroupManager).hasNoGroups();
+        verify(s3GroupManager).getS3GroupEntries();
+        verify(s3ObjectsForceFlushedCounter, never()).increment();
+    }
+
     private Collection<Record<Event>> generateRandomStringEventRecord() {
         return generateEventRecords(50);
     }