1919import io .micrometer .observation .ObservationRegistry ;
2020import io .micrometer .observation .contextpropagation .ObservationThreadLocalAccessor ;
2121
22+ import java .util .function .BiConsumer ;
23+
2224import org .apache .commons .logging .Log ;
2325import org .apache .commons .logging .LogFactory ;
2426import org .springframework .lang .Nullable ;
@@ -126,30 +128,43 @@ public void commandStarted(CommandStartedEvent event) {
126128 @ Override
127129 public void commandSucceeded (CommandSucceededEvent event ) {
128130
129- RequestContext requestContext = event .getRequestContext ();
131+ doInObservation ( event .getRequestContext (), ( observation , context ) -> {
130132
131- if (requestContext == null ) {
132- return ;
133- }
133+ context .setCommandSucceededEvent (event );
134134
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+ }
145138
146- observation .stop ();
139+ observation .stop ();
140+ });
147141 }
148142
149143 @ Override
150144 public void commandFailed (CommandFailedEvent event ) {
151145
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 ) {
153168
154169 if (requestContext == null ) {
155170 return ;
@@ -160,14 +175,7 @@ public void commandFailed(CommandFailedEvent event) {
160175 return ;
161176 }
162177
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 );
171179 }
172180
173181 /**
0 commit comments