From 0b79ed2859c3de3ef09d1c4b4110cb97aaa3b6ff Mon Sep 17 00:00:00 2001
From: gintarasm <gintaras.matulas@gmail.com>
Date: Mon, 8 Apr 2024 12:30:49 +0300
Subject: [PATCH] fix: retry when stream writer is closed

---
 .../sink/async/AsyncBigQuerySinkWriter.java   | 23 +++++++++++++++++++
 .../flink/bigquery/AsyncBigQuerySinkTest.java | 17 ++++++++++++++
 .../util/MockAsyncProtoClientProvider.java    | 16 +++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/src/main/java/com/vinted/flink/bigquery/sink/async/AsyncBigQuerySinkWriter.java b/src/main/java/com/vinted/flink/bigquery/sink/async/AsyncBigQuerySinkWriter.java
index 52c6604..0327ce4 100644
--- a/src/main/java/com/vinted/flink/bigquery/sink/async/AsyncBigQuerySinkWriter.java
+++ b/src/main/java/com/vinted/flink/bigquery/sink/async/AsyncBigQuerySinkWriter.java
@@ -128,6 +128,15 @@ protected void submitRequestEntries(List<StreamRequest> list, Consumer<List<Stre
                             logger.error("Trace-id {}, StreamWriter failed to append {}", traceId, t.getMessage());
                             var status = Status.fromThrowable(t);
                             switch (status.getCode()) {
+                                case ABORTED: {
+                                    if (isStreamWriterClosed(t)) {
+                                        logger.info("Trace-id {} steam writer closed. Retrying: {}", traceId, t.getMessage());
+                                        return retry(t, traceId, request);
+                                    } else {
+                                        logger.error("Trace-id {} Received error {} with status {}", traceId, t.getMessage(), status.getCode());
+                                        throw new AsyncWriterException(traceId, status.getCode(), t);
+                                    }
+                                }
                                 case UNAVAILABLE: {
                                     this.recreateStreamWriter(traceId, request.getStream(), writer.getWriterId(), request.getTable());
                                     return retry(t, traceId, request);
@@ -194,6 +203,20 @@ protected void submitRequestEntries(List<StreamRequest> list, Consumer<List<Stre
 
     }
 
+    private boolean isStreamWriterClosed(Throwable t) {
+        if (t == null) {
+            return false;
+        }
+
+        for (Throwable cause = t; cause != null; cause = cause.getCause()) {
+            if (cause instanceof Exceptions.StreamWriterClosedException) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     private boolean isMaximumRequestCallbackWaitTimeExceededException(Throwable t) {
         if (t == null) {
             return false;
diff --git a/src/test/java/com/vinted/flink/bigquery/AsyncBigQuerySinkTest.java b/src/test/java/com/vinted/flink/bigquery/AsyncBigQuerySinkTest.java
index f25d724..c686820 100644
--- a/src/test/java/com/vinted/flink/bigquery/AsyncBigQuerySinkTest.java
+++ b/src/test/java/com/vinted/flink/bigquery/AsyncBigQuerySinkTest.java
@@ -63,6 +63,23 @@ public void shouldRecreateWriterAndRetryFailingWithMaximumRequestCallbackWaitTim
         assertThat(mockClientProvider.getNumOfCreatedWriter()).isEqualTo(3);
     }
 
+    @Test
+    public void shouldRetryFailingWithAbortedStatus(@FlinkTest.FlinkParam FlinkTest.PipelineRunner runner, @FlinkTest.FlinkParam MockAsyncProtoClientProvider mockClientProvider) throws Exception {
+        mockClientProvider.givenStreamWriterClosed();
+        mockClientProvider.givenRetryCount(2);
+
+        assertThatThrownBy(() -> {
+            runner
+                    .withRetryCount(0)
+                    .runWithCustomSink(withBigQuerySink(mockClientProvider, pipeline(List.of(
+                            givenRow(1)
+                    ))));
+        }).isInstanceOf(JobExecutionException.class);
+
+
+        verify(mockClientProvider.getMockProtoWriter(), times(2)).append(any());
+    }
+
     @Test
     public void shouldFailAndNotRetryWhenUnknownErrorReceived(@FlinkTest.FlinkParam FlinkTest.PipelineRunner runner, @FlinkTest.FlinkParam MockAsyncProtoClientProvider mockClientProvider) throws Exception {
         mockClientProvider.givenFailingAppendWithStatus(Status.UNKNOWN);
diff --git a/src/test/java/com/vinted/flink/bigquery/util/MockAsyncProtoClientProvider.java b/src/test/java/com/vinted/flink/bigquery/util/MockAsyncProtoClientProvider.java
index 2619bd6..7a9f723 100644
--- a/src/test/java/com/vinted/flink/bigquery/util/MockAsyncProtoClientProvider.java
+++ b/src/test/java/com/vinted/flink/bigquery/util/MockAsyncProtoClientProvider.java
@@ -102,6 +102,14 @@ public void givenFailingAppendWithStatus(Status status) throws Descriptors.Descr
                 .thenReturn(createAppendRowsResponseError(new StatusException(status)));
     }
 
+    public void givenStreamWriterClosed() throws Descriptors.DescriptorValidationException, IOException {
+        var response =  createAppendRowsResponseError(
+                new StatusException(Status.ABORTED.withCause(createStreamWriterClosedException()))
+        );
+        Mockito.when(MockAsyncProtoClientProvider.protoWriter.append(Mockito.any()))
+                .thenReturn(response);
+    }
+
     public void givenTimeoutForAppend() throws Descriptors.DescriptorValidationException, IOException {
         Mockito.when(MockAsyncProtoClientProvider.protoWriter.append(Mockito.any()))
                 .thenReturn(createTimeoutAppendRowsResponse());
@@ -148,6 +156,14 @@ private static Exceptions.OffsetAlreadyExists createOffsetAlreadyExistsException
         return offsetMock;
     }
 
+    private static Exceptions.StreamWriterClosedException createStreamWriterClosedException() {
+        var offsetMock = Mockito.mock(Exceptions.StreamWriterClosedException.class);
+        Mockito.when(offsetMock.getStatus()).thenReturn(Status.ABORTED);
+        Mockito.when(offsetMock.getStreamName()).thenReturn("stream");
+        Mockito.when(offsetMock.getCause()).thenReturn(new RuntimeException());
+        return offsetMock;
+    }
+
     private static SettableApiFuture<AppendRowsResponse> createAppendRowsResponse() {
         SettableApiFuture<AppendRowsResponse> result = SettableApiFuture.create();
         result.set(AppendRowsResponse.newBuilder().buildPartial());