Fix unbounded async pipe sink retries (#18222) (#18225)

Author: Caideyipi

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtask.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.exception.pipe.PipeRuntimeException;
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkNonReportTimeConfigurableException;
 import org.apache.iotdb.commons.pipe.agent.task.connection.UnboundedBlockingPendingQueue;
 import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey;
 import org.apache.iotdb.commons.pipe.agent.task.subtask.PipeAbstractSinkSubtask;
@@ -170,6 +171,8 @@ private void transferHeartbeatEvent(final PipeHeartbeatEvent event) {
           })) {
         return;
       }
+    } catch (final PipeRuntimeSinkNonReportTimeConfigurableException e) {
+      throw e;
     } catch (final Exception e) {
       throw new PipeConnectionException(
           "PipeConnector: "

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java

@@ -23,6 +23,7 @@
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.client.ThriftClient;
 import org.apache.iotdb.commons.client.async.AsyncPipeDataTransferServiceClient;
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkNonReportTimeConfigurableException;
 import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey;
 import org.apache.iotdb.commons.pipe.config.PipeConfig;
 import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
@@ -85,7 +86,6 @@
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicLong;
 
 import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_ENABLE_SEND_TSFILE_LIMIT;
 import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_ENABLE_SEND_TSFILE_LIMIT_DEFAULT_VALUE;
@@ -539,6 +539,8 @@ private void logOnClientException(
    * @see PipeConnector#transfer(TsFileInsertionEvent) for more details.
    */
   private void transferQueuedEventsIfNecessary(final boolean forced) {
+    throwIfReceiverProbeIsDelayed();
+
     if ((retryEventQueue.isEmpty() && retryTsFileQueue.isEmpty())
         || (!forced
             && retryEventQueueEventCounter.getTabletInsertionEventCount()
@@ -600,6 +602,8 @@ private void transferQueuedEventsIfNecessary(final boolean forced) {
         }
       }
 
+      throwIfReceiverProbeIsDelayed();
+
       // Stop retrying if the execution time exceeds the threshold for better realtime performance
       if (System.currentTimeMillis() - retryStartTime
           > PipeConfig.getInstance().getPipeAsyncSinkMaxRetryExecutionTimeMsPerCall()) {
@@ -747,21 +751,58 @@ public void waitIfReceiverTemporarilyUnavailable(final TEndPoint endPoint) {
       return;
     }
 
-    while (!isClosed.get()) {
-      final long waitTimeInMs = backoff.getRemainingWaitTimeInMs();
-      if (waitTimeInMs <= 0) {
-        return;
+    while (!isClosed.get() && backoff.isActive()) {
+      if (backoff.isRetryMaxDurationExceeded()) {
+        final long probeDelayInMs = backoff.tryAcquireProbeAndGetDelayInMs();
+        if (probeDelayInMs <= 0) {
+          return;
+        }
+        throw createReceiverProbeDelayException(endPointKey, backoff);
       }
 
-      try {
-        Thread.sleep(waitTimeInMs);
-      } catch (final InterruptedException e) {
-        Thread.currentThread().interrupt();
-        return;
+      final long retryTimeInMs = backoff.reserveNextRetryTimeInMs();
+      while (!isClosed.get() && backoff.isActive()) {
+        if (backoff.isRetryMaxDurationExceeded()) {
+          break;
+        }
+
+        final long waitTimeInMs = retryTimeInMs - System.currentTimeMillis();
+        if (waitTimeInMs <= 0) {
+          return;
+        }
+
+        try {
+          Thread.sleep(Math.min(waitTimeInMs, 1000L));
+        } catch (final InterruptedException e) {
+          Thread.currentThread().interrupt();
+          return;
+        }
       }
     }
   }
 
+  private void throwIfReceiverProbeIsDelayed() {
+    for (final Map.Entry<String, ReceiverTemporaryUnavailableBackoff> entry :
+        receiverBackoffMap.entrySet()) {
+      final long probeDelayInMs = entry.getValue().getRemainingProbeDelayInMs();
+      if (probeDelayInMs <= 0) {
+        continue;
+      }
+
+      throw createReceiverProbeDelayException(entry.getKey(), entry.getValue());
+    }
+  }
+
+  private static PipeRuntimeSinkNonReportTimeConfigurableException
+      createReceiverProbeDelayException(
+          final String endPointKey, final ReceiverTemporaryUnavailableBackoff backoff) {
+    return new PipeRuntimeSinkNonReportTimeConfigurableException(
+        String.format(
+            "Receiver %s remained temporarily unavailable for more than %d ms, pause regular retries and probe every %d ms.",
+            endPointKey, backoff.getRetryMaxDurationInMs(), backoff.getRetryProbeIntervalInMs()),
+        Long.MAX_VALUE);
+  }
+
   public void recordReceiverStatus(final TEndPoint endPoint, final TSStatus status) {
     final String endPointKey = format(endPoint);
     if (Objects.isNull(endPointKey) || Objects.isNull(status)) {
@@ -781,10 +822,15 @@ public void recordReceiverStatus(final TEndPoint endPoint, final TSStatus status
             status);
       }
     } else if (isSuccess(status)) {
-      final ReceiverTemporaryUnavailableBackoff backoff = receiverBackoffMap.get(endPointKey);
-      if (Objects.nonNull(backoff) && backoff.getRemainingWaitTimeInMs() <= 0) {
-        receiverBackoffMap.remove(endPointKey, backoff);
-      }
+      receiverBackoffMap.computeIfPresent(
+          endPointKey,
+          (key, backoff) -> {
+            if (!backoff.shouldResetOnSuccess()) {
+              return backoff;
+            }
+            backoff.markAvailable();
+            return null;
+          });
     }
   }
 
@@ -994,23 +1040,90 @@ private static class ReceiverTemporaryUnavailableBackoff {
 
     private final long maxBackoffTimeInMs =
         Math.max(0, PipeConfig.getInstance().getPipeSinkSubtaskSleepIntervalMaxMs());
-    private final AtomicLong currentBackoffTimeInMs =
-        new AtomicLong(
-            Math.min(
-                Math.max(1, PipeConfig.getInstance().getPipeSinkSubtaskSleepIntervalInitMs()),
-                maxBackoffTimeInMs));
-    private final AtomicLong nextAvailableTimeInMs = new AtomicLong(0);
-
-    private long markTemporarilyUnavailable() {
-      final long backoffTimeInMs = currentBackoffTimeInMs.get();
-      nextAvailableTimeInMs.updateAndGet(
-          current -> Math.max(current, System.currentTimeMillis() + backoffTimeInMs));
-      currentBackoffTimeInMs.updateAndGet(this::getNextBackoffTimeInMs);
+    private final long initialBackoffTimeInMs =
+        Math.min(
+            Math.max(1, PipeConfig.getInstance().getPipeSinkSubtaskSleepIntervalInitMs()),
+            maxBackoffTimeInMs);
+    private final long retryMaxDurationInMs =
+        PipeConfig.getInstance().getPipeAsyncSinkRetryMaxDurationMs();
+    private final long retryProbeIntervalInMs =
+        Math.max(1, PipeConfig.getInstance().getPipeAsyncSinkRetryProbeIntervalMs());
+
+    private boolean active = false;
+    private long firstUnavailableTimeInMs = 0;
+    private long currentBackoffTimeInMs = initialBackoffTimeInMs;
+    private long failureBackoffUntilInMs = 0;
+    private long nextReservedRetryTimeInMs = 0;
+    private long nextProbeTimeInMs = 0;
+
+    private synchronized long markTemporarilyUnavailable() {
+      final long currentTimeInMs = System.currentTimeMillis();
+      if (!active) {
+        active = true;
+        firstUnavailableTimeInMs = currentTimeInMs;
+        currentBackoffTimeInMs = initialBackoffTimeInMs;
+        failureBackoffUntilInMs = 0;
+        nextReservedRetryTimeInMs = 0;
+        nextProbeTimeInMs = 0;
+      }
+
+      final long backoffTimeInMs = currentBackoffTimeInMs;
+      failureBackoffUntilInMs =
+          Math.max(failureBackoffUntilInMs, safeAdd(currentTimeInMs, backoffTimeInMs));
+      nextReservedRetryTimeInMs = Math.max(nextReservedRetryTimeInMs, failureBackoffUntilInMs);
+      currentBackoffTimeInMs = getNextBackoffTimeInMs(currentBackoffTimeInMs);
       return backoffTimeInMs;
     }
 
-    private long getRemainingWaitTimeInMs() {
-      return nextAvailableTimeInMs.get() - System.currentTimeMillis();
+    private synchronized boolean isActive() {
+      return active;
+    }
+
+    private synchronized boolean isRetryMaxDurationExceeded() {
+      return active
+          && retryMaxDurationInMs >= 0
+          && System.currentTimeMillis() - firstUnavailableTimeInMs >= retryMaxDurationInMs;
+    }
+
+    private synchronized long reserveNextRetryTimeInMs() {
+      final long currentTimeInMs = System.currentTimeMillis();
+      final long retryTimeInMs =
+          Math.max(currentTimeInMs, Math.max(failureBackoffUntilInMs, nextReservedRetryTimeInMs));
+      nextReservedRetryTimeInMs = safeAdd(retryTimeInMs, currentBackoffTimeInMs);
+      return retryTimeInMs;
+    }
+
+    private synchronized long tryAcquireProbeAndGetDelayInMs() {
+      final long currentTimeInMs = System.currentTimeMillis();
+      if (currentTimeInMs >= nextProbeTimeInMs) {
+        nextProbeTimeInMs = safeAdd(currentTimeInMs, retryProbeIntervalInMs);
+        return 0;
+      }
+      return nextProbeTimeInMs - currentTimeInMs;
+    }
+
+    private synchronized long getRemainingProbeDelayInMs() {
+      return isRetryMaxDurationExceeded()
+          ? Math.max(0, nextProbeTimeInMs - System.currentTimeMillis())
+          : 0;
+    }
+
+    private synchronized boolean shouldResetOnSuccess() {
+      return active
+          && (isRetryMaxDurationExceeded()
+              || failureBackoffUntilInMs - System.currentTimeMillis() <= 0);
+    }
+
+    private synchronized void markAvailable() {
+      active = false;
+    }
+
+    private long getRetryMaxDurationInMs() {
+      return retryMaxDurationInMs;
+    }
+
+    private long getRetryProbeIntervalInMs() {
+      return retryProbeIntervalInMs;
     }
 
     private long getNextBackoffTimeInMs(final long currentBackoffTimeInMs) {
@@ -1021,5 +1134,9 @@ private long getNextBackoffTimeInMs(final long currentBackoffTimeInMs) {
           ? maxBackoffTimeInMs
           : currentBackoffTimeInMs << 1;
     }
+
+    private static long safeAdd(final long left, final long right) {
+      return left >= Long.MAX_VALUE - right ? Long.MAX_VALUE : left + right;
+    }
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandler.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.commons.client.ThriftClient;
 import org.apache.iotdb.commons.client.async.AsyncPipeDataTransferServiceClient;
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkNonReportTimeConfigurableException;
 import org.apache.iotdb.commons.pipe.resource.log.PipeLogger;
 import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferSliceReqBuilder;
 import org.apache.iotdb.db.pipe.sink.protocol.thrift.async.IoTDBDataRegionAsyncSink;
@@ -92,7 +93,7 @@ public void onError(final Exception exception) {
    * @param client the client used for data transfer
    * @param req the request containing transfer details
    * @return {@code true} if the transfer was initiated successfully, {@code false} if the connector
-   *     is closed
+   *     is closed or the receiver probe is delayed
    * @throws TException if an error occurs during the transfer
    */
   protected boolean tryTransfer(
@@ -106,7 +107,13 @@ protected boolean tryTransfer(
     if (returnFalseIfSinkIsClosed(client)) {
       return false;
     }
-    sink.waitIfReceiverTemporarilyUnavailable(client.getEndPoint());
+    try {
+      sink.waitIfReceiverTemporarilyUnavailable(client.getEndPoint());
+    } catch (final PipeRuntimeSinkNonReportTimeConfigurableException e) {
+      returnClientToPool(client);
+      onError(e);
+      return false;
+    }
     if (returnFalseIfSinkIsClosed(client)) {
       return false;
     }
@@ -136,6 +143,21 @@ private boolean returnFalseIfSinkIsClosed(final AsyncPipeDataTransferServiceClie
     return true;
   }
 
+  private void returnClientToPool(final AsyncPipeDataTransferServiceClient client) {
+    client.setShouldReturnSelf(true);
+    client.returnSelf(
+        e -> {
+          if (e instanceof IllegalStateException) {
+            PipeLogger.log(
+                LOGGER::info,
+                "Illegal state when return the client to object pool, maybe the pool is already cleared. Will ignore.");
+            return true;
+          }
+          return false;
+        });
+    this.client = null;
+  }
+
   /**
    * @return {@code true} if all transmissions corresponding to the handler have been completed,
    *     {@code false} otherwise
@@ -276,6 +298,9 @@ private void fallbackToWholeRequest(
         return;
       }
       client.pipeTransfer(originalReq, this);
+    } catch (final PipeRuntimeSinkNonReportTimeConfigurableException e) {
+      returnClientToPool(client);
+      PipeTransferTrackableHandler.this.onError(e);
     } catch (final Exception e) {
       PipeTransferTrackableHandler.this.onError(e);
     }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskTest.java

@@ -20,9 +20,11 @@
 package org.apache.iotdb.db.pipe.agent.task.subtask.sink;
 
 import org.apache.iotdb.commons.conf.CommonDescriptor;
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkNonReportTimeConfigurableException;
 import org.apache.iotdb.commons.pipe.agent.task.connection.UnboundedBlockingPendingQueue;
 import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey;
 import org.apache.iotdb.commons.pipe.sink.protocol.PipeConnectorWithEventDiscard;
+import org.apache.iotdb.db.pipe.event.common.heartbeat.PipeHeartbeatEvent;
 import org.apache.iotdb.pipe.api.PipeConnector;
 import org.apache.iotdb.pipe.api.customizer.configuration.PipeConnectorRuntimeConfiguration;
 import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameterValidator;
@@ -39,9 +41,13 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.withSettings;
 
 public class PipeSinkSubtaskTest {
@@ -221,6 +227,46 @@ public void testCloseDoesNotWaitForeverForConnectorClose() throws Exception {
     }
   }
 
+  @Test
+  public void testHeartbeatPreservesReceiverProbeDelayException() throws Exception {
+    final long originalSleepIntervalInitMs =
+        CommonDescriptor.getInstance().getConfig().getPipeSinkSubtaskSleepIntervalInitMs();
+    final long originalSleepIntervalMaxMs =
+        CommonDescriptor.getInstance().getConfig().getPipeSinkSubtaskSleepIntervalMaxMs();
+    CommonDescriptor.getInstance().getConfig().setPipeSinkSubtaskSleepIntervalInitMs(1);
+    CommonDescriptor.getInstance().getConfig().setPipeSinkSubtaskSleepIntervalMaxMs(2);
+
+    final PipeConnector connector = mock(PipeConnector.class);
+    final UnboundedBlockingPendingQueue<Event> pendingQueue =
+        mock(UnboundedBlockingPendingQueue.class);
+    when(pendingQueue.waitedPoll()).thenReturn(new PipeHeartbeatEvent("1", false));
+    doThrow(new PipeRuntimeSinkNonReportTimeConfigurableException("probe delayed", Long.MAX_VALUE))
+        .when(connector)
+        .transfer(any(Event.class));
+
+    final PipeSinkSubtask subtask =
+        new PipeSinkSubtask(
+            "PipeSinkSubtaskTest",
+            System.currentTimeMillis(),
+            "data_test",
+            0,
+            pendingQueue,
+            connector);
+
+    try {
+      Assert.assertTrue(subtask.executeOnce());
+      verify(connector, never()).handshake();
+    } finally {
+      subtask.close();
+      CommonDescriptor.getInstance()
+          .getConfig()
+          .setPipeSinkSubtaskSleepIntervalInitMs(originalSleepIntervalInitMs);
+      CommonDescriptor.getInstance()
+          .getConfig()
+          .setPipeSinkSubtaskSleepIntervalMaxMs(originalSleepIntervalMaxMs);
+    }
+  }
+
   private static class BlockingHandshakeConnector
       implements PipeConnector, PipeConnectorWithEventDiscard {