InternalOperator automatic context propagation (#3625)

This change revisits InternalMonoOperator and InternalFluxOperator
implementations and wraps their async sources with a context restoring
Subscriber implementation.

This is a follow-up for #3549

reactor-core/src/main/java/reactor/core/publisher/ConnectableLift.java

@@ -69,7 +69,7 @@ public String stepName() {
 	@Override
 	public final CoreSubscriber<? super I> subscribeOrReturn(CoreSubscriber<? super O> actual) {
 		CoreSubscriber<? super I> input =
-				liftFunction.lifter.apply(source, actual);
+				liftFunction.lifter.apply(source, Operators.restoreContextOnSubscriberIfAutoCPEnabled(source, actual));
 
 		Objects.requireNonNull(input, "Lifted subscriber MUST NOT be null");
 

reactor-core/src/main/java/reactor/core/publisher/ConnectableLiftFuseable.java

@@ -71,7 +71,7 @@ public String stepName() {
 	@Override
 	public final CoreSubscriber<? super I> subscribeOrReturn(CoreSubscriber<? super O> actual) {
 		CoreSubscriber<? super I> input =
-				liftFunction.lifter.apply(source, actual);
+				liftFunction.lifter.apply(source, Operators.restoreContextOnSubscriberIfAutoCPEnabled(source, actual));
 
 		Objects.requireNonNull(input, "Lifted subscriber MUST NOT be null");
 

reactor-core/src/main/java/reactor/core/publisher/FluxBufferBoundary.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -53,7 +53,7 @@ final class FluxBufferBoundary<T, U, C extends Collection<? super T>>
 			Publisher<U> other,
 			Supplier<C> bufferSupplier) {
 		super(source);
-		this.other = Objects.requireNonNull(other, "other");
+		this.other = Operators.toFluxOrMono(Objects.requireNonNull(other, "other"));
 		this.bufferSupplier = Objects.requireNonNull(bufferSupplier, "bufferSupplier");
 	}
 

reactor-core/src/main/java/reactor/core/publisher/FluxBufferWhen.java

@@ -65,7 +65,7 @@ final class FluxBufferWhen<T, OPEN, CLOSE, BUFFER extends Collection<? super T>>
 			Supplier<BUFFER> bufferSupplier,
 			Supplier<? extends Queue<BUFFER>> queueSupplier) {
 		super(source);
-		this.start = Objects.requireNonNull(start, "start");
+		this.start = Operators.toFluxOrMono(Objects.requireNonNull(start, "start"));
 		this.end = Objects.requireNonNull(end, "end");
 		this.bufferSupplier = Objects.requireNonNull(bufferSupplier, "bufferSupplier");
 		this.queueSupplier = Objects.requireNonNull(queueSupplier, "queueSupplier");
@@ -360,7 +360,7 @@ void open(OPEN token) {
 
 			BufferWhenCloseSubscriber<T, BUFFER> bc = new BufferWhenCloseSubscriber<>(this, idx);
 			subscribers.add(bc);
-			p.subscribe(bc);
+			Operators.toFluxOrMono(p).subscribe(bc);
 		}
 
 		void openComplete(BufferWhenOpenSubscriber<OPEN> os) {

reactor-core/src/main/java/reactor/core/publisher/FluxConcatMap.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -448,7 +448,7 @@ void drain() {
 							}
 							else {
 								active = true;
-								p.subscribe(inner);
+								Operators.toFluxOrMono(p).subscribe(inner);
 							}
 						}
 					}
@@ -805,7 +805,7 @@ void drain() {
 							}
 							else {
 								active = true;
-								p.subscribe(inner);
+								Operators.toFluxOrMono(p).subscribe(inner);
 							}
 						}
 					}

reactor-core/src/main/java/reactor/core/publisher/FluxConcatMapNoPrefetch.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2020-2023 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.
@@ -203,7 +203,7 @@ public void onNext(T t) {
 					return;
 				}
 
-				p.subscribe(inner);
+				Operators.toFluxOrMono(p).subscribe(inner);
 			}
 			catch (Throwable e) {
 				Context ctx = actual.currentContext();

reactor-core/src/main/java/reactor/core/publisher/FluxDelaySubscription.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -40,7 +40,7 @@ final class FluxDelaySubscription<T, U> extends InternalFluxOperator<T, T>
 
 	FluxDelaySubscription(Flux<? extends T> source, Publisher<U> other) {
 		super(source);
-		this.other = Objects.requireNonNull(other, "other");
+		this.other = Operators.toFluxOrMono(Objects.requireNonNull(other, "other"));
 	}
 
 	@Override

reactor-core/src/main/java/reactor/core/publisher/FluxExpand.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2017-2023 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.
@@ -119,8 +119,8 @@ public void onNext(T t) {
 
 			Publisher<? extends T> p;
 			try {
-				p = Objects.requireNonNull(expander.apply(t),
-						"The expander returned a null Publisher");
+				p = Operators.toFluxOrMono(Objects.requireNonNull(expander.apply(t),
+						"The expander returned a null Publisher"));
 			}
 			catch (Throwable ex) {
 				Exceptions.throwIfFatal(ex);

reactor-core/src/main/java/reactor/core/publisher/FluxFilterWhen.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2017-2023 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.
@@ -279,7 +279,7 @@ void drain() {
 								FilterWhenInner inner = new FilterWhenInner(this, !(p instanceof Mono));
 								if (CURRENT.compareAndSet(this,null, inner)) {
 									state = STATE_RUNNING;
-									p.subscribe(inner);
+									Operators.toFluxOrMono(p).subscribe(inner);
 									break;
 								}
 							}

reactor-core/src/main/java/reactor/core/publisher/FluxGroupJoin.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -81,7 +81,7 @@ final class FluxGroupJoin<TLeft, TRight, TLeftEnd, TRightEnd, R>
 			Supplier<? extends Queue<Object>> queueSupplier,
 			Supplier<? extends Queue<TRight>> processorQueueSupplier) {
 		super(source);
-		this.other = Objects.requireNonNull(other, "other");
+		this.other = Operators.toFluxOrMono(Objects.requireNonNull(other, "other"));
 		this.leftEnd = Objects.requireNonNull(leftEnd, "leftEnd");
 		this.rightEnd = Objects.requireNonNull(rightEnd, "rightEnd");
 		this.processorQueueSupplier = Objects.requireNonNull(processorQueueSupplier, "processorQueueSupplier");
@@ -336,7 +336,7 @@ void drain() {
 								new LeftRightEndSubscriber(this, true, idx);
 						cancellations.add(end);
 
-						p.subscribe(end);
+						Operators.toFluxOrMono(p).subscribe(end);
 
 						ex = error;
 						if (ex != null) {
@@ -404,7 +404,7 @@ else if (mode == RIGHT_VALUE) {
 								new LeftRightEndSubscriber(this, false, idx);
 						cancellations.add(end);
 
-						p.subscribe(end);
+						Operators.toFluxOrMono(p).subscribe(end);
 
 						ex = error;
 						if (ex != null) {

reactor-core/src/main/java/reactor/core/publisher/FluxJoin.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -63,7 +63,7 @@ final class FluxJoin<TLeft, TRight, TLeftEnd, TRightEnd, R> extends
 			Function<? super TRight, ? extends Publisher<TRightEnd>> rightEnd,
 			BiFunction<? super TLeft, ? super TRight, ? extends R> resultSelector) {
 		super(source);
-		this.other = Objects.requireNonNull(other, "other");
+		this.other = Operators.toFluxOrMono(Objects.requireNonNull(other, "other"));
 		this.leftEnd = Objects.requireNonNull(leftEnd, "leftEnd");
 		this.rightEnd = Objects.requireNonNull(rightEnd, "rightEnd");
 		this.resultSelector = Objects.requireNonNull(resultSelector, "resultSelector");
@@ -289,7 +289,7 @@ void drain() {
 								new LeftRightEndSubscriber(this, true, idx);
 						cancellations.add(end);
 
-						p.subscribe(end);
+						Operators.toFluxOrMono(p).subscribe(end);
 
 						ex = error;
 						if (ex != null) {
@@ -366,7 +366,7 @@ else if (mode == RIGHT_VALUE) {
 								new LeftRightEndSubscriber(this, false, idx);
 						cancellations.add(end);
 
-						p.subscribe(end);
+						Operators.toFluxOrMono(p).subscribe(end);
 
 						ex = error;
 						if (ex != null) {

reactor-core/src/main/java/reactor/core/publisher/FluxLift.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -53,7 +53,7 @@ public Object scanUnsafe(Attr key) {
 	@Override
 	public CoreSubscriber<? super I> subscribeOrReturn(CoreSubscriber<? super O> actual) {
 		CoreSubscriber<? super I> input =
-				liftFunction.lifter.apply(source, actual);
+				liftFunction.lifter.apply(source, Operators.restoreContextOnSubscriberIfAutoCPEnabled(source, actual));
 
 		Objects.requireNonNull(input, "Lifted subscriber MUST NOT be null");
 

reactor-core/src/main/java/reactor/core/publisher/FluxLiftFuseable.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -55,7 +55,7 @@ public Object scanUnsafe(Attr key) {
 	@Override
 	public CoreSubscriber<? super I> subscribeOrReturn(CoreSubscriber<? super O> actual) {
 		CoreSubscriber<? super I> input =
-				liftFunction.lifter.apply(source, actual);
+				liftFunction.lifter.apply(source, Operators.restoreContextOnSubscriberIfAutoCPEnabled(source, actual));
 
 		Objects.requireNonNull(input, "Lifted subscriber MUST NOT be null");
 
@@ -65,6 +65,6 @@ public CoreSubscriber<? super I> subscribeOrReturn(CoreSubscriber<? super O> act
 			input = new FluxHide.SuppressFuseableSubscriber<>(input);
 		}
 		//otherwise QS is not required or user already made a compatible conversion
-		return input;
+		return Operators.restoreContextOnSubscriberIfAutoCPEnabled(this, input);
 	}
 }

reactor-core/src/main/java/reactor/core/publisher/FluxMergeSequential.java

@@ -205,7 +205,7 @@ public void onNext(T t) {
 			Publisher<? extends R> publisher;
 
 			try {
-				publisher = Objects.requireNonNull(mapper.apply(t), "publisher");
+				publisher = Operators.toFluxOrMono(Objects.requireNonNull(mapper.apply(t), "publisher"));
 			}
 			catch (Throwable ex) {
 				onError(Operators.onOperatorError(s, ex, t, actual.currentContext()));

reactor-core/src/main/java/reactor/core/publisher/FluxOnErrorResume.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -91,8 +91,8 @@ public void onError(Throwable t) {
 				Publisher<? extends T> p;
 
 				try {
-					p = Objects.requireNonNull(nextFactory.apply(t),
-					"The nextFactory returned a null Publisher");
+					p = Operators.toFluxOrMono(Objects.requireNonNull(nextFactory.apply(t),
+					"The nextFactory returned a null Publisher"));
 				}
 				catch (Throwable e) {
 					Throwable _e = Operators.onOperatorError(e, actual.currentContext());

reactor-core/src/main/java/reactor/core/publisher/FluxPublishMulticast.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -78,8 +78,9 @@ public CoreSubscriber<? super T> subscribeOrReturn(CoreSubscriber<? super R> act
 				queueSupplier,
 				actual.currentContext());
 
-		Publisher<? extends R> out = Objects.requireNonNull(transform.apply(multicast),
-				"The transform returned a null Publisher");
+		Publisher<? extends R> out = Operators.toFluxOrMono(Objects.requireNonNull(
+				transform.apply(multicast),
+				"The transform returned a null Publisher"));
 
 		if (out instanceof Fuseable) {
 			out.subscribe(new CancelFuseableMulticaster<>(actual, multicast));

reactor-core/src/main/java/reactor/core/publisher/FluxRepeatWhen.java

@@ -76,7 +76,7 @@ public CoreSubscriber<? super T> subscribeOrReturn(CoreSubscriber<? super T> act
 			return null;
 		}
 
-		p.subscribe(other);
+		Operators.toFluxOrMono(p).subscribe(other);
 
 		if (!main.cancelled) {
 			return main;

reactor-core/src/main/java/reactor/core/publisher/FluxRetry.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -18,6 +18,7 @@
 
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 
+import org.reactivestreams.Publisher;
 import reactor.core.CorePublisher;
 import reactor.core.CoreSubscriber;
 
@@ -114,7 +115,9 @@ void resubscribe() {
 						produced(c);
 					}
 
-					source.subscribe(this);
+					// Not wrapping source, but just the subscriber due to requirements
+					// in TailCallSubscribeTest#retry
+					source.subscribe(Operators.restoreContextOnSubscriberIfPublisherNonInternal(source, this));
 
 				} while (WIP.decrementAndGet(this) != 0);
 			}

reactor-core/src/main/java/reactor/core/publisher/FluxSample.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -52,7 +52,7 @@ final class FluxSample<T, U> extends InternalFluxOperator<T, T> {
 
 	FluxSample(Flux<? extends T> source, Publisher<U> other) {
 		super(source);
-		this.other = Objects.requireNonNull(other, "other");
+		this.other = Operators.toFluxOrMono(Objects.requireNonNull(other, "other"));
 	}
 
 	@Override

reactor-core/src/main/java/reactor/core/publisher/FluxSampleFirst.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -191,7 +191,7 @@ public void onNext(T t) {
 				SampleFirstOther<U> other = new SampleFirstOther<>(this);
 
 				if (Operators.replace(OTHER, this, other)) {
-					p.subscribe(other);
+					Operators.toFluxOrMono(p).subscribe(other);
 				}
 			}
 			else {

reactor-core/src/main/java/reactor/core/publisher/FluxSampleTimeout.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -207,7 +207,7 @@ public void onNext(T t) {
 			SampleTimeoutOther<T, U> os = new SampleTimeoutOther<>(this, t, idx);
 
 			if (Operators.replace(OTHER, this, os)) {
-				p.subscribe(os);
+				Operators.toFluxOrMono(p).subscribe(os);
 			}
 		}
 

reactor-core/src/main/java/reactor/core/publisher/FluxSkipUntilOther.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 VMware Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2016-2023 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.
@@ -42,7 +42,7 @@ final class FluxSkipUntilOther<T, U> extends InternalFluxOperator<T, T> {
 
 	FluxSkipUntilOther(Flux<? extends T> source, Publisher<U> other) {
 		super(source);
-		this.other = Objects.requireNonNull(other, "other");
+		this.other = Operators.toFluxOrMono(Objects.requireNonNull(other, "other"));
 	}
 
 	@Override