diff --git a/lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandService.java b/lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandService.java index 84ba84e9..8a8ae03e 100644 --- a/lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandService.java +++ b/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. + * + *

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. + * + *

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()}. + * + *

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(); } diff --git a/lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandServiceImpl.java b/lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandServiceImpl.java index 698b1e11..c0d94007 100644 --- a/lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandServiceImpl.java +++ b/lib-cmd-queue-redis/src/main/java/io/seqera/data/command/CommandServiceImpl.java @@ -16,6 +16,7 @@ */ package io.seqera.data.command; +import java.time.Duration; import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; @@ -23,6 +24,7 @@ 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. + * + *

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

String submit(Command

command) { // Create submitted state with params object directly (serialized via @JsonTypeInfo) @@ -331,8 +391,18 @@ private boolean processCommandWithHandler( * @throws RuntimeException if the handler throws an exception */ private CommandResult executeWithTimeout(CommandHandler handler, Command

command) { - // Submit handler execution to thread pool for async execution - final Future> 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> future = executor.submit(() -> { + try { + return handler.execute(command); + } + finally { + inflight.decrementAndGet(); + } + }); try { // Block until result is available or timeout expires diff --git a/lib-cmd-queue-redis/src/test/groovy/io/seqera/data/command/CommandServiceDrainTest.groovy b/lib-cmd-queue-redis/src/test/groovy/io/seqera/data/command/CommandServiceDrainTest.groovy new file mode 100644 index 00000000..01c11fd9 --- /dev/null +++ b/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 + */ +// 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 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 { + 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 execute(Command 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 checkStatus(Command command, CommandState state) { + return CommandResult.running() + } + } + +} diff --git a/lib-data-stream-redis/src/main/java/io/seqera/data/stream/AbstractMessageStream.java b/lib-data-stream-redis/src/main/java/io/seqera/data/stream/AbstractMessageStream.java index 693f5bbf..3e32803f 100644 --- a/lib-data-stream-redis/src/main/java/io/seqera/data/stream/AbstractMessageStream.java +++ b/lib-data-stream-redis/src/main/java/io/seqera/data/stream/AbstractMessageStream.java @@ -101,6 +101,12 @@ public abstract class AbstractMessageStream implements Closeable { private static final AtomicInteger count = new AtomicInteger(); + /** + * Granularity at which an in-loop pause re-checks {@link #closing}, so a cooperative + * shutdown is not held up for a whole poll interval or backoff delay. + */ + private static final long PAUSE_SLICE_MILLIS = 50; + private final Map> listeners = new ConcurrentHashMap<>(); private final ExponentialAttempt attempt = new ExponentialAttempt(); @@ -111,7 +117,14 @@ public abstract class AbstractMessageStream implements Closeable { private final StreamMetrics metrics; - private Thread thread; + private volatile Thread thread; + + /** + * Set by {@link #awaitQuiescent(Duration)} to stop the dispatcher from claiming further + * messages. The dispatcher observes it at the head of its loop and at every pause slice, + * so it exits at a safe point rather than being interrupted mid-call. + */ + private volatile boolean closing; private final String name0; @@ -276,7 +289,10 @@ private void consumeOne(String streamId, MessageConsumer consumer, AtomicInte */ protected void processMessages() { log.trace("Message stream - starting listener thread"); - while (!Thread.currentThread().isInterrupted()) { + // `closing` is checked first so a cooperative shutdown claims no further message; the + // cycle already in progress below always runs to completion, which is what lets a + // consumer finish its work (and its database writes) before the context tears down. + while (!closing && !Thread.currentThread().isInterrupted()) { try { final var count = new AtomicInteger(); for (Map.Entry> entry : listeners.entrySet()) { @@ -289,34 +305,93 @@ protected void processMessages() { // if no message was sent, sleep for a while before retrying if (count.get() == 0) { log.trace("Message stream - await before checking for new messages"); - Thread.sleep(pollInterval().toMillis()); + pause(pollInterval().toMillis()); } } - catch (InterruptedException e) { - log.debug("Message streaming interrupt exception - cause: {}", e.getMessage()); - Thread.currentThread().interrupt(); - break; - } catch (Throwable e) { + // A forced stop (close() fallback) surfaces as an interrupt, possibly wrapped by + // the underlying client. Treat it as "exit now", not as a stream error to retry: + // logging it at ERROR with a backoff would turn every hard shutdown into noise. + if (e instanceof InterruptedException || Thread.currentThread().isInterrupted()) { + log.debug("Message stream {} interrupted - exiting listener thread", name0); + Thread.currentThread().interrupt(); + break; + } final var d0 = attempt.delay(); log.error("Unexpected error on message stream {} (await: {}) - cause: {}", name0, d0, e.getMessage(), e); - sleep(d0.toMillis()); + pause(d0.toMillis()); } } log.trace("Message stream - exiting listener thread"); } /** - * Shutdown orderly the stream + * Sleep up to {@code millis}, returning early once {@link #closing} is set or the thread is + * interrupted. Used instead of a single long sleep so neither the poll interval nor an + * exponential backoff delay can hold up a cooperative shutdown. + */ + private void pause(long millis) { + final long deadline = System.currentTimeMillis() + millis; + long remaining; + while (!closing + && !Thread.currentThread().isInterrupted() + && (remaining = deadline - System.currentTimeMillis()) > 0) { + sleep(Math.min(PAUSE_SLICE_MILLIS, remaining)); + } + } + + /** + * Stop claiming new messages and wait for the dispatcher to finish the cycle it is running. + * + *

This is the cooperative half of {@link #close()}, exposed separately so a caller can + * drain the stream while its collaborators — a database connection pool, for instance — are + * still usable, and only then release resources. + * + *

Safe to call more than once, and safe to call before any consumer was registered. + * + * @param timeout + * how long to wait for the dispatcher to exit + * @return + * {@code true} if the dispatcher stopped within the timeout, {@code false} if it is + * still running, in which case the caller decides whether to force a stop + */ + public boolean awaitQuiescent(Duration timeout) { + closing = true; + final Thread t = thread; + if (t == null) { + return true; + } + try { + t.join(Math.max(1, timeout.toMillis())); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return !t.isAlive(); + } + + /** + * Shutdown orderly the stream. + * + *

Cooperative first: {@link #awaitQuiescent(Duration)} lets the dispatcher finish the + * message it is holding and leave the loop at a safe point. Interrupting a thread parked in a + * Redis read can hand a RESP-desynced connection back to the pool (libseqera#92), so the + * interrupt below is a fallback for a dispatcher that overran {@link #closeTimeout()}, not the + * normal path. + * + *

Callers that need the drain to complete while other beans are still alive should call + * {@link #awaitQuiescent(Duration)} themselves, ahead of this method. */ @Override public void close() { if (thread == null) { return; } - // interrupt the thread + if (awaitQuiescent(closeTimeout())) { + return; + } + log.warn("Message stream {} did not stop within {} - forcing interrupt", name0, closeTimeout()); thread.interrupt(); - // wait for the termination try { thread.join(1_000); } @@ -325,6 +400,16 @@ public void close() { } } + /** + * How long {@link #close()} waits for a cooperative stop before interrupting the dispatcher. + * Subclasses may override to align with an application-level shutdown budget. + * + * @return the cooperative close timeout, {@code 10s} by default + */ + protected Duration closeTimeout() { + return Duration.ofSeconds(10); + } + public int length(String streamId) { return stream.length(streamId); } diff --git a/lib-data-stream-redis/src/test/groovy/io/seqera/data/stream/AbstractMessageStreamDrainTest.groovy b/lib-data-stream-redis/src/test/groovy/io/seqera/data/stream/AbstractMessageStreamDrainTest.groovy new file mode 100644 index 00000000..d5f655da --- /dev/null +++ b/lib-data-stream-redis/src/test/groovy/io/seqera/data/stream/AbstractMessageStreamDrainTest.groovy @@ -0,0 +1,171 @@ +/* + * 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.stream + +import java.time.Duration +import java.util.concurrent.CountDownLatch +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicInteger + +import io.micronaut.test.extensions.spock.annotation.MicronautTest +import io.seqera.data.stream.impl.LocalMessageStream +import io.seqera.random.LongRndKey +import jakarta.inject.Inject +import spock.lang.Specification +/** + * Covers the cooperative shutdown contract: a consumer already running must be allowed to + * finish, because at that point it may be mid-way through work against resources the caller + * is about to tear down. + * + * @author Paolo Di Tommaso + */ +@MicronautTest(environments = ['test']) +class AbstractMessageStreamDrainTest extends Specification { + + @Inject + LocalMessageStream target + + def 'awaitQuiescent should let an in-progress consumer finish without interrupting it'() { + given: + def id = "stream-${LongRndKey.rndHex()}" + def stream = new TestPlainStream(target) + and: 'a consumer that is slow enough to still be running when the drain starts' + def entered = new CountDownLatch(1) + def completed = new AtomicBoolean(false) + def interrupted = new AtomicBoolean(false) + stream.addConsumer(id, { msg -> + entered.countDown() + try { + Thread.sleep(500) + completed.set(true) + } + catch (InterruptedException e) { + interrupted.set(true) + Thread.currentThread().interrupt() + } + return true + }) + + when: 'a message is picked up and the drain begins while the consumer is still inside it' + stream.offer(id, 'one') + entered.await(5, TimeUnit.SECONDS) + def quiesced = stream.awaitQuiescent(Duration.ofSeconds(10)) + + then: 'the drain waits for it rather than cutting it short' + quiesced + completed.get() + !interrupted.get() + + cleanup: + stream.close() + } + + def 'awaitQuiescent should stop the dispatcher claiming further messages'() { + given: + def id = "stream-${LongRndKey.rndHex()}" + def stream = new TestPlainStream(target) + def seen = new AtomicInteger() + def entered = new CountDownLatch(1) + stream.addConsumer(id, { msg -> + seen.incrementAndGet() + entered.countDown() + Thread.sleep(300) + return true + }) + + when: 'two messages are queued but the drain starts during the first' + stream.offer(id, 'one') + entered.await(5, TimeUnit.SECONDS) + stream.offer(id, 'two') + stream.awaitQuiescent(Duration.ofSeconds(10)) + and: 'well past the poll interval, so a live dispatcher would have taken the second' + Thread.sleep(1_500) + + then: 'only the message already claimed was delivered' + seen.get() == 1 + + cleanup: + stream.close() + } + + def 'awaitQuiescent should report false when the consumer outlives the timeout'() { + given: + def id = "stream-${LongRndKey.rndHex()}" + def stream = new TestPlainStream(target) + def entered = new CountDownLatch(1) + stream.addConsumer(id, { msg -> + entered.countDown() + Thread.sleep(2_000) + return true + }) + + when: + stream.offer(id, 'one') + entered.await(5, TimeUnit.SECONDS) + def quiesced = stream.awaitQuiescent(Duration.ofMillis(200)) + + then: 'the caller is told the drain did not complete, and decides what to do next' + !quiesced + + cleanup: + stream.close() + } + + def 'close should drain cooperatively instead of interrupting the consumer'() { + given: 'this is the behaviour change - close() used to interrupt the dispatcher first' + def id = "stream-${LongRndKey.rndHex()}" + def stream = new TestPlainStream(target) + def entered = new CountDownLatch(1) + def completed = new AtomicBoolean(false) + def interrupted = new AtomicBoolean(false) + stream.addConsumer(id, { msg -> + entered.countDown() + try { + Thread.sleep(500) + completed.set(true) + } + catch (InterruptedException e) { + interrupted.set(true) + Thread.currentThread().interrupt() + } + return true + }) + + when: + stream.offer(id, 'one') + entered.await(5, TimeUnit.SECONDS) + stream.close() + + then: 'the consumer ran to completion and was never interrupted' + completed.get() + !interrupted.get() + } + + def 'awaitQuiescent should be a no-op when no consumer was ever registered'() { + given: + def stream = new TestPlainStream(target) + + expect: + stream.awaitQuiescent(Duration.ofSeconds(1)) + + cleanup: + stream.close() + } + +}