Server adapters release buffers on error/cancel

Review and update Servlet and Undertow adapters to release any data
buffers they be holding on to at the time of error or cancellation.

Also remove onDiscard hooks from Reactor and Undertow request body.
For Reactor we expect it to be handled. For Undertow there isn't
any Reactor Core upstream for the callback to be useful.

Issue: SPR-17410

Author: rstoyanchev

spring-web/src/main/java/org/springframework/http/server/reactive/AbstractListenerReadPublisher.java

@@ -163,6 +163,14 @@ public final void onError(Throwable ex) {
 	 */
 	protected abstract void readingPaused();
 
+	/**
+	 * Invoked after an I/O read error from the underlying server or after a
+	 * cancellation signal from the downstream consumer to allow sub-classes
+	 * to discard any current cached data they might have.
+	 * @since 5.1.2
+	 */
+	protected abstract void discardData();
+
 
 	// Private methods for use in State...
 
@@ -416,7 +424,10 @@ <T> void request(AbstractListenerReadPublisher<T> publisher, long n) {
 		}
 
 		<T> void cancel(AbstractListenerReadPublisher<T> publisher) {
-			if (!publisher.changeState(this, COMPLETED)) {
+			if (publisher.changeState(this, COMPLETED)) {
+				publisher.discardData();
+			}
+			else {
 				publisher.state.get().cancel(publisher);
 			}
 		}
@@ -439,6 +450,7 @@ <T> void onAllDataRead(AbstractListenerReadPublisher<T> publisher) {
 
 		<T> void onError(AbstractListenerReadPublisher<T> publisher, Throwable t) {
 			if (publisher.changeState(this, COMPLETED)) {
+				publisher.discardData();
 				Subscriber<? super T> s = publisher.subscriber;
 				if (s != null) {
 					s.onError(t);

spring-web/src/main/java/org/springframework/http/server/reactive/AbstractListenerWriteProcessor.java

@@ -158,6 +158,9 @@ public void cancel() {
 
 	@Override
 	public final void subscribe(Subscriber<? super Void> subscriber) {
+		// Technically, cancellation from the result subscriber should be propagated
+		// to the upstream subscription. In practice, HttpHandler server adapters
+		// don't have a reason to cancel the result subscription.
 		this.resultPublisher.subscribe(subscriber);
 	}
 
@@ -176,8 +179,14 @@ public final void subscribe(Subscriber<? super Void> subscriber) {
 	 * data item for writing once that is possible.
 	 */
 	protected void dataReceived(T data) {
-		if (this.currentData != null) {
-			throw new IllegalStateException("Current data not processed yet: " + this.currentData);
+		T prev = this.currentData;
+		if (prev != null) {
+			// This shouldn't happen:
+			//   1. dataReceived can only be called from REQUESTED state
+			//   2. currentData is cleared before requesting
+			discardData(data);
+			cancel();
+			onError(new IllegalStateException("Received new data while current not processed yet."));
 		}
 		this.currentData = data;
 	}
@@ -226,6 +235,16 @@ protected void writingComplete() {
 	protected void writingFailed(Throwable ex) {
 	}
 
+	/**
+	 * Invoked after any error (either from the upstream write Publisher, or
+	 * from I/O operations to the underlying server) and cancellation
+	 * to discard in-flight data that was in
+	 * the process of being written when the error took place.
+	 * @param data the data to be released
+	 * @since 5.1.2
+	 */
+	protected abstract void discardData(T data);
+
 
 	// Private methods for use from State's...
 
@@ -245,6 +264,7 @@ private void changeStateToReceived(State oldState) {
 
 	private void changeStateToComplete(State oldState) {
 		if (changeState(oldState, State.COMPLETED)) {
+			discardCurrentData();
 			writingComplete();
 			this.resultPublisher.publishComplete();
 		}
@@ -263,6 +283,14 @@ private void writeIfPossible() {
 		}
 	}
 
+	private void discardCurrentData() {
+		T data = this.currentData;
+		this.currentData = null;
+		if (data != null) {
+			discardData(data);
+		}
+	}
+
 
 	/**
 	 * Represents a state for the {@link Processor} to be in.
@@ -378,11 +406,14 @@ public <T> void onSubscribe(AbstractListenerWriteProcessor<T> processor, Subscri
 		}
 
 		public <T> void onNext(AbstractListenerWriteProcessor<T> processor, T data) {
-			throw new IllegalStateException(toString());
+			processor.discardData(data);
+			processor.cancel();
+			processor.onError(new IllegalStateException("Illegal onNext without demand"));
 		}
 
 		public <T> void onError(AbstractListenerWriteProcessor<T> processor, Throwable ex) {
 			if (processor.changeState(this, COMPLETED)) {
+				processor.discardCurrentData();
 				processor.writingComplete();
 				processor.resultPublisher.publishError(ex);
 			}

spring-web/src/main/java/org/springframework/http/server/reactive/ReactorServerHttpRequest.java

@@ -29,9 +29,7 @@
 import reactor.netty.http.server.HttpServerRequest;
 
 import org.springframework.core.io.buffer.DataBuffer;
-import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
-import org.springframework.core.io.buffer.PooledDataBuffer;
 import org.springframework.http.HttpCookie;
 import org.springframework.http.HttpHeaders;
 import org.springframework.lang.Nullable;
@@ -165,8 +163,7 @@ protected SslInfo initSslInfo() {
 
 	@Override
 	public Flux<DataBuffer> getBody() {
-		Flux<DataBuffer> body = this.request.receive().retain().map(this.bufferFactory::wrap);
-		return body.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
+		return this.request.receive().retain().map(this.bufferFactory::wrap);
 	}
 
 	@SuppressWarnings("unchecked")

spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpRequest.java

@@ -302,6 +302,11 @@ protected void readingPaused() {
 			// no-op
 		}
 
+		@Override
+		protected void discardData() {
+			// Nothing to discard since we pass data buffers on immediately..
+		}
+
 
 		private class RequestBodyPublisherReadListener implements ReadListener {
 

spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpResponse.java

@@ -334,6 +334,7 @@ protected boolean write(DataBuffer dataBuffer) throws IOException {
 			boolean ready = ServletServerHttpResponse.this.isWritePossible();
 			int remaining = dataBuffer.readableByteCount();
 			if (ready && remaining > 0) {
+				// In case of IOException, onError handling should call discardData(DataBuffer)..
 				int written = writeToOutputStream(dataBuffer);
 				if (logger.isTraceEnabled()) {
 					logger.trace(getLogPrefix() + "Wrote " + written + " of " + remaining + " bytes");
@@ -359,6 +360,11 @@ else if (rsWriteLogger.isTraceEnabled()) {
 		protected void writingComplete() {
 			bodyProcessor = null;
 		}
+
+		@Override
+		protected void discardData(DataBuffer dataBuffer) {
+			DataBufferUtils.release(dataBuffer);
+		}
 	}
 
 }

spring-web/src/main/java/org/springframework/http/server/reactive/UndertowServerHttpRequest.java

@@ -116,8 +116,7 @@ protected SslInfo initSslInfo() {
 
 	@Override
 	public Flux<DataBuffer> getBody() {
-		return Flux.from(this.body)
-				.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
+		return Flux.from(this.body);
 	}
 
 	@SuppressWarnings("unchecked")
@@ -201,6 +200,10 @@ else if (read == -1) {
 			}
 		}
 
+		@Override
+		protected void discardData() {
+			// Nothing to discard since we pass data buffers on immediately..
+		}
 	}
 
 	private static class UndertowDataBuffer implements PooledDataBuffer {

spring-web/src/main/java/org/springframework/http/server/reactive/UndertowServerHttpResponse.java

@@ -181,6 +181,7 @@ protected boolean write(DataBuffer dataBuffer) throws IOException {
 			// Track write listener calls from here on..
 			this.writePossible = false;
 
+			// In case of IOException, onError handling should call discardData(DataBuffer)..
 			int total = buffer.remaining();
 			int written = writeByteBuffer(buffer);
 
@@ -235,6 +236,11 @@ protected void writingFailed(Throwable ex) {
 			cancel();
 			onError(ex);
 		}
+
+		@Override
+		protected void discardData(DataBuffer dataBuffer) {
+			DataBufferUtils.release(dataBuffer);
+		}
 	}
 
 

spring-web/src/test/java/org/springframework/http/server/reactive/ListenerReadPublisherTests.java

@@ -16,58 +16,95 @@
 
 package org.springframework.http.server.reactive;
 
-import java.io.IOException;
-
+import org.junit.Before;
 import org.junit.Test;
-import org.mockito.invocation.InvocationOnMock;
-import org.mockito.stubbing.Answer;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 
 import org.springframework.core.io.buffer.DataBuffer;
 
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.isA;
-import static org.mockito.Mockito.mock;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
 
 /**
- * Unit tests for {@link AbstractListenerReadPublisher}
+ * Unit tests for {@link AbstractListenerReadPublisher}.
  *
  * @author Violeta Georgieva
- * @since 5.0
+ * @author Rossen Stoyanchev
  */
 public class ListenerReadPublisherTests {
 
+	private final TestListenerReadPublisher publisher = new TestListenerReadPublisher();
+
+	private final TestSubscriber subscriber = new TestSubscriber();
+
+
+	@Before
+	public void setup() {
+		this.publisher.subscribe(this.subscriber);
+	}
+
+
 	@Test
-	@SuppressWarnings("unchecked")
-	public void testReceiveTwoRequestCallsWhenOnSubscribe() {
-		Subscriber<DataBuffer> subscriber = mock(Subscriber.class);
-		doAnswer(new SubscriptionAnswer()).when(subscriber).onSubscribe(isA(Subscription.class));
+	public void twoReads() {
+
+		this.subscriber.getSubscription().request(2);
+		this.publisher.onDataAvailable();
+
+		assertEquals(2, this.publisher.getReadCalls());
+	}
+
+	@Test // SPR-17410
+	public void discardDataOnError() {
 
-		TestListenerReadPublisher publisher = new TestListenerReadPublisher();
-		publisher.subscribe(subscriber);
-		publisher.onDataAvailable();
+		this.subscriber.getSubscription().request(2);
+		this.publisher.onDataAvailable();
+		this.publisher.onError(new IllegalStateException());
 
-		assertTrue(publisher.getReadCalls() == 2);
+		assertEquals(2, this.publisher.getReadCalls());
+		assertEquals(1, this.publisher.getDiscardCalls());
 	}
 
-	private static final class TestListenerReadPublisher extends AbstractListenerReadPublisher {
+	@Test // SPR-17410
+	public void discardDataOnCancel() {
+
+		this.subscriber.getSubscription().request(2);
+		this.subscriber.setCancelOnNext(true);
+		this.publisher.onDataAvailable();
+
+		assertEquals(1, this.publisher.getReadCalls());
+		assertEquals(1, this.publisher.getDiscardCalls());
+	}
+
+
+	private static final class TestListenerReadPublisher extends AbstractListenerReadPublisher<DataBuffer> {
 
 		private int readCalls = 0;
 
+		private int discardCalls = 0;
+
+
 		public TestListenerReadPublisher() {
 			super("");
 		}
 
+
+		public int getReadCalls() {
+			return this.readCalls;
+		}
+
+		public int getDiscardCalls() {
+			return this.discardCalls;
+		}
+
 		@Override
 		protected void checkOnDataAvailable() {
 			// no-op
 		}
 
 		@Override
-		protected DataBuffer read() throws IOException {
-			readCalls++;
+		protected DataBuffer read() {
+			this.readCalls++;
 			return mock(DataBuffer.class);
 		}
 
@@ -76,22 +113,48 @@ protected void readingPaused() {
 			// No-op
 		}
 
-		public int getReadCalls() {
-			return this.readCalls;
+		@Override
+		protected void discardData() {
+			this.discardCalls++;
 		}
-
 	}
 
-	private static final class SubscriptionAnswer implements Answer<Subscription> {
+
+	private static final class TestSubscriber implements Subscriber<DataBuffer> {
+
+		private Subscription subscription;
+
+		private boolean cancelOnNext;
+
+
+		public Subscription getSubscription() {
+			return this.subscription;
+		}
+
+		public void setCancelOnNext(boolean cancelOnNext) {
+			this.cancelOnNext = cancelOnNext;
+		}
+
+
+		@Override
+		public void onSubscribe(Subscription subscription) {
+			this.subscription = subscription;
+		}
 
 		@Override
-		public Subscription answer(InvocationOnMock invocation) throws Throwable {
-			Subscription arg = (Subscription) invocation.getArguments()[0];
-			arg.request(1);
-			arg.request(1);
-			return arg;
+		public void onNext(DataBuffer dataBuffer) {
+			if (this.cancelOnNext) {
+				this.subscription.cancel();
+			}
 		}
 
+		@Override
+		public void onError(Throwable t) {
+		}
+
+		@Override
+		public void onComplete() {
+		}
 	}
 
 }