Drain in-flight work before releasing the queue

lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandService.java

@@ -16,6 +16,7 @@
  */
 package io.seqera.data.command;
 
+import java.time.Duration;
 import java.util.Optional;
 
 /**
@@ -97,6 +98,34 @@ public interface CommandService {
     /**
      * Stop consuming commands from the queue.
      * Called during shutdown to gracefully stop processing.
+     *
+     * <p>This releases the queue immediately and does not wait for handler executions already in
+     * progress. Prefer {@link #drain(Duration)} when those executions depend on resources — a
+     * database connection pool, for instance — that are about to be torn down.
      */
     void stop();
+
+    /**
+     * Stop claiming new commands and wait for the ones already being handled to finish.
+     *
+     * <p>Intended to be called while the rest of the application is still alive, so that a handler
+     * mid-execution can complete its work and record its outcome instead of failing against
+     * resources that have already been closed. On return the queue is released, as with
+     * {@link #stop()}.
+     *
+     * <p>Deliberately framework-agnostic: the caller decides what triggers it and how the timeout
+     * relates to any container-level shutdown budget.
+     *
+     * @param timeout
+     *      how long to wait for in-flight commands to finish
+     * @return
+     *      {@code true} if nothing was left running, {@code false} if the timeout was reached with
+     *      commands still in flight — the caller may then log, report, or proceed regardless
+     */
+    boolean drain(Duration timeout);
+
+    /**
+     * @return the number of command handler executions currently in progress
+     */
+    int activeCommands();
 }

lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandServiceImpl.java

@@ -16,13 +16,15 @@
  */
 package io.seqera.data.command;
 
+import java.time.Duration;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import io.micronaut.scheduling.TaskExecutors;
 import io.seqera.data.command.store.CommandStateStore;
@@ -74,6 +76,22 @@ public class CommandServiceImpl implements CommandService {
 
     private volatile boolean started = false;
 
+    /**
+     * Handler invocations submitted to {@link #executor} that have not returned yet.
+     *
+     * <p>This is deliberately not derived from {@link #executor}: that pool is shared and
+     * container-managed, so its queue says nothing about this service. It also has to be counted
+     * around the submitted task rather than around {@link Future#get}, because
+     * {@link #executeWithTimeout} abandons the future on overrun while the handler keeps running —
+     * which is exactly the work a drain must wait for.
+     */
+    private final AtomicInteger inflight = new AtomicInteger();
+
+    /**
+     * Granularity at which {@link #drain(Duration)} re-checks {@link #inflight}.
+     */
+    private static final long DRAIN_POLL_MILLIS = 50;
+
     @Override
     public void start() {
         if (started) {
@@ -95,6 +113,48 @@ public void stop() {
         log.info("Command service stopped");
     }
 
+    @Override
+    public boolean drain(Duration timeout) {
+        if (!started) {
+            return activeCommands() == 0;
+        }
+        started = false;
+        final long deadline = System.currentTimeMillis() + Math.max(0, timeout.toMillis());
+
+        // 1. Stop claiming new commands and let the dispatcher finish the message it holds. This
+        //    does not release the stream, so an in-progress handler can still acknowledge.
+        queue.awaitQuiescent(timeout);
+
+        // 2. Wait for handler executions abandoned by executeWithTimeout() that are still running.
+        //    These are the ones that matter: they are mid-flight against the database, and letting
+        //    them finish here is the whole point of draining before the context tears down.
+        while (inflight.get() > 0 && System.currentTimeMillis() < deadline) {
+            try {
+                Thread.sleep(Math.min(DRAIN_POLL_MILLIS, Math.max(1, deadline - System.currentTimeMillis())));
+            }
+            catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                break;
+            }
+        }
+
+        final int remaining = inflight.get();
+        // 3. Release the stream. Done last so steps 1-2 ran with every collaborator still usable.
+        queue.close();
+
+        if (remaining > 0) {
+            log.warn("Command service drained with {} command(s) still running after {}", remaining, timeout);
+            return false;
+        }
+        log.info("Command service drained - no command left in flight");
+        return true;
+    }
+
+    @Override
+    public int activeCommands() {
+        return inflight.get();
+    }
+
     @Override
     public <P> String submit(Command<P> command) {
         // Create submitted state with params object directly (serialized via @JsonTypeInfo)
@@ -331,8 +391,18 @@ private <P, R> boolean processCommandWithHandler(
      * @throws RuntimeException if the handler throws an exception
      */
     private <P, R> CommandResult<R> executeWithTimeout(CommandHandler<P, R> handler, Command<P> command) {
-        // Submit handler execution to thread pool for async execution
-        final Future<CommandResult<R>> future = executor.submit(() -> handler.execute(command));
+        // Submit handler execution to thread pool for async execution. The counter is decremented
+        // inside the task, not after future.get(), so an execution abandoned on timeout below is
+        // still counted for as long as it actually runs — see the inflight field javadoc.
+        inflight.incrementAndGet();
+        final Future<CommandResult<R>> future = executor.submit(() -> {
+            try {
+                return handler.execute(command);
+            }
+            finally {
+                inflight.decrementAndGet();
+            }
+        });
 
         try {
             // Block until result is available or timeout expires

lib-cmd-queue-redis/src/test/groovy/io/seqera/data/command/CommandServiceDrainTest.groovy

@@ -0,0 +1,147 @@
+/*
+ * Copyright 2026, Seqera Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.seqera.data.command
+
+import java.time.Duration
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicBoolean
+
+import com.github.f4b6a3.tsid.TsidCreator
+import io.micronaut.test.extensions.spock.annotation.MicronautTest
+import io.micronaut.test.support.TestPropertyProvider
+import io.seqera.data.command.store.CommandStateStore
+import jakarta.inject.Inject
+
+import spock.lang.Specification
+/**
+ * Covers {@code drain()}: an execution abandoned by the execute-timeout keeps running in the
+ * background, and a shutdown must wait for it. That window is where a handler is still writing
+ * to a database whose connection pool is about to be closed.
+ *
+ * @author Paolo Di Tommaso <paolo.ditommaso@gmail.com>
+ */
+// rebuildContext: drain() releases the queue for good, so each feature needs its own
+// CommandService rather than a context shared across the class
+@MicronautTest(packages = ["io.seqera.data.stream"], transactional = false, rebuildContext = true)
+class CommandServiceDrainTest extends Specification implements TestPropertyProvider {
+
+    @Inject
+    CommandService commandService
+
+    @Inject
+    CommandStateStore store
+
+    @Override
+    Map<String, String> getProperties() {
+        return [
+            'command-queue.poll-interval'  : '100ms',
+            // deliberately shorter than the handler below, so execute() is abandoned mid-flight
+            // exactly as it is in production, where batches outlive the 1s default
+            'command-queue.execute-timeout': '200ms'
+        ]
+    }
+
+    def 'drain should wait for a handler abandoned by the execute timeout'() {
+        given:
+        def handler = new SlowHandler(runFor: 1_500)
+        commandService.registerHandler(handler)
+        commandService.start()
+
+        and:
+        def command = new TestCommand(TsidCreator.getTsid().toLowerCase(), 'slow-drain', new TestParams(1, 'x'))
+
+        when: 'the command is picked up and overruns the execute timeout'
+        commandService.submit(command)
+        handler.entered.await(10, TimeUnit.SECONDS)
+
+        then: 'it is counted as in flight even though execute() already returned to the dispatcher'
+        commandService.activeCommands() == 1
+
+        when:
+        def drained = commandService.drain(Duration.ofSeconds(10))
+
+        then: 'drain blocked until the handler finished, and it was never interrupted'
+        drained
+        handler.completed.get()
+        !handler.interrupted.get()
+        commandService.activeCommands() == 0
+    }
+
+    def 'drain should report false and leave the count visible when the handler outlives the timeout'() {
+        given:
+        def handler = new SlowHandler(runFor: 3_000)
+        commandService.registerHandler(handler)
+        commandService.start()
+
+        and:
+        def command = new TestCommand(TsidCreator.getTsid().toLowerCase(), 'slow-drain', new TestParams(2, 'x'))
+
+        when:
+        commandService.submit(command)
+        handler.entered.await(10, TimeUnit.SECONDS)
+        def drained = commandService.drain(Duration.ofMillis(300))
+
+        then: 'the caller is told the drain was incomplete rather than silently proceeding'
+        !drained
+        commandService.activeCommands() == 1
+
+        cleanup: 'let the abandoned handler finish so it does not outlive the test'
+        handler.finished.await(10, TimeUnit.SECONDS)
+    }
+
+    def 'drain should be a no-op when the service was never started'() {
+        expect:
+        commandService.drain(Duration.ofSeconds(1))
+        commandService.activeCommands() == 0
+    }
+
+    static class SlowHandler implements CommandHandler<TestParams, TestResult> {
+        long runFor
+        final CountDownLatch entered = new CountDownLatch(1)
+        final CountDownLatch finished = new CountDownLatch(1)
+        final AtomicBoolean completed = new AtomicBoolean(false)
+        final AtomicBoolean interrupted = new AtomicBoolean(false)
+
+        @Override
+        String type() { 'slow-drain' }
+
+        @Override
+        CommandResult<TestResult> execute(Command<TestParams> command) {
+            entered.countDown()
+            try {
+                Thread.sleep(runFor)
+                completed.set(true)
+            }
+            catch (InterruptedException e) {
+                interrupted.set(true)
+                Thread.currentThread().interrupt()
+            }
+            finally {
+                finished.countDown()
+            }
+            return CommandResult.success(new TestResult('done', command.params().value))
+        }
+
+        @Override
+        CommandResult<TestResult> checkStatus(Command<TestParams> command, CommandState state) {
+            return CommandResult.running()
+        }
+    }
+
+}