fix: avoid blocking alarm pool during stream reconnects (#1400)

tmdiep · web-flow · commit 3197cc0bc377 · 2023-05-15T07:33:08.000-04:00
* fix: avoid blocking alarm pool during stream reconnects Do not process background requests in alarm pool while publish and subscribe streams are reconnecting. Avoid grabbing the same lock as reinitialize, which will be held until the stream is successfully reconnected. Also reverts #1394, which can cause thread explosion. * Remove unintended change
diff --git a/google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/AlarmFactory.java b/google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/AlarmFactory.java
@@ -47,11 +47,4 @@ static AlarmFactory create(Duration duration) {
                 duration.toNanos(),
                 NANOSECONDS);
   }
-
-  /** Runnable is executed by an unbounded pool, although the alarm pool is bounded. */
-  static AlarmFactory createUnbounded(Duration duration) {
-    AlarmFactory underlying = create(duration);
-    return runnable ->
-        underlying.newAlarm(() -> SystemExecutors.getFuturesExecutor().execute(runnable));
-  }
 }
diff --git a/google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/wire/PublisherImpl.java b/google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/wire/PublisherImpl.java
@@ -117,6 +117,12 @@ void failBatch(int startIdx, CheckedApiException e) {
   @GuardedBy("monitor.monitor")
   private final Queue<InFlightBatch> batchesInFlight = new ArrayDeque<>();
 
+  // reconnectingMonitor is always acquired after monitor.monitor when both are held.
+  private final CloseableMonitor reconnectingMonitor = new CloseableMonitor();
+
+  @GuardedBy("reconnectingMonitor.monitor")
+  private boolean reconnecting = false;
+
   @VisibleForTesting
   PublisherImpl(
       PublishStreamFactory streamFactory,
@@ -146,7 +152,7 @@ public PublisherImpl(
     this(
         streamFactory,
         new BatchPublisherImpl.Factory(),
-        AlarmFactory.createUnbounded(
+        AlarmFactory.create(
             Duration.ofNanos(
                 Objects.requireNonNull(batchingSettings.getDelayThreshold()).toNanos())),
         initialRequest,
@@ -189,6 +195,9 @@ private void rebatchForRestart() {
   @Override
   public void triggerReinitialize(CheckedApiException streamError) {
     try (CloseableMonitor.Hold h = monitor.enter()) {
+      try (CloseableMonitor.Hold rh = reconnectingMonitor.enter()) {
+        reconnecting = true;
+      }
       connection.reinitialize(initialRequest);
       rebatchForRestart();
       Collection<InFlightBatch> batches = batchesInFlight;
@@ -201,6 +210,9 @@ public void triggerReinitialize(CheckedApiException streamError) {
                         .get()
                         .publish(batch.messagesToSend(), batch.firstSequenceNumber()));
           });
+      try (CloseableMonitor.Hold rh = reconnectingMonitor.enter()) {
+        reconnecting = false;
+      }
     } catch (CheckedApiException e) {
       onPermanentError(e);
     }
@@ -219,7 +231,7 @@ protected void handlePermanentError(CheckedApiException error) {
   @Override
   protected void start() {
     try (CloseableMonitor.Hold h = monitor.enter()) {
-      this.alarmFuture = Optional.of(alarmFactory.newAlarm(this::flushToStream));
+      this.alarmFuture = Optional.of(alarmFactory.newAlarm(this::backgroundFlushToStream));
     }
   }
 
@@ -276,6 +288,15 @@ public void cancelOutstandingPublishes() {
     }
   }
 
+  private void backgroundFlushToStream() {
+    try (CloseableMonitor.Hold h = reconnectingMonitor.enter()) {
+      if (reconnecting) {
+        return;
+      }
+    }
+    flushToStream();
+  }
+
   private void flushToStream() {
     try (CloseableMonitor.Hold h = monitor.enter()) {
       if (shutdown) return;
diff --git a/google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/wire/SubscriberImpl.java b/google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/wire/SubscriberImpl.java
@@ -82,6 +82,12 @@ public class SubscriberImpl extends ProxyService
   @GuardedBy("monitor.monitor")
   private boolean shutdown = false;
 
+  // reconnectingMonitor is always acquired after monitor.monitor when both are held.
+  private final CloseableMonitor reconnectingMonitor = new CloseableMonitor();
+
+  @GuardedBy("reconnectingMonitor.monitor")
+  private boolean reconnecting = false;
+
   @VisibleForTesting
   SubscriberImpl(
       SubscribeStreamFactory streamFactory,
@@ -205,6 +211,9 @@ public void triggerReinitialize(CheckedApiException streamError) {
 
     try (CloseableMonitor.Hold h = monitor.enter()) {
       if (shutdown) return;
+      try (CloseableMonitor.Hold rh = reconnectingMonitor.enter()) {
+        reconnecting = true;
+      }
       connection.reinitialize(getInitialRequest());
       connection.modifyConnection(
           connectedSubscriber -> {
@@ -214,6 +223,9 @@ public void triggerReinitialize(CheckedApiException streamError) {
                 .requestForRestart()
                 .ifPresent(request -> connectedSubscriber.get().allowFlow(request));
           });
+      try (CloseableMonitor.Hold rh = reconnectingMonitor.enter()) {
+        reconnecting = false;
+      }
     } catch (CheckedApiException e) {
       onPermanentError(e);
     }
@@ -239,6 +251,11 @@ public void onClientResponse(List<SequencedMessage> messages) throws CheckedApiE
   }
 
   private void processBatchFlowRequest() {
+    try (CloseableMonitor.Hold h = reconnectingMonitor.enter()) {
+      if (reconnecting) {
+        return;
+      }
+    }
     try (CloseableMonitor.Hold h = monitor.enter()) {
       if (shutdown) return;
       connection.modifyConnection(

Original file line number	Diff line number	Diff line change
`@@ -47,11 +47,4 @@ static AlarmFactory create(Duration duration) {`
`47`	`47`	`duration.toNanos(),`
`48`	`48`	`NANOSECONDS);`
`49`	`49`	`}`
`50`		`-`
`51`		`- /** Runnable is executed by an unbounded pool, although the alarm pool is bounded. */`
`52`		`- static AlarmFactory createUnbounded(Duration duration) {`
`53`		`- AlarmFactory underlying = create(duration);`
`54`		`- return runnable ->`
`55`		`- underlying.newAlarm(() -> SystemExecutors.getFuturesExecutor().execute(runnable));`
`56`		`- }`
`57`	`50`	`}`