(celeborn) 01/06: [CELEBORN-1315] Manually close the RocksDB/LevelDB instance when checkVersion throw Exception

[ni...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

nicholasjiang pushed a commit to branch branch-0.4
in repository https://gitbox.apache.org/repos/asf/celeborn.git

commit b67e2389f4cf0aa1746a5337bec8146b54d1cb36
Author: SteNicholas <pr...@163.com>
AuthorDate: Fri Mar 8 15:03:57 2024 +0800

    [CELEBORN-1315] Manually close the RocksDB/LevelDB instance when checkVersion throw Exception
    
    ### What changes were proposed in this pull request?
    
    Should close the `RocksDB`/`LevelDB` instance when `checkVersion` throw Exception.
    
    Backport [[SPARK-46389][CORE] Manually close the RocksDB/LevelDB instance when checkVersion throw Exception](https://github.com/apache/spark/pull/44327).
    
    ### Why are the changes needed?
    
    In the process of initializing the DB in `RocksDBProvider`/`LevelDBProvider`, there is a `checkVersion` step that may throw an exception. After the exception is thrown, the upper-level caller cannot hold the already opened RockDB/LevelDB instance, so it cannot perform resource cleanup, which poses a potential risk of handle leakage. So this PR manually closes the `RocksDB`/`LevelDB` instance when `checkVersion` throws an exception.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    CI.
    
    Closes #2369 from SteNicholas/CELEBORN-1315.
    
    Authored-by: SteNicholas <pr...@163.com>
    Signed-off-by: mingji <fe...@alibaba-inc.com>
---
 .../celeborn/service/deploy/worker/shuffledb/LevelDBProvider.java  | 7 ++++++-
 .../celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java  | 4 ++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDBProvider.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDBProvider.java
index 27734dca5..7bac20415 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDBProvider.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDBProvider.java
@@ -83,7 +83,12 @@ public class LevelDBProvider {
         }
       }
       // if there is a version mismatch, we throw an exception, which means the service is unusable
-      checkVersion(tmpDb, version);
+      try {
+        checkVersion(tmpDb, version);
+      } catch (IOException ioe) {
+        tmpDb.close();
+        throw ioe;
+      }
     }
     return tmpDb;
   }
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java
index 13d65a86e..4229008c0 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java
@@ -104,7 +104,11 @@ public class RocksDBProvider {
         // is unusable
         checkVersion(tmpDb, version);
       } catch (RocksDBException e) {
+        tmpDb.close();
         throw new IOException(e.getMessage(), e);
+      } catch (IOException ioe) {
+        tmpDb.close();
+        throw ioe;
       }
     }
     return tmpDb;