Fix leaking resources from guice modules

trinodb · Mar 18, 2024 · d67576e · d67576e
1 parent fbaff02
commit d67576e
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 29 deletions.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcModule.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcModule.java
@@ -31,14 +31,14 @@
 import io.trino.spi.connector.ConnectorSplitManager;
 import io.trino.spi.function.table.ConnectorTableFunction;
 import io.trino.spi.procedure.Procedure;
-import jakarta.annotation.PreDestroy;
 
 import java.util.concurrent.ExecutorService;
 
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.configuration.ConditionalModule.conditionalModule;
 import static io.airlift.configuration.ConfigBinder.configBinder;
+import static io.trino.plugin.base.ClosingBinder.closingBinder;
 import static org.weakref.jmx.guice.ExportBinder.newExporter;
 
 public class JdbcModule
@@ -103,6 +103,9 @@ public void setup(Binder binder)
                 .in(Scopes.SINGLETON);
 
         newSetBinder(binder, JdbcQueryEventListener.class);
+
+        closingBinder(binder)
+                .registerExecutor(ExecutorService.class, ForRecordCursor.class);
     }
 
     public static Multibinder<SessionPropertiesProvider> sessionPropertiesProviderBinder(Binder binder)
@@ -134,10 +137,4 @@ public static void bindTablePropertiesProvider(Binder binder, Class<? extends Ta
     {
         tablePropertiesProviderBinder(binder).addBinding().to(type).in(Scopes.SINGLETON);
     }
-
-    @PreDestroy
-    public void shutdownRecordCursorExecutor(@ForRecordCursor ExecutorService executor)
-    {
-        executor.shutdownNow();
-    }
 }
diff --git a/...trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreModule.java b/...trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreModule.java
@@ -26,13 +26,13 @@
 import io.trino.plugin.hive.ForHiveMetastore;
 import io.trino.plugin.hive.metastore.HiveMetastoreFactory;
 import io.trino.plugin.hive.metastore.RawHiveMetastoreFactory;
-import jakarta.annotation.PreDestroy;
 
 import java.util.concurrent.ExecutorService;
 
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.concurrent.Threads.threadsNamed;
 import static io.airlift.configuration.ConfigBinder.configBinder;
+import static io.trino.plugin.base.ClosingBinder.closingBinder;
 import static io.trino.plugin.base.security.UserNameProvider.SIMPLE_USER_NAME_PROVIDER;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.Executors.newFixedThreadPool;
@@ -91,12 +91,9 @@ protected void setup(Binder binder)
                 .setDefault()
                 .toInstance(SIMPLE_USER_NAME_PROVIDER);
         binder.bind(Key.get(boolean.class, AllowHiveTableRename.class)).toInstance(true);
-    }
 
-    @PreDestroy
-    public void shutdownsWriteStatisticExecutor(@ThriftHiveWriteStatisticsExecutor ExecutorService executor)
-    {
-        executor.shutdownNow();
+        closingBinder(binder)
+                .registerExecutor(ExecutorService.class, ThriftHiveWriteStatisticsExecutor.class);
     }
 
     private static class ThriftHiveMetastoreStatisticExecutorProvider

diff --git a/plugin/trino-kafka/src/main/java/io/trino/plugin/kafka/schema/confluent/ConfluentModule.java b/plugin/trino-kafka/src/main/java/io/trino/plugin/kafka/schema/confluent/ConfluentModule.java
@@ -57,9 +57,7 @@
 import io.trino.spi.HostAddress;
 import io.trino.spi.TrinoException;
 import io.trino.spi.type.TypeManager;
-import jakarta.annotation.PreDestroy;
 
-import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -75,6 +73,7 @@
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static io.airlift.configuration.ConditionalModule.conditionalModule;
 import static io.airlift.configuration.ConfigBinder.configBinder;
+import static io.trino.plugin.base.ClosingBinder.closingBinder;
 import static io.trino.plugin.kafka.encoder.EncoderModule.encoderFactory;
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.util.Objects.requireNonNull;
@@ -107,6 +106,9 @@ protected void setup(Binder binder)
         binder.bind(TableDescriptionSupplier.class).toProvider(ConfluentSchemaRegistryTableDescriptionSupplier.Factory.class).in(Scopes.SINGLETON);
         newMapBinder(binder, String.class, SchemaParser.class).addBinding("AVRO").to(AvroSchemaParser.class).in(Scopes.SINGLETON);
         newMapBinder(binder, String.class, SchemaParser.class).addBinding("PROTOBUF").to(LazyLoadedProtobufSchemaParser.class).in(Scopes.SINGLETON);
+
+        closingBinder(binder)
+                .registerCloseable(SchemaRegistryClient.class);
     }
 
     @Provides
@@ -138,13 +140,6 @@ public static SchemaRegistryClient createSchemaRegistryClient(
                 classLoader);
     }
 
-    @PreDestroy
-    public void destroy(SchemaRegistryClient client)
-            throws IOException
-    {
-        client.close();
-    }
-
     private class ConfluentDecoderModule
             implements Module
     {

diff --git a/plugin/trino-phoenix5/src/main/java/io/trino/plugin/phoenix5/PhoenixClientModule.java b/plugin/trino-phoenix5/src/main/java/io/trino/plugin/phoenix5/PhoenixClientModule.java
@@ -59,7 +59,6 @@
 import io.trino.spi.connector.ConnectorPageSinkProvider;
 import io.trino.spi.connector.ConnectorPageSourceProvider;
 import io.trino.spi.connector.ConnectorSplitManager;
-import jakarta.annotation.PreDestroy;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
 import org.apache.phoenix.jdbc.PhoenixDriver;
@@ -75,6 +74,7 @@
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.configuration.ConditionalModule.conditionalModule;
 import static io.airlift.configuration.ConfigBinder.configBinder;
+import static io.trino.plugin.base.ClosingBinder.closingBinder;
 import static io.trino.plugin.jdbc.JdbcModule.bindSessionPropertiesProvider;
 import static io.trino.plugin.jdbc.JdbcModule.bindTablePropertiesProvider;
 import static io.trino.plugin.phoenix5.ConfigurationInstantiator.newEmptyConfiguration;
@@ -147,6 +147,9 @@ protected void setup(Binder binder)
         install(new JdbcDiagnosticModule());
         install(new IdentifierMappingModule());
         install(new DecimalModule());
+
+        closingBinder(binder)
+                .registerExecutor(ExecutorService.class, ForRecordCursor.class);
     }
 
     private void checkConfiguration(String connectionUrl)
@@ -211,10 +214,4 @@ public ExecutorService createRecordCursorExecutor()
     {
         return newDirectExecutorService();
     }
-
-    @PreDestroy
-    public void shutdownRecordCursorExecutor(@ForRecordCursor ExecutorService executor)
-    {
-        executor.shutdownNow();
-    }
 }