Merge pull request #1275 from data-integrations/bugfix/PLUGIN-1660-Add-retry-for-snapshot-creation

PLUGIN-1660 - Add retry for create and delete snapshot.

Author: greeshmaswaminathan

src/main/java/io/cdap/plugin/gcp/publisher/source/PubSubDirectDStream.java

@@ -55,6 +55,7 @@ public class PubSubDirectDStream<T> extends InputDStream<T> implements Streaming
 
   private static final Logger LOG = LoggerFactory.getLogger(PubSubDirectDStream.class);
   private static final String CDAP_PIPELINE = "cdap_pipeline";
+  private static final int MAX_SNAPSHOT_ATTEMPTS = 3;
 
   private final Credentials credentials;
   private final PubSubSubscriberConfig config;
@@ -64,6 +65,7 @@ public class PubSubDirectDStream<T> extends InputDStream<T> implements Streaming
   private final SerializableFunction<PubSubMessage, T> mappingFn;
   private final StreamingContext streamingContext;
   private final String pipeline;
+  private final BackoffConfig backoffConfig;
 
   private SubscriptionAdminClient subscriptionAdminClient;
   private ProjectSnapshotName currentSnapshotName;
@@ -82,6 +84,7 @@ public PubSubDirectDStream(io.cdap.cdap.etl.api.streaming.StreamingContext conte
     this.pipeline = context.getPipelineName();
     this.credentials = PubSubSubscriberUtil.createCredentials(config.getServiceAccount(),
                                                               config.isServiceAccountFilePath());
+    backoffConfig = BackoffConfig.defaultInstance();
   }
 
   @Override
@@ -195,7 +198,7 @@ private String generateName(String subscription) {
   }
 
   private void createSubscriptionIfNotPresent() throws IOException, InterruptedException {
-    PubSubSubscriberUtil.createSubscription(() -> true, BackoffConfig.defaultInstance(),
+    PubSubSubscriberUtil.createSubscription(() -> true, backoffConfig,
                                             ProjectSubscriptionName.format(config.getProject(),
                                                                            config.getSubscription()),
                                             TopicName.format(config.getProject(), config.getTopic()),
@@ -243,15 +246,21 @@ private ProjectSnapshotName fetchSnapShot(String subscriptionId,
 
   private Snapshot createSnapshot(ProjectSnapshotName projectSnapshotName,
                                   ProjectSubscriptionName projectSubscriptionName) {
-    // Creation takes around 3.5 s .
     LOG.debug("Creating snapshot {} for subscription {} in Pub/Sub .", projectSnapshotName.toString(),
               projectSubscriptionName.toString());
-    CreateSnapshotRequest request = CreateSnapshotRequest.newBuilder()
-      .setName(projectSnapshotName.toString())
-      .setSubscription(projectSubscriptionName.toString())
-      .putAllLabels(Collections.singletonMap(CDAP_PIPELINE, getLabelValue(pipeline)))
-      .build();
-    return subscriptionAdminClient.createSnapshot(request);
+    try {
+      return PubSubSubscriberUtil.callWithRetry(() -> {
+        // Creation takes around 3.5 s .
+        CreateSnapshotRequest request = CreateSnapshotRequest.newBuilder()
+          .setName(projectSnapshotName.toString())
+          .setSubscription(projectSubscriptionName.toString())
+          .putAllLabels(Collections.singletonMap(CDAP_PIPELINE, getLabelValue(pipeline)))
+          .build();
+        return subscriptionAdminClient.createSnapshot(request);
+      }, backoffConfig, MAX_SNAPSHOT_ATTEMPTS);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
   }
 
   @VisibleForTesting
@@ -268,9 +277,16 @@ static String getLabelValue(String pipeline) {
   }
 
   private void deleteSnapshot(ProjectSnapshotName projectSnapshotName) {
-    // Deletion takes around 2.5 s .
-    // TODO - Consider making this asynchronous
-    subscriptionAdminClient.deleteSnapshot(projectSnapshotName);
+    try {
+      PubSubSubscriberUtil.callWithRetry(() -> {
+        // Deletion takes around 2.5 s .
+        // TODO - Consider making this asynchronous
+        subscriptionAdminClient.deleteSnapshot(projectSnapshotName);
+        return null;
+      }, backoffConfig, MAX_SNAPSHOT_ATTEMPTS);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
   }
 
   private void saveSnapshotAsState(Snapshot snapshot, String subscription,

src/main/java/io/cdap/plugin/gcp/publisher/source/PubSubSubscriberUtil.java

@@ -227,6 +227,36 @@ public static Credentials createCredentials(String serviceAccount, boolean servi
     }
   }
 
+  /**
+   * Call provided supplier with retries.
+   *
+   * @param supplier      The supplier to be invoked
+   * @param backoffConfig {@link BackoffConfig} for the retries
+   * @param maxAttempts   Integer indicating max number of attempts
+   * @param <T>
+   * @return Value returned by supplier
+   * @throws Exception Any exception that is not retryable or exceeded retry
+   */
+  public static <T> T callWithRetry(Supplier<T> supplier,
+                                    BackoffConfig backoffConfig, int maxAttempts) throws Exception {
+    int backoff = backoffConfig.getInitialBackoffMs();
+    ApiException lastApiException = null;
+    while (maxAttempts-- > 0) {
+      try {
+        return supplier.get();
+      } catch (ApiException ae) {
+        lastApiException = ae;
+        //Retry if the exception is retryable.
+        if (PubSubSubscriberUtil.isApiExceptionRetryable(ae)) {
+          backoff = PubSubSubscriberUtil.sleepAndIncreaseBackoff(() -> true, backoff, backoffConfig);
+          continue;
+        }
+        throw ae;
+      }
+    }
+    throw new RuntimeException(lastApiException);
+  }
+
   private static SubscriptionAdminClient buildSubscriptionAdminClient(Credentials credentials) throws IOException {
     SubscriptionAdminSettings.Builder builder = SubscriptionAdminSettings.newBuilder();
     if (credentials != null) {

src/test/java/io/cdap/plugin/gcp/publisher/source/PubSubSubscriberUtilTest.java

@@ -0,0 +1,46 @@
+package io.cdap.plugin.gcp.publisher.source;
+
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.api.gax.rpc.InternalException;
+import com.google.api.gax.rpc.NotFoundException;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.util.function.Supplier;
+
+/**
+ * Tests for {@link PubSubSubscriberUtil}
+ */
+public class PubSubSubscriberUtilTest {
+  @Test(expected = NotFoundException.class)
+  public void testCallWithRetryNonRetryable() throws Exception {
+    PubSubSubscriberUtil.callWithRetry(() -> {
+      throw new NotFoundException(new RuntimeException("subscription not found"),
+                                  GrpcStatusCode.of(Status.Code.NOT_FOUND), false);
+    }, BackoffConfig.defaultInstance(), 3);
+  }
+
+  @Test(expected = RuntimeException.class)
+  public void testCallWithRetryMaxRetry() throws Exception {
+    PubSubSubscriberUtil.callWithRetry(() -> {
+      throw new InternalException(new StatusRuntimeException(Status.INTERNAL),
+                                  GrpcStatusCode.of(Status.Code.INTERNAL), false);
+    }, BackoffConfig.defaultInstance(), 3);
+  }
+
+  @Test
+  public void testCallWithRetrySuccess() throws Exception {
+    InternalException internalException = new InternalException(new StatusRuntimeException(Status.INTERNAL),
+                                                                GrpcStatusCode.of(Status.Code.INTERNAL), false);
+    Supplier<String> testSupplier = Mockito.mock(Supplier.class);
+    String returnValue = "success";
+    Mockito.when(testSupplier.get()).thenThrow(internalException).thenThrow(internalException)
+      .thenReturn(returnValue);
+
+    String result = PubSubSubscriberUtil.callWithRetry(testSupplier, BackoffConfig.defaultInstance(), 3);
+    Assert.assertSame(returnValue, result);
+  }
+}