Ensuring only one task is executed at a time and it always schedules the next one

Author: LikeTheSalad

opamp-client/src/main/java/io/opentelemetry/opamp/client/internal/request/service/HttpRequestService.java

@@ -27,6 +27,8 @@
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -43,7 +45,8 @@ public final class HttpRequestService implements RequestService {
   private final AtomicBoolean isRunning = new AtomicBoolean(false);
   private final AtomicBoolean hasStopped = new AtomicBoolean(false);
   private final AtomicReference<PeriodicDelay> currentDelay;
-  private final AtomicReference<ScheduledFuture<?>> currentTask = new AtomicReference<>();
+  private final AtomicReference<ScheduledFuture<?>> scheduledTask = new AtomicReference<>();
+  private final Lock sendLock = new ReentrantLock();
   private final RetryAfterParser retryAfterParser;
   @Nullable private Callback callback;
   @Nullable private Supplier<Request> requestSupplier;
@@ -100,26 +103,12 @@ public void start(Callback callback, Supplier<Request> requestSupplier) {
     if (isRunning.compareAndSet(false, true)) {
       this.callback = callback;
       this.requestSupplier = requestSupplier;
-      currentTask.set(
-          executorService.schedule(
-              this::periodicSend, getNextDelay().toNanos(), TimeUnit.NANOSECONDS));
+      scheduleNextExecution();
     } else {
       throw new IllegalStateException("HttpRequestService is already running");
     }
   }
 
-  private void periodicSend() {
-    doSendRequest();
-    // schedule the next execution
-    currentTask.set(
-        executorService.schedule(
-            this::periodicSend, getNextDelay().toNanos(), TimeUnit.NANOSECONDS));
-  }
-
-  private void sendOnce() {
-    executorService.execute(this::doSendRequest);
-  }
-
   private Duration getNextDelay() {
     return Objects.requireNonNull(currentDelay.get()).getNextDelay();
   }
@@ -138,7 +127,43 @@ public void sendRequest() {
       throw new IllegalStateException("HttpRequestService is not running");
     }
 
-    sendOnce();
+    executorService.execute(this::requestOnDemand);
+  }
+
+  private void requestOnDemand() {
+    if (sendLock.tryLock()) {
+      try {
+        ScheduledFuture<?> scheduledFuture = scheduledTask.get();
+        if (scheduledFuture != null) {
+          // Cancel future task
+          scheduledFuture.cancel(false);
+        }
+        sendAndScheduleNext();
+      } finally {
+        sendLock.unlock();
+      }
+    }
+  }
+
+  private void periodicRequest() {
+    if (sendLock.tryLock()) {
+      try {
+        sendAndScheduleNext();
+      } finally {
+        sendLock.unlock();
+      }
+    }
+  }
+
+  private void sendAndScheduleNext() {
+    doSendRequest();
+    scheduleNextExecution();
+  }
+
+  private void scheduleNextExecution() {
+    scheduledTask.set(
+        executorService.schedule(
+            this::periodicRequest, getNextDelay().toNanos(), TimeUnit.NANOSECONDS));
   }
 
   private void doSendRequest() {
@@ -150,7 +175,7 @@ private void doSendRequest() {
     try (HttpSender.Response response = future.get(30, TimeUnit.SECONDS)) {
       getCallback().onConnectionSuccess();
       if (isSuccessful(response)) {
-        handleSuccessResponse(
+        handleHttpSuccess(
             Response.create(ServerToAgent.ADAPTER.decode(response.bodyInputStream())));
       } else {
         handleHttpError(response);
@@ -187,7 +212,7 @@ private static boolean isSuccessful(HttpSender.Response response) {
     return response.statusCode() >= 200 && response.statusCode() < 300;
   }
 
-  private void handleSuccessResponse(Response response) {
+  private void handleHttpSuccess(Response response) {
     useRegularDelay();
     ServerToAgent serverToAgent = response.getServerToAgent();
 
@@ -211,28 +236,19 @@ private void handleErrorResponse(ServerErrorResponse errorResponse) {
 
   private void useRegularDelay() {
     if (currentDelay.compareAndSet(periodicRetryDelay, periodicRequestDelay)) {
-      cancelCurrentTask();
       periodicRequestDelay.reset();
     }
   }
 
   private void useRetryDelay(@Nullable Duration retryAfter) {
     if (currentDelay.compareAndSet(periodicRequestDelay, periodicRetryDelay)) {
-      cancelCurrentTask();
       periodicRetryDelay.reset();
       if (retryAfter != null && periodicRetryDelay instanceof AcceptsDelaySuggestion) {
         ((AcceptsDelaySuggestion) periodicRetryDelay).suggestDelay(retryAfter);
       }
     }
   }
 
-  private void cancelCurrentTask() {
-    ScheduledFuture<?> future = currentTask.get();
-    if (future != null) {
-      future.cancel(false);
-    }
-  }
-
   private Callback getCallback() {
     return Objects.requireNonNull(callback);
   }

opamp-client/src/test/java/io/opentelemetry/opamp/client/internal/request/service/HttpRequestServiceTest.java

@@ -126,6 +126,19 @@ void verifySendingRequest_happyPath() {
     verify(callback).onConnectionSuccess();
   }
 
+  @Test
+  void verifyWhenSendingOnDemandRequest_andDelayChanges() {
+    // Initial state
+    assertThat(assertAndGetSingleCurrentTask().delay).isEqualTo(REGULAR_DELAY);
+
+    // Trigger delay strategy change
+    requestSender.enqueueResponse(createFailedResponse(503));
+    httpRequestService.sendRequest();
+
+    // Expected state
+    assertThat(assertAndGetSingleCurrentTask().delay).isEqualTo(RETRY_DELAY);
+  }
+
   @Test
   void verifySendingRequest_whenTheresAParsingError() {
     HttpSender.Response httpResponse = createSuccessfulResponse(new byte[] {1, 2, 3});
@@ -240,14 +253,14 @@ void verifySendingRequest_duringRegularMode() {
   private void verifyRetryDelayOnError(
       HttpSender.Response errorResponse, Duration expectedRetryDelay) {
     requestSender.enqueueResponse(errorResponse);
-    ScheduledTask previousTask = getCurrentScheduledTask();
+    ScheduledTask previousTask = assertAndGetSingleCurrentTask();
 
     previousTask.run();
 
     verifySingleRequestSent();
     verify(periodicRetryDelay).reset();
     verify(callback).onRequestFailed(any());
-    ScheduledTask retryTask = getCurrentScheduledTask();
+    ScheduledTask retryTask = assertAndGetSingleCurrentTask();
     assertThat(retryTask.delay).isEqualTo(expectedRetryDelay);
 
     // Retry with another error
@@ -258,7 +271,7 @@ private void verifyRetryDelayOnError(
 
     verifySingleRequestSent();
     verify(callback).onRequestFailed(any());
-    ScheduledTask retryTask2 = getCurrentScheduledTask();
+    ScheduledTask retryTask2 = assertAndGetSingleCurrentTask();
     assertThat(retryTask2.delay).isEqualTo(expectedRetryDelay);
 
     // Retry with a success
@@ -271,7 +284,7 @@ private void verifyRetryDelayOnError(
     verify(periodicRequestDelay).reset();
     verifySingleRequestSent();
     verifyRequestSuccessCallback(serverToAgent);
-    assertThat(getCurrentScheduledTask().delay).isEqualTo(REGULAR_DELAY);
+    assertThat(assertAndGetSingleCurrentTask().delay).isEqualTo(REGULAR_DELAY);
   }
 
   private Request createRequestSupplier() {
@@ -280,7 +293,7 @@ private Request createRequestSupplier() {
     return Request.create(agentToServer);
   }
 
-  private ScheduledTask getCurrentScheduledTask() {
+  private ScheduledTask assertAndGetSingleCurrentTask() {
     assertThat(scheduledTasks).hasSize(1);
     return scheduledTasks.get(0);
   }