Add wait for container stop on cancellation

Author: JiHongKim98

src/main/java/org/springframework/data/redis/stream/StreamPollTask.java

@@ -41,6 +41,7 @@
  * {@link Task} that invokes a {@link BiFunction read function} to poll on a Redis Stream.
  *
  * @author Mark Paluch
+ * @author JiHongKim98
  * @see 2.2
  */
 class StreamPollTask<K, V extends Record<K, ?>> implements Task {
@@ -54,6 +55,8 @@ class StreamPollTask<K, V extends Record<K, ?>> implements Task {
 	private final PollState pollState;
 	private final TypeDescriptor targetType;
 
+	private final CountDownLatch cancelLatch = new CountDownLatch(1);
+
 	private volatile boolean isInEventLoop = false;
 
 	StreamPollTask(StreamReadRequest<K> streamRequest, StreamListener<K, V> listener, ErrorHandler errorHandler,
@@ -83,6 +86,12 @@ private static PollState createPollState(StreamReadRequest<?> streamRequest) {
 	@Override
 	public void cancel() throws DataAccessResourceFailureException {
 		this.pollState.cancel();
+
+		try {
+			cancelLatch.await();
+		} catch (InterruptedException e) {
+			Thread.currentThread().interrupt();
+		}
 	}
 
 	@Override
@@ -112,6 +121,7 @@ public void run() {
 			doLoop();
 		} finally {
 			isInEventLoop = false;
+			cancelLatch.countDown();
 		}
 	}
 
@@ -134,6 +144,7 @@ private void doLoop() {
 			} catch (RuntimeException ex) {
 
 				if (cancelSubscriptionOnError.test(ex)) {
+					cancelLatch.countDown();
 					cancel();
 				}
 

src/test/java/org/springframework/data/redis/stream/AbstractStreamMessageListenerContainerIntegrationTests.java

@@ -384,6 +384,27 @@ void containerRestartShouldRestartSubscription() throws InterruptedException {
 		cancelAwait(subscription);
 	}
 
+	@Test // GH-2261
+	void containerShouldStopGracefully() throws InterruptedException {
+		StreamMessageListenerContainer<String, MapRecord<String, String, String>> container = StreamMessageListenerContainer
+			.create(connectionFactory, containerOptions);
+
+		BlockingQueue<MapRecord<String, String, String>> queue = new LinkedBlockingQueue<>();
+		container.start();
+		Subscription subscription = container.receive(StreamOffset.create("my-stream", ReadOffset.from("0-0")), r -> {
+			try {
+				Thread.sleep(1500);
+			} catch (InterruptedException e) {
+				// ignore
+			}
+			queue.add(r);
+		});
+		redisTemplate.opsForStream().add("my-stream", Collections.singletonMap("key", "value1"));
+		subscription.await(DEFAULT_TIMEOUT);
+		container.stop();
+		assertThat(queue.poll(500, TimeUnit.MILLISECONDS)).isNotNull();
+	}
+
 	private static void cancelAwait(Subscription subscription) {
 
 		subscription.cancel();