Add operator for deferred writes

2015-12-23 16:38:40 -05:00
parent 9789db23d2
commit 8c89b478d9
2 changed files with 411 additions and 0 deletions
--- a/spring-web-reactive/src/main/java/org/springframework/http/server/reactive/WriteWithOperator.java
+++ b/spring-web-reactive/src/main/java/org/springframework/http/server/reactive/WriteWithOperator.java
@@ -0,0 +1,218 @@
+/*
+ * Copyright 2002-2015 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://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.springframework.http.server.reactive;
+
+import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import reactor.core.subscriber.SubscriberBarrier;
+import reactor.core.support.Assert;
+import reactor.fn.Function;
+
+
+/**
+ * Given a write function that accepts a source {@code Publisher<T>} to write
+ * with and returns {@code Publisher<Void>} for the result, this operator helps
+ * to defer the invocation of the write function, until we know if the source
+ * publisher will begin publishing without an error. If the first emission is
+ * an error, the write function is bypassed, and the error is sent directly
+ * through the result publisher. Otherwise the write function is invoked.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class WriteWithOperator<T> implements Function<Subscriber<? super Void>, Subscriber<? super T>> {
+
+	private final java.util.function.Function<Publisher<T>, Publisher<Void>> writeFunction;
+
+
+	public WriteWithOperator(java.util.function.Function<Publisher<T>, Publisher<Void>> writeFunction) {
+		this.writeFunction = writeFunction;
+	}
+
+	@Override
+	public Subscriber<? super T> apply(Subscriber<? super Void> subscriber) {
+		return new WriteWithBarrier(subscriber);
+	}
+
+
+	private class WriteWithBarrier extends SubscriberBarrier<T, Void> implements Publisher<T> {
+
+		/**
+		 * We've at at least one emission, we've called the write function, the write
+		 * subscriber has subscribed and cached signals have been emitted to it.
+		 * We're now simply passing data through to the write subscriber.
+		 **/
+		private boolean readyToWrite = false;
+
+		/** No emission from upstream yet */
+		private boolean beforeFirstEmission = true;
+
+		/** Cached signal before readyToWrite */
+		private T item;
+
+		/** Cached 1st/2nd signal before readyToWrite */
+		private Throwable error;
+
+		/** Cached 1st/2nd signal before readyToWrite */
+		private boolean completed = false;
+
+		/** The actual writeSubscriber vs the downstream completion subscriber */
+		private Subscriber<? super T> writeSubscriber;
+
+
+		public WriteWithBarrier(Subscriber<? super Void> subscriber) {
+			super(subscriber);
+		}
+
+
+		@Override
+		protected void doOnSubscribe(Subscription subscription) {
+			super.doOnSubscribe(subscription);
+			((Subscription) super.upstream()).request(1); // bypass doRequest
+		}
+
+		@Override
+		public void doNext(T item) {
+			if (this.readyToWrite) {
+				this.writeSubscriber.onNext(item);
+				return;
+			}
+			synchronized (this) {
+				if (this.readyToWrite) {
+					this.writeSubscriber.onNext(item);
+				}
+				else if (this.beforeFirstEmission) {
+					this.item = item;
+					this.beforeFirstEmission = false;
+					writeFunction.apply(this).subscribe(downstream());
+				}
+				else {
+					subscription.cancel();
+					downstream().onError(new IllegalStateException("Unexpected item."));
+				}
+			}
+		}
+
+		@Override
+		public void doError(Throwable ex) {
+			if (this.readyToWrite) {
+				this.writeSubscriber.onError(ex);
+				return;
+			}
+			synchronized (this) {
+				if (this.readyToWrite) {
+					this.writeSubscriber.onError(ex);
+				}
+				else if (this.beforeFirstEmission) {
+					this.beforeFirstEmission = false;
+					downstream().onError(ex);
+				}
+				else {
+					this.error = ex;
+				}
+			}
+		}
+
+		@Override
+		public void doComplete() {
+			if (this.readyToWrite) {
+				this.writeSubscriber.onComplete();
+				return;
+			}
+			synchronized (this) {
+				if (this.readyToWrite) {
+					this.writeSubscriber.onComplete();
+				}
+				else if (this.beforeFirstEmission) {
+					this.completed = true;
+					this.beforeFirstEmission = false;
+					writeFunction.apply(this).subscribe(downstream());
+				}
+				else {
+					this.completed = true;
+				}
+			}
+		}
+
+		@Override
+		public void subscribe(Subscriber<? super T> subscriber) {
+			synchronized (this) {
+				Assert.isNull(this.writeSubscriber, "Only one writeSubscriber supported.");
+				this.writeSubscriber = subscriber;
+
+				if (this.error != null || this.completed) {
+					this.writeSubscriber.onSubscribe(NO_OP_SUBSCRIPTION);
+					emitCachedSignals();
+				}
+				else {
+					this.writeSubscriber.onSubscribe(this);
+				}
+			}
+		}
+
+		/**
+		 * Emit cached signals to the write subscriber.
+		 * @return true if no more signals expected
+		 */
+		private boolean emitCachedSignals() {
+			if (this.item != null) {
+				this.writeSubscriber.onNext(this.item);
+			}
+			if (this.error != null) {
+				this.writeSubscriber.onError(this.error);
+				return true;
+			}
+			if (this.completed) {
+				this.writeSubscriber.onComplete();
+				return true;
+			}
+			return false;
+		}
+
+		@Override
+		protected void doRequest(long n) {
+			if (this.readyToWrite) {
+				super.doRequest(n);
+				return;
+			}
+			synchronized (this) {
+				if (this.writeSubscriber != null) {
+					readyToWrite = true;
+					if (emitCachedSignals()) {
+						return;
+					}
+					n--;
+					if (n == 0) {
+						return;
+					}
+					super.doRequest(n);
+				}
+			}
+		}
+	}
+
+	private final static Subscription NO_OP_SUBSCRIPTION = new Subscription() {
+
+		@Override
+		public void request(long n) {
+		}
+
+		@Override
+		public void cancel() {
+		}
+	};
+
+}
--- a/spring-web-reactive/src/test/java/org/springframework/http/server/reactive/WriteWithOperatorTests.java
+++ b/spring-web-reactive/src/test/java/org/springframework/http/server/reactive/WriteWithOperatorTests.java
@@ -0,0 +1,193 @@
+/*
+ * Copyright 2002-2015 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://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.springframework.http.server.reactive;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import reactor.Publishers;
+import reactor.core.publisher.PublisherFactory;
+import reactor.core.subscriber.SubscriberBarrier;
+import reactor.rx.Streams;
+import reactor.rx.action.Signal;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author Rossen Stoyanchev
+ */
+@SuppressWarnings("ThrowableResultOfMethodCallIgnored")
+public class WriteWithOperatorTests {
+
+	private OneByOneAsyncWriter writer;
+
+	private WriteWithOperator<String> operator;
+
+
+	@Before
+	public void setUp() throws Exception {
+		this.writer = new OneByOneAsyncWriter();
+		this.operator = new WriteWithOperator<>(this.writer::writeWith);
+	}
+
+	@Test
+	public void errorBeforeFirstItem() throws Exception {
+		IllegalStateException error = new IllegalStateException("boo");
+		Publisher<Void> completion = Publishers.lift(Publishers.error(error), this.operator);
+		List<Signal<Void>> signals = Streams.wrap(completion).materialize().toList().await(5, TimeUnit.SECONDS);
+
+		assertEquals(1, signals.size());
+		assertSame("Unexpected signal: " + signals.get(0), error, signals.get(0).getThrowable());
+	}
+
+	@Test
+	public void completionBeforeFirstItem() throws Exception {
+		Publisher<Void> completion = Publishers.lift(Publishers.empty(), this.operator);
+		List<Signal<Void>> signals = Streams.wrap(completion).materialize().toList().await(5, TimeUnit.SECONDS);
+
+		assertEquals(1, signals.size());
+		assertTrue("Unexpected signal: " + signals.get(0), signals.get(0).isOnComplete());
+
+		assertEquals(0, this.writer.items.size());
+		assertTrue(this.writer.completed);
+	}
+
+	@Test
+	public void writeOneItem() throws Exception {
+		Publisher<Void> completion = Publishers.lift(Publishers.just("one"), this.operator);
+		List<Signal<Void>> signals = Streams.wrap(completion).materialize().toList().await(5, TimeUnit.SECONDS);
+
+		assertEquals(1, signals.size());
+		assertTrue("Unexpected signal: " + signals.get(0), signals.get(0).isOnComplete());
+
+		assertEquals(1, this.writer.items.size());
+		assertEquals("one", this.writer.items.get(0));
+		assertTrue(this.writer.completed);
+	}
+
+
+	@Test
+	public void writeMultipleItems() throws Exception {
+		List<String> items = Arrays.asList("one", "two", "three");
+		Publisher<Void> completion = Publishers.lift(Publishers.from(items), this.operator);
+		List<Signal<Void>> signals = Streams.wrap(completion).materialize().consumeAsList().await(5, TimeUnit.SECONDS);
+
+		assertEquals(1, signals.size());
+		assertTrue("Unexpected signal: " + signals.get(0), signals.get(0).isOnComplete());
+
+		assertEquals(3, this.writer.items.size());
+		assertEquals("one", this.writer.items.get(0));
+		assertEquals("two", this.writer.items.get(1));
+		assertEquals("three", this.writer.items.get(2));
+		assertTrue(this.writer.completed);
+	}
+
+	@Test
+	public void errorAfterMultipleItems() throws Exception {
+		IllegalStateException error = new IllegalStateException("boo");
+		Publisher<String> publisher = PublisherFactory.create(subscriber -> {
+			int i = subscriber.context().incrementAndGet();
+			subscriber.onNext(String.valueOf(i));
+			if (i == 3) {
+				subscriber.onError(error);
+			}
+		}, subscriber -> new AtomicInteger());
+		Publisher<Void> completion = Publishers.lift(publisher, this.operator);
+		List<Signal<Void>> signals = Streams.wrap(completion).materialize().toList().await(5, TimeUnit.SECONDS);
+
+		assertEquals(1, signals.size());
+		assertSame("Unexpected signal: " + signals.get(0), error, signals.get(0).getThrowable());
+
+		assertEquals(3, this.writer.items.size());
+		assertEquals("1", this.writer.items.get(0));
+		assertEquals("2", this.writer.items.get(1));
+		assertEquals("3", this.writer.items.get(2));
+		assertSame(error, this.writer.error);
+	}
+
+
+	private static class OneByOneAsyncWriter {
+
+		private List<String> items = new ArrayList<>();
+
+		private boolean completed = false;
+
+		private Throwable error;
+
+
+		public Publisher<Void> writeWith(Publisher<String> publisher) {
+			return subscriber -> {
+				Executors.newSingleThreadScheduledExecutor().schedule(
+						(Runnable) () -> publisher.subscribe(new WriteSubscriber(subscriber)),
+						50, TimeUnit.MILLISECONDS);
+			};
+		}
+
+		private class WriteSubscriber extends SubscriberBarrier<String, Void> {
+
+			public WriteSubscriber(Subscriber<? super Void> subscriber) {
+				super(subscriber);
+			}
+
+			@Override
+			protected void doOnSubscribe(Subscription subscription) {
+				subscription.request(1);
+			}
+
+			@Override
+			public void doNext(String item) {
+				items.add(item);
+				this.subscription.request(1);
+			}
+
+			@Override
+			public void doError(Throwable ex) {
+				error = ex;
+				this.subscriber.onError(ex);
+			}
+
+			@Override
+			public void doComplete() {
+				completed = true;
+				this.subscriber.onComplete();
+			}
+		}
+	}
+
+	private final static Subscription NO_OP_SUBSCRIPTION = new Subscription() {
+
+		@Override
+		public void request(long n) {
+		}
+
+		@Override
+		public void cancel() {
+		}
+	};
+
+}