[Iceberg]Fix jmx metrics overwrite each other between nodes in query runner

Author: hantangwangd

presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnectorFactory.java

@@ -18,6 +18,8 @@
 import com.facebook.presto.spi.connector.ConnectorContext;
 import com.facebook.presto.spi.connector.ConnectorFactory;
 
+import javax.management.MBeanServer;
+
 import java.lang.reflect.InvocationTargetException;
 import java.util.Map;
 import java.util.Optional;
@@ -27,6 +29,13 @@
 public class IcebergConnectorFactory
         implements ConnectorFactory
 {
+    private final MBeanServer mBeanServer;
+
+    public IcebergConnectorFactory(MBeanServer mBeanServer)
+    {
+        this.mBeanServer = mBeanServer;
+    }
+
     @Override
     public String getName()
     {
@@ -45,8 +54,8 @@ public Connector create(String catalogName, Map<String, String> config, Connecto
         ClassLoader classLoader = IcebergConnectorFactory.class.getClassLoader();
         try {
             return (Connector) classLoader.loadClass(InternalIcebergConnectorFactory.class.getName())
-                    .getMethod("createConnector", String.class, Map.class, ConnectorContext.class, Optional.class)
-                    .invoke(null, catalogName, config, context, Optional.empty());
+                    .getMethod("createConnector", String.class, Map.class, ConnectorContext.class, Optional.class, MBeanServer.class)
+                    .invoke(null, catalogName, config, context, Optional.empty(), mBeanServer);
         }
         catch (InvocationTargetException e) {
             Throwable targetException = e.getTargetException();

presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergPlugin.java

@@ -19,15 +19,30 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 
+import javax.management.MBeanServer;
+
+import java.lang.management.ManagementFactory;
 import java.util.Set;
 
 public class IcebergPlugin
         implements Plugin
 {
+    private final MBeanServer mBeanServer;
+
+    public IcebergPlugin()
+    {
+        this(ManagementFactory.getPlatformMBeanServer());
+    }
+
+    public IcebergPlugin(MBeanServer mBeanServer)
+    {
+        this.mBeanServer = mBeanServer;
+    }
+
     @Override
     public Iterable<ConnectorFactory> getConnectorFactories()
     {
-        return ImmutableList.of(new IcebergConnectorFactory());
+        return ImmutableList.of(new IcebergConnectorFactory(mBeanServer));
     }
 
     @Override

presto-iceberg/src/main/java/com/facebook/presto/iceberg/InternalIcebergConnectorFactory.java

@@ -57,7 +57,6 @@
 
 import javax.management.MBeanServer;
 
-import java.lang.management.ManagementFactory;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -72,7 +71,8 @@ public static Connector createConnector(
             String catalogName,
             Map<String, String> config,
             ConnectorContext context,
-            Optional<ExtendedHiveMetastore> metastore)
+            Optional<ExtendedHiveMetastore> metastore,
+            MBeanServer mBeanServer)
     {
         ClassLoader classLoader = InternalIcebergConnectorFactory.class.getClassLoader();
         try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
@@ -88,8 +88,7 @@ public static Connector createConnector(
                     new CachingModule(),
                     new HiveCommonModule(),
                     binder -> {
-                        MBeanServer platformMBeanServer = ManagementFactory.getPlatformMBeanServer();
-                        binder.bind(MBeanServer.class).toInstance(new RebindSafeMBeanServer(platformMBeanServer));
+                        binder.bind(MBeanServer.class).toInstance(new RebindSafeMBeanServer(mBeanServer));
                         binder.bind(NodeVersion.class).toInstance(new NodeVersion(context.getNodeManager().getCurrentNode().getVersion()));
                         binder.bind(NodeManager.class).toInstance(context.getNodeManager());
                         binder.bind(TypeManager.class).toInstance(context.getTypeManager());

presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergQueryRunner.java

@@ -38,6 +38,9 @@
 import com.google.common.collect.ImmutableSet;
 import io.airlift.tpch.TpchTable;
 
+import javax.management.MBeanServer;
+import javax.management.MBeanServerFactory;
+
 import java.io.IOException;
 import java.net.URI;
 import java.nio.file.Files;
@@ -197,7 +200,13 @@ public static DistributedQueryRunner createIcebergQueryRunner(
         queryRunner.installPlugin(new TpcdsPlugin());
         queryRunner.createCatalog("tpcds", "tpcds");
 
-        queryRunner.installPlugin(new IcebergPlugin());
+        queryRunner.getServers().forEach(server -> {
+            MBeanServer mBeanServer = MBeanServerFactory.newMBeanServer();
+            server.installPlugin(new IcebergPlugin(mBeanServer));
+            if (addJmxPlugin) {
+                server.installPlugin(new JmxPlugin(mBeanServer));
+            }
+        });
 
         String catalogType = extraConnectorProperties.getOrDefault("iceberg.catalog.type", HIVE.name());
         Path icebergDataDirectory = getIcebergDataDirectoryPath(queryRunner.getCoordinator().getDataDirectory(), catalogType, format, addStorageFormatToPath);
@@ -211,7 +220,6 @@ public static DistributedQueryRunner createIcebergQueryRunner(
         queryRunner.createCatalog(ICEBERG_CATALOG, "iceberg", icebergProperties);
 
         if (addJmxPlugin) {
-            queryRunner.installPlugin(new JmxPlugin());
             queryRunner.createCatalog("jmx", "jmx");
         }
 

presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergConnectorFactory.java

@@ -18,6 +18,7 @@
 import com.google.common.collect.ImmutableMap;
 import org.testng.annotations.Test;
 
+import java.lang.management.ManagementFactory;
 import java.util.Map;
 
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -38,7 +39,7 @@ public void testCachingHiveMetastore()
 
     private static void createConnector(Map<String, String> config)
     {
-        ConnectorFactory factory = new IcebergConnectorFactory();
+        ConnectorFactory factory = new IcebergConnectorFactory(ManagementFactory.getPlatformMBeanServer());
         factory.create("iceberg-test", config, new TestingConnectorContext())
                 .shutdown();
     }

presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergParquetMetadataCaching.java

@@ -45,7 +45,7 @@ protected QueryRunner createQueryRunner()
                 PARQUET,
                 false,
                 true,
-                OptionalInt.of(1),
+                OptionalInt.of(2),
                 Optional.empty());
     }