From ad2fb7cd365dd3bd8acc5f5b637a8f39197432b8 Mon Sep 17 00:00:00 2001
From: "gabry.wu" <gabrywu@apache.org>
Date: Tue, 19 Dec 2023 15:58:12 +0800
Subject: [PATCH] Core: Shutdown scheduler in Lock manager (#9150)

---
 .../org/apache/iceberg/util/LockManagers.java | 19 ++++++++++++++++++-
 .../iceberg/util/TestInMemoryLockManager.java |  2 +-
 2 files changed, 19 insertions(+), 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
index 2dc7406d92b6..96622cb57f83 100644
--- 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
index 7f7a39706517..b4d9296387ee 100644
--- 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();
   }