pq: add support for Batch#unread

Author: yaauie

logstash-core/src/main/java/org/logstash/ackedqueue/Batch.java

@@ -65,6 +65,12 @@ public void close() throws IOException {
         }
     }
 
+    public void unread() throws IOException {
+        if (firstSeqNum >= 0L) {
+            this.queue.unread(firstSeqNum, elements.size());
+        }
+    }
+
     public int size() {
         return elements.size();
     }
@@ -77,6 +83,14 @@ public Queue getQueue() {
         return queue;
     }
 
+    long firstSeqNum() {
+        return firstSeqNum;
+    }
+
+    long lastSeqNum() {
+        return firstSeqNum + elements.size();
+    }
+
     /**
      *
      * @param serialized Collection of serialized elements

logstash-core/src/main/java/org/logstash/ackedqueue/Page.java

@@ -36,7 +36,6 @@ public final class Page implements Closeable {
     final int pageNum;
     private long minSeqNum; // TODO: see if we can make it final?
     private int elementCount;
-    private long firstUnreadSeqNum;
     private final Queue queue;
     private final PageIO pageIO;
     private boolean writable;
@@ -45,41 +44,62 @@ public final class Page implements Closeable {
     // TODO: go steal LocalCheckpointService in feature/seq_no from ES
     // TODO: https://github.com/elastic/elasticsearch/blob/feature/seq_no/core/src/main/java/org/elasticsearch/index/seqno/LocalCheckpointService.java
     private final BitSet ackedSeqNums;
+    private final BitSet readSeqNums;
     private Checkpoint lastCheckpoint;
 
+    // Use {@link PageFactory}
     Page(int pageNum, Queue queue, long minSeqNum, int elementCount, long firstUnreadSeqNum, BitSet ackedSeqNums, @NotNull PageIO pageIO, boolean writable) {
         this.pageNum = pageNum;
         this.queue = queue;
 
         this.minSeqNum = minSeqNum;
         this.elementCount = elementCount;
-        this.firstUnreadSeqNum = firstUnreadSeqNum;
         this.ackedSeqNums = ackedSeqNums;
+        this.readSeqNums = readSeqNums(minSeqNum, firstUnreadSeqNum, elementCount);
         this.lastCheckpoint = new Checkpoint(0, 0, 0, 0, 0);
         this.pageIO = pageIO;
         this.writable = writable;
 
         assert this.pageIO != null : "invalid null pageIO";
     }
 
+    private static BitSet readSeqNums(final long minSeqNum, final long firstUnreadSeqNum, final int elementCount) {
+        final BitSet bitSet = new BitSet(elementCount);
+        if (firstUnreadSeqNum > minSeqNum) {
+            bitSet.set(0, (int) (firstUnreadSeqNum - minSeqNum));
+        }
+        return bitSet;
+    }
+
     public String toString() {
-        return "pageNum=" + this.pageNum + ", minSeqNum=" + this.minSeqNum + ", elementCount=" + this.elementCount + ", firstUnreadSeqNum=" + this.firstUnreadSeqNum;
+        return "pageNum=" + this.pageNum +
+                ", minSeqNum=" + this.minSeqNum +
+                ", elementCount=" + this.elementCount +
+                ", firstUnreadSeqNum=" + (this.minSeqNum + this.readSeqNums.nextClearBit(0));
     }
 
     /**
-     * @param limit the maximum number of elements to read, actual number readcan be smaller
+     * @param limit the maximum number of elements to read, actual number read can be smaller
      * @return {@link SequencedList} collection of serialized elements read
      * @throws IOException if an IO error occurs
      */
     public SequencedList<byte[]> read(int limit) throws IOException {
         // first make sure this page is activated, activating previously activated is harmless
         this.pageIO.activate();
 
-        SequencedList<byte[]> serialized = this.pageIO.read(this.firstUnreadSeqNum, limit);
-        assert serialized.getSeqNums().get(0) == this.firstUnreadSeqNum :
-            String.format("firstUnreadSeqNum=%d != first result seqNum=%d", this.firstUnreadSeqNum, serialized.getSeqNums().get(0));
+        // determine where to begin our read
+        int firstUnreadOffset = this.readSeqNums.nextClearBit(0);
+        long firstUnreadSeqNum = this.minSeqNum + firstUnreadOffset;
+
+        // determine how many contiguous events we can read without re-emitting any already-read events
+        int nextReadOffset = this.readSeqNums.nextSetBit(firstUnreadOffset);
+        int effectiveLimit = (nextReadOffset < 0) ? limit : Math.min(limit, nextReadOffset - firstUnreadOffset);
+
+        SequencedList<byte[]> serialized = this.pageIO.read(firstUnreadSeqNum, effectiveLimit);
+        assert serialized.getSeqNums().get(0) == firstUnreadSeqNum :
+            String.format("firstUnreadSeqNum=%d != first result seqNum=%d", firstUnreadSeqNum, serialized.getSeqNums().get(0));
 
-        this.firstUnreadSeqNum += serialized.getElements().size();
+        this.readSeqNums.set(firstUnreadOffset, firstUnreadOffset + serialized.getElements().size());
 
         return serialized;
     }
@@ -93,7 +113,6 @@ public void write(byte[] bytes, long seqNum, int checkpointMaxWrites) throws IOE
 
         if (this.minSeqNum <= 0) {
             this.minSeqNum = seqNum;
-            this.firstUnreadSeqNum = seqNum;
         }
         this.elementCount++;
 
@@ -136,7 +155,11 @@ public boolean isFullyAcked() {
     }
 
     public long unreadCount() {
-        return this.elementCount <= 0 ? 0 : Math.max(0, (maxSeqNum() - this.firstUnreadSeqNum) + 1);
+        if (this.elementCount <= 0) {
+            return 0;
+        }
+
+        return this.elementCount - this.readSeqNums.cardinality();
     }
 
     /**
@@ -187,6 +210,14 @@ public boolean ack(long firstSeqNum, int count, int checkpointMaxAcks) throws IO
         return done;
     }
 
+    public boolean unread(long firstSeqNum, int count) {
+        final boolean wasFullyRead = isFullyRead();
+        int unreadChunkFirstOffset = Ints.checkedCast(Math.subtractExact(firstSeqNum, this.minSeqNum));
+        this.readSeqNums.clear(unreadChunkFirstOffset, unreadChunkFirstOffset+count);
+
+        return wasFullyRead;
+    }
+
     public void checkpoint() throws IOException {
         if (this.writable) {
             headPageCheckpoint();

logstash-core/src/main/java/org/logstash/ackedqueue/Queue.java

@@ -758,6 +758,25 @@ public void ack(final long firstAckSeqNum, final int ackCount) throws IOExceptio
         }
     }
 
+    public void unread(final long firstUnreadSeqNum, final int unreadCount) throws IOException {
+        lock.lock();
+        try {
+            if (containsSeq(headPage, firstUnreadSeqNum)) {
+                this.headPage.unread(firstUnreadSeqNum, unreadCount);
+            } else {
+                final int resultIndex = binaryFindPageForSeqnum(firstUnreadSeqNum);
+                final Page relevantTailPage = this.tailPages.get(resultIndex);
+                if (relevantTailPage.unread(firstUnreadSeqNum, unreadCount)) {
+                    this.unreadTailPages.add(0, relevantTailPage);
+                }
+                notEmpty.signalAll();
+            }
+            this.unreadCount -= unreadCount;
+        } finally {
+            lock.unlock();
+        }
+    }
+
     public CheckpointIO getCheckpointIO() {
         return this.checkpointIO;
     }
@@ -870,6 +889,10 @@ public long getUnackedCount() {
         return sumPages(Page::getUnackedCount);
     }
 
+    long getPageCount() {
+        return sumPages((page) -> 1L);
+    }
+
     private long sumPages(final ToLongFunction<Page> valueGetter) {
         lock.lock();
         try {

logstash-core/src/test/java/org/logstash/ackedqueue/QueueTest.java

@@ -39,6 +39,8 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
 
 import org.junit.After;
 import org.junit.Before;
@@ -168,6 +170,52 @@ public void writeWhenPageEqualsQueueSize() throws IOException {
         }
     }
 
+    @Test
+    public void unreadEvents() throws IOException {
+        final Settings queueSettings = TestSettings.fileSettingsBuilder(dataPath)
+                .compressionCodecFactory(CompressionCodec.fromConfigValue("disabled"))
+                .capacity(1024)
+                .build();
+        try (Queue q = new Queue(TestSettings.fileSettingsBuilder(dataPath).capacity(1024).build())) {
+            q.open();
+
+            final AtomicLong atomicLong = new AtomicLong();
+            Supplier<String> tenByteStringSupplier = () -> String.format("%010d", atomicLong.incrementAndGet());
+
+            // insert enough events to fill more than two pages
+            while (q.getPageCount() < 3) {
+                q.write(new StringElement(tenByteStringSupplier.get()));
+            }
+
+            // read batches until there is only one page in the quick-pick unreadTailPages
+            final ArrayList<Batch> batches = new ArrayList<>();
+            while (q.unreadTailPages.size() > 1) {
+                batches.add(q.readBatch(2, 500L));
+            }
+
+            // ack the first batch
+            final Batch ackBatch = batches.remove(0);
+            ackBatch.close();
+
+            // unread the next batch, validate that the page got put back on unreadTailPages
+            final Batch unreadBatch = batches.remove(0);
+            unreadBatch.unread();
+            assertThat(q.unreadTailPages.size(), is(2));
+
+            // read another batch, ensure it is the same as the one we previously unread
+            final long unreadBatchFirstSeqNum = unreadBatch.firstSeqNum();
+            final Batch nextReadBatch = q.readBatch(3, 500L); // ask for more than the known 2-event unread region
+            assertThat(q.unreadTailPages.size(), is(1));
+
+            assertThat(nextReadBatch.firstSeqNum(), is(unreadBatchFirstSeqNum));
+            assertThat(nextReadBatch.size(), is(2)); // even though we asked for more, we are getting the entire contiguously-unread region
+
+            // now read a new batch, and make sure it's newer than the previous batches
+            final Batch lastBatch = q.readBatch(2, 500L);
+            assertThat(lastBatch.firstSeqNum(), is(1 + batches.get(batches.size() - 1).firstSeqNum()));
+        }
+    }
+
     @Test
     public void singleWriteMultiRead() throws IOException {
         try (Queue q = new Queue(TestSettings.persistedQueueSettings(100, dataPath))) {

logstash-core/src/test/java/org/logstash/ackedqueue/TestSettings.java

@@ -31,4 +31,8 @@ public static Settings persistedQueueSettings(int capacity, long size, String fo
         return SettingsImpl.fileSettingsBuilder(folder).capacity(capacity)
             .queueMaxBytes(size).elementClass(StringElement.class).build();
     }
+
+    public static Settings.Builder fileSettingsBuilder(String folder) {
+        return SettingsImpl.fileSettingsBuilder(folder).elementClass(StringElement.class);
+    }
 }