fix getting last account address

tronprotocol · Nov 15, 2023 · fbb7fac · fbb7fac
1 parent 7b08d38
commit fbb7fac
Show file tree

Hide file tree

Showing 5 changed files with 216 additions and 63 deletions.
diff --git a/chainbase/src/main/java/org/tron/core/db/common/iterator/DBIterator.java b/chainbase/src/main/java/org/tron/core/db/common/iterator/DBIterator.java
@@ -6,6 +6,7 @@
 import java.io.Closeable;
 import java.util.Iterator;
 import java.util.Map.Entry;
+import java.util.NoSuchElementException;
 
 public interface DBIterator extends Iterator<Entry<byte[], byte[]>>, AutoCloseable, Closeable {
 
@@ -20,4 +21,57 @@ public interface DBIterator extends Iterator<Entry<byte[], byte[]>>, AutoCloseab
     this.seek(afterThat == null ? key : afterThat);
     return Iterators.filter(this, entry -> Bytes.indexOf(entry.getKey(), key) == 0);
   }
+
+  /**
+   * An iterator is either positioned at a key/value pair, or
+   * not valid.  This method returns true iff the iterator is valid.
+   *
+   * REQUIRES: iterator not closed
+   *
+   * @throws IllegalStateException if the iterator is closed.
+   * @return an iterator is either positioned at a key/value pair
+   */
+  boolean valid();
+
+  /**
+   * The underlying storage for
+   * the returned slice is valid only until the next modification of
+   * the iterator.
+   *
+   * REQUIRES: valid() && !closed
+   *
+   * @throws IllegalStateException if the iterator is closed.
+   * @throws NoSuchElementException if the iterator is not valid.
+   *
+   * @return the key for the current entry
+   */
+  byte[] getKey();
+
+  /**
+   * The underlying storage for
+   * the returned slice is valid only until the next modification of
+   * the iterator.
+   *
+   * REQUIRES: valid() && !closed
+   *
+   * @throws IllegalStateException if the iterator is closed.
+   * @throws NoSuchElementException if the iterator is not valid.
+   *
+   * @return the value for the current entry
+   */
+  byte[] getValue();
+
+  /**
+   * @throws IllegalStateException if the iterator is closed.
+   */
+  void checkState();
+
+  /**
+   * @throws NoSuchElementException if the iterator is not valid.
+   */
+  default void checkValid() {
+    if (!valid()) {
+      throw new NoSuchElementException();
+    }
+  }
 }
diff --git a/chainbase/src/main/java/org/tron/core/db/common/iterator/RockStoreIterator.java b/chainbase/src/main/java/org/tron/core/db/common/iterator/RockStoreIterator.java
@@ -3,6 +3,7 @@
 import java.io.IOException;
 import java.util.Map.Entry;
 import java.util.NoSuchElementException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import lombok.extern.slf4j.Slf4j;
 import org.rocksdb.RocksIterator;
 
@@ -13,23 +14,22 @@ public final class RockStoreIterator implements DBIterator {
   private final RocksIterator dbIterator;
   private boolean first = true;
 
-  private boolean valid = true;
+  private final AtomicBoolean close = new AtomicBoolean(false);
 
   public RockStoreIterator(RocksIterator dbIterator) {
     this.dbIterator = dbIterator;
   }
 
   @Override
   public void close() throws IOException {
-    if (valid) {
+    if (close.compareAndSet(false, true)) {
       dbIterator.close();
-      valid = false;
     }
   }
 
   @Override
   public boolean hasNext() {
-    if (!valid) {
+    if (close.get()) {
       return false;
     }
     boolean hasNext = false;
@@ -40,13 +40,12 @@ public boolean hasNext() {
         first = false;
       }
       if (!(hasNext = dbIterator.isValid())) { // false is last item
-        dbIterator.close();
-        valid = false;
+        close();
       }
     } catch (Exception e) {
       logger.error(e.getMessage(), e);
       try {
-        dbIterator.close();
+        close();
       } catch (Exception e1) {
         logger.error(e.getMessage(), e);
       }
@@ -56,7 +55,7 @@ public boolean hasNext() {
 
   @Override
   public Entry<byte[], byte[]> next() {
-    if (!valid) {
+    if (close.get()) {
       throw new NoSuchElementException();
     }
     byte[] key = dbIterator.key();
@@ -82,16 +81,49 @@ public byte[] setValue(byte[] value) {
 
   @Override
   public void seek(byte[] key) {
+    checkState();
     dbIterator.seek(key);
+    this.first = false;
   }
 
   @Override
   public void seekToFirst() {
+    checkState();
     dbIterator.seekToFirst();
+    this.first = false;
   }
 
   @Override
   public void seekToLast() {
+    checkState();
     dbIterator.seekToLast();
+    this.first = false;
   }
-}
+
+  @Override
+  public boolean valid() {
+    checkState();
+    return dbIterator.isValid();
+  }
+
+  @Override
+  public byte[] getKey() {
+    checkState();
+    checkValid();
+    return dbIterator.key();
+  }
+
+  @Override
+  public byte[] getValue() {
+    checkState();
+    checkValid();
+    return dbIterator.value();
+  }
+
+  @Override
+  public void checkState() {
+    if (close.get()) {
+      throw new IllegalStateException("iterator has been closed");
+    }
+  }
+}
diff --git a/chainbase/src/main/java/org/tron/core/db/common/iterator/StoreIterator.java b/chainbase/src/main/java/org/tron/core/db/common/iterator/StoreIterator.java
@@ -3,6 +3,7 @@
 import java.io.IOException;
 import java.util.Map.Entry;
 import java.util.NoSuchElementException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import lombok.extern.slf4j.Slf4j;
 import org.iq80.leveldb.DBIterator;
 
@@ -13,23 +14,22 @@ public final class StoreIterator implements org.tron.core.db.common.iterator.DBI
   private final DBIterator dbIterator;
   private boolean first = true;
 
-  private boolean valid = true;
+  private final AtomicBoolean close = new AtomicBoolean(false);
 
   public StoreIterator(DBIterator dbIterator) {
     this.dbIterator = dbIterator;
   }
 
   @Override
   public void close() throws IOException {
-    if (valid) {
+    if (close.compareAndSet(false, true)) {
       dbIterator.close();
-      valid = false;
     }
   }
 
   @Override
   public boolean hasNext() {
-    if (!valid) {
+    if (close.get()) {
       return false;
     }
 
@@ -42,8 +42,7 @@ public boolean hasNext() {
       }
 
       if (!(hasNext = dbIterator.hasNext())) { // false is last item
-        dbIterator.close();
-        valid = false;
+        close();
       }
     } catch (Exception e) {
       logger.error(e.getMessage(), e);
@@ -54,7 +53,7 @@ public boolean hasNext() {
 
   @Override
   public Entry<byte[], byte[]> next() {
-    if (!valid) {
+    if (close.get()) {
       throw new NoSuchElementException();
     }
     return dbIterator.next();
@@ -67,16 +66,50 @@ public void remove() {
 
   @Override
   public void seek(byte[] key) {
+    checkState();
     dbIterator.seek(key);
+    this.first = false;
   }
 
   @Override
   public void seekToFirst() {
+    checkState();
     dbIterator.seekToFirst();
+    this.first = false;
   }
 
   @Override
   public void seekToLast() {
+    checkState();
     dbIterator.seekToLast();
+    this.first = false;
+  }
+
+  @Override
+  public boolean valid() {
+    checkState();
+    return dbIterator.hasNext();
+  }
+
+  @Override
+  public byte[] getKey() {
+    checkState();
+    checkValid();
+    return dbIterator.peekNext().getKey();
+  }
+
+  @Override
+  public byte[] getValue() {
+    checkState();
+    checkValid();
+    return dbIterator.peekNext().getValue();
+  }
+
+  @Override
+  public void checkState() {
+    if (close.get()) {
+      throw new IllegalStateException("iterator has been closed");
+    }
   }
 }
+
diff --git a/chainbase/src/main/java/org/tron/core/service/RewardCalService.java b/chainbase/src/main/java/org/tron/core/service/RewardCalService.java
@@ -75,13 +75,7 @@ private void destroy() {
   }
 
   public void calReward() throws IOException {
-    try (DBIterator iterator = rewardCacheStore.iterator()) {
-      iterator.seekToLast();
-      if (iterator.hasNext()) {
-        byte[] key  = iterator.next().getKey();
-        System.arraycopy(key, 0, lastAccount, 0, ADDRESS_SIZE);
-      }
-    }
+    initLastAccount();
     es.submit(this::startRewardCal);
   }
 
@@ -93,14 +87,18 @@ public void calRewardForTest() throws IOException {
       return;
     }
     accountIterator = (DBIterator) accountStore.getDb().iterator();
+    initLastAccount();
+    startRewardCal();
+  }
+
+  private void initLastAccount() throws IOException {
     try (DBIterator iterator = rewardCacheStore.iterator()) {
       iterator.seekToLast();
-      if (iterator.hasNext()) {
-        byte[] key  = iterator.next().getKey();
+      if (iterator.valid()) {
+        byte[] key  = iterator.getKey();
         System.arraycopy(key, 0, lastAccount, 0, ADDRESS_SIZE);
       }
     }
-    startRewardCal();
   }