Merge pull request #4669 from halibobo1205/fix/db_iterator

fix(db): db iterator

Author: halibobo1205

chainbase/src/main/java/org/tron/core/db/common/iterator/RockStoreIterator.java

@@ -5,15 +5,16 @@
 import java.util.NoSuchElementException;
 import lombok.extern.slf4j.Slf4j;
 import org.rocksdb.RocksIterator;
-import org.tron.common.error.TronDBException;
 
 
 @Slf4j(topic = "DB")
 public final class RockStoreIterator implements DBIterator {
 
-  private RocksIterator dbIterator;
+  private final RocksIterator dbIterator;
   private boolean first = true;
 
+  private boolean valid = true;
+
   public RockStoreIterator(RocksIterator dbIterator) {
     this.dbIterator = dbIterator;
   }
@@ -25,6 +26,9 @@ public void close() throws IOException {
 
   @Override
   public boolean hasNext() {
+    if (!valid) {
+      return false;
+    }
     boolean hasNext = false;
     // true is first item
     try {
@@ -34,6 +38,7 @@ public boolean hasNext() {
       }
       if (!(hasNext = dbIterator.isValid())) { // false is last item
         dbIterator.close();
+        valid = false;
       }
     } catch (Exception e) {
       logger.error(e.getMessage(), e);
@@ -48,7 +53,7 @@ public boolean hasNext() {
 
   @Override
   public Entry<byte[], byte[]> next() {
-    if (!dbIterator.isValid()) {
+    if (!valid) {
       throw new NoSuchElementException();
     }
     byte[] key = dbIterator.key();

chainbase/src/main/java/org/tron/core/db/common/iterator/StoreIterator.java

@@ -2,16 +2,19 @@
 
 import java.io.IOException;
 import java.util.Map.Entry;
+import java.util.NoSuchElementException;
 import lombok.extern.slf4j.Slf4j;
 import org.iq80.leveldb.DBIterator;
-import org.tron.common.error.TronDBException;
+
 
 @Slf4j(topic = "DB")
 public final class StoreIterator implements org.tron.core.db.common.iterator.DBIterator {
 
-  private DBIterator dbIterator;
+  private final DBIterator dbIterator;
   private boolean first = true;
 
+  private boolean valid = true;
+
   public StoreIterator(DBIterator dbIterator) {
     this.dbIterator = dbIterator;
   }
@@ -23,6 +26,10 @@ public void close() throws IOException {
 
   @Override
   public boolean hasNext() {
+    if (!valid) {
+      return false;
+    }
+
     boolean hasNext = false;
     // true is first item
     try {
@@ -33,6 +40,7 @@ public boolean hasNext() {
 
       if (!(hasNext = dbIterator.hasNext())) { // false is last item
         dbIterator.close();
+        valid = false;
       }
     } catch (Exception e) {
       logger.error(e.getMessage(), e);
@@ -43,6 +51,9 @@ public boolean hasNext() {
 
   @Override
   public Entry<byte[], byte[]> next() {
+    if (!valid) {
+      throw new NoSuchElementException();
+    }
     return dbIterator.next();
   }
 

framework/src/test/java/org/tron/core/db/DBIteratorTest.java

@@ -0,0 +1,90 @@
+package org.tron.core.db;
+
+import static org.fusesource.leveldbjni.JniDBFactory.factory;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Paths;
+import java.util.NoSuchElementException;
+import org.iq80.leveldb.DB;
+import org.iq80.leveldb.Options;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.rocksdb.RocksDB;
+import org.rocksdb.RocksDBException;
+import org.tron.common.utils.FileUtil;
+import org.tron.core.db.common.iterator.RockStoreIterator;
+import org.tron.core.db.common.iterator.StoreIterator;
+
+public class DBIteratorTest {
+
+  @BeforeClass
+  public static void init() {
+    File file = Paths.get("database-iterator").toFile();
+    if (!file.exists()) {
+      file.mkdirs();
+    }
+  }
+
+  @AfterClass
+  public static void clear() {
+    File file = Paths.get("database-iterator").toFile();
+    if (file.exists()) {
+      FileUtil.deleteDir(Paths.get("database-iterator").toFile());
+    }
+  }
+
+
+  @Test
+  public void testLevelDb() throws IOException {
+    File file = new File("database-iterator/testLevelDb");
+    try {
+      DB db = factory.open(file, new Options().createIfMissing(true));
+      db.put("1".getBytes(StandardCharsets.UTF_8), "1".getBytes(StandardCharsets.UTF_8));
+      db.put("2".getBytes(StandardCharsets.UTF_8), "2".getBytes(StandardCharsets.UTF_8));
+      StoreIterator  iterator = new StoreIterator(db.iterator());
+      while (iterator.hasNext()) {
+        iterator.next();
+      }
+      Assert.assertFalse(iterator.hasNext());
+      try {
+        iterator.next();
+      } catch (Exception e) {
+        Assert.assertTrue(e instanceof  NoSuchElementException);
+      }
+      db.close();
+    } finally {
+      factory.destroy(file, new Options());
+    }
+
+
+  }
+
+  @Test
+  public void testRocksDb() throws RocksDBException {
+    File file = new File("database-iterator/testRocksDb");
+    try (org.rocksdb.Options options = new org.rocksdb.Options().setCreateIfMissing(true)) {
+      RocksDB db = RocksDB.open(options, file.toString());
+      db.put("1".getBytes(StandardCharsets.UTF_8), "1".getBytes(StandardCharsets.UTF_8));
+      db.put("2".getBytes(StandardCharsets.UTF_8), "2".getBytes(StandardCharsets.UTF_8));
+      RockStoreIterator iterator = new RockStoreIterator(db.newIterator());
+      while (iterator.hasNext()) {
+        iterator.next();
+      }
+      Assert.assertFalse(iterator.hasNext());
+      try {
+        iterator.next();
+      } catch (Exception e) {
+        Assert.assertTrue(e instanceof  NoSuchElementException);
+      }
+      db.close();
+    } finally {
+      RocksDB.destroyDB(file.toString(), new org.rocksdb.Options());
+    }
+
+  }
+
+}