Improve cancel handling in AbstractListenerReadPublisher

Closes gh-30393
spring-projects · Jan 17, 2024 · c4a34fa · c4a34fa
1 parent e3f185a
commit c4a34fa
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/...src/main/java/org/springframework/http/server/reactive/AbstractListenerReadPublisher.java b/...src/main/java/org/springframework/http/server/reactive/AbstractListenerReadPublisher.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -123,7 +123,7 @@ public final void onDataAvailable() {
 	}
 
 	/**
-	 * Subclasses can call this method to delegate a contain notification when
+	 * Subclasses can call this method to delegate a container notification when
 	 * all data has been read.
 	 */
 	public void onAllDataRead() {
@@ -362,6 +362,12 @@ <T> void onError(AbstractListenerReadPublisher<T> publisher, Throwable ex) {
 				publisher.errorPending = ex;
 				publisher.handlePendingCompletionOrError();
 			}
+
+			@Override
+			<T> void cancel(AbstractListenerReadPublisher<T> publisher) {
+				publisher.completionPending = true;
+				publisher.handlePendingCompletionOrError();
+			}
 		},
 
 		NO_DEMAND {
@@ -435,6 +441,13 @@ <T> void onError(AbstractListenerReadPublisher<T> publisher, Throwable ex) {
 				publisher.errorPending = ex;
 				publisher.handlePendingCompletionOrError();
 			}
+
+			@Override
+			<T> void cancel(AbstractListenerReadPublisher<T> publisher) {
+				publisher.discardData();
+				publisher.completionPending = true;
+				publisher.handlePendingCompletionOrError();
+			}
 		},
 
 		COMPLETED {

diff --git a/...eb/src/test/java/org/springframework/http/server/reactive/ListenerReadPublisherTests.java b/...eb/src/test/java/org/springframework/http/server/reactive/ListenerReadPublisherTests.java
@@ -104,6 +104,9 @@ protected void checkOnDataAvailable() {
 
 		@Override
 		protected DataBuffer read() {
+			if (this.discardCalls != 0) {
+				return null;
+			}
 			this.readCalls++;
 			return mock();
 		}

diff --git a/...ava/org/springframework/web/reactive/socket/adapter/AbstractListenerWebSocketSession.java b/...ava/org/springframework/web/reactive/socket/adapter/AbstractListenerWebSocketSession.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -290,8 +290,10 @@ else if (rsReadLogger.isTraceEnabled()) {
 
 		@Override
 		protected void discardData() {
+			Queue<Object> queue = this.pendingMessages;
+			this.pendingMessages = Queues.empty().get(); // prevent further reading
 			while (true) {
-				WebSocketMessage message = (WebSocketMessage) this.pendingMessages.poll();
+				WebSocketMessage message = (WebSocketMessage) queue.poll();
 				if (message == null) {
 					return;
 				}