How can I make a Flux emit an extra element if no element has been emitted for a given time? - project-reactor

Question

I am implementing a Heartbeat for a WebFlux SSE endpoint. To avoid a timeout in the client, I want to make sure that an element is emitted at least every, say, 10 seconds.
I came up with the following solution that emits a heartbeat element every 10 seconds regardless of whether a real element has been emitted or not:
originalFlux.mergeWith(Flux.interval(Duration.ofSeconds(10), Duration.ofSeconds(10)).map(ignored -> "heartbeat")
This is probably good enough for my use case but still I wonder if it is possible to emit the heartbeat only if no real element has been emitted in the last 10 seconds. I played around with the timeout operator which implements exactly the timing behavior I am looking for, but that emits an error and cancels the originalFlux instead of just emitting an extra element.
The following code using timeout passes my test but looks too complicated and as far as I understand could lose elements from the originalFlux if they are emitted between cancelling and re-subscribing to it:
ConnectableFlux<String> sharedOriginalFlux = originalFlux.publish();
CompletableFuture<Disposable> eventualSubscription = new CompletableFuture<>();
return addHeartbeat(sharedOriginalFlux)
.doOnSubscribe(ignored -> eventualSubscription.complete(sharedOriginalFlux.connect()))
.doFinally(ignored -> eventualSubscription.thenAccept(Disposable::dispose))
private Flux<String> addHeartbeat(Flux<String> sharedOriginalFlux) {
return sharedOriginalFlux.timeout(
Duration.ofSeconds(10),
Flux.mergeSequential(
Mono.just("heartbeat"),
Flux.defer(() -> addHeartbeat(sharedOriginalFlux))));
}
Is there a simple and safe way to do this?

Answer 1

It's not necessarily simpler, but another option could be to create a separate processor that can wrap the original Flux to provide a heartbeat (which shouldn't miss any elements):
public class HeartbeatProcessor<T> {
private final FluxProcessor<T, T> processor;
private final FluxSink<T> sink;
private final T heartbeatValue;
private final Duration heartbeatPeriod;
private Disposable d;
public HeartbeatProcessor(Flux<T> orig, T heartbeatValue, Duration heartbeatPeriod) {
this.heartbeatValue = heartbeatValue;
this.heartbeatPeriod = heartbeatPeriod;
this.processor = DirectProcessor.<T>create().serialize();
this.sink = processor.sink();
this.d = Mono.just(heartbeatValue).delayElement(heartbeatPeriod).subscribe(this::emit);
orig.subscribe(this::emit);
}
private void emit(T val) {
sink.next(val);
d.dispose();
this.d = Mono.just(heartbeatValue).delayElement(heartbeatPeriod).subscribe(this::emit);
}
public Flux<T> getFlux() {
return processor;
}
}
You could then call it as follows:
new HeartbeatProcessor<>(elements, "heartbeat", Duration.ofSeconds(10))
.getFlux()
.subscribe(System.out::println);

Answer 2

Flux.switchMap is a good candidate for this job: It switches to a new Publisher (and cancels the previous one) whenever the original Flux emits an item. In your case, the new Publisher is your heartbeat Flux.interval, prepended with the original item T:
public static Flux<String> addHeartbeat(Flux<String> originalFlux) {
return originalFlux
.startWith("heartbeat")
.materialize()
.switchMap(signal -> switch (signal.getType()) {
case ON_NEXT -> Flux.interval(Duration.ofSeconds(10))
.map(ignored -> "heartbeat")
.startWith(signal.get());
case ON_COMPLETE -> Mono.empty();
case ON_ERROR -> Mono.error(signal.getThrowable());
default -> Mono.error(new IllegalStateException());
});
}
Flux.switchMap is almost fit for the job, but it differs on two points from your desired solution:
It will only emit elements once the first element is received.
This means you have no heartbeat before the first item. This is solved by adding Flux.startWith("heartbeat"), which will emit "heartbeat" immediately on subscription, which then is processed by the switchMap into a heartbeat every 10 seconds.
The Publisher of the last element is never cancelled.
Since every generated Publisher is a Flux.interval that never completes, the onComplete signal will never reach the user. This is solved by transforming the onComplete signal into an emitted Signal item using Flux.materialize(), then map the onComplete Signal into an empty Publisher just to cancel the previous Publisher. This also creates onNext and onError Signals, which we have to handle each:
a. Signal.ON_NEXT can be processed as usual, retrieving the original with Signal.get()
b. Signal.ON_COMPLETE is mapped to an empty Mono that immediately completes.
c. Signal.ON_ERROR should relay the error downstream using Mono.error(Throwable).
d. The Signal enum contains more values, but they are not produced by Flux.materialize().
Here is the test to test this solution:
#Test
public void shouldAddHeartbeat() {
Flux<String> originalFlux = Flux.just(25, 15, 7, 5)
.concatMap(delay -> Mono.delay(Duration.ofSeconds(delay)).thenReturn(delay + " seconds delay"));
Flux<String> withHeartbeat = addHeartbeat(originalFlux);
StepVerifier.withVirtualTime(() -> withHeartbeat)
.expectNext("heartbeat")
.thenAwait(Duration.ofSeconds(10)).expectNext("heartbeat")
.thenAwait(Duration.ofSeconds(10)).expectNext("heartbeat")
.thenAwait(Duration.ofSeconds(5)).expectNext("25 seconds delay")
.thenAwait(Duration.ofSeconds(10)).expectNext("heartbeat")
.thenAwait(Duration.ofSeconds(5)).expectNext("15 seconds delay")
.thenAwait(Duration.ofSeconds(7)).expectNext("7 seconds delay")
.thenAwait(Duration.ofSeconds(5)).expectNext("5 seconds delay")
.verifyComplete();
}