Skip to content

Commit a6bc58f

Browse files
ozangunalpozangunalp
committed
Framework-level graceful shutdown based on pausable channels
Pause channels and drains in-flight messages during graceful shutdown and consumer seek (Kafka) operations.
1 parent a2ad4fc commit a6bc58f

16 files changed

Lines changed: 429 additions & 128 deletions

File tree

‎FOLLOW-UP-PAUSABLE-DRAIN.md‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
# PausableChannel drain — implementation plan
2+
3+
## Current state
4+
5+
- `PausableChannel` has `pauseAndDrain()` returning `Uni<Void>` — pauses and waits for in-flight messages to complete
6+
- `PausableChannelDecorator.PauserChannel` implements WIP tracking via message wrapping (intercepts ack/nack)
7+
- WIP tracking happens AFTER the buffer, so drain works with any buffer strategy
8+
- `ConsumerSeekTest` uses `pauseAndDrain()` for the seek workflow
9+
- Channels require `pausable=true` to get a `PausableChannel` registered
10+
- WIP tracking is always active on pausable channels
11+
12+
## Remaining steps
13+
14+
### 1. Default `pausable=true`
15+
16+
Change the default of `pausable` from `false` to `true`. Channels that don't want it can opt out with `pausable=false`.
17+
18+
**Files:**
19+
- `smallrye-reactive-messaging-provider/.../ConfiguredChannelFactory.java` — change the condition that checks for `pausable` config to default to `true`
20+
21+
### 2. Gate WIP tracking on `graceful-shutdown=true`
22+
23+
The message wrapping in `trackMessage()` adds overhead on every message (ack/nack interception). This should only be active when `graceful-shutdown=true` (which is the default).
24+
25+
When `graceful-shutdown=false`:
26+
- `PauserChannel` is still registered (pause/resume/clearBuffer work)
27+
- `pauseAndDrain()` pauses but completes immediately (no WIP tracking)
28+
- No message wrapping overhead
29+
30+
**Files:**
31+
- `PausableChannelDecorator` — accept `gracefulShutdown` flag per channel, conditionally apply `trackMessage` transform
32+
- `ConfiguredChannelFactory` — pass `graceful-shutdown` config value to the decorator
33+
34+
### 3. Use `pauseAndDrain` in connector graceful shutdown
35+
36+
Replace the private `waitForProcessing()` polling loops in Kafka commit handlers with `pauseAndDrain()` on the `PausableChannel`. The connector shutdown path becomes:
37+
38+
```java
39+
pausableChannel.pauseAndDrain().await().atMost(Duration.ofMillis(timeout));
40+
commitAllAndAwait();
41+
```
42+
43+
**Files:**
44+
- `KafkaThrottledLatestProcessedCommit.terminate()` — remove `waitForProcessing()` loop
45+
- `KafkaCheckpointCommit.terminate()` — remove `waitForProcessing()` loop
46+
- `KafkaShareGroupCommit.terminate()` — remove `waitForProcessing()` loop
47+
- These handlers need access to the `PausableChannel` — pass via factory or lookup from `ChannelRegistry`
48+
49+
### 4. Remove Kafka-specific `drainProcessing` from `KafkaCommitHandler`
50+
51+
Once graceful shutdown uses `PausableChannel.pauseAndDrain()`, the `drainProcessing()` default method on `KafkaCommitHandler` and its override in `KafkaLatestCommit` are no longer needed.
52+
53+
**Files:**
54+
- `KafkaCommitHandler.java` — remove `drainProcessing()` default method
55+
- `KafkaLatestCommit.java` — remove `drainProcessing()` override, `received()` override, `seekPendingPartitions`, WIP counter
56+
57+
### 5. Tests
58+
59+
- Update existing `PausableChannelTest` to verify `pauseAndDrain()` behavior
60+
- Add a test without `pausable=true` config to verify auto-registration
61+
- Verify `graceful-shutdown=false` skips WIP tracking

‎api/src/main/java/io/smallrye/reactive/messaging/ChannelRegistry.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,4 +47,8 @@ Subscriber<? extends Message<?>> register(String name,
4747

4848
PausableChannel getPausable(String name);
4949

50+
default Map<String, PausableChannel> getPausableChannels() {
51+
return Map.of();
52+
}
53+
5054
}

‎api/src/main/java/io/smallrye/reactive/messaging/PausableChannel.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
package io.smallrye.reactive.messaging;
22

3+
import java.time.Duration;
4+
5+
import io.smallrye.mutiny.Uni;
6+
37
/**
48
* A channel that can be paused and resumed.
59
*/
@@ -28,4 +32,31 @@ public interface PausableChannel {
2832
*/
2933
default void clearBuffer() {
3034
}
35+
36+
/**
37+
* Pauses the channel if not already paused, then waits for all in-flight messages
38+
* (delivered to the consumer but not yet acked or nacked) to complete.
39+
* <p>
40+
* This is useful before performing operations that require no in-flight processing,
41+
* such as seeking a Kafka consumer to a new position or during graceful shutdown.
42+
*
43+
* @return a {@link Uni} that completes when all in-flight messages
44+
* have been acknowledged or negatively acknowledged
45+
*/
46+
default Uni<Void> pauseAndDrain() {
47+
return Uni.createFrom().voidItem();
48+
}
49+
50+
/**
51+
* Returns the maximum duration to wait for in-flight messages to drain
52+
* during graceful shutdown.
53+
* <p>
54+
* Configurable via the {@code graceful-shutdown.drain-timeout} channel property.
55+
* Defaults to 10 seconds.
56+
*
57+
* @return the drain timeout duration
58+
*/
59+
default Duration getDrainTimeout() {
60+
return Duration.ofSeconds(10);
61+
}
3162
}

‎smallrye-reactive-messaging-amqp/src/test/java/io/smallrye/reactive/messaging/amqp/GracefulShutdownAmqpTest.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
package io.smallrye.reactive.messaging.amqp;
2+
3+
import static org.assertj.core.api.Assertions.assertThat;
4+
import static org.awaitility.Awaitility.await;
5+
6+
import java.util.List;
7+
import java.util.UUID;
8+
import java.util.concurrent.CopyOnWriteArrayList;
9+
import java.util.concurrent.TimeUnit;
10+
import java.util.concurrent.atomic.AtomicInteger;
11+
12+
import jakarta.enterprise.context.ApplicationScoped;
13+
14+
import org.eclipse.microprofile.reactive.messaging.Incoming;
15+
import org.jboss.weld.environment.se.Weld;
16+
import org.jboss.weld.environment.se.WeldContainer;
17+
import org.junit.jupiter.api.AfterEach;
18+
import org.junit.jupiter.api.Test;
19+
20+
import io.smallrye.reactive.messaging.annotations.Blocking;
21+
import io.smallrye.reactive.messaging.test.common.config.MapBasedConfig;
22+
import io.smallrye.reactive.messaging.test.common.config.SmallRyeConfigTestUtil;
23+
24+
public class GracefulShutdownAmqpTest extends AmqpBrokerTestBase {
25+
26+
private WeldContainer container;
27+
28+
@AfterEach
29+
public void cleanup() {
30+
if (container != null) {
31+
container.shutdown();
32+
}
33+
MapBasedConfig.cleanup();
34+
SmallRyeConfigTestUtil.releaseConfig();
35+
}
36+
37+
@Test
38+
public void testGracefulShutdownDrainsInFlightMessages() {
39+
String address = UUID.randomUUID().toString();
40+
41+
new MapBasedConfig()
42+
.with("mp.messaging.incoming.data.address", address)
43+
.with("mp.messaging.incoming.data.connector", AmqpConnector.CONNECTOR_NAME)
44+
.with("mp.messaging.incoming.data.host", host)
45+
.with("mp.messaging.incoming.data.port", port)
46+
.with("amqp-username", username)
47+
.with("amqp-password", password)
48+
.with("mp.messaging.incoming.data.tracing-enabled", false)
49+
.with("mp.messaging.incoming.data.graceful-shutdown", true)
50+
.write();
51+
52+
Weld weld = new Weld();
53+
weld.addBeanClass(SlowAmqpConsumerBean.class);
54+
container = initializeContainer(weld);
55+
56+
await().until(() -> isAmqpConnectorAlive(container));
57+
await().until(() -> isAmqpConnectorReady(container));
58+
59+
SlowAmqpConsumerBean bean = container.select(SlowAmqpConsumerBean.class).get();
60+
List<Integer> received = bean.getReceived();
61+
62+
AtomicInteger counter = new AtomicInteger();
63+
usage.produceTenIntegers(address, counter::getAndIncrement);
64+
65+
// Wait for some messages to be consumed
66+
await().atMost(30, TimeUnit.SECONDS).until(() -> received.size() >= 3);
67+
68+
int countBeforeShutdown = received.size();
69+
70+
// Shutdown triggers GracefulShutdownController at Priority 40,
71+
// which drains in-flight messages before AmqpConnector.terminate at Priority 50
72+
container.shutdown();
73+
container = null;
74+
75+
// After shutdown, verify in-flight messages completed
76+
int countAfterShutdown = received.size();
77+
assertThat(countAfterShutdown).isGreaterThanOrEqualTo(countBeforeShutdown);
78+
}
79+
80+
@ApplicationScoped
81+
public static class SlowAmqpConsumerBean {
82+
83+
private final List<Integer> received = new CopyOnWriteArrayList<>();
84+
85+
@Incoming("data")
86+
@Blocking
87+
public void consume(int payload) throws InterruptedException {
88+
Thread.sleep(100);
89+
received.add(payload);
90+
}
91+
92+
public List<Integer> getReceived() {
93+
return received;
94+
}
95+
}
96+
}

‎smallrye-reactive-messaging-kafka/src/main/java/io/smallrye/reactive/messaging/kafka/commit/KafkaCheckpointCommit.java‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -228,39 +228,13 @@ public <K, V> Uni<Void> handle(IncomingKafkaRecord<K, V> record) {
228228
*/
229229
@Override
230230
public void terminate(boolean graceful) {
231-
if (graceful) {
232-
long stillUnprocessed = waitForProcessing();
233-
if (stillUnprocessed > 0) {
234-
log.messageStillUnprocessedAfterTimeout(stillUnprocessed);
235-
}
236-
}
237-
238231
removeFromState(checkpointStateMap.keySet())
239232
.chain(this::persistProcessingState)
240233
.runSubscriptionOn(this::runOnContext)
241234
.await().atMost(Duration.ofMillis(getTimeoutInMillis()));
242235
stateStore.close();
243236
}
244237

245-
private long waitForProcessing() {
246-
int attempt = autoCommitInterval / 100;
247-
for (int i = 0; i < attempt; i++) {
248-
long sum = checkpointStateMap.values().stream().map(CheckpointState::getUnprocessedRecords).mapToLong(l -> l).sum();
249-
if (sum == 0) {
250-
return sum;
251-
}
252-
log.waitingForMessageProcessing(sum);
253-
try {
254-
Thread.sleep(100);
255-
} catch (InterruptedException e) {
256-
Thread.currentThread().interrupt();
257-
break;
258-
}
259-
}
260-
261-
return checkpointStateMap.values().stream().map(CheckpointState::getUnprocessedRecords).mapToLong(l -> l).sum();
262-
}
263-
264238
/**
265239
*
266240
* @param partitions assigned partitions

‎smallrye-reactive-messaging-kafka/src/main/java/io/smallrye/reactive/messaging/kafka/commit/KafkaLatestCommit.java‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -79,21 +79,40 @@ public void partitionsSeeked(Collection<TopicPartition> partitions) {
7979
partitionsRevoked(partitions);
8080
}
8181

82+
@Override
83+
public void terminate(boolean graceful) {
84+
if (graceful) {
85+
Map<TopicPartition, OffsetAndMetadata> toCommit = runOnContextAndAwait(() -> {
86+
Map<TopicPartition, OffsetAndMetadata> map = new HashMap<>();
87+
for (Map.Entry<TopicPartition, Long> entry : offsets.entrySet()) {
88+
map.put(entry.getKey(), new OffsetAndMetadata(entry.getValue()));
89+
}
90+
offsets.clear();
91+
return map;
92+
});
93+
if (!toCommit.isEmpty()) {
94+
try {
95+
consumer.unwrap().commitSync(toCommit);
96+
} catch (Exception e) {
97+
log.failedToCommit(toCommit, e);
98+
}
99+
}
100+
}
101+
}
102+
82103
@Override
83104
public <K, V> Uni<Void> handle(IncomingKafkaRecord<K, V> record) {
84-
runOnContext(() -> {
105+
return Uni.createFrom().voidItem().invoke(() -> {
85106
Map<TopicPartition, OffsetAndMetadata> map = new HashMap<>();
86107
TopicPartition key = TopicPartitions.getTopicPartition(record);
87108
Long last = offsets.get(key);
88-
// Verify that the latest committed offset before this one.
89109
if (last == null || last < record.getOffset() + 1) {
90110
offsets.put(key, record.getOffset() + 1);
91111
map.put(key, new OffsetAndMetadata(record.getOffset() + 1, null));
92112
consumer.commitAsync(map)
93113
.subscribe().with(ignored -> {
94114
}, throwable -> log.failedToCommitAsync(key, record.getOffset() + 1, throwable));
95115
}
96-
});
97-
return Uni.createFrom().voidItem().runSubscriptionOn(record::runOnMessageContext);
116+
}).runSubscriptionOn(record::runOnMessageContext);
98117
}
99118
}

‎smallrye-reactive-messaging-kafka/src/main/java/io/smallrye/reactive/messaging/kafka/commit/KafkaShareGroupCommit.java‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import static io.smallrye.reactive.messaging.kafka.i18n.KafkaLogging.log;
44

5-
import java.time.Duration;
65
import java.util.ArrayList;
76
import java.util.List;
87
import java.util.Map;
@@ -152,32 +151,6 @@ private void stopRenewTimer() {
152151
@Override
153152
public void terminate(boolean graceful) {
154153
stopRenewTimer();
155-
if (graceful) {
156-
waitForProcessing();
157-
}
158-
}
159-
160-
private void waitForProcessing() {
161-
int attempts = renewInterval / 100;
162-
for (int i = 0; i < attempts; i++) {
163-
long pending = inProgress.values().stream().mapToLong(Map::size).sum();
164-
if (pending == 0) {
165-
return;
166-
}
167-
log.shareGroupPendingOffsets("all", pending + " records still in progress");
168-
try {
169-
Thread.sleep(100);
170-
} catch (InterruptedException e) {
171-
Thread.currentThread().interrupt();
172-
break;
173-
}
174-
// commitSyncAndAwait
175-
try {
176-
consumer.commit().await().atMost(Duration.ofMillis(getTimeoutInMillis()));
177-
} catch (Exception e) {
178-
log.shareGroupAcknowledgementCommitFailed(e);
179-
}
180-
}
181154
}
182155

183156
/**

‎smallrye-reactive-messaging-kafka/src/main/java/io/smallrye/reactive/messaging/kafka/commit/KafkaThrottledLatestProcessedCommit.java‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -488,13 +488,6 @@ public TooManyMessagesWithoutAckException(TopicPartition topic, long offset, lon
488488

489489
@Override
490490
public void terminate(boolean graceful) {
491-
if (graceful) {
492-
long stillUnprocessed = waitForProcessing();
493-
if (stillUnprocessed > 0) {
494-
log.messageStillUnprocessedAfterTimeout(stillUnprocessed);
495-
}
496-
}
497-
498491
commitAllAndAwait();
499492
runOnContextAndAwait(() -> {
500493
offsetStores.clear();
@@ -503,26 +496,6 @@ public void terminate(boolean graceful) {
503496
});
504497
}
505498

506-
private long waitForProcessing() {
507-
int attempt = autoCommitInterval / 100;
508-
for (int i = 0; i < attempt; i++) {
509-
long sum = offsetStores.values().stream().map(OffsetStore::getUnprocessedCount).mapToLong(l -> l).sum();
510-
if (sum == 0) {
511-
return sum;
512-
}
513-
log.waitingForMessageProcessing(sum);
514-
try {
515-
Thread.sleep(100);
516-
} catch (InterruptedException e) {
517-
Thread.currentThread().interrupt();
518-
break;
519-
}
520-
}
521-
522-
return offsetStores.values().stream().map(OffsetStore::getUnprocessedCount).mapToLong(l -> l).sum();
523-
524-
}
525-
526499
private void commitAllAndAwait() {
527500
Map<TopicPartition, Long> offsetsMapping = runOnContextAndAwait(
528501
this::clearLesserSequentiallyProcessedOffsetsAndReturnLargestOffsetMapping);

0 commit comments

Comments
 (0)