`ObservationRequestTracker` remains open for observed Cassandra statements

[BenEfrati]

Bug description

When Cassandra request observation is enabled, RequestIdGenerator appears to break the observation wrapper chain.
The problem is that RequestIdGenerator.getDecoratedStatement(Statement<?> statement, String requestId) calls:

return statement.setCustomPayload(unmodifiableMap);

https://github.com/apache/cassandra-java-driver/blob/c7fe73fe23149c10fe365402b55a6dc88b90b0d4/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java#L82

Since Statement<?> is immutable, setCustomPayload(...) returns a new statement instance.

If the incoming statement was already wrapped by Spring Data Cassandra observation infrastructure (for example an ObservationStatement created from CqlSessionObservationInterceptor / tracked by ObservationRequestTracker

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

Line 169 in d41b7ea

Statement<?> observableStatement = ObservationStatement.createProxy(startObservation(statement, false, methodName),

), the returned statement is no longer the wrapper/proxy instance. This drops the observation wrapper and breaks the observation lifecycle, so the corresponding long task timer is never stopped.

In production, this leads to a buildup of open long task timers and memory growth.

Expected behavior

Observation should be stopped on call to onSucess

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

Lines 53 to 66 in d41b7ea

	public void onSuccess(Request request, long latencyNanos, DriverExecutionProfile executionProfile, Node node,
	String requestLogPrefix) {
	if (request instanceof CassandraObservationSupplier supplier) {
	Observation observation = supplier.getObservation();
	if (log.isDebugEnabled()) {
	log.debug("Closing observation [" + observation + "]");
	}
	observation.stop();
	}
	}

Actual behavior

Applying the request id via setCustomPayload(...) creates a new statement instance and loses the observation wrapper. The request still executes, but ObservationRequestTracker never closes the corresponding long task timer.

Confirmed root cause

The exact failure flow is:

1. Spring Data Cassandra wraps the statement for observation.
2. W3CContextRequestIdGenerator adds the request id via statement.setCustomPayload(...).
3. Because the statement is immutable, a new plain statement instance is returned.
4. The observation wrapper is lost.
5. Request completes without the original tracked wrapper.
6. ObservationRequestTracker does not stop the long task timer.

The relevant method shape is effectively:

default Statement<?> getDecoratedStatement(
    Statement<?> statement, String requestId) {

  Map<String, ByteBuffer> existing = new HashMap<>(statement.getCustomPayload());
  String key = getCustomPayloadKey();
  existing.put(key, ByteBuffer.wrap(requestId.getBytes(StandardCharsets.UTF_8)));

  Map<String, ByteBuffer> unmodifiableMap = Collections.unmodifiableMap(existing);

  return statement.setCustomPayload(unmodifiableMap);
}

Environment

• Spring Data Cassandra: 5.0.5
• Spring Boot: 4.0.6
• Micrometer: 1.16.5
• Java: 25
• Cassandra driver: 4.19.2

Possible fix

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

Lines 105 to 110 in d41b7ea

	Object result = invocation.proceed();
	if (result instanceof Statement<?>) {
	this.delegate = (Statement<?>) result;
	}

Object result = invocation.proceed(); 
if (result instanceof Statement<?> statement) { 
    this.delegate = statement
    return createProxy(this.observation, statement); ;
} 
return result;

[BenEfrati]

This test failed with existing implementation.

@Testcontainers
class CassandraObservationRegressionTest {

  @Container
  private static final CassandraContainer cassandra = new CassandraContainer("cassandra:4.1");

  private final ObservationRegistry observationRegistry = ObservationRegistry.create();

  @Test
  void requestIdDecorationMustPreserveObservation() throws Exception {
    SimpleMeterRegistry meterRegistry = new SimpleMeterRegistry();
    observationRegistry.observationConfig().observationHandler(new DefaultMeterObservationHandler(meterRegistry));
    CqlSessionBuilder builder = CqlSession.builder()
        .addContactPoint(cassandra.getContactPoint())
        .withLocalDatacenter(cassandra.getLocalDatacenter())
        .withRequestIdGenerator(new W3CContextRequestIdGenerator("payload-key"))
        .withAuthCredentials(cassandra.getUsername(), cassandra.getPassword());

    ObservableCqlSessionFactoryBean observableCqlSessionFactoryBean =
        new ObservableCqlSessionFactoryBean(builder, observationRegistry);
    observableCqlSessionFactoryBean.afterPropertiesSet();
    try (CqlSession session = observableCqlSessionFactoryBean.getObject()) {
      session.execute("SELECT now() FROM system.local");
      session.executeAsync("SELECT now() FROM system.local").toCompletableFuture().get();
    }
    assertThat(meterRegistry.find("execute")
        .timers()
        .stream()
        .mapToDouble(Timer::count)
        .sum())
        .isOne();
    assertThat(meterRegistry.find("execute.active")
        .longTaskTimers()
        .stream()
        .mapToDouble(LongTaskTimer::activeTasks)
        .sum())
        .isZero();
    assertThat(meterRegistry.find("execute.cassandra.node.success")
        .counters()
        .stream()
        .mapToDouble(Counter::count)
        .sum())
        .isOne();

    assertThat(meterRegistry.find("executeAsync")
        .timers()
        .stream()
        .mapToDouble(Timer::count)
        .sum())
        .isOne();
    assertThat(meterRegistry.find("executeAsync.active")
        .longTaskTimers()
        .stream()
        .mapToDouble(LongTaskTimer::activeTasks)
        .sum())
        .isZero();
    assertThat(meterRegistry.find("executeAsync.cassandra.node.success")
        .counters()
        .stream()
        .mapToDouble(Counter::count)
        .sum())
        .isOne();
  }

}

With this fix, metrics are created and active task is zero.

Object result = invocation.proceed(); 
if (result instanceof Statement<?> statement) { 
    this.delegate = statement
    return createProxy(this.observation, statement); ;
} 
return result;