HDDS-15288. Use SequenceIdType in StateManagerImpl.

hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/ha/SequenceIdGenerator.java

@@ -222,7 +222,7 @@ default RequestType getType() {
   static final class StateManagerImpl implements StateManager {
     private Table<String, Long> sequenceIdTable;
     private final DBTransactionBuffer transactionBuffer;
-    private final Map<String, Long> sequenceIdToLastIdMap;
+    private final Map<SequenceIdType, Long> sequenceIdToLastIdMap;
 
     private StateManagerImpl(Table<String, Long> sequenceIdTable,
                                DBTransactionBuffer trxBuffer) {
@@ -235,10 +235,11 @@ private StateManagerImpl(Table<String, Long> sequenceIdTable,
     @Override
     public Boolean allocateBatch(String sequenceIdName,
                                  Long expectedLastId, Long newLastId) {
-      Long lastId = sequenceIdToLastIdMap.computeIfAbsent(sequenceIdName,
+      SequenceIdType idType = SequenceIdType.valueOf(sequenceIdName);
+      Long lastId = sequenceIdToLastIdMap.computeIfAbsent(idType,
           key -> {
             try {
-              Long idInDb = this.sequenceIdTable.get(key);
+              Long idInDb = this.sequenceIdTable.get(key.name());
               return idInDb != null ? idInDb : INVALID_SEQUENCE_ID;
             } catch (IOException ioe) {
               throw new RuntimeException("Failed to get lastId from db", ioe);
@@ -253,18 +254,18 @@ public Boolean allocateBatch(String sequenceIdName,
 
       try {
         transactionBuffer
-            .addToBuffer(sequenceIdTable, sequenceIdName, newLastId);
+            .addToBuffer(sequenceIdTable, idType.name(), newLastId);
       } catch (IOException ioe) {
         throw new RuntimeException("Failed to put lastId to Batch", ioe);
       }
 
-      sequenceIdToLastIdMap.put(sequenceIdName, newLastId);
+      sequenceIdToLastIdMap.put(idType, newLastId);
       return true;
     }
 
     @Override
     public Long getLastId(String sequenceIdName) {
-      return sequenceIdToLastIdMap.get(sequenceIdName);
+      return sequenceIdToLastIdMap.get(SequenceIdType.valueOf(sequenceIdName));
     }
 
     @Override
@@ -286,7 +287,7 @@ private void initialize() throws IOException {
               "sequenceIdName should not be null");
           Objects.requireNonNull(lastId,
               "lastId should not be null");
-          sequenceIdToLastIdMap.put(sequenceIdName, lastId);
+          sequenceIdToLastIdMap.put(SequenceIdType.valueOf(sequenceIdName), lastId);
         }
       }
     }

hadoop-hdds/server-scm/src/test/java/org/apache/hadoop/hdds/scm/ha/TestSequenceIDGenerator.java

@@ -19,6 +19,8 @@
 
 import static org.apache.hadoop.hdds.scm.ScmConfigKeys.OZONE_SCM_SEQUENCE_ID_BATCH_SIZE;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -34,6 +36,7 @@
 import org.apache.hadoop.hdds.scm.metadata.SCMMetadataStoreImpl;
 import org.apache.hadoop.hdds.utils.db.Table;
 import org.apache.hadoop.ozone.container.common.SCMTestUtils;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
@@ -180,4 +183,92 @@ public StateManager createStateManager(
       }
     }
   }
+
+  @Test
+  public void testAllocateBatchFromDBWhenMissingInMap() throws Exception {
+    OzoneConfiguration conf = SCMTestUtils.getConf(testDir);
+    SCMMetadataStore scmMetadataStore = new SCMMetadataStoreImpl(conf);
+    scmMetadataStore.start(conf);
+    SCMHAManager scmHAManager = SCMHAManagerStub.getInstance(true);
+
+    // Create the StateManager directly using its Builder
+    SequenceIdGenerator.StateManager stateManager =
+        new SequenceIdGenerator.StateManagerImpl.Builder()
+            .setRatisServer(scmHAManager.getRatisServer())
+            .setDBTransactionBuffer(scmHAManager.getDBTransactionBuffer())
+            .setSequenceIdTable(scmMetadataStore.getSequenceIdTable())
+            .build();
+
+    String localIdKey = SequenceIdType.localId.name();
+
+    // Verify initial state from empty DB
+    Assertions.assertNull(stateManager.getLastId(localIdKey));
+
+    // Allocate a new batch, which puts 100L into the sequenceIdToLastIdMap map
+    assertTrue(stateManager.allocateBatch(localIdKey, 0L, 100L));
+    // Verify the map was updated
+    assertEquals(100L, stateManager.getLastId(localIdKey));
+
+    // Allocate a new batch, which puts 100L into the sequenceIdToLastIdMap map
+    assertTrue(stateManager.allocateBatch(localIdKey, 100L, 200L));
+    // Verify the map was updated
+    assertEquals(200L, stateManager.getLastId(localIdKey));
+
+    // This call should fail because expectedLastId in db should be (200L)
+    // But we are passing 0L
+    assertFalse(stateManager.allocateBatch(localIdKey, 0L, 100L));
+  }
+
+  @Test
+  public void testReinitializePopulatesSequenceIdMapFromDB() throws Exception {
+    OzoneConfiguration conf = SCMTestUtils.getConf(testDir);
+    SCMMetadataStore scmMetadataStore = new SCMMetadataStoreImpl(conf);
+    scmMetadataStore.start(conf);
+    SCMHAManager scmHAManager = SCMHAManagerStub.getInstance(true);
+
+    String containerId = SequenceIdType.containerId.name();
+    // Simulate an SCM restart by writing a raw String directly to the database.
+    scmMetadataStore.getSequenceIdTable().put(containerId, 100L);
+
+    // Create the StateManager directly using its Builder
+    SequenceIdGenerator.StateManager stateManager =
+        new SequenceIdGenerator.StateManagerImpl.Builder()
+            .setRatisServer(scmHAManager.getRatisServer())
+            .setDBTransactionBuffer(scmHAManager.getDBTransactionBuffer())
+            .setSequenceIdTable(scmMetadataStore.getSequenceIdTable())
+            .build();
+
+    // Check if reinitialize() correctly converts DB key into SequenceIdType Enums
+    // for the sequenceIdToLastIdMap used.
+    stateManager.reinitialize(scmMetadataStore.getSequenceIdTable());
+
+    assertEquals(100L, stateManager.getLastId(containerId));
+    assertTrue(stateManager.allocateBatch(containerId, 100L, 1100L));
+    assertEquals(1100L, stateManager.getLastId(containerId));
+  }
+
+  @Test
+  public void testAllocateBatchFailsOnUnknownSequenceId() throws Exception {
+    OzoneConfiguration conf = SCMTestUtils.getConf(testDir);
+    SCMMetadataStore scmMetadataStore = new SCMMetadataStoreImpl(conf);
+    scmMetadataStore.start(conf);
+    SCMHAManager scmHAManager = SCMHAManagerStub.getInstance(true);
+
+    // Create the StateManager directly using its Builder
+    SequenceIdGenerator.StateManager stateManager =
+        new SequenceIdGenerator.StateManagerImpl.Builder()
+            .setRatisServer(scmHAManager.getRatisServer())
+            .setDBTransactionBuffer(scmHAManager.getDBTransactionBuffer())
+            .setSequenceIdTable(scmMetadataStore.getSequenceIdTable())
+            .build();
+
+    try {
+      // sequenceIdName string must match one of the predefined Enums.
+      // Passing an invalid string should immediately throw an exception before hitting the db.
+      stateManager.allocateBatch("unknownSequenceId", 0L, 1L);
+      fail("Expected allocateBatch to reject an unknown sequence id");
+    } catch (Exception e) {
+      // ignore
+    }
+  }
 }