fix: Every time use KeyValueStorage.Batch, close it in finally (#4578)

Co-authored-by: nuolin <[email protected]>

Author: StevenLuMT

bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/storage/ldb/SingleDirectoryDbLedgerStorage.java

@@ -827,23 +827,24 @@ public void checkpoint(Checkpoint checkpoint) throws IOException {
             // Write all the pending entries into the entry logger and collect the offset
             // position for each entry
 
-            Batch batch = entryLocationIndex.newBatch();
-            writeCacheBeingFlushed.forEach((ledgerId, entryId, entry) -> {
-                long location = entryLogger.addEntry(ledgerId, entry);
-                entryLocationIndex.addLocation(batch, ledgerId, entryId, location);
-            });
+            try (Batch batch = entryLocationIndex.newBatch()) {
+                writeCacheBeingFlushed.forEach((ledgerId, entryId, entry) -> {
+                    long location = entryLogger.addEntry(ledgerId, entry);
+                    entryLocationIndex.addLocation(batch, ledgerId, entryId, location);
+                });
 
-            long entryLoggerStart = MathUtils.nowInNano();
-            entryLogger.flush();
-            recordSuccessfulEvent(dbLedgerStorageStats.getFlushEntryLogStats(), entryLoggerStart);
+                long entryLoggerStart = MathUtils.nowInNano();
+                entryLogger.flush();
+                recordSuccessfulEvent(dbLedgerStorageStats.getFlushEntryLogStats(), entryLoggerStart);
 
-            long batchFlushStartTime = MathUtils.nowInNano();
-            batch.flush();
-            batch.close();
-            recordSuccessfulEvent(dbLedgerStorageStats.getFlushLocationIndexStats(), batchFlushStartTime);
-            if (log.isDebugEnabled()) {
-                log.debug("DB batch flushed time : {} s",
-                        MathUtils.elapsedNanos(batchFlushStartTime) / (double) TimeUnit.SECONDS.toNanos(1));
+                long batchFlushStartTime = MathUtils.nowInNano();
+                batch.flush();
+
+                recordSuccessfulEvent(dbLedgerStorageStats.getFlushLocationIndexStats(), batchFlushStartTime);
+                if (log.isDebugEnabled()) {
+                    log.debug("DB batch flushed time : {} s",
+                            MathUtils.elapsedNanos(batchFlushStartTime) / (double) TimeUnit.SECONDS.toNanos(1));
+                }
             }
 
             long ledgerIndexStartTime = MathUtils.nowInNano();
@@ -1095,20 +1096,20 @@ public long addLedgerToIndex(long ledgerId, boolean isFenced, byte[] masterKey,
         MutableLong numberOfEntries = new MutableLong();
 
         // Iterate over all the entries pages
-        Batch batch = entryLocationIndex.newBatch();
-        for (LedgerCache.PageEntries page: pages) {
-            try (LedgerEntryPage lep = page.getLEP()) {
-                lep.getEntries((entryId, location) -> {
-                    entryLocationIndex.addLocation(batch, ledgerId, entryId, location);
-                    numberOfEntries.increment();
-                    return true;
-                });
+        try (Batch batch = entryLocationIndex.newBatch()) {
+            for (LedgerCache.PageEntries page : pages) {
+                try (LedgerEntryPage lep = page.getLEP()) {
+                    lep.getEntries((entryId, location) -> {
+                        entryLocationIndex.addLocation(batch, ledgerId, entryId, location);
+                        numberOfEntries.increment();
+                        return true;
+                    });
+                }
             }
-        }
 
-        ledgerIndex.flush();
-        batch.flush();
-        batch.close();
+            ledgerIndex.flush();
+            batch.flush();
+        }
 
         return numberOfEntries.longValue();
     }

bookkeeper-server/src/test/java/org/apache/bookkeeper/bookie/storage/ldb/EntryLocationIndexTest.java

@@ -93,16 +93,15 @@ public void deleteBatchLedgersTest() throws Exception {
         int numEntriesPerLedger = 100;
 
         int location = 0;
-        KeyValueStorage.Batch batch = idx.newBatch();
-        for (int entryId = 0; entryId < numEntriesPerLedger; ++entryId) {
-            for (int ledgerId = 0; ledgerId < numLedgers; ++ledgerId) {
-                idx.addLocation(batch, ledgerId, entryId, location);
-                location++;
+        try (KeyValueStorage.Batch batch = idx.newBatch()) {
+            for (int entryId = 0; entryId < numEntriesPerLedger; ++entryId) {
+                for (int ledgerId = 0; ledgerId < numLedgers; ++ledgerId) {
+                    idx.addLocation(batch, ledgerId, entryId, location);
+                    location++;
+                }
             }
+            batch.flush();
         }
-        batch.flush();
-        batch.close();
-
 
         int expectedLocation = 0;
         for (int entryId = 0; entryId < numEntriesPerLedger; ++entryId) {

bookkeeper-server/src/test/java/org/apache/bookkeeper/bookie/storage/ldb/KeyValueStorageTest.java

@@ -156,16 +156,16 @@ public void simple() throws Exception {
         assertEquals(null, db.get(toArray(10)));
         assertTrue(db.count() > 0);
 
-        Batch batch = db.newBatch();
-        batch.remove(toArray(11));
-        batch.remove(toArray(12));
-        batch.remove(toArray(13));
-        batch.flush();
-        assertEquals(null, db.get(toArray(11)));
-        assertEquals(null, db.get(toArray(12)));
-        assertEquals(null, db.get(toArray(13)));
-        assertEquals(14L, fromArray(db.get(toArray(14))));
-        batch.close();
+        try (Batch batch = db.newBatch()) {
+            batch.remove(toArray(11));
+            batch.remove(toArray(12));
+            batch.remove(toArray(13));
+            batch.flush();
+            assertEquals(null, db.get(toArray(11)));
+            assertEquals(null, db.get(toArray(12)));
+            assertEquals(null, db.get(toArray(13)));
+            assertEquals(14L, fromArray(db.get(toArray(14))));
+        }
 
         db.close();
         FileUtils.deleteDirectory(tmpDir);