HDDS-12368. Seek to correct start key in KeyManagerImpl#getTableEntries

hadoop-hdds/framework/src/test/java/org/apache/hadoop/hdds/utils/MapBackedTableIterator.java

@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.hdds.utils;
+
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.hadoop.hdds.utils.db.Table;
+import org.apache.hadoop.hdds.utils.db.TableIterator;
+
+/**
+ * Generic Table Iterator implementation that can be used for unit tests to reduce redundant mocking in tests.
+ */
+public class MapBackedTableIterator<V> implements TableIterator<String, Table.KeyValue<String, V>> {
+
+  private Iterator<Table.KeyValue<String, V>> itr;
+  private final String prefix;
+  private final TreeMap<String, V> values;
+
+  public MapBackedTableIterator(TreeMap<String, V> values, String prefix) {
+    this.prefix = prefix;
+    this.values = values;
+    this.seekToFirst();
+  }
+
+  @Override
+  public void seekToFirst() {
+    this.itr = this.values.entrySet().stream()
+        .filter(e -> prefix == null || e.getKey().startsWith(prefix))
+        .map(e -> Table.newKeyValue(e.getKey(), e.getValue())).iterator();
+  }
+
+  @Override
+  public void seekToLast() {
+
+  }
+
+  @Override
+  public Table.KeyValue<String, V> seek(String s) throws IOException {
+    this.itr = this.values.entrySet().stream()
+        .filter(e -> prefix == null || e.getKey().startsWith(prefix))
+        .filter(e -> e.getKey().compareTo(s) >= 0)
+        .map(e -> Table.newKeyValue(e.getKey(), e.getValue())).iterator();
+    Map.Entry<String, V> firstEntry = values.ceilingEntry(s);
+    return firstEntry == null ? null : Table.newKeyValue(firstEntry.getKey(), firstEntry.getValue());
+  }
+
+  @Override
+  public void removeFromDB() throws IOException {
+
+  }
+
+  @Override
+  public void close() throws IOException {
+
+  }
+
+  @Override
+  public boolean hasNext() {
+    return this.itr.hasNext();
+  }
+
+  @Override
+  public Table.KeyValue<String, V> next() {
+    return itr.next();
+  }
+}

hadoop-ozone/ozone-manager/pom.xml

@@ -362,6 +362,12 @@
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.ozone</groupId>
+      <artifactId>hdds-server-framework</artifactId>
+      <type>test-jar</type>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>org.apache.ozone</groupId>
       <artifactId>hdds-server-scm</artifactId>

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/KeyManagerImpl.java

@@ -696,6 +696,7 @@ private <V, R> List<Table.KeyValue<String, R>> getTableEntries(String startKey,
     */
     if (startKey != null) {
       tableIterator.seek(startKey);
+    } else {
       tableIterator.seekToFirst();
     }
     int currentCount = 0;

hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/TestKeyManagerImpl.java

@@ -0,0 +1,214 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.ozone.om;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.hadoop.hdds.conf.OzoneConfiguration;
+import org.apache.hadoop.hdds.utils.MapBackedTableIterator;
+import org.apache.hadoop.hdds.utils.db.Table;
+import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
+import org.apache.hadoop.ozone.om.helpers.RepeatedOmKeyInfo;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.mockito.Mockito;
+
+/**
+ * Test class for unit tests KeyManagerImpl.
+ */
+public class TestKeyManagerImpl {
+  private static Stream<Arguments> getTableIteratorParameters() {
+    return Stream.of(
+        Arguments.argumentSet("Fetch first 50 entries for volume 0, bucket 0",
+            5, 10, 100, 0, 0, 0, 0, 0, 50, null),
+        Arguments.argumentSet("Fetch first 50 entries for any volume/bucket", 5, 10, 100, null, null, 0, 0, 0, 50,
+            null),
+        Arguments.argumentSet("Fetch first 30 entries for volume 1, bucket 1", 5, 10, 100, 1, 1, 0, 0, 0, 30, null),
+        Arguments.argumentSet("Fetch 20 entries from offset (2,2,10) for volume 2, bucket 2", 5, 10, 100, 2, 2, 2, 2,
+            10, 20, null),
+        Arguments.argumentSet("Fetch 40 entries from offset (2,2,50) for volume 3, bucket 3", 5, 10, 100, 3, 3, 3, 3,
+            50, 40, null),
+        Arguments.argumentSet("Fetch 200 entries from the very beginning (null start offsets)", 5, 10, 100, null,
+            null, null, null, null, 200, null),
+        Arguments.argumentSet("Fetch 200 entries starting from bucket 3, key 50, spanning 3 buckets", 5, 10, 100,
+            null, null, 0, 3, 50, 200, null),
+        Arguments.argumentSet("Invalid: bucket is set but volume is null", 5, 10, 100, null, 1, 0, 0, 0, 10,
+            IOException.class),
+        Arguments.argumentSet("Invalid: volume is set but bucket is null", 5, 10, 100, 1, null, 0, 0, 0, 10,
+            IOException.class),
+        Arguments.argumentSet("Fetch 50 entries from volume 2, bucket 5, but only 31 exist", 5, 10, 100, 2, 5, 2, 5,
+            70, 50, null),
+        Arguments.argumentSet("Start from last volume (4), second-last bucket (8), key 80 but only 131 entries exist",
+            5, 10, 100, null, null, 4, 8, 80, 200, null)
+    );
+  }
+
+  @SuppressWarnings({"checkstyle:ParameterNumber"})
+  private <V> List<Table.KeyValue<String, V>> mockTableIterator(
+      Class<V> valueClass, Table<String, V> table, int numberOfVolumes, int numberOfBucketsPerVolume,
+      int numberOfKeysPerBucket, String volumeNamePrefix, String bucketNamePrefix, String keyPrefix,
+      Integer volumeNumberFilter, Integer bucketNumberFilter, Integer startVolumeNumber, Integer startBucketNumber,
+      Integer startKeyNumber, int numberOfEntries) throws IOException {
+    TreeMap<String, V> values = new TreeMap<>();
+    List<Table.KeyValue<String, V>> keyValues = new ArrayList<>();
+    String startKey = startVolumeNumber == null || startBucketNumber == null || startKeyNumber == null ? null
+        : (String.format("/%s%010d/%s%010d/%s%010d", volumeNamePrefix, startVolumeNumber, bucketNamePrefix,
+        startBucketNumber, keyPrefix, startKeyNumber));
+    for (int i = 0; i < numberOfVolumes; i++) {
+      for (int j = 0; j < numberOfBucketsPerVolume; j++) {
+        for (int k = 0; k < numberOfKeysPerBucket; k++) {
+          String key = String.format("/%s%010d/%s%010d/%s%010d", volumeNamePrefix, i, bucketNamePrefix, j,
+              keyPrefix, k);
+          V value = valueClass == String.class ? (V) key : Mockito.mock(valueClass);
+          values.put(key, value);
+
+          if ((volumeNumberFilter == null || i == volumeNumberFilter) &&
+              (bucketNumberFilter == null || j == bucketNumberFilter) &&
+              (startKey == null || startKey.compareTo(key) <= 0)) {
+            keyValues.add(Table.newKeyValue(key, value));
+          }
+        }
+      }
+    }
+
+    when(table.iterator(anyString())).thenAnswer(i -> new MapBackedTableIterator<>(values, i.getArgument(0)));
+    return keyValues.subList(0, Math.min(numberOfEntries, keyValues.size()));
+  }
+
+  @ParameterizedTest
+  @MethodSource("getTableIteratorParameters")
+  @SuppressWarnings({"checkstyle:ParameterNumber"})
+  public void testGetDeletedKeyEntries(int numberOfVolumes, int numberOfBucketsPerVolume, int numberOfKeysPerBucket,
+                                       Integer volumeNumber, Integer bucketNumber,
+                                       Integer startVolumeNumber, Integer startBucketNumber, Integer startKeyNumber,
+                                       int numberOfEntries, Class<? extends Exception> expectedException)
+      throws IOException {
+    String volumeNamePrefix = "volume";
+    String bucketNamePrefix = "bucket";
+    String keyPrefix = "key";
+    OzoneConfiguration configuration = new OzoneConfiguration();
+    OMMetadataManager metadataManager = Mockito.mock(OMMetadataManager.class);
+    when(metadataManager.getBucketKeyPrefix(anyString(), anyString())).thenAnswer(i ->
+        "/" + i.getArguments()[0] + "/" + i.getArguments()[1] + "/");
+    KeyManagerImpl km = new KeyManagerImpl(null, null, metadataManager, configuration, null, null, null);
+    Table<String, RepeatedOmKeyInfo> mockedDeletedTable = Mockito.mock(Table.class);
+    when(metadataManager.getDeletedTable()).thenReturn(mockedDeletedTable);
+    List<Table.KeyValue<String, List<OmKeyInfo>>> expectedEntries = mockTableIterator(
+        RepeatedOmKeyInfo.class, mockedDeletedTable, numberOfVolumes, numberOfBucketsPerVolume, numberOfKeysPerBucket,
+        volumeNamePrefix, bucketNamePrefix, keyPrefix, volumeNumber, bucketNumber, startVolumeNumber, startBucketNumber,
+        startKeyNumber, numberOfEntries).stream()
+        .map(kv -> {
+          try {
+            String key = kv.getKey();
+            RepeatedOmKeyInfo value = kv.getValue();
+            List<OmKeyInfo> omKeyInfos = Collections.singletonList(Mockito.mock(OmKeyInfo.class));
+            when(value.cloneOmKeyInfoList()).thenReturn(omKeyInfos);
+            return Table.newKeyValue(key, omKeyInfos);
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        }).collect(Collectors.toList());
+    String volumeName = volumeNumber == null ? null : (String.format("%s%010d", volumeNamePrefix, volumeNumber));
+    String bucketName = bucketNumber == null ? null : (String.format("%s%010d", bucketNamePrefix, bucketNumber));
+    String startKey = startVolumeNumber == null || startBucketNumber == null || startKeyNumber == null ? null
+        : (String.format("/%s%010d/%s%010d/%s%010d", volumeNamePrefix, startVolumeNumber, bucketNamePrefix,
+        startBucketNumber, keyPrefix, startKeyNumber));
+    if (expectedException != null) {
+      assertThrows(expectedException, () -> km.getDeletedKeyEntries(volumeName, bucketName, startKey, numberOfEntries));
+    } else {
+      assertEquals(expectedEntries, km.getDeletedKeyEntries(volumeName, bucketName, startKey, numberOfEntries));
+    }
+  }
+
+  @ParameterizedTest
+  @MethodSource("getTableIteratorParameters")
+  @SuppressWarnings({"checkstyle:ParameterNumber"})
+  public void testGetRenameKeyEntries(int numberOfVolumes, int numberOfBucketsPerVolume, int numberOfKeysPerBucket,
+                                      Integer volumeNumber, Integer bucketNumber,
+                                      Integer startVolumeNumber, Integer startBucketNumber, Integer startKeyNumber,
+                                      int numberOfEntries, Class<? extends Exception> expectedException)
+      throws IOException {
+    String volumeNamePrefix = "volume";
+    String bucketNamePrefix = "bucket";
+    String keyPrefix = "";
+    OzoneConfiguration configuration = new OzoneConfiguration();
+    OMMetadataManager metadataManager = Mockito.mock(OMMetadataManager.class);
+    when(metadataManager.getBucketKeyPrefix(anyString(), anyString())).thenAnswer(i ->
+        "/" + i.getArguments()[0] + "/" + i.getArguments()[1] + "/");
+    KeyManagerImpl km = new KeyManagerImpl(null, null, metadataManager, configuration, null, null, null);
+    Table<String, String> mockedRenameTable = Mockito.mock(Table.class);
+    when(metadataManager.getSnapshotRenamedTable()).thenReturn(mockedRenameTable);
+    List<Table.KeyValue<String, String>> expectedEntries = mockTableIterator(
+        String.class, mockedRenameTable, numberOfVolumes, numberOfBucketsPerVolume, numberOfKeysPerBucket,
+        volumeNamePrefix, bucketNamePrefix, keyPrefix, volumeNumber, bucketNumber, startVolumeNumber, startBucketNumber,
+        startKeyNumber, numberOfEntries);
+    String volumeName = volumeNumber == null ? null : (String.format("%s%010d", volumeNamePrefix, volumeNumber));
+    String bucketName = bucketNumber == null ? null : (String.format("%s%010d", bucketNamePrefix, bucketNumber));
+    String startKey = startVolumeNumber == null || startBucketNumber == null || startKeyNumber == null ? null
+        : (String.format("/%s%010d/%s%010d/%s%010d", volumeNamePrefix, startVolumeNumber, bucketNamePrefix,
+        startBucketNumber, keyPrefix, startKeyNumber));
+    if (expectedException != null) {
+      assertThrows(expectedException, () -> km.getRenamesKeyEntries(volumeName, bucketName, startKey, numberOfEntries));
+    } else {
+      assertEquals(expectedEntries, km.getRenamesKeyEntries(volumeName, bucketName, startKey, numberOfEntries));
+    }
+  }
+
+  @ParameterizedTest
+  @MethodSource("getTableIteratorParameters")
+  @SuppressWarnings({"checkstyle:ParameterNumber"})
+  public void testGetDeletedDirEntries(int numberOfVolumes, int numberOfBucketsPerVolume, int numberOfKeysPerBucket,
+                                       Integer volumeNumber, Integer bucketNumber,
+                                       Integer startVolumeNumber, Integer startBucketNumber, Integer startKeyNumber,
+                                       int numberOfEntries, Class<? extends Exception> expectedException)
+      throws IOException {
+    String volumeNamePrefix = "";
+    String bucketNamePrefix = "";
+    String keyPrefix = "key";
+    startVolumeNumber = null;
+    OzoneConfiguration configuration = new OzoneConfiguration();
+    OMMetadataManager metadataManager = Mockito.mock(OMMetadataManager.class);
+    when(metadataManager.getBucketKeyPrefixFSO(anyString(), anyString())).thenAnswer(i ->
+        "/" + i.getArguments()[0] + "/" + i.getArguments()[1] + "/");
+    KeyManagerImpl km = new KeyManagerImpl(null, null, metadataManager, configuration, null, null, null);
+    Table<String, OmKeyInfo> mockedDeletedDirTable = Mockito.mock(Table.class);
+    when(metadataManager.getDeletedDirTable()).thenReturn(mockedDeletedDirTable);
+    List<Table.KeyValue<String, OmKeyInfo>> expectedEntries = mockTableIterator(
+        OmKeyInfo.class, mockedDeletedDirTable, numberOfVolumes, numberOfBucketsPerVolume, numberOfKeysPerBucket,
+        volumeNamePrefix, bucketNamePrefix, keyPrefix, volumeNumber, bucketNumber, startVolumeNumber, startBucketNumber,
+        startKeyNumber, numberOfEntries);
+    String volumeName = volumeNumber == null ? null : (String.format("%s%010d", volumeNamePrefix, volumeNumber));
+    String bucketName = bucketNumber == null ? null : (String.format("%s%010d", bucketNamePrefix, bucketNumber));
+    if (expectedException != null) {
+      assertThrows(expectedException, () -> km.getDeletedDirEntries(volumeName, bucketName, numberOfEntries));
+    } else {
+      assertEquals(expectedEntries, km.getDeletedDirEntries(volumeName, bucketName, numberOfEntries));
+    }
+  }
+}