From 389238f6229eccfd9b8a56e580762c521990bc44 Mon Sep 17 00:00:00 2001
From: Juergen Hoeller <hoellerj@vmware.com>
Date: Mon, 14 Aug 2023 15:14:22 +0200
Subject: [PATCH] Add registerReactiveTypeOverride method to
 ReactiveAdapterRegistry

Closes gh-31047
---
 .../core/ReactiveAdapterRegistry.java         | 41 ++++++++++++++++---
 .../core/ReactiveAdapterRegistryTests.java    | 31 ++++++++++++++
 2 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/spring-core/src/main/java/org/springframework/core/ReactiveAdapterRegistry.java b/spring-core/src/main/java/org/springframework/core/ReactiveAdapterRegistry.java
index c4614bb39055..e1aa6db3aaca 100644
--- a/spring-core/src/main/java/org/springframework/core/ReactiveAdapterRegistry.java
+++ b/spring-core/src/main/java/org/springframework/core/ReactiveAdapterRegistry.java
@@ -105,16 +105,45 @@ public ReactiveAdapterRegistry() {
 	 * Register a reactive type along with functions to adapt to and from a
 	 * Reactive Streams {@link Publisher}. The function arguments assume that
 	 * their input is neither {@code null} nor {@link Optional}.
+	 * <p>This variant registers the new adapter after existing adapters.
+	 * It will be matched for the exact reactive type if no earlier adapter was
+	 * registered for the specific type, and it will be matched for assignability
+	 * in a second pass if no earlier adapter had an assignable type before.
+	 * @see #registerReactiveTypeOverride
+	 * @see #getAdapter
 	 */
 	public void registerReactiveType(ReactiveTypeDescriptor descriptor,
 			Function<Object, Publisher<?>> toAdapter, Function<Publisher<?>, Object> fromAdapter) {
 
-		if (reactorPresent) {
-			this.adapters.add(new ReactorAdapter(descriptor, toAdapter, fromAdapter));
-		}
-		else {
-			this.adapters.add(new ReactiveAdapter(descriptor, toAdapter, fromAdapter));
-		}
+		this.adapters.add(buildAdapter(descriptor, toAdapter, fromAdapter));
+	}
+
+	/**
+	 * Register a reactive type along with functions to adapt to and from a
+	 * Reactive Streams {@link Publisher}. The function arguments assume that
+	 * their input is neither {@code null} nor {@link Optional}.
+	 * <p>This variant registers the new adapter first, effectively overriding
+	 * any previously registered adapters for the same reactive type. This allows
+	 * for overriding existing adapters, in particular default adapters.
+	 * <p>Note that existing adapters for specific types will still match before
+	 * an assignability match with the new adapter. In order to override all
+	 * existing matches, a new reactive type adapter needs to be registered
+	 * for every specific type, not relying on subtype assignability matches.
+	 * @since 5.3.30
+	 * @see #registerReactiveType
+	 * @see #getAdapter
+	 */
+	public void registerReactiveTypeOverride(ReactiveTypeDescriptor descriptor,
+			Function<Object, Publisher<?>> toAdapter, Function<Publisher<?>, Object> fromAdapter) {
+
+		this.adapters.add(0, buildAdapter(descriptor, toAdapter, fromAdapter));
+	}
+
+	private ReactiveAdapter buildAdapter(ReactiveTypeDescriptor descriptor,
+			Function<Object, Publisher<?>> toAdapter, Function<Publisher<?>, Object> fromAdapter) {
+
+		return (reactorPresent ? new ReactorAdapter(descriptor, toAdapter, fromAdapter) :
+				new ReactiveAdapter(descriptor, toAdapter, fromAdapter));
 	}
 
 	/**
diff --git a/spring-core/src/test/java/org/springframework/core/ReactiveAdapterRegistryTests.java b/spring-core/src/test/java/org/springframework/core/ReactiveAdapterRegistryTests.java
index a969a482f631..87b750dfbc92 100644
--- a/spring-core/src/test/java/org/springframework/core/ReactiveAdapterRegistryTests.java
+++ b/spring-core/src/test/java/org/springframework/core/ReactiveAdapterRegistryTests.java
@@ -39,6 +39,7 @@
  * Unit tests for {@link ReactiveAdapterRegistry}.
  *
  * @author Rossen Stoyanchev
+ * @author Juergen Hoeller
  */
 @SuppressWarnings("unchecked")
 class ReactiveAdapterRegistryTests {
@@ -54,14 +55,40 @@ void getAdapterForReactiveSubType() {
 		ReactiveAdapter adapter2 = getAdapter(ExtendedFlux.class);
 		assertThat(adapter2).isSameAs(adapter1);
 
+		// Register regular reactive type (after existing adapters)
 		this.registry.registerReactiveType(
 				ReactiveTypeDescriptor.multiValue(ExtendedFlux.class, ExtendedFlux::empty),
 				o -> (ExtendedFlux<?>) o,
 				ExtendedFlux::from);
 
+		// Matches for ExtendedFlux itself
 		ReactiveAdapter adapter3 = getAdapter(ExtendedFlux.class);
 		assertThat(adapter3).isNotNull();
 		assertThat(adapter3).isNotSameAs(adapter1);
+
+		// Does not match for ExtendedFlux subclass since the default Flux adapter
+		// is being assignability-checked first when no specific match was found
+		ReactiveAdapter adapter4 = getAdapter(ExtendedExtendedFlux.class);
+		assertThat(adapter4).isSameAs(adapter1);
+
+		// Register reactive type override (before existing adapters)
+		this.registry.registerReactiveTypeOverride(
+				ReactiveTypeDescriptor.multiValue(Flux.class, ExtendedFlux::empty),
+				o -> (ExtendedFlux<?>) o,
+				ExtendedFlux::from);
+
+		// Override match for Flux
+		ReactiveAdapter adapter5 = getAdapter(Flux.class);
+		assertThat(adapter5).isNotNull();
+		assertThat(adapter5).isNotSameAs(adapter1);
+
+		// Initially registered adapter specifically matches for ExtendedFlux
+		ReactiveAdapter adapter6 = getAdapter(ExtendedFlux.class);
+		assertThat(adapter6).isSameAs(adapter3);
+
+		// Override match for ExtendedFlux subclass
+		ReactiveAdapter adapter7 = getAdapter(ExtendedExtendedFlux.class);
+		assertThat(adapter7).isSameAs(adapter5);
 	}
 
 
@@ -81,6 +108,10 @@ public void subscribe(CoreSubscriber<? super T> actual) {
 	}
 
 
+	private static class ExtendedExtendedFlux<T> extends ExtendedFlux<T> {
+	}
+
+
 	@Nested
 	class Reactor {