add-max-retry-property-support (#2062)

AntonZarutsky · pongad · commit 97f6164476e9 · 2017-05-18T10:01:20.000+10:00
diff --git a/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/spi/v1/Publisher.java b/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/spi/v1/Publisher.java
@@ -344,8 +344,7 @@ public void onSuccess(PublishResponse result) {
                                 + "the expected %s results. Please contact Cloud Pub/Sub support "
                                 + "if this frequently occurs",
                             result.getMessageIdsCount(), outstandingBatch.size()));
-                for (OutstandingPublish oustandingMessage :
-                    outstandingBatch.outstandingPublishes) {
+                for (OutstandingPublish oustandingMessage : outstandingBatch.outstandingPublishes) {
                   oustandingMessage.publishResult.setException(t);
                 }
                 return;
@@ -368,9 +367,10 @@ public void onFailure(Throwable t) {
                 computeNextBackoffDelayMs(outstandingBatch, retrySettings, longRandom);
 
             if (!isRetryable(t)
+                || retrySettings.getMaxAttempts() > 0
+                    && outstandingBatch.getAttempt() > retrySettings.getMaxAttempts()
                 || System.currentTimeMillis() + nextBackoffDelay
-                    > outstandingBatch.creationTime
-                        + retrySettings.getTotalTimeout().toMillis()) {
+                    > outstandingBatch.creationTime + retrySettings.getTotalTimeout().toMillis()) {
               try {
                 for (OutstandingPublish outstandingPublish :
                     outstandingBatch.outstandingPublishes) {
@@ -408,6 +408,10 @@ private static final class OutstandingBatch {
       this.batchSizeBytes = batchSizeBytes;
     }
 
+    public int getAttempt() {
+      return attempt;
+    }
+
     public int size() {
       return outstandingPublishes.size();
     }
@@ -506,7 +510,7 @@ interface LongRandom {
    * Constructs a new {@link Builder} using the given topic.
    *
    * <p>Example of creating a {@code Publisher}.
-   * <pre> {@code
+   * <pre>{@code
    * String projectName = "my_project";
    * String topicName = "my_topic";
    * TopicName topic = TopicName.create(projectName, topicName);
diff --git a/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/spi/v1/PublisherImplTest.java b/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/spi/v1/PublisherImplTest.java
@@ -16,13 +16,8 @@
 
 package com.google.cloud.pubsub.spi.v1;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
-
-import com.google.api.gax.batching.BatchingSettings;
 import com.google.api.core.ApiFuture;
+import com.google.api.gax.batching.BatchingSettings;
 import com.google.api.gax.batching.FlowControlSettings;
 import com.google.api.gax.batching.FlowController.LimitExceededBehavior;
 import com.google.api.gax.grpc.ChannelProvider;
@@ -40,14 +35,17 @@
 import io.grpc.inprocess.InProcessChannelBuilder;
 import io.grpc.inprocess.InProcessServerBuilder;
 import io.grpc.internal.ServerImpl;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Executor;
-import org.threeten.bp.Duration;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.threeten.bp.Duration;
+
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+
+import static org.junit.Assert.*;
 
 @RunWith(JUnit4.class)
 public class PublisherImplTest {
@@ -279,6 +277,81 @@ public void testPublishFailureRetries() throws Exception {
     publisher.shutdown();
   }
 
+  @Test(expected = ExecutionException.class)
+  public void testPublishFailureRetries_retriesDisabled() throws Exception {
+    Publisher publisher =
+        getTestPublisherBuilder()
+            .setExecutorProvider(SINGLE_THREAD_EXECUTOR)
+            .setRetrySettings(
+                Publisher.Builder.DEFAULT_RETRY_SETTINGS
+                    .toBuilder()
+                    .setTotalTimeout(Duration.ofSeconds(10))
+                    .setMaxAttempts(1)
+                    .build())
+            .build();
+
+    testPublisherServiceImpl.addPublishError(new Throwable("Transiently failing"));
+
+    ApiFuture<String> publishFuture1 = sendTestMessage(publisher, "A");
+
+    try {
+      publishFuture1.get();
+    } finally {
+      assertSame(testPublisherServiceImpl.getCapturedRequests().size(), 1);
+      publisher.shutdown();
+    }
+  }
+
+  @Test
+  public void testPublishFailureRetries_maxRetriesSetup() throws Exception {
+    Publisher publisher =
+        getTestPublisherBuilder()
+            .setExecutorProvider(SINGLE_THREAD_EXECUTOR)
+            .setRetrySettings(
+                Publisher.Builder.DEFAULT_RETRY_SETTINGS
+                    .toBuilder()
+                    .setTotalTimeout(Duration.ofSeconds(10))
+                    .setMaxAttempts(3)
+                    .build())
+            .build();
+
+    testPublisherServiceImpl.addPublishError(new Throwable("Transiently failing"));
+    testPublisherServiceImpl.addPublishError(new Throwable("Transiently failing"));
+    testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1"));
+
+    ApiFuture<String> publishFuture1 = sendTestMessage(publisher, "A");
+
+    assertEquals("1", publishFuture1.get());
+
+    assertEquals(3, testPublisherServiceImpl.getCapturedRequests().size());
+    publisher.shutdown();
+  }
+
+  @Test
+  public void testPublishFailureRetries_maxRetriesSetUnlimited() throws Exception {
+    Publisher publisher =
+        getTestPublisherBuilder()
+            .setExecutorProvider(SINGLE_THREAD_EXECUTOR)
+            .setRetrySettings(
+                Publisher.Builder.DEFAULT_RETRY_SETTINGS
+                    .toBuilder()
+                    .setTotalTimeout(Duration.ofSeconds(10))
+                    .setMaxAttempts(0)
+                    .build())
+            .build();
+
+    testPublisherServiceImpl.addPublishError(new Throwable("Transiently failing"));
+    testPublisherServiceImpl.addPublishError(new Throwable("Transiently failing"));
+    testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1"));
+
+    ApiFuture<String> publishFuture1 = sendTestMessage(publisher, "A");
+
+    assertEquals("1", publishFuture1.get());
+
+    assertEquals(3, testPublisherServiceImpl.getCapturedRequests().size());
+    publisher.shutdown();
+  }
+
   public void testPublishFailureRetries_exceededsRetryDuration() throws Exception {
     Publisher publisher =
         getTestPublisherBuilder()
