Track active subscriptions as Runnable in the processor

Author: simonbasle

spring-context/src/main/java/org/springframework/scheduling/annotation/ScheduledAnnotationBeanPostProcessor.java

@@ -28,6 +28,7 @@
 import java.util.Set;
 import java.util.TimeZone;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
@@ -144,6 +145,8 @@ public class ScheduledAnnotationBeanPostProcessor
 
 	private final Map<Object, Set<ScheduledTask>> scheduledTasks = new IdentityHashMap<>(16);
 
+	private final Map<Object, List<Runnable>> reactiveSubscriptions = new IdentityHashMap<>(16);
+
 
 	/**
 	 * Create a default {@code ScheduledAnnotationBeanPostProcessor}.
@@ -573,8 +576,8 @@ protected void processScheduledSync(Scheduled scheduled, Method method, Object b
 	protected void processScheduledAsync(Scheduled scheduled, Method method, Object bean) {
 		Runnable task;
 		try {
-			boolean isFixedDelaySpecialCase = scheduled.fixedDelay() > 0 || StringUtils.hasText(scheduled.fixedDelayString());
-			task = ScheduledAnnotationReactiveSupport.createSubscriptionRunnable(method, bean, isFixedDelaySpecialCase);
+			task = ScheduledAnnotationReactiveSupport.createSubscriptionRunnable(method, bean, scheduled,
+					this.reactiveSubscriptions.computeIfAbsent(bean, k -> new CopyOnWriteArrayList<>()));
 		}
 		catch (IllegalArgumentException ex) {
 			throw new IllegalStateException("Could not create recurring task for @Scheduled method '" + method.getName() + "': " + ex.getMessage());
@@ -620,7 +623,8 @@ private static boolean isP(char ch) {
 	/**
 	 * Return all currently scheduled tasks, from {@link Scheduled} methods
 	 * as well as from programmatic {@link SchedulingConfigurer} interaction.
-	 * <p>Note this doesn't include scheduled reactive methods.
+	 * <p>Note this includes upcoming scheduled subscriptions for reactive methods,
+	 * but doesn't cover any currently active subscription for such methods.
 	 * @since 5.0.2
 	 */
 	@Override
@@ -639,20 +643,27 @@ public Set<ScheduledTask> getScheduledTasks() {
 	@Override
 	public void postProcessBeforeDestruction(Object bean, String beanName) {
 		Set<ScheduledTask> tasks;
+		List<Runnable> liveSubscriptions;
 		synchronized (this.scheduledTasks) {
 			tasks = this.scheduledTasks.remove(bean);
+			liveSubscriptions = this.reactiveSubscriptions.remove(bean);
 		}
 		if (tasks != null) {
 			for (ScheduledTask task : tasks) {
 				task.cancel();
 			}
 		}
+		if (liveSubscriptions != null) {
+			for (Runnable subscription : liveSubscriptions) {
+				subscription.run(); // equivalent to cancelling the subscription
+			}
+		}
 	}
 
 	@Override
 	public boolean requiresDestruction(Object bean) {
 		synchronized (this.scheduledTasks) {
-			return this.scheduledTasks.containsKey(bean);
+			return this.scheduledTasks.containsKey(bean) || this.reactiveSubscriptions.containsKey(bean);
 		}
 	}
 
@@ -666,6 +677,12 @@ public void destroy() {
 				}
 			}
 			this.scheduledTasks.clear();
+			Collection<List<Runnable>> allLiveSubscriptions = this.reactiveSubscriptions.values();
+			for (List<Runnable> liveSubscriptions : allLiveSubscriptions) {
+				for (Runnable liveSubscription : liveSubscriptions) {
+					liveSubscription.run(); //equivalent to cancelling the subscription
+				}
+			}
 		}
 		this.registrar.destroy();
 	}

spring-context/src/main/java/org/springframework/scheduling/annotation/ScheduledAnnotationReactiveSupport.java

@@ -18,6 +18,7 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 
 import org.apache.commons.logging.Log;
@@ -32,9 +33,11 @@
 import org.springframework.core.KotlinDetector;
 import org.springframework.core.ReactiveAdapter;
 import org.springframework.core.ReactiveAdapterRegistry;
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.ReflectionUtils;
+import org.springframework.util.StringUtils;
 
 /**
  * Helper class for @{@link ScheduledAnnotationBeanPostProcessor} to support reactive cases
@@ -148,33 +151,16 @@ static Publisher<?> getPublisherFor(Method method, Object bean) {
 	 * (i.e. the task blocks until completion of the Publisher, then the delay is applied
 	 * until next iteration).
 	 */
-	static Runnable createSubscriptionRunnable(Method method, Object targetBean, boolean isFixedDelaySpecialCase) {
+	static Runnable createSubscriptionRunnable(Method method, Object targetBean, Scheduled scheduled,
+			List<Runnable> subscriptionTrackerRegistry) {
+		boolean shouldBlock = scheduled.fixedDelay() > 0 || StringUtils.hasText(scheduled.fixedDelayString());
 		final Publisher<?> publisher = getPublisherFor(method, targetBean);
-		if (isFixedDelaySpecialCase) {
+		if (shouldBlock) {
 			return () -> {
 				final CountDownLatch latch = new CountDownLatch(1);
-				publisher.subscribe(new Subscriber<Object>() {
-					@Override
-					public void onSubscribe(Subscription s) {
-						s.request(Integer.MAX_VALUE);
-					}
-
-					@Override
-					public void onNext(Object o) {
-						// NO-OP
-					}
-
-					@Override
-					public void onError(Throwable ex) {
-						LOGGER.warn("Unexpected error occurred in scheduled reactive task", ex);
-						latch.countDown();
-					}
-
-					@Override
-					public void onComplete() {
-						latch.countDown();
-					}
-				});
+				TrackingSubscriber subscriber = new TrackingSubscriber(subscriptionTrackerRegistry, latch);
+				subscriptionTrackerRegistry.add(subscriber);
+				publisher.subscribe(subscriber);
 				try {
 					latch.await();
 				}
@@ -183,27 +169,80 @@ public void onComplete() {
 				}
 			};
 		}
-		return () -> publisher.subscribe(new Subscriber<Object>() {
-			@Override
-			public void onSubscribe(Subscription s) {
-				s.request(Integer.MAX_VALUE);
-			}
+		return () -> {
+			final TrackingSubscriber subscriber = new TrackingSubscriber(subscriptionTrackerRegistry);
+			subscriptionTrackerRegistry.add(subscriber);
+			publisher.subscribe(subscriber);
+		};
+	}
+
+	/**
+	 * A {@code Subscriber} which keeps track of its {@code Subscription} and exposes the
+	 * capacity to cancel the subscription as a {@code Runnable}. Can optionally support
+	 * blocking if a {@code CountDownLatch} is passed at construction.
+	 */
+	private static final class TrackingSubscriber implements Subscriber<Object>, Runnable {
+
+		private final List<Runnable> subscriptionTrackerRegistry;
+
+		@Nullable
+		private final CountDownLatch blockingLatch;
+
+		/*
+		Implementation note: since this is created last minute when subscribing,
+		there shouldn't be a way to cancel the tracker externally from the
+		ScheduledAnnotationBeanProcessor before the #setSubscription(Subscription)
+		 method is called.
+		 */
+		@Nullable
+		private Subscription s;
+
+		TrackingSubscriber(List<Runnable> subscriptionTrackerRegistry) {
+			this(subscriptionTrackerRegistry, null);
+		}
 
-			@Override
-			public void onNext(Object o) {
-				// NO-OP
+		TrackingSubscriber(List<Runnable> subscriptionTrackerRegistry, @Nullable CountDownLatch latch) {
+			this.subscriptionTrackerRegistry = subscriptionTrackerRegistry;
+			this.blockingLatch = latch;
+		}
+
+		@Override
+		public void run() {
+			if (this.s != null) {
+				this.s.cancel();
+			}
+			if (this.blockingLatch != null) {
+				this.blockingLatch.countDown();
 			}
+		}
+
+		@Override
+		public void onSubscribe(Subscription s) {
+			this.s = s;
+			s.request(Integer.MAX_VALUE);
+		}
+
+		@Override
+		public void onNext(Object o) {
+			// NO-OP
+		}
 
-			@Override
-			public void onError(Throwable ex) {
-				LOGGER.warn("Unexpected error occurred in scheduled reactive task", ex);
+		@Override
+		public void onError(Throwable ex) {
+			this.subscriptionTrackerRegistry.remove(this);
+			LOGGER.warn("Unexpected error occurred in scheduled reactive task", ex);
+			if (this.blockingLatch != null) {
+				this.blockingLatch.countDown();
 			}
+		}
 
-			@Override
-			public void onComplete() {
-				// NO-OP
+		@Override
+		public void onComplete() {
+			this.subscriptionTrackerRegistry.remove(this);
+			if (this.blockingLatch != null) {
+				this.blockingLatch.countDown();
 			}
-		});
+		}
 	}
 
 }