Refactor ProcessScheduledJobsServiceCommand to dispose of processing jobs on shutdown

service/src/test/java/org/whispersystems/textsecuregcm/scheduler/JobSchedulerTest.java

@@ -1,14 +1,10 @@
 package org.whispersystems.textsecuregcm.scheduler;
 
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.RegisterExtension;
-import org.whispersystems.textsecuregcm.storage.DynamoDbExtension;
-import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema;
-import org.whispersystems.textsecuregcm.util.TestClock;
-import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
-import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
-import javax.annotation.Nullable;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
@@ -16,8 +12,14 @@ import java.time.ZoneId;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.atomic.AtomicInteger;
-
-import static org.junit.jupiter.api.Assertions.*;
+import javax.annotation.Nullable;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.whispersystems.textsecuregcm.storage.DynamoDbExtension;
+import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema;
+import org.whispersystems.textsecuregcm.util.TestClock;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
+import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
 
 class JobSchedulerTest {
 
@@ -80,16 +82,16 @@ class JobSchedulerTest {
     // Clock time is before scheduled job time
     testClock.pin(CURRENT_TIME.minusMillis(1));
 
-    scheduler.processAvailableJobs().join();
+    scheduler.processAvailableJobs().block();
     assertEquals(0, scheduler.jobsProcessed.get());
 
     // Clock time is after scheduled job time
     testClock.pin(CURRENT_TIME.plusMillis(1));
 
-    scheduler.processAvailableJobs().join();
+    scheduler.processAvailableJobs().block();
     assertEquals(1, scheduler.jobsProcessed.get());
 
-    scheduler.processAvailableJobs().join();
+    scheduler.processAvailableJobs().block();
     assertEquals(1, scheduler.jobsProcessed.get(),
         "Jobs should be cleared after successful processing; job counter should not increment on second run");
   }
@@ -112,10 +114,10 @@ class JobSchedulerTest {
 
     scheduler.scheduleJob(scheduler.buildRunAtAttribute(CURRENT_TIME, 0L), CURRENT_TIME, null).join();
 
-    scheduler.processAvailableJobs().join();
+    scheduler.processAvailableJobs().block();
     assertEquals(1, jobsEncountered.get());
 
-    scheduler.processAvailableJobs().join();
+    scheduler.processAvailableJobs().block();
     assertEquals(2, jobsEncountered.get(),
         "Jobs should not be cleared after failed processing; encountered job counter should increment on second run");
   }

service/src/test/java/org/whispersystems/textsecuregcm/workers/ProcessScheduledJobsServiceCommandTest.java

@@ -0,0 +1,123 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.workers;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.whispersystems.textsecuregcm.scheduler.JobScheduler;
+import reactor.core.publisher.Mono;
+import reactor.test.publisher.TestPublisher;
+
+@Timeout(value = 5, threadMode = Timeout.ThreadMode.SEPARATE_THREAD)
+class ProcessScheduledJobsServiceCommandTest {
+
+  private ScheduledExecutorService scheduledExecutorService;
+
+  @BeforeEach
+  void setUp() {
+    scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
+  }
+
+  @Test
+  void testDisposeOnStopCancels() throws Exception {
+    // This test publisher will never emit any values or intentionally complete
+    final TestPublisher<Integer> testPublisher = TestPublisher.create();
+    final TestJobScheduler testJobScheduler = new TestJobScheduler(testPublisher);
+
+    final ProcessScheduledJobsServiceCommand.ScheduledJobProcessor scheduledJobProcessor =
+        new ProcessScheduledJobsServiceCommand.ScheduledJobProcessor(testJobScheduler, scheduledExecutorService, 60);
+
+    scheduledJobProcessor.start();
+    testJobScheduler.getStartLatch().await();
+
+    scheduledJobProcessor.stop();
+    testJobScheduler.getEndLatch().await();
+
+    scheduledExecutorService.shutdown();
+    assertTrue(scheduledExecutorService.awaitTermination(1, TimeUnit.SECONDS), "The submitted task should complete");
+
+    testPublisher.assertCancelled();
+  }
+
+  @Test
+  void testCompletedPublisher() throws Exception {
+    final TestPublisher<Integer> testPublisher = TestPublisher.create();
+    testPublisher.complete();
+
+    final TestJobScheduler testJobScheduler = new TestJobScheduler(testPublisher);
+
+    final ProcessScheduledJobsServiceCommand.ScheduledJobProcessor scheduledJobProcessor =
+        new ProcessScheduledJobsServiceCommand.ScheduledJobProcessor(testJobScheduler, scheduledExecutorService, 60);
+
+    scheduledJobProcessor.start();
+    testJobScheduler.getStartLatch().await();
+
+    scheduledJobProcessor.stop();
+    testJobScheduler.getEndLatch().await();
+
+    scheduledExecutorService.shutdown();
+    assertTrue(scheduledExecutorService.awaitTermination(1, TimeUnit.SECONDS), "The submitted task should complete");
+
+    testPublisher.assertNotCancelled();
+  }
+
+  private static class TestJobScheduler extends JobScheduler {
+
+    private final TestPublisher<Integer> testPublisher;
+    private final CountDownLatch startLatch = new CountDownLatch(1);
+    private final CountDownLatch endLatch = new CountDownLatch(1);
+
+    protected TestJobScheduler(TestPublisher<Integer> testPublisher) {
+      super(null, null, null, null);
+      this.testPublisher = testPublisher;
+    }
+
+    /**
+     * A {@link CountDownLatch} indicating whether the {@link Mono} returned by {@link #processAvailableJobs()} has been
+     * subscribed to.
+     */
+    public CountDownLatch getStartLatch() {
+      return startLatch;
+    }
+
+    /**
+     * A {@link CountDownLatch} indicating whether the {@link Mono} returned by {@link #processAvailableJobs()} has
+     * terminated or been canceled.
+     */
+    public CountDownLatch getEndLatch() {
+      return endLatch;
+    }
+
+    @Override
+    public String getSchedulerName() {
+      return "test";
+    }
+
+    @Override
+    protected CompletableFuture<String> processJob(@Nullable byte[] jobData) {
+      return CompletableFuture.failedFuture(new IllegalStateException("Not implemented"));
+    }
+
+    @Override
+    public Mono<Void> processAvailableJobs() {
+      return testPublisher.flux()
+          .then()
+          .doOnSubscribe(ignored -> startLatch.countDown())
+          .doOnTerminate(endLatch::countDown)
+          .doOnCancel(endLatch::countDown);
+    }
+  }
+
+}