Add untilStarted & untilStopped to ReactiveMessagePipeline (#214)

Author: lhotari

gradle/libs.versions.toml

@@ -52,6 +52,7 @@ log4j-slf4j2-impl = { module = "org.apache.logging.log4j:log4j-slf4j2-impl", ver
 mockito-core = { module = "org.mockito:mockito-core", version.ref = "mockito" }
 pulsar-client-api = { module = "org.apache.pulsar:pulsar-client-api", version.ref = "pulsar" }
 pulsar-client-shaded = { module = "org.apache.pulsar:pulsar-client", version.ref = "pulsar" }
+pulsar-client-all = { module = "org.apache.pulsar:pulsar-client-all", version.ref = "pulsar" }
 rat-gradle = { module = "org.nosphere.apache:creadur-rat-gradle", version.ref = "rat-gradle" }
 reactor-core = { module = "io.projectreactor:reactor-core", version.ref = "reactor" }
 reactor-test = { module = "io.projectreactor:reactor-test", version.ref = "reactor" }

pulsar-client-reactive-adapter/build.gradle

@@ -35,6 +35,7 @@ dependencies {
 	testImplementation libs.bundles.log4j
 	testImplementation libs.mockito.core
 
+	intTestImplementation libs.pulsar.client.all
 	intTestImplementation project(':pulsar-client-reactive-producer-cache-caffeine')
 	intTestImplementation project(path: ':pulsar-client-reactive-producer-cache-caffeine-shaded', configuration: 'shadow')
 	intTestImplementation libs.junit.jupiter

pulsar-client-reactive-adapter/src/intTest/java/org/apache/pulsar/reactive/client/adapter/ReactiveMessagePipelineE2ETests.java

@@ -36,8 +36,11 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import org.apache.pulsar.client.admin.PulsarAdmin;
 import org.apache.pulsar.client.api.PulsarClient;
 import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.common.policies.data.SubscriptionStats;
+import org.apache.pulsar.common.policies.data.TopicStats;
 import org.apache.pulsar.reactive.client.api.MessageSpec;
 import org.apache.pulsar.reactive.client.api.MessageSpecBuilder;
 import org.apache.pulsar.reactive.client.api.ReactiveMessagePipeline;
@@ -94,6 +97,38 @@ void shouldConsumeMessages() throws Exception {
 		}
 	}
 
+	@Test
+	void shouldSupportWaitingForConsumingToStartAndStop() throws Exception {
+		try (PulsarClient pulsarClient = SingletonPulsarContainer.createPulsarClient();
+				PulsarAdmin pulsarAdmin = SingletonPulsarContainer.createPulsarAdmin()) {
+			String topicName = "test" + UUID.randomUUID();
+			ReactivePulsarClient reactivePulsarClient = AdaptedReactivePulsarClientFactory.create(pulsarClient);
+			ReactiveMessagePipeline pipeline = reactivePulsarClient.messageConsumer(Schema.STRING)
+				.subscriptionName("sub")
+				.topic(topicName)
+				.build()
+				.messagePipeline()
+				.messageHandler((message) -> Mono.empty())
+				.build()
+				.start();
+
+			// wait for consuming to start
+			pipeline.untilStarted().block(Duration.ofSeconds(5));
+			// there should be an existing subscription
+			List<String> subscriptions = pulsarAdmin.topics().getSubscriptions(topicName);
+			assertThat(subscriptions).as("subscription should be created").contains("sub");
+
+			// stop the pipeline
+			pipeline.stop();
+			// and wait for it to stop
+			pipeline.untilStopped().block(Duration.ofSeconds(5));
+			// there should be no consumers
+			TopicStats topicStats = pulsarAdmin.topics().getStats(topicName);
+			SubscriptionStats subStats = topicStats.getSubscriptions().get("sub");
+			assertThat(subStats.getConsumers()).isEmpty();
+		}
+	}
+
 	@ParameterizedTest
 	@EnumSource(MessageOrderScenario.class)
 	void shouldRetainMessageOrder(MessageOrderScenario messageOrderScenario) throws Exception {

pulsar-client-reactive-adapter/src/intTest/java/org/apache/pulsar/reactive/client/adapter/SingletonPulsarContainer.java

@@ -19,6 +19,7 @@
 
 package org.apache.pulsar.reactive.client.adapter;
 
+import org.apache.pulsar.client.admin.PulsarAdmin;
 import org.apache.pulsar.client.api.PulsarClient;
 import org.apache.pulsar.client.api.PulsarClientException;
 import org.testcontainers.containers.PulsarContainer;
@@ -44,6 +45,12 @@ static PulsarClient createPulsarClient() throws PulsarClientException {
 			.build();
 	}
 
+	static PulsarAdmin createPulsarAdmin() throws PulsarClientException {
+		return PulsarAdmin.builder()
+			.serviceHttpUrl(SingletonPulsarContainer.PULSAR_CONTAINER.getHttpServiceUrl())
+			.build();
+	}
+
 	static DockerImageName getPulsarImage() {
 		return DockerImageName.parse("apachepulsar/pulsar:4.0.4");
 	}

pulsar-client-reactive-adapter/src/main/java/org/apache/pulsar/reactive/client/internal/adapter/ReactiveConsumerAdapter.java

@@ -25,6 +25,7 @@
 import org.apache.pulsar.client.api.Consumer;
 import org.apache.pulsar.client.api.ConsumerBuilder;
 import org.apache.pulsar.client.api.PulsarClient;
+import org.apache.pulsar.reactive.client.internal.api.InternalConsumerListener;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
@@ -45,12 +46,20 @@ class ReactiveConsumerAdapter<T> {
 	}
 
 	private Mono<Consumer<T>> createConsumerMono() {
-		return AdapterImplementationFactory.adaptPulsarFuture(
-				() -> this.consumerBuilderFactory.apply(this.pulsarClientSupplier.get()).subscribeAsync());
+		return Mono.deferContextual((contextView) -> AdapterImplementationFactory
+			.adaptPulsarFuture(
+					() -> this.consumerBuilderFactory.apply(this.pulsarClientSupplier.get()).subscribeAsync())
+			.doOnSuccess((consumer) -> contextView.<InternalConsumerListener>getOrEmpty(InternalConsumerListener.class)
+				.ifPresent((listener) -> listener.onConsumerCreated(consumer))));
 	}
 
 	private Mono<Void> closeConsumer(Consumer<?> consumer) {
-		return Mono.fromFuture(consumer::closeAsync).doOnSuccess((__) -> this.LOG.info("Consumer closed {}", consumer));
+		return Mono.deferContextual((contextView) -> Mono.fromFuture(consumer::closeAsync).doFinally((signalType) -> {
+			this.LOG.info("Consumer closed {}", consumer);
+			contextView.<InternalConsumerListener>getOrEmpty(InternalConsumerListener.class)
+				.ifPresent((listener) -> listener.onConsumerClosed(consumer));
+		}));
+
 	}
 
 	<R> Mono<R> usingConsumer(Function<Consumer<T>, Mono<R>> usingConsumerAction) {

pulsar-client-reactive-api/src/main/java/org/apache/pulsar/reactive/client/api/ReactiveMessagePipeline.java

@@ -19,20 +19,26 @@
 
 package org.apache.pulsar.reactive.client.api;
 
+import reactor.core.publisher.Mono;
+
 /**
  * Reactive message pipeline interface.
  */
 public interface ReactiveMessagePipeline extends AutoCloseable {
 
 	/**
-	 * Starts the reactive pipeline.
-	 * @return the pipeline
+	 * Starts the reactive pipeline asynchronously.
+	 * @return the pipeline instance
+	 * @see #untilStarted() For returning a reactive publisher (Mono) that completes after
+	 * the pipeline has actually started.
 	 */
 	ReactiveMessagePipeline start();
 
 	/**
-	 * Stops the reactive pipeline.
-	 * @return the reactive pipeline
+	 * Stops the reactive pipeline asynchronously.
+	 * @return the pipeline instance
+	 * @see #untilStopped() For returning a reactive publisher (Mono) that completes after
+	 * the pipeline has actually stopped.
 	 */
 	ReactiveMessagePipeline stop();
 
@@ -43,11 +49,54 @@ public interface ReactiveMessagePipeline extends AutoCloseable {
 	boolean isRunning();
 
 	/**
-	 * Closes the reactive pipeline.
+	 * Closes the reactive pipeline asynchronously without waiting for shutdown
+	 * completion.
 	 * @throws Exception if an error occurs
 	 */
 	default void close() throws Exception {
 		stop();
 	}
 
+	/**
+	 * <p>
+	 * Returns a reactive publisher (Mono) that completes after the pipeline has
+	 * successfully subscribed to the input topic(s) and started consuming messages for
+	 * the first time after pipeline creation. This method is not intended to be used
+	 * after a pipeline restarts following failure. Use this method to wait for consumer
+	 * and Pulsar subscription creation. This helps avoid race conditions when sending
+	 * messages immediately after the pipeline starts.
+	 * </p>
+	 * <p>
+	 * The {@link #start()} method must be called before invoking this method.
+	 * </p>
+	 * <p>
+	 * To wait for the operation to complete synchronously, it is necessary to call
+	 * {@link Mono#block()} on the returned Mono.
+	 * </p>
+	 * @return a Mono that completes after the pipeline has created its underlying Pulsar
+	 * consumer
+	 */
+	default Mono<Void> untilStarted() {
+		return Mono.empty();
+	}
+
+	/**
+	 * <p>
+	 * Returns a reactive publisher (Mono) that completes after the pipeline has closed
+	 * the underlying Pulsar consumer and stopped consuming new messages.
+	 * </p>
+	 * <p>
+	 * The {@link #stop()} method must be called before invoking this method.
+	 * </p>
+	 * <p>
+	 * To wait for the operation to complete synchronously, it is necessary to call
+	 * {@link Mono#block()} on the returned Mono.
+	 * </p>
+	 * @return a Mono that completes when the pipeline has closed the underlying Pulsar
+	 * consumer
+	 */
+	default Mono<Void> untilStopped() {
+		return Mono.empty();
+	}
+
 }

pulsar-client-reactive-api/src/main/java/org/apache/pulsar/reactive/client/internal/api/DefaultReactiveMessagePipeline.java

@@ -21,6 +21,7 @@
 
 import java.time.Duration;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
@@ -67,6 +68,10 @@ class DefaultReactiveMessagePipeline<T> implements ReactiveMessagePipeline {
 
 	private final MessageGroupingFunction groupingFunction;
 
+	private final AtomicReference<InternalConsumerListenerImpl> consumerListener = new AtomicReference<>();
+
+	private final AtomicReference<CompletableFuture<Void>> pipelineStoppedFuture = new AtomicReference<>();
+
 	DefaultReactiveMessagePipeline(ReactiveMessageConsumer<T> messageConsumer,
 			Function<Message<T>, Publisher<Void>> messageHandler, BiConsumer<Message<T>, Throwable> errorLogger,
 			Retry pipelineRetrySpec, Duration handlingTimeout, Function<Mono<Void>, Publisher<Void>> transformer,
@@ -83,7 +88,14 @@ class DefaultReactiveMessagePipeline<T> implements ReactiveMessagePipeline {
 		this.pipeline = messageConsumer.consumeMany(this::createMessageConsumer)
 			.then()
 			.transform(transformer)
-			.transform(this::decoratePipeline);
+			.transform(this::decoratePipeline)
+			.doFinally((signalType) -> {
+				CompletableFuture<Void> f = this.pipelineStoppedFuture.get();
+				if (f != null) {
+					f.complete(null);
+				}
+			})
+			.doFirst(() -> this.pipelineStoppedFuture.set(new CompletableFuture<>()));
 	}
 
 	private Mono<Void> decorateMessageHandler(Mono<Void> messageHandler) {
@@ -168,14 +180,26 @@ public ReactiveMessagePipeline start() {
 		if (this.killSwitch.get() != null) {
 			throw new IllegalStateException("Message handler is already running.");
 		}
-		Disposable disposable = this.pipeline.subscribe(null, this::logError, this::logUnexpectedCompletion);
+		InternalConsumerListenerImpl consumerListener = new InternalConsumerListenerImpl();
+		Disposable disposable = this.pipeline.contextWrite(Context.of(InternalConsumerListener.class, consumerListener))
+			.subscribe(null, this::logError, this::logUnexpectedCompletion);
 		if (!this.killSwitch.compareAndSet(null, disposable)) {
 			disposable.dispose();
 			throw new IllegalStateException("Message handler was already running.");
 		}
+		this.consumerListener.set(consumerListener);
 		return this;
 	}
 
+	@Override
+	public Mono<Void> untilStarted() {
+		if (!isRunning()) {
+			throw new IllegalStateException("Pipeline isn't running. Call start first.");
+		}
+		InternalConsumerListenerImpl internalConsumerListener = this.consumerListener.get();
+		return internalConsumerListener.waitForConsumerCreated();
+	}
+
 	private void logError(Throwable throwable) {
 		LOG.error("ReactiveMessageHandler was unexpectedly terminated.", throwable);
 	}
@@ -195,9 +219,44 @@ public ReactiveMessagePipeline stop() {
 		return this;
 	}
 
+	@Override
+	public Mono<Void> untilStopped() {
+		if (isRunning()) {
+			throw new IllegalStateException("Pipeline is running. Call stop first.");
+		}
+		CompletableFuture<Void> f = this.pipelineStoppedFuture.get();
+		if (f != null) {
+			return Mono.fromFuture(f, true);
+		}
+		else {
+			return Mono.empty();
+		}
+	}
+
 	@Override
 	public boolean isRunning() {
 		return this.killSwitch.get() != null;
 	}
 
+	private static final class InternalConsumerListenerImpl implements InternalConsumerListener {
+
+		private final CompletableFuture<Void> createdFuture;
+
+		private InternalConsumerListenerImpl() {
+			this.createdFuture = new CompletableFuture<>();
+		}
+
+		@Override
+		public void onConsumerCreated(Object nativeConsumer) {
+			if (!this.createdFuture.isDone()) {
+				this.createdFuture.complete(null);
+			}
+		}
+
+		Mono<Void> waitForConsumerCreated() {
+			return Mono.fromFuture(this.createdFuture, true);
+		}
+
+	}
+
 }

pulsar-client-reactive-api/src/main/java/org/apache/pulsar/reactive/client/internal/api/InternalConsumerListener.java

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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
+ *
+ *   https://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 org.apache.pulsar.reactive.client.internal.api;
+
+/**
+ * Internal interface to signal the creation and closing of a native consumer. This is not
+ * to be intended to be used by applications.
+ */
+public interface InternalConsumerListener {
+
+	/**
+	 * Called when a new native consumer is created. This is called each time a new
+	 * consumer is created initially or as a result of a reactive pipeline retry.
+	 * @param nativeConsumer the native consumer instance
+	 */
+	default void onConsumerCreated(Object nativeConsumer) {
+		// no-op
+	}
+
+	/**
+	 * Called when a native consumer is closed. This is called each time a consumer is
+	 * closed as a result of a reactive pipeline retry or when the pipeline is closed.
+	 * @param nativeConsumer the native consumer instance
+	 */
+	default void onConsumerClosed(Object nativeConsumer) {
+		// no-op
+	}
+
+}