Merge pull request #50 from bellofreedom/key_iterator

Add key iterator for HaloDB

src/main/java/com/oath/halodb/HaloDB.java

@@ -80,6 +80,10 @@ public HaloDBIterator newIterator() throws HaloDBException {
         return new HaloDBIterator(dbInternal);
     }
 
+    public HaloDBKeyIterator newKeyIterator() {
+        return new HaloDBKeyIterator(dbInternal);
+    }
+
     public void pauseCompaction() throws HaloDBException {
         try {
             dbInternal.pauseCompaction();

src/main/java/com/oath/halodb/HaloDBKeyIterator.java

@@ -0,0 +1,116 @@
+package com.oath.halodb;
+
+import java.io.IOException;
+import java.nio.channels.ClosedChannelException;
+import java.util.*;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class HaloDBKeyIterator implements  Iterator<RecordKey>{
+    private static final Logger logger = LoggerFactory.getLogger(HaloDBIterator.class);
+
+    private Iterator<Integer> outer;
+    private Iterator<IndexFileEntry> inner;
+    private HaloDBFile currentFile;
+
+    private RecordKey next;
+
+    private final HaloDBInternal dbInternal;
+
+    HaloDBKeyIterator(HaloDBInternal dbInternal) {
+        this.dbInternal = dbInternal;
+        outer = dbInternal.listDataFileIds().iterator();
+    }
+
+    @Override
+    public boolean hasNext() {
+        if (next != null) {
+            return true;
+        }
+
+        try {
+            // inner == null means this is the first time hasNext() is called.
+            // use moveToNextFile() to move to the first file.
+            if (inner == null && !moveToNextFile()) {
+                return false;
+            }
+
+            do {
+                if (readNextRecord()) {
+                    return true;
+                }
+            } while (moveToNextFile());
+
+            return false;
+
+        } catch (IOException e) {
+            logger.error("Error in Iterator", e);
+            return false;
+        }
+    }
+
+    @Override
+    public RecordKey next() {
+        if (hasNext()) {
+            RecordKey key = next;
+            next = null;
+            return key;
+        }
+        throw new NoSuchElementException();
+    }
+
+    private boolean moveToNextFile() throws IOException {
+        while (outer.hasNext()) {
+            int fileId = outer.next();
+            currentFile = dbInternal.getHaloDBFile(fileId);
+            if (currentFile != null) {
+                try {
+                    inner = currentFile.getIndexFile().newIterator();
+                    return true;
+                } catch (ClosedChannelException e) {
+                    if (dbInternal.isClosing()) {
+                        //TODO: define custom Exception classes for HaloDB.
+                        throw new RuntimeException("DB is closing");
+                    }
+                    logger.debug("Index file {} closed, probably by compaction thread. Skipping to next one", fileId);
+                }
+            }
+            logger.debug("Data file {} deleted, probably by compaction thread. Skipping to next one", fileId);
+        }
+
+        return false;
+    }
+
+    private boolean readNextRecord() {
+        while (inner.hasNext()) {
+            IndexFileEntry entry = inner.next();
+            try {
+                try {
+                    next = readValidRecordKey(entry);
+                    if (next != null) {
+                        return true;
+                    }
+                } catch (ClosedChannelException e) {
+                    if (dbInternal.isClosing()) {
+                        throw new RuntimeException("DB is closing");
+                    }
+                    logger.debug("Data file {} closed, probably by compaction thread. Skipping to next one", currentFile.getFileId());
+                    break;
+                }
+            } catch (IOException e) {
+                logger.info("Error in iterator", e);
+                break;
+            }
+        }
+        return false;
+    }
+
+    private RecordKey readValidRecordKey(IndexFileEntry entry) throws IOException {
+        InMemoryIndexMetaData meta = Utils.getMetaData(entry, currentFile.getFileId());
+        RecordKey key = null;
+        if (dbInternal.isRecordFresh(entry.getKey(), meta)) {
+            key = new RecordKey(entry.getKey());
+        }
+        return key;
+    }
+}

src/main/java/com/oath/halodb/RecordKey.java

@@ -0,0 +1,31 @@
+package com.oath.halodb;
+
+import java.util.*;
+
+public class RecordKey {
+    final byte[] key;
+    public RecordKey(byte[] key) {
+        this.key = key;
+    }
+
+    public byte[] getBytes() {
+        return key;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        // to be used in tests as we don't check if the headers are the same.
+
+        if (this == obj) {
+            return true;
+        }
+        if (!(obj instanceof RecordKey)) {
+            return false;
+        }
+
+        RecordKey recordKey = (RecordKey)obj;
+        return Arrays.equals(this.key, recordKey.getBytes());
+    }
+
+
+}

src/test/java/com/oath/halodb/HaloDBKeyIteratorTest.java

@@ -0,0 +1,97 @@
+package com.oath.halodb;
+
+import java.io.IOException;
+import java.nio.channels.ClosedChannelException;
+import java.util.*;
+import mockit.Invocation;
+import mockit.Mock;
+import mockit.MockUp;
+import org.hamcrest.MatcherAssert;
+import org.hamcrest.Matchers;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class HaloDBKeyIteratorTest extends TestBase {
+
+    @Test(expectedExceptions = NoSuchElementException.class, dataProvider = "Options")
+    public void testWithEmptyDB(HaloDBOptions options) throws HaloDBException {
+        String directory = TestUtils.getTestDirectory("HaloDBKeyIteratorTest", "testWithEmptyDB");
+
+        HaloDB db = getTestDB(directory, options);
+        HaloDBKeyIterator iterator = db.newKeyIterator();
+        Assert.assertFalse(iterator.hasNext());
+        iterator.next();
+    }
+
+    @Test(dataProvider = "Options")
+    public void testWithDelete(HaloDBOptions options) throws HaloDBException {
+        String directory =  TestUtils.getTestDirectory("HaloDBKeyIteratorTest", "testWithEmptyDB");
+
+        options.setCompactionDisabled(true);
+
+        HaloDB db = getTestDB(directory, options);
+        int noOfRecords = 10_000;
+        List<Record> records = TestUtils.insertRandomRecords(db, noOfRecords);
+
+        // delete all records.
+        for (Record r : records) {
+            db.delete(r.getKey());
+        }
+
+        HaloDBKeyIterator iterator = db.newKeyIterator();
+        Assert.assertFalse(iterator.hasNext());
+
+        // close and open the db again.
+        db.close();
+        db = getTestDBWithoutDeletingFiles(directory, options);
+        iterator = db.newKeyIterator();
+        Assert.assertFalse(iterator.hasNext());
+    }
+
+    @Test(dataProvider = "Options")
+    public void testPutAndGetDB(HaloDBOptions options) throws HaloDBException {
+        String directory = TestUtils.getTestDirectory("HaloDBKeyIteratorTest", "testPutAndGetDB");
+
+        options.setCompactionDisabled(true);
+        options.setMaxFileSize(10 * 1024);
+
+        HaloDB db = getTestDB(directory, options);
+
+        int noOfRecords = 10_000;
+        List<Record> records = TestUtils.insertRandomRecords(db, noOfRecords);
+
+        List<RecordKey> keys = new LinkedList<>();
+        for (Record record : records) {
+            keys.add(new RecordKey(record.getKey()));
+        }
+
+        List<RecordKey> actual = new ArrayList<>();
+        db.newKeyIterator().forEachRemaining(actual::add);
+
+        MatcherAssert.assertThat(actual, Matchers.containsInAnyOrder(keys.toArray()));
+    }
+
+    @Test(dataProvider = "Options")
+    public void testPutUpdateAndGetDB(HaloDBOptions options) throws HaloDBException {
+        String directory = TestUtils.getTestDirectory("HaloDBKeyIteratorTest", "testPutUpdateAndGetDB");
+
+        options.setCompactionDisabled(true);
+        options.setMaxFileSize(10 * 1024);
+
+        HaloDB db = getTestDB(directory, options);
+
+        int noOfRecords = 10_000;
+        List<Record> records = TestUtils.insertRandomRecords(db, noOfRecords);
+
+        List<Record> updated = TestUtils.updateRecords(db, records);
+
+        List<RecordKey> keys = new LinkedList<>();
+        for (Record record : updated) {
+            keys.add(new RecordKey(record.getKey()));
+        }
+
+        List<RecordKey> actual = new ArrayList<>();
+        db.newKeyIterator().forEachRemaining(actual::add);
+        MatcherAssert.assertThat(actual, Matchers.containsInAnyOrder(keys.toArray()));
+    }
+}