Add disposeGracefully method to Scheduler

reactor-core/src/jcstress/java/reactor/core/scheduler/SingleSchedulerStressTest.java

@@ -0,0 +1,138 @@
+/*
+ * Copyright (c) 2022 VMware Inc. or its affiliates, All Rights Reserved.
+ *
+ * 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
+ *
+ *   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 reactor.core.scheduler;
+
+import java.time.Duration;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import org.openjdk.jcstress.annotations.Actor;
+import org.openjdk.jcstress.annotations.Arbiter;
+import org.openjdk.jcstress.annotations.Expect;
+import org.openjdk.jcstress.annotations.JCStressTest;
+import org.openjdk.jcstress.annotations.Outcome;
+import org.openjdk.jcstress.annotations.State;
+import org.openjdk.jcstress.infra.results.ZZZ_Result;
+import org.openjdk.jcstress.infra.results.Z_Result;
+
+public abstract class SingleSchedulerStressTest {
+
+	@JCStressTest
+	@Outcome(id = {"true"}, expect = Expect.ACCEPTABLE, desc = "Task scheduled after racing restart")
+	@State
+	public static class StartDisposeStressTest {
+
+		final SingleScheduler scheduler = new SingleScheduler(Thread::new);
+		{
+			scheduler.start();
+		}
+
+		@Actor
+		public void restart1() {
+			scheduler.disposeGracefully(Duration.ofMillis(100)).block(Duration.ofMillis(100));
+			scheduler.start();
+		}
+
+		@Actor
+		public void restart2() {
+			scheduler.disposeGracefully(Duration.ofMillis(100)).block(Duration.ofMillis(100));
+			scheduler.start();
+		}
+
+		@Arbiter
+		public void arbiter(Z_Result r) {
+			// At this stage, at least one actor called scheduler.start(),
+			// so we should be able to execute a task.
+			final CountDownLatch latch = new CountDownLatch(1);
+			if (scheduler.isDisposed()) {
+				return;
+			}
+			scheduler.schedule(() -> {
+				r.r1 = true;
+				latch.countDown();
+			});
+			try {
+				r.r1 = latch.await(100, TimeUnit.MILLISECONDS);
+			} catch (InterruptedException ignored) {
+			}
+			scheduler.dispose();
+		}
+	}
+
+	@JCStressTest
+	@Outcome(id = {"true, true, true"}, expect = Expect.ACCEPTABLE,
+			desc = "Both time out, task gets rejected, scheduler disposed eventually")
+	@State
+	public static class DisposeGracefullyStressTest {
+		final SingleScheduler scheduler = new SingleScheduler(Thread::new);
+		final CountDownLatch arbiterLatch = new CountDownLatch(1);
+		{
+			scheduler.start();
+			// Schedule a task that disallows graceful closure until the arbiter kicks in
+			// to make sure that actors fail while waiting.
+			scheduler.schedule(() -> {
+				while (true) {
+					try {
+						if (arbiterLatch.await(20, TimeUnit.MILLISECONDS)) {
+							return;
+						}
+					}
+					catch (InterruptedException ignored) {
+					}
+				}
+			});
+		}
+
+		@Actor
+		public void disposeGracefully1(ZZZ_Result r) {
+			long start = System.nanoTime();
+			try {
+				scheduler.disposeGracefully(Duration.ofMillis(20)).block();
+			} catch (Exception e) {
+				long duration = System.nanoTime() - start;
+				// Validate that the wait took non-zero time.
+				r.r1 = (e.getCause() instanceof TimeoutException) && Duration.ofNanos(duration).toMillis() > 15;
+			}
+		}
+
+		@Actor
+		public void disposeGracefully2(ZZZ_Result r) {
+			long start = System.nanoTime();
+			try {
+				scheduler.disposeGracefully(Duration.ofMillis(20)).block();
+			} catch (Exception e) {
+				long duration = System.nanoTime() - start;
+				// Validate that the wait took non-zero time.
+				r.r2 = (e.getCause() instanceof TimeoutException) && Duration.ofNanos(duration).toMillis() > 15;
+			}
+		}
+
+		@Arbiter
+		public void arbiter(ZZZ_Result r) {
+			// Release the task blocking graceful closure.
+			arbiterLatch.countDown();
+			try {
+				scheduler.schedule(() -> {});
+			} catch (RejectedExecutionException e) {
+				scheduler.disposeGracefully(Duration.ofMillis(20)).block();
+				r.r3 = scheduler.isDisposed() && scheduler.state.executor.isTerminated();
+			}
+		}
+	}
+}

reactor-core/src/main/java/reactor/core/Disposable.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2015-2022 VMware Inc. or its affiliates, All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,9 +16,11 @@
 
 package reactor.core;
 
+import java.time.Duration;
 import java.util.Collection;
 import java.util.function.Supplier;
 
+import reactor.core.publisher.Mono;
 import reactor.util.annotation.Nullable;
 
 /**
@@ -165,4 +167,11 @@ default boolean addAll(Collection<? extends Disposable> ds) {
 		 */
 		int size();
 	}
+
+	// TODO(dj): add javadoc
+	interface Graceful extends Disposable {
+		default Mono<Void> disposeGracefully(Duration gracePeriod) {
+			return Mono.fromRunnable(this::dispose);
+		}
+	}
 }

reactor-core/src/main/java/reactor/core/scheduler/BoundedElasticScheduler.java

@@ -17,6 +17,7 @@
 package reactor.core.scheduler;
 
 import java.time.Clock;
+import java.time.Duration;
 import java.time.Instant;
 import java.time.ZoneId;
 import java.util.ArrayList;
@@ -49,6 +50,9 @@
 import reactor.core.Disposables;
 import reactor.core.Exceptions;
 import reactor.core.Scannable;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
 import reactor.util.Logger;
 import reactor.util.Loggers;
 import reactor.util.annotation.Nullable;
@@ -102,6 +106,7 @@ public String toString() {
 	final Clock         clock;
 	final ThreadFactory factory;
 	final long          ttlMillis;
+	final Sinks.Empty<Void> disposedNotifier;
 
 	volatile BoundedServices boundedServices;
 	static final AtomicReferenceFieldUpdater<BoundedElasticScheduler, BoundedServices> BOUNDED_SERVICES =
@@ -131,6 +136,7 @@ public String toString() {
 		this.factory = threadFactory;
 		this.clock = Objects.requireNonNull(clock, "A Clock must be provided");
 		this.ttlMillis = ttlMillis;
+		this.disposedNotifier = Sinks.empty();
 
 		this.boundedServices = SHUTDOWN; //initially disposed, EVICTOR is also null
 	}
@@ -204,6 +210,25 @@ public void dispose() {
 		}
 	}
 
+	@Override
+	public Mono<Void> disposeGracefully(Duration gracePeriod) {
+		BoundedServices services = BOUNDED_SERVICES.get(this);
+		if (services != SHUTDOWN && BOUNDED_SERVICES.compareAndSet(this, services, SHUTDOWN)) {
+			ScheduledExecutorService e = EVICTOR.getAndSet(this, null);
+			Mono<Void> evictorDone;
+			if (e != null) {
+				Sinks.Empty<Void> done = Sinks.empty();
+				Schedulers.shutdownAndAwait(e, gracePeriod, done);
+				evictorDone = done.asMono();
+			} else {
+				evictorDone = Mono.empty();
+			}
+			Mono.whenDelayError(services.disposeGracefully(gracePeriod), evictorDone)
+			    .subscribe(v -> {}, disposedNotifier::tryEmitError, disposedNotifier::tryEmitEmpty);
+		}
+		return disposedNotifier.asMono();
+	}
+
 	@Override
 	public Disposable schedule(Runnable task) {
 		//tasks running once will call dispose on the BoundedState, decreasing its usage by one
@@ -321,7 +346,7 @@ public Worker createWorker() {
 	}
 
 
-	static final class BoundedServices extends AtomicInteger implements Disposable {
+	static final class BoundedServices extends AtomicInteger implements Disposable.Graceful {
 
 		/**
 		 * Constant for this counter of live executors to reflect the whole pool has been
@@ -342,6 +367,7 @@ static final class BoundedServices extends AtomicInteger implements Disposable {
 		//duplicated Clock field from parent so that SHUTDOWN can be instantiated and partially used
 		final Clock                               clock;
 		final Deque<BoundedState>                 idleQueue;
+		final Sinks.Empty<Void>                   disposedNotifier = Sinks.empty();
 
 		volatile BoundedState[]                                                   busyArray;
 		static final AtomicReferenceFieldUpdater<BoundedServices, BoundedState[]> BUSY_ARRAY =
@@ -526,6 +552,20 @@ public void dispose() {
 			for (int i = 0; i < arr.length; i++) {
 				arr[i].shutdown();
 			}
+			disposedNotifier.tryEmitEmpty();
+		}
+
+		@Override
+		public Mono<Void> disposeGracefully(Duration gracePeriod) {
+			int disposed = getAndSet(DISPOSED);
+			if (disposed >= 0) {
+				BoundedState[] arr = BUSY_ARRAY.getAndSet(this, ALL_SHUTDOWN);
+				Mono.whenDelayError(
+						Flux.fromIterable(idleQueue).flatMap(bs -> bs.shutdownGracefully(gracePeriod)),
+						Flux.fromArray(arr).flatMap(bs -> bs.shutdownGracefully(gracePeriod))
+				).subscribe(v -> {}, disposedNotifier::tryEmitError, disposedNotifier::tryEmitEmpty);
+			}
+			return disposedNotifier.asMono();
 		}
 	}
 
@@ -543,6 +583,7 @@ static class BoundedState implements Disposable, Scannable {
 
 		final BoundedServices          parent;
 		final ScheduledExecutorService executor;
+		final Sinks.Empty<Void>        disposedNotifier;
 
 		long idleSinceTimestamp = -1L;
 
@@ -552,6 +593,7 @@ static class BoundedState implements Disposable, Scannable {
 		BoundedState(BoundedServices parent, ScheduledExecutorService executor) {
 			this.parent = parent;
 			this.executor = executor;
+			this.disposedNotifier = Sinks.empty();
 		}
 
 		/**
@@ -599,6 +641,7 @@ boolean tryEvict(long evictionTimestamp, long ttlMillis) {
 			if (elapsed >= ttlMillis) {
 				if (MARK_COUNT.compareAndSet(this, 0, EVICTED)) {
 					executor.shutdownNow();
+					disposedNotifier.tryEmitEmpty();
 					return true;
 				}
 			}
@@ -612,6 +655,7 @@ boolean tryEvict(long evictionTimestamp, long ttlMillis) {
 		 * to this method (for APIs that take a {@link Disposable}).
 		 *
 		 * @see #shutdown()
+		 * @see #shutdownGracefully(Duration)
 		 * @see #dispose()
 		 */
 		void release() {
@@ -642,6 +686,19 @@ void shutdown() {
 			this.idleSinceTimestamp = -1L;
 			MARK_COUNT.set(this, EVICTED);
 			this.executor.shutdownNow();
+			this.disposedNotifier.tryEmitEmpty();
+		}
+
+		Mono<Void> shutdownGracefully(Duration gracePeriod) {
+			this.idleSinceTimestamp = -1L;
+			int inUse = MARK_COUNT.getAndSet(this, EVICTED);
+			if (inUse >= 0) {
+				// TODO(dj): consecutive calls to dispose could simply get an empty mono
+				// but they'd lose the ability to get notified once the shutdown is done
+				// so we need an field for the disposed signal
+				Schedulers.shutdownAndAwait(this.executor, gracePeriod, this.disposedNotifier);
+			}
+			return this.disposedNotifier.asMono();
 		}
 
 		/**