Close AWS SDK client in Glue v1

trinodb · May 10, 2024 · 3207453 · 3207453
1 parent f55c040
commit 3207453
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/ClosingBinder.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/ClosingBinder.java
@@ -17,7 +17,9 @@
 import com.google.inject.Binder;
 import com.google.inject.BindingAnnotation;
 import com.google.inject.Inject;
+import com.google.inject.Injector;
 import com.google.inject.Key;
+import com.google.inject.Provider;
 import com.google.inject.multibindings.Multibinder;
 import io.trino.plugin.base.util.AutoCloseableCloser;
 import jakarta.annotation.PreDestroy;
@@ -26,7 +28,9 @@
 import java.lang.annotation.Target;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
+import java.util.function.Consumer;
 
+import static com.google.common.base.Verify.verifyNotNull;
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static java.lang.annotation.ElementType.FIELD;
 import static java.lang.annotation.ElementType.METHOD;
@@ -71,6 +75,45 @@ public void registerCloseable(Key<? extends AutoCloseable> key)
         closeables.addBinding().to(key);
     }
 
+    public <T> void registerResource(Class<T> type, Consumer<? super T> close)
+    {
+        registerResource(Key.get(type), close);
+    }
+
+    public <T> void registerResource(Key<T> key, Consumer<? super T> close)
+    {
+        closeables.addBinding().toProvider(new ResourceCloser<T>(key, close));
+    }
+
+    private static class ResourceCloser<T>
+            implements Provider<AutoCloseable>
+    {
+        private final Key<T> key;
+        private final Consumer<? super T> close;
+        private Injector injector;
+
+        private ResourceCloser(Key<T> key, Consumer<? super T> close)
+        {
+            this.key = requireNonNull(key, "key is null");
+            this.close = requireNonNull(close, "close is null");
+        }
+
+        @Inject
+        public void setInjector(Injector injector)
+        {
+            this.injector = injector;
+        }
+
+        @Override
+        public AutoCloseable get()
+        {
+            T object = injector.getInstance(key);
+            verifyNotNull(object, "null at key %s", key);
+            Consumer<? super T> close = this.close;
+            return () -> close.accept(object);
+        }
+    }
+
     private record Cleanup(
             @ForCleanup Set<ExecutorService> executors,
             @ForCleanup Set<AutoCloseable> closeables)

diff --git a/.../trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/v1/GlueMetastoreModule.java b/.../trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/v1/GlueMetastoreModule.java
@@ -78,6 +78,7 @@ protected void setup(Binder binder)
         binder.bind(GlueHiveMetastoreFactory.class).in(Scopes.SINGLETON);
         binder.bind(Key.get(GlueMetastoreStats.class, ForGlueHiveMetastore.class)).toInstance(new GlueMetastoreStats());
         binder.bind(AWSGlueAsync.class).toProvider(HiveGlueClientProvider.class).in(Scopes.SINGLETON);
+        closingBinder(binder).registerResource(AWSGlueAsync.class, AWSGlueAsync::shutdown);
         newExporter(binder).export(GlueHiveMetastore.class).withGeneratedName();
 
         binder.bind(Key.get(boolean.class, AllowHiveTableRename.class)).toInstance(false);