Merge branch 'grumpyjames-master'

whtblktone · Apr 10, 2017 · 534f06c · 534f06c
2 parents 2687e16 + a85ca47
commit 534f06c
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 1 deletion.
diff --git a/src/main/java/com/lmax/disruptor/BatchEventProcessor.java b/src/main/java/com/lmax/disruptor/BatchEventProcessor.java
@@ -37,6 +37,7 @@ public final class BatchEventProcessor<T>
     private final EventHandler<? super T> eventHandler;
     private final Sequence sequence = new Sequence(Sequencer.INITIAL_CURSOR_VALUE);
     private final TimeoutHandler timeoutHandler;
+    private final BatchStartAware batchStartAware;
 
     /**
      * Construct a {@link EventProcessor} that will automatically track the progress by updating its sequence when
@@ -60,7 +61,10 @@ public BatchEventProcessor(
             ((SequenceReportingEventHandler<?>) eventHandler).setSequenceCallback(sequence);
         }
 
-        timeoutHandler = (eventHandler instanceof TimeoutHandler) ? (TimeoutHandler) eventHandler : null;
+        batchStartAware =
+                (eventHandler instanceof BatchStartAware) ? (BatchStartAware) eventHandler : null;
+        timeoutHandler =
+                (eventHandler instanceof TimeoutHandler) ? (TimeoutHandler) eventHandler : null;
     }
 
     @Override
@@ -122,6 +126,10 @@ public void run()
                 try
                 {
                     final long availableSequence = sequenceBarrier.waitFor(nextSequence);
+                    if (batchStartAware != null)
+                    {
+                        batchStartAware.onBatchStart(availableSequence - nextSequence + 1);
+                    }
 
                     while (nextSequence <= availableSequence)
                     {

diff --git a/src/main/java/com/lmax/disruptor/BatchStartAware.java b/src/main/java/com/lmax/disruptor/BatchStartAware.java
@@ -0,0 +1,6 @@
+package com.lmax.disruptor;
+
+public interface BatchStartAware
+{
+    void onBatchStart(final long batchSize);
+}
diff --git a/src/test/java/com/lmax/disruptor/BatchEventProcessorTest.java b/src/test/java/com/lmax/disruptor/BatchEventProcessorTest.java
@@ -18,10 +18,14 @@
 import com.lmax.disruptor.support.StubEvent;
 import org.junit.Test;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
 import static com.lmax.disruptor.RingBuffer.createMultiProducer;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 public final class BatchEventProcessorTest
@@ -137,4 +141,50 @@ public void onEvent(StubEvent event, long sequence, boolean endOfBatch) throws E
         }
     }
 
+    @Test
+    public void reportAccurateBatchSizesAtBatchStartTime()
+        throws Exception
+    {
+        final List<Long> batchSizes = new ArrayList<Long>();
+        final CountDownLatch eventLatch = new CountDownLatch(6);
+
+        final class LoopbackEventHandler
+            implements EventHandler<StubEvent>, BatchStartAware
+        {
+
+            @Override
+            public void onBatchStart(long batchSize)
+            {
+                batchSizes.add(batchSize);
+            }
+
+            @Override
+            public void onEvent(StubEvent event, long sequence, boolean endOfBatch)
+                throws Exception
+            {
+                if (!endOfBatch)
+                {
+                    ringBuffer.publish(ringBuffer.next());
+                }
+                eventLatch.countDown();
+            }
+        }
+
+        final BatchEventProcessor<StubEvent> batchEventProcessor =
+            new BatchEventProcessor<StubEvent>(
+                ringBuffer, sequenceBarrier, new LoopbackEventHandler());
+
+        ringBuffer.publish(ringBuffer.next());
+        ringBuffer.publish(ringBuffer.next());
+        ringBuffer.publish(ringBuffer.next());
+
+        Thread thread = new Thread(batchEventProcessor);
+        thread.start();
+        eventLatch.await();
+
+        batchEventProcessor.halt();
+        thread.join();
+
+        assertEquals(Arrays.asList(3L, 2L, 1L), batchSizes);
+    }
 }