Fix block location iteration with rocksdb

`setPrefixSameAsStart` does not work when there is a composite key, and both components need to be specified (2 longs). Therefore, use `setIterateUpperBound` instead. Fixes #11753 pr-link: #11752 change-id: cid-5c7c45fa69a5c0f0aa078d86baf229c4256d6a50
Alluxio · Jul 11, 2020 · 0aff938 · 0aff938
1 parent fec5421
commit 0aff938
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/core/server/master/src/main/java/alluxio/master/metastore/rocks/RocksBlockStore.java b/core/server/master/src/main/java/alluxio/master/metastore/rocks/RocksBlockStore.java
@@ -30,6 +30,7 @@
 import org.rocksdb.RocksDB;
 import org.rocksdb.RocksDBException;
 import org.rocksdb.RocksIterator;
+import org.rocksdb.Slice;
 import org.rocksdb.WriteOptions;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -147,9 +148,12 @@ public void close() {
 
   @Override
   public List<BlockLocation> getLocations(long id) {
+    byte[] startKey = RocksUtils.toByteArray(id, 0);
+    byte[] endKey = RocksUtils.toByteArray(id, Long.MAX_VALUE);
+
     try (RocksIterator iter = db().newIterator(mBlockLocationsColumn.get(),
-        new ReadOptions().setPrefixSameAsStart(true))) {
-      iter.seek(Longs.toByteArray(id));
+        new ReadOptions().setIterateUpperBound(new Slice(endKey)))) {
+      iter.seek(startKey);
       List<BlockLocation> locations = new ArrayList<>();
       for (; iter.isValid(); iter.next()) {
         try {

diff --git a/core/server/master/src/test/java/alluxio/master/metastore/rocks/RocksBlockStoreTest.java b/core/server/master/src/test/java/alluxio/master/metastore/rocks/RocksBlockStoreTest.java
@@ -23,6 +23,7 @@
 import org.junit.rules.TemporaryFolder;
 
 import java.util.Iterator;
+import java.util.List;
 
 public class RocksBlockStoreTest {
   @Rule
@@ -59,4 +60,24 @@ public void testIterator() throws Exception {
     }
     assertFalse(iter.hasNext());
   }
+
+  @Test
+  public void blockLocations() throws Exception {
+    final int blockCount = 5;
+    final int workerIdStart = 100000;
+    RocksBlockStore blockStore = new RocksBlockStore(mFolder.newFolder().getAbsolutePath());
+    // create blocks and locations
+    for (int i = 0; i < blockCount; i++) {
+      blockStore.putBlock(i, Block.BlockMeta.newBuilder().setLength(i).build());
+      blockStore
+          .addLocation(i, Block.BlockLocation.newBuilder().setWorkerId(workerIdStart + i).build());
+    }
+
+    // validate locations
+    for (int i = 0; i < blockCount; i++) {
+      List<Block.BlockLocation> locations = blockStore.getLocations(i);
+      assertEquals(1, locations.size());
+      assertEquals(workerIdStart + i, locations.get(0).getWorkerId());
+    }
+  }
 }