From e5aff2645bf856963c64b7428eddcac0fd95f4d3 Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@vmware.com>
Date: Thu, 24 Aug 2023 13:57:29 +0200
Subject: [PATCH] Polishing.

Refactor duplicate code into callback.

See #4481
---
 .../MongoObservationCommandListener.java      | 56 +++++++++++--------
 1 file changed, 32 insertions(+), 24 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java
index 2cbe4f03d..9c525e9e4 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java
@@ -19,6 +19,8 @@ import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
 
+import java.util.function.BiConsumer;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.lang.Nullable;
@@ -126,30 +128,43 @@ public class MongoObservationCommandListener implements CommandListener {
 	@Override
 	public void commandSucceeded(CommandSucceededEvent event) {
 
-		RequestContext requestContext = event.getRequestContext();
+		doInObservation(event.getRequestContext(), (observation, context) -> {
 
-		if (requestContext == null) {
-			return;
-		}
+			context.setCommandSucceededEvent(event);
 
-		Observation observation = requestContext.getOrDefault(ObservationThreadLocalAccessor.KEY, null);
-		if (observation == null || !(observation.getContext()instanceof MongoHandlerContext context)) {
-			return;
-		}
+			if (log.isDebugEnabled()) {
+				log.debug("Command succeeded - will stop observation [" + observation + "]");
+			}
 
-		context.setCommandSucceededEvent(event);
-
-		if (log.isDebugEnabled()) {
-			log.debug("Command succeeded - will stop observation [" + observation + "]");
-		}
-
-		observation.stop();
+			observation.stop();
+		});
 	}
 
 	@Override
 	public void commandFailed(CommandFailedEvent event) {
 
-		RequestContext requestContext = event.getRequestContext();
+		doInObservation(event.getRequestContext(), (observation, context) -> {
+
+			context.setCommandFailedEvent(event);
+
+			if (log.isDebugEnabled()) {
+				log.debug("Command failed - will stop observation [" + observation + "]");
+			}
+
+			observation.error(event.getThrowable());
+			observation.stop();
+		});
+	}
+
+	/**
+	 * Performs the given action for the {@link Observation} and {@link MongoHandlerContext} if there is an ongoing Mongo
+	 * Observation. Exceptions thrown by the action are relayed to the caller.
+	 *
+	 * @param requestContext the context to extract the Observation from.
+	 * @param action the action to invoke.
+	 */
+	private void doInObservation(@Nullable RequestContext requestContext,
+			BiConsumer<Observation, MongoHandlerContext> action) {
 
 		if (requestContext == null) {
 			return;
@@ -160,14 +175,7 @@ public class MongoObservationCommandListener implements CommandListener {
 			return;
 		}
 
-		context.setCommandFailedEvent(event);
-
-		if (log.isDebugEnabled()) {
-			log.debug("Command failed - will stop observation [" + observation + "]");
-		}
-
-		observation.error(event.getThrowable());
-		observation.stop();
+		action.accept(observation, context);
 	}
 
 	/**