Ensure JobClusterActor's logs always contain the clusterName

mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/jobcluster/JobClusterActor.java

@@ -167,10 +167,13 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.slf4j.MDC;
 import rx.Observable;
 import rx.functions.Action1;
 import rx.schedulers.Schedulers;
 import rx.subjects.BehaviorSubject;
+import scala.PartialFunction;
+import scala.runtime.BoxedUnit;
 
 
 /**
@@ -186,6 +189,8 @@ public class JobClusterActor extends AbstractActorWithTimers implements IJobClus
     private static final Integer DEFAULT_LIMIT = 100;
     private static final Integer DEFAULT_ACTIVE_JOB_LIMIT = 5000;
 
+    private static final String MDC_KEY_CLUSTER_NAME = "clusterName";
+
     private final Logger logger = LoggerFactory.getLogger(JobClusterActor.class);
 
     private static final String CHECK_EXPIRED_TIMER_KEY = "EXPIRE_OLD_JOBS";
@@ -673,6 +678,24 @@ MetricGroupId getMetricGroupId(String name) {
         return new MetricGroupId("JobClusterActor", new BasicTag("jobCluster", name));
     }
 
+    /**
+     * Wraps every message this actor processes so that the cluster name is available via
+     * SLF4J's MDC for the duration of that message's handling. This lets every log line
+     * emitted while handling a message automatically include the cluster name (for logging
+     * backends that render MDC values, e.g. structured/JSON encoders, or a pattern layout
+     * with %X{clusterName}), instead of requiring each of this class's log call sites to pass
+     * `name` explicitly.
+     */
+    @Override
+    public void aroundReceive(PartialFunction<Object, BoxedUnit> receive, Object msg) {
+        MDC.put(MDC_KEY_CLUSTER_NAME, name);
+        try {
+            super.aroundReceive(receive, msg);
+        } finally {
+            MDC.remove(MDC_KEY_CLUSTER_NAME);
+        }
+    }
+
     @Override
     public void preStart() throws Exception {
         logger.info("JobClusterActor {} started", name);

mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/jobcluster/JobClusterAkkaTest.java

@@ -64,8 +64,11 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import akka.actor.AbstractActor.Receive;
 import akka.actor.ActorRef;
 import akka.actor.ActorSystem;
+import akka.japi.pf.ReceiveBuilder;
+import akka.testkit.TestActorRef;
 import akka.testkit.javadsl.TestKit;
 import com.netflix.mantis.master.scheduler.TestHelpers;
 import com.typesafe.config.Config;
@@ -141,6 +144,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.BeforeClass;
@@ -150,6 +154,7 @@
 import org.junit.rules.TemporaryFolder;
 import org.mockito.Mockito;
 import org.mockito.stubbing.Answer;
+import org.slf4j.MDC;
 import rx.schedulers.Schedulers;
 import rx.subjects.BehaviorSubject;
 
@@ -331,6 +336,34 @@ private JobDefinition createJob(String name2) throws InvalidJobException {
         return createJob(name2, 0, MantisJobDurationType.Perpetual, null);
     }
 
+    // LOGGING TESTS ////////////////////////////////////////////////////////////////////////////////
+    @Test
+    public void testAroundReceivePutsClusterNameInMdcDuringMessageHandlingAndClearsItAfter() {
+        String clusterName = "mdcTestCluster";
+        MantisSchedulerFactory schedulerMockFactory = mock(MantisSchedulerFactory.class);
+
+        TestActorRef<JobClusterActor> actorRef = TestActorRef.create(
+            system,
+            props(clusterName, jobStore, schedulerMockFactory, eventPublisher, costsCalculator, 0));
+        JobClusterActor actor = actorRef.underlyingActor();
+
+        // Call aroundReceive directly (bypassing the mailbox/dispatcher) with a receive
+        // handler that captures the MDC value synchronously while "processing" the message,
+        // since MDC is thread-local and reading it from the test thread after the fact
+        // (post message-send) wouldn't observe the value set-and-cleared during handling.
+        AtomicReference<String> observedDuringProcessing = new AtomicReference<>();
+        Receive captureMdcReceive = ReceiveBuilder.create()
+            .matchAny(msg -> observedDuringProcessing.set(MDC.get("clusterName")))
+            .build();
+
+        actor.aroundReceive(captureMdcReceive.onMessage(), "any-test-message");
+
+        assertEquals("MDC should contain the cluster name while the actor is handling a message",
+            clusterName, observedDuringProcessing.get());
+        assertNull("MDC should be cleared once message handling completes",
+            MDC.get("clusterName"));
+    }
+
     // CLUSTER CRUD TESTS ///////////////////////////////////////////////////////////////////////////
     @Test
     public void testJobClusterCreate() throws Exception  {