Reuse ManagedLedgerFactory instances across SQL queries (#4813)

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/PulsarConnector.java

@@ -86,11 +86,6 @@ public class PulsarConnector implements Connector {
         } catch (Exception e) {
             log.error(e, "Failed to close pulsar connector");
         }
-        try {
-            PulsarConnectorCache.shutdown();
-        } catch (Exception e) {
-            log.error("Failed to shutdown pulsar connector cache");
-        }
         try {
             lifeCycleManager.stop();
         } catch (Exception e) {

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/PulsarConnectorCache.java

@@ -39,11 +39,14 @@ import java.util.Map;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.common.annotations.VisibleForTesting;
+
 public class PulsarConnectorCache {
 
     private static final Logger log = Logger.get(PulsarConnectorCache.class);
 
-    private static PulsarConnectorCache instance;
+    @VisibleForTesting
+    static PulsarConnectorCache instance;
 
     private final ManagedLedgerFactory managedLedgerFactory;
 

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/PulsarSplitManager.java

@@ -133,16 +133,6 @@ public class PulsarSplitManager implements ConnectorSplitManager {
         return new FixedSplitSource(splits);
     }
 
-    @VisibleForTesting
-    ManagedLedgerFactory getManagedLedgerFactory() throws Exception {
-        ClientConfiguration bkClientConfiguration = new ClientConfiguration()
-                .setZkServers(this.pulsarConnectorConfig.getZookeeperUri())
-                .setClientTcpNoDelay(false)
-                .setStickyReadsEnabled(false)
-                .setUseV2WireProtocol(true);
-        return new ManagedLedgerFactoryImpl(bkClientConfiguration);
-    }
-
     @VisibleForTesting
     Collection<PulsarSplit> getSplitsPartitionedTopic(int numSplits, TopicName topicName, PulsarTableHandle
             tableHandle, SchemaInfo schemaInfo, TupleDomain<ColumnHandle> tupleDomain) throws Exception {
@@ -165,59 +155,40 @@ public class PulsarSplitManager implements ConnectorSplitManager {
 
         int splitRemainder = actualNumSplits % numPartitions;
 
-        ManagedLedgerFactory managedLedgerFactory = getManagedLedgerFactory();
-
-        try {
-            List<PulsarSplit> splits = new LinkedList<>();
-            for (int i = 0; i < numPartitions; i++) {
-
-                int splitsForThisPartition = (splitRemainder > i) ? splitsPerPartition + 1 : splitsPerPartition;
-                splits.addAll(
-                        getSplitsForTopic(
-                                topicName.getPartition(i).getPersistenceNamingEncoding(),
-                                managedLedgerFactory,
-                                splitsForThisPartition,
-                                tableHandle,
-                                schemaInfo,
-                                topicName.getPartition(i).getLocalName(),
-                                tupleDomain)
-                );
-            }
-            return splits;
-        } finally {
-            if (managedLedgerFactory != null) {
-                try {
-                    managedLedgerFactory.shutdown();
-                } catch (Exception e) {
-                    log.error(e);
-                }
-            }
+        ManagedLedgerFactory managedLedgerFactory = PulsarConnectorCache.getConnectorCache(pulsarConnectorConfig)
+                .getManagedLedgerFactory();
+
+        List<PulsarSplit> splits = new LinkedList<>();
+        for (int i = 0; i < numPartitions; i++) {
+
+            int splitsForThisPartition = (splitRemainder > i) ? splitsPerPartition + 1 : splitsPerPartition;
+            splits.addAll(
+                    getSplitsForTopic(
+                            topicName.getPartition(i).getPersistenceNamingEncoding(),
+                            managedLedgerFactory,
+                            splitsForThisPartition,
+                            tableHandle,
+                            schemaInfo,
+                            topicName.getPartition(i).getLocalName(),
+                            tupleDomain));
         }
+        return splits;
     }
 
     @VisibleForTesting
-    Collection<PulsarSplit> getSplitsNonPartitionedTopic(int numSplits, TopicName topicName, PulsarTableHandle
-            tableHandle, SchemaInfo schemaInfo, TupleDomain<ColumnHandle> tupleDomain) throws Exception {
-        ManagedLedgerFactory managedLedgerFactory = null;
-        try {
-            managedLedgerFactory = getManagedLedgerFactory();
-
-            return getSplitsForTopic(
-                    topicName.getPersistenceNamingEncoding(),
-                    managedLedgerFactory,
-                    numSplits,
-                    tableHandle,
-                    schemaInfo,
-                    tableHandle.getTableName(), tupleDomain);
-        } finally {
-            if (managedLedgerFactory != null) {
-                try {
-                    managedLedgerFactory.shutdown();
-                } catch (Exception e) {
-                    log.error(e);
-                }
-            }
-        }
+    Collection<PulsarSplit> getSplitsNonPartitionedTopic(int numSplits, TopicName topicName,
+            PulsarTableHandle tableHandle, SchemaInfo schemaInfo, TupleDomain<ColumnHandle> tupleDomain)
+            throws Exception {
+        ManagedLedgerFactory managedLedgerFactory = PulsarConnectorCache.getConnectorCache(pulsarConnectorConfig)
+                .getManagedLedgerFactory();
+
+        return getSplitsForTopic(
+                topicName.getPersistenceNamingEncoding(),
+                managedLedgerFactory,
+                numSplits,
+                tableHandle,
+                schemaInfo,
+                tableHandle.getTableName(), tupleDomain);
     }
 
     @VisibleForTesting

pulsar-sql/presto-pulsar/src/test/java/org/apache/pulsar/sql/presto/TestPulsarConnector.java

@@ -920,7 +920,8 @@ public abstract class TestPulsarConnector {
             }
         });
 
-        doReturn(managedLedgerFactory).when(this.pulsarSplitManager).getManagedLedgerFactory();
+        PulsarConnectorCache.instance = mock(PulsarConnectorCache.class);
+        when(PulsarConnectorCache.instance.getManagedLedgerFactory()).thenReturn(managedLedgerFactory);
 
         for (Map.Entry<TopicName, PulsarSplit> split : splits.entrySet()) {