[WIP]: Add OpenTelemetry metrics support for ClusterInfo

d2/src/main/java/com/linkedin/d2/jmx/ClusterInfoJmx.java

@@ -18,14 +18,22 @@
 
 import com.linkedin.d2.balancer.LoadBalancerStateItem;
 import com.linkedin.d2.balancer.simple.ClusterInfoItem;
+import com.linkedin.d2.jmx.ClusterInfoOtelMetricsProvider;
 
 
 public class ClusterInfoJmx implements ClusterInfoJmxMBean
 {
   private final ClusterInfoItem _clusterInfoItem;
+  private final ClusterInfoOtelMetricsProvider _clusterInfoOtelMetricsProvider;
 
-  public ClusterInfoJmx(ClusterInfoItem clusterInfoItem) {
+  public ClusterInfoJmx(ClusterInfoItem clusterInfoItem)
+  {
+    this(clusterInfoItem, new NoOpClusterInfoOtelMetricsProvider());
+  }
+
+  public ClusterInfoJmx(ClusterInfoItem clusterInfoItem, ClusterInfoOtelMetricsProvider clusterInfoOtelMetricsProvider) {
     _clusterInfoItem = clusterInfoItem;
+    _clusterInfoOtelMetricsProvider = clusterInfoOtelMetricsProvider;
   }
 
   @Override
@@ -37,11 +45,19 @@ public ClusterInfoItem getClusterInfoItem()
   @Override
   public int getCanaryDistributionPolicy()
   {
+    String clusterName = _clusterInfoItem.getClusterPropertiesItem().getProperty().getClusterName();
+
     switch (_clusterInfoItem.getClusterPropertiesItem().getDistribution())
     {
-      case STABLE: return 0;
-      case CANARY: return 1;
-      default: return -1;
+      case STABLE: 
+        _clusterInfoOtelMetricsProvider.recordCanaryDistributionPolicy(clusterName, 0);
+        return 0;
+      case CANARY: 
+        _clusterInfoOtelMetricsProvider.recordCanaryDistributionPolicy(clusterName, 1);
+        return 1;
+      default: 
+        _clusterInfoOtelMetricsProvider.recordCanaryDistributionPolicy(clusterName, -1);
+        return -1;
     }
   }
 }

d2/src/main/java/com/linkedin/d2/jmx/ClusterInfoOtelMetricsProvider.java

@@ -0,0 +1,17 @@
+package com.linkedin.d2.jmx;
+
+/**
+ * Interface for OpenTelemetry metrics collection for ClusterInfo.
+ */
+public interface ClusterInfoOtelMetricsProvider {
+
+    /**
+   * Records the current canary distribution policy for the cluster 
+   * (0=STABLE, 1=CANARY, -1=UNSPECIFIED)
+   * 
+   * @param clusterName the name of the cluster
+   * @param policy the canary distribution policy
+   */
+    void recordCanaryDistributionPolicy(String clusterName, int policy);
+
+}

d2/src/main/java/com/linkedin/d2/jmx/JmxManager.java

@@ -50,10 +50,17 @@ public class JmxManager
 
   private final MBeanServer   _server;
   private final Set<ObjectName> _registeredNames = new HashSet<>();
+  private ClusterInfoOtelMetricsProvider _clusterInfoOtelMetricsProvider;
 
   public JmxManager()
+  {
+    this(new NoOpClusterInfoOtelMetricsProvider());
+  }
+
+  public JmxManager(ClusterInfoOtelMetricsProvider clusterInfoOtelMetricsProvider)
   {
     _server = ManagementFactory.getPlatformMBeanServer();
+    _clusterInfoOtelMetricsProvider = clusterInfoOtelMetricsProvider;
   }
 
   MBeanServer getMBeanServer()
@@ -118,7 +125,7 @@ public synchronized JmxManager registerServicePropertiesJmxBean(String name, Ser
 
   public synchronized JmxManager registerClusterInfo(String name, ClusterInfoItem clusterInfo)
   {
-    checkReg(new ClusterInfoJmx(clusterInfo), name);
+    checkReg(new ClusterInfoJmx(clusterInfo, _clusterInfoOtelMetricsProvider), name);
 
     return this;
   }

d2/src/main/java/com/linkedin/d2/jmx/NoOpClusterInfoOtelMetricsProvider.java

@@ -0,0 +1,10 @@
+package com.linkedin.d2.jmx;
+
+public class NoOpClusterInfoOtelMetricsProvider implements ClusterInfoOtelMetricsProvider
+{
+  @Override
+  public void recordCanaryDistributionPolicy(String clusterName, int policy)
+  {
+    // No-op
+  }
+}

d2/src/test/java/com/linkedin/d2/jmx/ClusterInfoJmxTest.java

@@ -31,6 +31,8 @@
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
+import static org.mockito.Mockito.*;
+
 
 public class ClusterInfoJmxTest
 {
@@ -71,4 +73,28 @@ public int getPartitionId(URI uri) {
     );
     Assert.assertEquals(clusterInfoJmx.getCanaryDistributionPolicy(), expectedValue);
   }
+
+  @Test(dataProvider = "getCanaryDistributionPoliciesTestData")
+  public void testGetCanaryDistributionPolicyWithMetricsProvider(CanaryDistributionProvider.Distribution distribution, int expectedValue)
+  {
+    ClusterInfoOtelMetricsProvider mockProvider = mock(ClusterInfoOtelMetricsProvider.class);
+    ClusterInfoItem clusterInfoItem = new ClusterInfoItem(_mockedSimpleBalancerState, new ClusterProperties("Foo"), new PartitionAccessor() {
+      @Override
+      public int getMaxPartitionId() {
+        return 0;
+      }
+
+      @Override
+      public int getPartitionId(URI uri) {
+        return 0;
+      }
+    }, distribution);
+
+    ClusterInfoJmx clusterInfoJmx = new ClusterInfoJmx(clusterInfoItem, mockProvider);
+
+    int actualValue = clusterInfoJmx.getCanaryDistributionPolicy();
+
+    Assert.assertEquals(actualValue, expectedValue);
+    verify(mockProvider, times(1)).recordCanaryDistributionPolicy("Foo", expectedValue);
+  }
 }