Core: Shutdown scheduler in Lock manager (apache#9150)

cdouglas · Apr 22, 2024 · ad2fb7c · ad2fb7c
1 parent 8e52c83
commit ad2fb7c
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/core/src/main/java/org/apache/iceberg/util/LockManagers.java b/core/src/main/java/org/apache/iceberg/util/LockManagers.java
@@ -18,9 +18,11 @@
  */
 package org.apache.iceberg.util;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
@@ -154,6 +156,20 @@ public void initialize(Map<String, String> properties) {
               CatalogProperties.LOCK_HEARTBEAT_THREADS,
               CatalogProperties.LOCK_HEARTBEAT_THREADS_DEFAULT);
     }
+
+    @Override
+    public void close() throws Exception {
+      if (scheduler != null) {
+        List<Runnable> tasks = scheduler.shutdownNow();
+        tasks.forEach(
+            task -> {
+              if (task instanceof Future) {
+                ((Future<?>) task).cancel(true);
+              }
+            });
+        scheduler = null;
+      }
+    }
   }
 
   /**
@@ -260,10 +276,11 @@ public boolean release(String entityId, String ownerId) {
     }
 
     @Override
-    public void close() {
+    public void close() throws Exception {
       HEARTBEATS.values().forEach(future -> future.cancel(false));
       HEARTBEATS.clear();
       LOCKS.clear();
+      super.close();
     }
   }
 

diff --git a/core/src/test/java/org/apache/iceberg/util/TestInMemoryLockManager.java b/core/src/test/java/org/apache/iceberg/util/TestInMemoryLockManager.java
@@ -49,7 +49,7 @@ public void before() {
   }
 
   @AfterEach
-  public void after() {
+  public void after() throws Exception {
     lockManager.close();
   }