19
19
import io .micrometer .observation .ObservationRegistry ;
20
20
import io .micrometer .observation .contextpropagation .ObservationThreadLocalAccessor ;
21
21
22
+ import java .util .function .BiConsumer ;
23
+
22
24
import org .apache .commons .logging .Log ;
23
25
import org .apache .commons .logging .LogFactory ;
24
26
import org .springframework .lang .Nullable ;
@@ -126,30 +128,43 @@ public void commandStarted(CommandStartedEvent event) {
126
128
@ Override
127
129
public void commandSucceeded (CommandSucceededEvent event ) {
128
130
129
- RequestContext requestContext = event .getRequestContext ();
131
+ doInObservation ( event .getRequestContext (), ( observation , context ) -> {
130
132
131
- if (requestContext == null ) {
132
- return ;
133
- }
133
+ context .setCommandSucceededEvent (event );
134
134
135
- Observation observation = requestContext .getOrDefault (ObservationThreadLocalAccessor .KEY , null );
136
- if (observation == null || !(observation .getContext ()instanceof MongoHandlerContext context )) {
137
- return ;
138
- }
139
-
140
- context .setCommandSucceededEvent (event );
141
-
142
- if (log .isDebugEnabled ()) {
143
- log .debug ("Command succeeded - will stop observation [" + observation + "]" );
144
- }
135
+ if (log .isDebugEnabled ()) {
136
+ log .debug ("Command succeeded - will stop observation [" + observation + "]" );
137
+ }
145
138
146
- observation .stop ();
139
+ observation .stop ();
140
+ });
147
141
}
148
142
149
143
@ Override
150
144
public void commandFailed (CommandFailedEvent event ) {
151
145
152
- RequestContext requestContext = event .getRequestContext ();
146
+ doInObservation (event .getRequestContext (), (observation , context ) -> {
147
+
148
+ context .setCommandFailedEvent (event );
149
+
150
+ if (log .isDebugEnabled ()) {
151
+ log .debug ("Command failed - will stop observation [" + observation + "]" );
152
+ }
153
+
154
+ observation .error (event .getThrowable ());
155
+ observation .stop ();
156
+ });
157
+ }
158
+
159
+ /**
160
+ * Performs the given action for the {@link Observation} and {@link MongoHandlerContext} if there is an ongoing Mongo
161
+ * Observation. Exceptions thrown by the action are relayed to the caller.
162
+ *
163
+ * @param requestContext the context to extract the Observation from.
164
+ * @param action the action to invoke.
165
+ */
166
+ private void doInObservation (@ Nullable RequestContext requestContext ,
167
+ BiConsumer <Observation , MongoHandlerContext > action ) {
153
168
154
169
if (requestContext == null ) {
155
170
return ;
@@ -160,14 +175,7 @@ public void commandFailed(CommandFailedEvent event) {
160
175
return ;
161
176
}
162
177
163
- context .setCommandFailedEvent (event );
164
-
165
- if (log .isDebugEnabled ()) {
166
- log .debug ("Command failed - will stop observation [" + observation + "]" );
167
- }
168
-
169
- observation .error (event .getThrowable ());
170
- observation .stop ();
178
+ action .accept (observation , context );
171
179
}
172
180
173
181
/**
0 commit comments