Accept CqlSessionBuilder and unwrap instrumented session for reactive observability.

We now accept `CqlSessionBuilder` in ObservableReactiveSessionFactoryBean to use the session directly to avoid duplicate instrumentation.

Also, DefaultBridgedReactiveSession tries to unwrap any observability-proxied sessions.

Closes #1366

Author: mp911de

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/aot/CassandraRuntimeHints.java

@@ -31,7 +31,7 @@
 import org.springframework.data.cassandra.observability.CassandraObservationSupplier;
 import org.springframework.data.cassandra.repository.support.SimpleCassandraRepository;
 import org.springframework.data.cassandra.repository.support.SimpleReactiveCassandraRepository;
-import org.springframework.data.repository.util.ReactiveWrappers;
+import org.springframework.data.util.ReactiveWrappers;
 import org.springframework.lang.Nullable;
 import org.springframework.util.ClassUtils;
 
@@ -91,6 +91,19 @@ public void registerHints(org.springframework.aot.hint.RuntimeHints hints, @Null
 			}
 
 			hints.proxies().registerJdkProxy(CqlSession.class, SpringProxy.class, Advised.class, DecoratingProxy.class);
+			Class<?> observationDecorated;
+			try {
+				observationDecorated = Class.forName(
+						"org.springframework.data.cassandra.observability.CqlSessionObservationInterceptor.ObservationDecoratedProxy",
+						false, classLoader);
+			} catch (Exception e) {
+				observationDecorated = null;
+			}
+
+			if (observationDecorated != null) {
+				hints.proxies().registerJdkProxy(CqlSession.class, SpringProxy.class, Advised.class, DecoratingProxy.class,
+						observationDecorated);
+			}
 		}
 	}
 }

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/cql/session/DefaultBridgedReactiveSession.java

@@ -27,6 +27,8 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.springframework.aop.TargetSource;
+import org.springframework.aop.framework.AopProxyUtils;
 import org.springframework.data.cassandra.ReactiveResultSet;
 import org.springframework.data.cassandra.ReactiveSession;
 import org.springframework.util.Assert;
@@ -54,7 +56,7 @@
  * Calls are deferred until a subscriber subscribes to the resulting {@link org.reactivestreams.Publisher}. The calls
  * are executed by subscribing to {@link CompletionStage} and returning the result as calls complete.
  * <p>
- * Elements are emitted on netty EventLoop threads. {@link AsyncResultSet} allows {@link AsyncResultSet#fetchNextPage()}
+ * Elements are emitted on netty EventLoop threads. {@link AsyncResultSet} allows {@link AsyncResultSet#fetchNextPage()
  * asynchronous requesting} of subsequent pages. The next page is requested after emitting all elements of the previous
  * page. However, this is an intermediate solution until Datastax can provide a fully reactive driver.
  * <p>
@@ -85,6 +87,18 @@ public DefaultBridgedReactiveSession(CqlSession session) {
 
 		Assert.notNull(session, "Session must not be null");
 
+		// potentially unwrap a ObservationDecoratedProxy as reactive observability
+		// requires its own approach to span creation. We do not want to participate in
+		// async API spans but rather drive our own spans.
+		if (session instanceof TargetSource) {
+			Class<?>[] interfaces = session.getClass().getInterfaces();
+			for (Class<?> anInterface : interfaces) {
+				if (anInterface.getName().endsWith("ObservationDecoratedProxy")) {
+					session = (CqlSession) AopProxyUtils.getSingletonTarget(session);
+				}
+			}
+		}
+
 		this.session = session;
 	}
 
@@ -239,7 +253,6 @@ public Flux<Row> availableRows() {
 			return Flux.fromIterable(resultSet.currentPage());
 		}
 
-
 		@Override
 		public ColumnDefinitions getColumnDefinitions() {
 			return this.resultSet.getColumnDefinitions();

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/observability/CqlSessionObservationInterceptor.java

@@ -25,6 +25,7 @@
 
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
+import org.springframework.aop.TargetSource;
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.CqlSession;
@@ -183,4 +184,11 @@ private Observation startObservation(Statement<?> statement, boolean prepare, St
 		return observation.start();
 	}
 
+	/**
+	 * Marker interface for components that want to participate in observation but do not want to work with a
+	 * {@code CqlSession} that is already decorated for observation.
+	 */
+	public interface ObservationDecoratedProxy extends TargetSource {
+
+	}
 }

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/observability/ObservableCqlSessionFactory.java

@@ -17,7 +17,10 @@
 
 import io.micrometer.observation.ObservationRegistry;
 
+import org.springframework.aop.RawTargetAccess;
+import org.springframework.aop.TargetSource;
 import org.springframework.aop.framework.ProxyFactory;
+import org.springframework.data.cassandra.observability.CqlSessionObservationInterceptor.ObservationDecoratedProxy;
 import org.springframework.util.Assert;
 
 import com.datastax.oss.driver.api.core.CqlSession;
@@ -65,7 +68,10 @@ public static CqlSession wrap(CqlSession session, String remoteServiceName, Obse
 		proxyFactory.setTarget(session);
 		proxyFactory.addAdvice(new CqlSessionObservationInterceptor(session, remoteServiceName, observationRegistry));
 		proxyFactory.addInterface(CqlSession.class);
+		proxyFactory.addInterface(ObservationDecoratedProxy.class);
 
 		return (CqlSession) proxyFactory.getProxy();
 	}
+
+
 }

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/observability/ObservableReactiveSession.java

@@ -180,4 +180,5 @@ private Observation startObservation(@Nullable Observation parent, Statement<?>
 	private static Observation getParentObservation(ContextView contextView) {
 		return contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null);
 	}
+
 }

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/observability/ObservableReactiveSessionFactoryBean.java

@@ -17,6 +17,8 @@
 
 import io.micrometer.observation.ObservationRegistry;
 
+import org.springframework.aop.TargetSource;
+import org.springframework.aop.framework.AopProxyUtils;
 import org.springframework.beans.factory.config.AbstractFactoryBean;
 import org.springframework.data.cassandra.ReactiveSession;
 import org.springframework.data.cassandra.core.cql.session.DefaultBridgedReactiveSession;
@@ -25,6 +27,7 @@
 import org.springframework.util.ObjectUtils;
 
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.CqlSessionBuilder;
 
 /**
  * Factory bean to construct a {@link ReactiveSession} integrated with given {@link ObservationRegistry}. The required
@@ -41,10 +44,30 @@ public class ObservableReactiveSessionFactoryBean extends AbstractFactoryBean<Re
 
 	private final CqlSession cqlSession;
 
+	private final boolean requiresDestroy;
+
 	private final ObservationRegistry observationRegistry;
 
 	private @Nullable String remoteServiceName;
 
+	/**
+	 * Construct a new {@link ObservableReactiveSessionFactoryBean}.
+	 *
+	 * @param cqlSessionBuilder must not be {@literal null}.
+	 * @param observationRegistry must not be {@literal null}.
+	 * @since 4.0.5
+	 */
+	public ObservableReactiveSessionFactoryBean(CqlSessionBuilder cqlSessionBuilder,
+			ObservationRegistry observationRegistry) {
+
+		Assert.notNull(cqlSessionBuilder, "CqlSessionBuilder must not be null");
+		Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
+
+		this.cqlSession = cqlSessionBuilder.build();
+		this.requiresDestroy = true;
+		this.observationRegistry = observationRegistry;
+	}
+
 	/**
 	 * Construct a new {@link ObservableReactiveSessionFactoryBean}.
 	 *
@@ -56,7 +79,9 @@ public ObservableReactiveSessionFactoryBean(CqlSession cqlSession, ObservationRe
 		Assert.notNull(cqlSession, "CqlSession must not be null");
 		Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
 
-		this.cqlSession = cqlSession;
+		this.cqlSession = cqlSession instanceof TargetSource c ? (CqlSession) AopProxyUtils.getSingletonTarget(c)
+				: cqlSession;
+		this.requiresDestroy = false;
 		this.observationRegistry = observationRegistry;
 	}
 
@@ -81,6 +106,14 @@ public String getRemoteServiceName() {
 		return remoteServiceName;
 	}
 
+	@Override
+	public void destroy() {
+
+		if (requiresDestroy) {
+			cqlSession.close();
+		}
+	}
+
 	/**
 	 * Set the remote service name.
 	 *

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/observability/ObservableReactiveSessionFactoryBeanUnitTests.java

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2023 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
+ *
+ *      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 org.springframework.data.cassandra.observability;
+
+import static org.assertj.core.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import io.micrometer.observation.ObservationRegistry;
+
+import org.junit.jupiter.api.Test;
+import org.springframework.data.cassandra.ReactiveSession;
+import org.springframework.data.cassandra.core.cql.session.DefaultBridgedReactiveSession;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.CqlSessionBuilder;
+
+/**
+ * Unit tests for {@link ObservableReactiveSessionFactoryBean}.
+ *
+ * @author Mark Paluch
+ */
+class ObservableReactiveSessionFactoryBeanUnitTests {
+
+	@Test // GH-1366
+	void sessionFactoryBeanUnwrapsObservationProxy() throws Exception {
+
+		CqlSession sessionMock = mock(CqlSession.class);
+		ObservationRegistry registry = ObservationRegistry.NOOP;
+
+		CqlSession wrapped = ObservableCqlSessionFactory.wrap(sessionMock, registry);
+
+		ObservableReactiveSessionFactoryBean bean = new ObservableReactiveSessionFactoryBean(wrapped, registry);
+		bean.afterPropertiesSet();
+
+		ReactiveSession object = bean.getObject();
+		Object cqlSession = ReflectionTestUtils.getField(ReflectionTestUtils.getField(object, "delegate"), "session");
+
+		assertThat(cqlSession).isSameAs(sessionMock);
+	}
+
+	@Test // GH-1366
+	void reactiveSessionDelegateBeanUnwrapsObservationProxy() throws Exception {
+
+		CqlSession sessionMock = mock(CqlSession.class);
+		ObservationRegistry registry = ObservationRegistry.NOOP;
+
+		CqlSession wrapped = ObservableCqlSessionFactory.wrap(sessionMock, registry);
+
+		ReactiveSession object = new DefaultBridgedReactiveSession(wrapped);
+		Object cqlSession = ReflectionTestUtils.getField(object, "session");
+
+		assertThat(cqlSession).isSameAs(sessionMock);
+	}
+
+	@Test // GH-1366
+	void closesCqlSessionViaBuilder() throws Exception {
+
+		CqlSessionBuilder builderMock = mock(CqlSessionBuilder.class);
+		CqlSession sessionMock = mock(CqlSession.class);
+
+		doReturn(sessionMock).when(builderMock).build();
+
+		ObservationRegistry registry = ObservationRegistry.NOOP;
+
+		ObservableReactiveSessionFactoryBean bean = new ObservableReactiveSessionFactoryBean(builderMock, registry);
+		bean.afterPropertiesSet();
+		bean.destroy();
+
+		verify(sessionMock).close();
+	}
+
+	@Test // GH-1366
+	void doesNotCloseCqlSession() throws Exception {
+
+		CqlSession sessionMock = mock(CqlSession.class);
+
+		ObservationRegistry registry = ObservationRegistry.NOOP;
+
+		ObservableReactiveSessionFactoryBean bean = new ObservableReactiveSessionFactoryBean(sessionMock, registry);
+		bean.afterPropertiesSet();
+		bean.destroy();
+
+		verifyNoInteractions(sessionMock);
+	}
+}