Skip to content

Commit 0b151aa

Browse files
mp911demp911de
committed
Polishing.
Extract ifContextPresent(…) utility method to apply contextual actions if the context matches the given type. See #1541
1 parent 74fca38 commit 0b151aa

File tree

2 files changed

+154
-28
lines changed

2 files changed

+154
-28
lines changed

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

Lines changed: 38 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,11 @@
1919
import io.micrometer.observation.Observation.Context;
2020
import io.micrometer.observation.Observation.Event;
2121

22+
import java.util.function.Consumer;
23+
2224
import org.apache.commons.logging.Log;
2325
import org.apache.commons.logging.LogFactory;
26+
2427
import org.springframework.data.cassandra.observability.CassandraObservation.Events;
2528
import org.springframework.data.cassandra.observability.CassandraObservation.HighCardinalityKeyNames;
2629
import org.springframework.lang.Nullable;
@@ -50,9 +53,9 @@ public enum ObservationRequestTracker implements RequestTracker {
5053
public void onSuccess(Request request, long latencyNanos, DriverExecutionProfile executionProfile, Node node,
5154
String requestLogPrefix) {
5255

53-
if (request instanceof CassandraObservationSupplier) {
56+
if (request instanceof CassandraObservationSupplier supplier) {
5457

55-
Observation observation = ((CassandraObservationSupplier) request).getObservation();
58+
Observation observation = supplier.getObservation();
5659

5760
if (log.isDebugEnabled()) {
5861
log.debug("Closing observation [" + observation + "]");
@@ -66,9 +69,9 @@ public void onSuccess(Request request, long latencyNanos, DriverExecutionProfile
6669
public void onError(Request request, Throwable error, long latencyNanos, DriverExecutionProfile executionProfile,
6770
@Nullable Node node, String requestLogPrefix) {
6871

69-
if (request instanceof CassandraObservationSupplier) {
72+
if (request instanceof CassandraObservationSupplier supplier) {
7073

71-
Observation observation = ((CassandraObservationSupplier) request).getObservation();
74+
Observation observation = supplier.getObservation();
7275
observation.error(error);
7376

7477
if (log.isDebugEnabled()) {
@@ -83,22 +86,17 @@ public void onError(Request request, Throwable error, long latencyNanos, DriverE
8386
public void onNodeError(Request request, Throwable error, long latencyNanos, DriverExecutionProfile executionProfile,
8487
Node node, String requestLogPrefix) {
8588

86-
if (request instanceof CassandraObservationSupplier) {
87-
88-
Observation observation = ((CassandraObservationSupplier) request).getObservation();
89-
Context context = observation.getContext();
90-
91-
if (context instanceof CassandraObservationContext) {
89+
if (request instanceof CassandraObservationSupplier supplier) {
9290

93-
((CassandraObservationContext) context).setNode(node);
91+
Observation observation = supplier.getObservation();
92+
ifContextPresent(observation, CassandraObservationContext.class, context -> context.setNode(node));
9493

95-
observation.highCardinalityKeyValue(
96-
String.format(HighCardinalityKeyNames.NODE_ERROR_TAG.asString(), node.getEndPoint()), error.toString());
97-
observation.event(Event.of(Events.NODE_ERROR.getValue()));
94+
observation.highCardinalityKeyValue(
95+
String.format(HighCardinalityKeyNames.NODE_ERROR_TAG.asString(), node.getEndPoint()), error.toString());
96+
observation.event(Event.of(Events.NODE_ERROR.getValue()));
9897

99-
if (log.isDebugEnabled()) {
100-
log.debug("Marking node error for [" + observation + "]");
101-
}
98+
if (log.isDebugEnabled()) {
99+
log.debug("Marking node error for [" + observation + "]");
102100
}
103101
}
104102
}
@@ -107,20 +105,15 @@ public void onNodeError(Request request, Throwable error, long latencyNanos, Dri
107105
public void onNodeSuccess(Request request, long latencyNanos, DriverExecutionProfile executionProfile, Node node,
108106
String requestLogPrefix) {
109107

110-
if (request instanceof CassandraObservationSupplier) {
111-
112-
Observation observation = ((CassandraObservationSupplier) request).getObservation();
113-
Context context = observation.getContext();
108+
if (request instanceof CassandraObservationSupplier supplier) {
114109

115-
if (context instanceof CassandraObservationContext) {
110+
Observation observation = supplier.getObservation();
111+
ifContextPresent(observation, CassandraObservationContext.class, context -> context.setNode(node));
116112

117-
((CassandraObservationContext) context).setNode(node);
113+
observation.event(Event.of(Events.NODE_SUCCESS.getValue()));
118114

119-
observation.event(Event.of(Events.NODE_SUCCESS.getValue()));
120-
121-
if (log.isDebugEnabled()) {
122-
log.debug("Marking node success for [" + observation + "]");
123-
}
115+
if (log.isDebugEnabled()) {
116+
log.debug("Marking node success for [" + observation + "]");
124117
}
125118
}
126119
}
@@ -130,4 +123,21 @@ public void close() throws Exception {
130123

131124
}
132125

126+
/**
127+
* If the {@link Observation} is a real observation (i.e. not no-op) and the context is of the given type, apply the
128+
* consumer function to the context.
129+
*/
130+
static <T extends Context> void ifContextPresent(Observation observation, Class<T> contextType,
131+
Consumer<T> contextConsumer) {
132+
133+
if (observation.isNoop()) {
134+
return;
135+
}
136+
137+
Context context = observation.getContext();
138+
if (contextType.isInstance(context)) {
139+
contextConsumer.accept(contextType.cast(context));
140+
}
141+
}
142+
133143
}

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

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Copyright 2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.cassandra.observability;
17+
18+
import static org.assertj.core.api.Assertions.*;
19+
import static org.mockito.Mockito.*;
20+
21+
import io.micrometer.observation.Observation;
22+
23+
import java.net.InetSocketAddress;
24+
import java.util.function.Consumer;
25+
26+
import org.junit.jupiter.api.Test;
27+
import org.mockito.Answers;
28+
29+
import org.springframework.lang.Nullable;
30+
31+
import com.datastax.oss.driver.api.core.session.Request;
32+
import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
33+
import com.datastax.oss.driver.internal.core.metadata.DefaultEndPoint;
34+
import com.datastax.oss.driver.internal.core.metadata.DefaultNode;
35+
36+
/**
37+
* Unit tests for {@link ObservationRequestTracker}.
38+
*
39+
* @author Mark Paluch
40+
*/
41+
class ObservationRequestTrackerUnitTests {
42+
43+
@Test // GH-1541
44+
void shouldStopObservation() {
45+
46+
Request request = mockRequest(null);
47+
48+
ObservationRequestTracker.INSTANCE.onSuccess(request, 0, null, null, "");
49+
50+
verify(((CassandraObservationSupplier) request).getObservation()).stop();
51+
}
52+
53+
@Test // GH-1541
54+
void shouldAssociateNodeWithContext() {
55+
56+
CassandraObservationContext context = new CassandraObservationContext(null, "foo", false, "foo", "foo", "bar");
57+
58+
Request request = mockRequest(context);
59+
InternalDriverContext driverContext = mock(InternalDriverContext.class, Answers.RETURNS_MOCKS);
60+
61+
DefaultNode node = new DefaultNode(new DefaultEndPoint(InetSocketAddress.createUnresolved("localhost", 1234)),
62+
driverContext);
63+
ObservationRequestTracker.INSTANCE.onNodeSuccess(request, 0, null, node, "");
64+
65+
assertThat(context.getNode()).isEqualTo(node);
66+
}
67+
68+
@Test // GH-1541
69+
void noOpObservationShouldNotAssociateContext() {
70+
71+
CassandraObservationContext context = new CassandraObservationContext(null, "foo", false, "foo", "foo", "bar");
72+
Request request = mockRequest(context, observation -> {
73+
when(observation.isNoop()).thenReturn(true);
74+
});
75+
InternalDriverContext driverContext = mock(InternalDriverContext.class, Answers.RETURNS_MOCKS);
76+
77+
DefaultNode node = new DefaultNode(new DefaultEndPoint(InetSocketAddress.createUnresolved("localhost", 1234)),
78+
driverContext);
79+
ObservationRequestTracker.INSTANCE.onNodeSuccess(request, 0, null, node, "");
80+
81+
assertThat(context.getNode()).isNull();
82+
}
83+
84+
@Test // GH-1541
85+
void observationWithOtherContextShouldNotAssociateContext() {
86+
87+
Request request = mockRequest(mock(Observation.Context.class));
88+
InternalDriverContext driverContext = mock(InternalDriverContext.class, Answers.RETURNS_MOCKS);
89+
90+
DefaultNode node = new DefaultNode(new DefaultEndPoint(InetSocketAddress.createUnresolved("localhost", 1234)),
91+
driverContext);
92+
93+
assertThatNoException().isThrownBy(() -> {
94+
ObservationRequestTracker.INSTANCE.onNodeSuccess(request, 0, null, node, "");
95+
});
96+
}
97+
98+
private static Request mockRequest(@Nullable Observation.Context context) {
99+
return mockRequest(context, observation -> {});
100+
}
101+
102+
private static Request mockRequest(@Nullable Observation.Context context,
103+
Consumer<Observation> observationCustomizer) {
104+
105+
Request request = mock(Request.class, withSettings().extraInterfaces(CassandraObservationSupplier.class));
106+
107+
Observation observation = mock(Observation.class);
108+
CassandraObservationSupplier supplier = (CassandraObservationSupplier) request;
109+
when(supplier.getObservation()).thenReturn(observation);
110+
when(observation.getContext()).thenReturn(context);
111+
112+
observationCustomizer.accept(observation);
113+
114+
return request;
115+
}
116+
}

0 commit comments

Comments
 (0)