Decoded 'pipeline.batch.metrics' setting into BatchSizeSamplingType and spread around to reach in memory queue client and control the batch size metrics. Covered with tests the readBatch code to verify the effectiveness of the flag

Author: andsel

logstash-core/src/main/java/org/logstash/execution/QueueReadClientBase.java

@@ -48,6 +48,18 @@
 @JRubyClass(name = "QueueReadClientBase")
 public abstract class QueueReadClientBase extends RubyObject implements QueueReadClient {
 
+    public enum BatchSizeSamplingType {
+        NONE, FULL;
+
+        public static BatchSizeSamplingType decode(String type) {
+            return switch (type) {
+                case "false" -> NONE;
+                case "true" -> FULL;
+                default -> throw new IllegalArgumentException("Invalid batch size type: " + type);
+            };
+        }
+    }
+
     private static final long serialVersionUID = 1L;
 
     protected int batchSize = 125;
@@ -62,6 +74,7 @@ public abstract class QueueReadClientBase extends RubyObject implements QueueRea
     private transient LongCounter pipelineMetricFiltered;
     private transient TimerMetric pipelineMetricTime;
     private transient HistogramMetric pipelineMetricBatch;
+    protected BatchSizeSamplingType batchSizeSamplingType = BatchSizeSamplingType.FULL;
 
     protected QueueReadClientBase(final Ruby runtime, final RubyClass metaClass) {
         super(runtime, metaClass);
@@ -196,7 +209,9 @@ public void startMetrics(QueueBatch batch) {
         // JTODO getId has been deprecated in JDK 19, when JDK 21 is the target version use threadId() instead
         long threadId = Thread.currentThread().getId();
         inflightBatches.put(threadId, batch);
-        pipelineMetricBatch.update(batch.filteredSize());
+        if (batchSizeSamplingType == BatchSizeSamplingType.FULL) {
+            pipelineMetricBatch.update(batch.filteredSize());
+        }
     }
 
     @Override

logstash-core/src/main/java/org/logstash/ext/JrubyMemoryReadClientExt.java

@@ -47,9 +47,11 @@ public JrubyMemoryReadClientExt(final Ruby runtime, final RubyClass metaClass) {
 
     @SuppressWarnings("rawtypes")
     private JrubyMemoryReadClientExt(final Ruby runtime, final RubyClass metaClass,
-                                     BlockingQueue queue, int batchSize, int waitForMillis) {
+                                     BlockingQueue queue, int batchSize, int waitForMillis,
+                                     BatchSizeSamplingType batchSizeSamplingType) {
         super(runtime, metaClass);
         this.queue = queue;
+        this.batchSizeSamplingType = batchSizeSamplingType;
         this.batchSize = batchSize;
         this.waitForNanos = TimeUnit.NANOSECONDS.convert(waitForMillis, TimeUnit.MILLISECONDS);
         this.waitForMillis = waitForMillis;
@@ -58,8 +60,14 @@ private JrubyMemoryReadClientExt(final Ruby runtime, final RubyClass metaClass,
     @SuppressWarnings("rawtypes")
     public static JrubyMemoryReadClientExt create(BlockingQueue queue, int batchSize,
                                                   int waitForMillis) {
+        return create(queue, batchSize, waitForMillis, BatchSizeSamplingType.FULL);
+    }
+
+    @SuppressWarnings("rawtypes")
+    public static JrubyMemoryReadClientExt create(BlockingQueue queue, int batchSize, int waitForMillis,
+                                                  BatchSizeSamplingType batchSizeSamplingType) {
         return new JrubyMemoryReadClientExt(RubyUtil.RUBY,
-                RubyUtil.MEMORY_READ_CLIENT_CLASS, queue, batchSize, waitForMillis);
+                RubyUtil.MEMORY_READ_CLIENT_CLASS, queue, batchSize, waitForMillis, batchSizeSamplingType);
     }
 
     @Override

logstash-core/src/main/java/org/logstash/ext/JrubyWrappedSynchronousQueueExt.java

@@ -26,7 +26,6 @@
 import org.jruby.Ruby;
 import org.jruby.RubyClass;
 import org.jruby.RubyNumeric;
-import org.jruby.RubyString;
 import org.jruby.anno.JRubyClass;
 import org.jruby.anno.JRubyMethod;
 import org.jruby.runtime.ThreadContext;
@@ -43,6 +42,7 @@ public final class JrubyWrappedSynchronousQueueExt extends AbstractWrappedQueueE
     private static final long serialVersionUID = 1L;
 
     private transient BlockingQueue<JrubyEventExtLibrary.RubyEvent> queue;
+    private QueueReadClientBase.BatchSizeSamplingType batchMetricsSamplingType;
 
     public JrubyWrappedSynchronousQueueExt(final Ruby runtime, final RubyClass metaClass) {
         super(runtime, metaClass);
@@ -54,7 +54,8 @@ public JrubyWrappedSynchronousQueueExt initialize(final ThreadContext context,
                                                       IRubyObject size,
                                                       IRubyObject batchMetricsSampling) {
         int typedSize = ((RubyNumeric)size).getIntValue();
-        String batchMetricsSamplingType = ((RubyString) batchMetricsSampling).asJavaString();
+        this.batchMetricsSamplingType = QueueReadClientBase.BatchSizeSamplingType.decode(batchMetricsSampling.asJavaString());
+
         this.queue = new ArrayBlockingQueue<>(typedSize);
         return this;
     }
@@ -68,7 +69,7 @@ protected JRubyAbstractQueueWriteClientExt getWriteClient(final ThreadContext co
     protected QueueReadClientBase getReadClient() {
         // batch size and timeout are currently hard-coded to 125 and 50ms as values observed
         // to be reasonable tradeoffs between latency and throughput per PR #8707
-        return JrubyMemoryReadClientExt.create(queue, 125, 50);
+        return JrubyMemoryReadClientExt.create(queue, 125, 50, batchMetricsSamplingType);
     }
 
     @Override

logstash-core/src/test/java/org/logstash/ext/JrubyMemoryReadClientExtTest.java

@@ -30,6 +30,7 @@
 import org.logstash.RubyTestBase;
 import org.logstash.RubyUtil;
 import org.logstash.execution.QueueBatch;
+import org.logstash.execution.QueueReadClientBase.BatchSizeSamplingType;
 import org.logstash.instrument.metrics.AbstractNamespacedMetricExt;
 import org.logstash.instrument.metrics.MetricKeys;
 import org.logstash.instrument.metrics.MockNamespacedMetric;
@@ -49,7 +50,7 @@ public final class JrubyMemoryReadClientExtTest extends RubyTestBase {
     @SuppressWarnings("deprecation")
     public void testInflightBatchesTracking() throws InterruptedException, IOException {
         final BlockingQueue<JrubyEventExtLibrary.RubyEvent> queue = new ArrayBlockingQueue<>(10);
-        final JrubyMemoryReadClientExt client = JrubyMemoryReadClientExt.create(queue, 5, 50);
+        final JrubyMemoryReadClientExt client = JrubyMemoryReadClientExt.create(queue, 5, 50, BatchSizeSamplingType.NONE);
 
         final ThreadContext context = client.getRuntime().getCurrentContext();
 
@@ -71,7 +72,7 @@ public void givenNonEmptyQueueAndEnabledBatchSizeMetricThenHistogramContainsData
         final BlockingQueue<JrubyEventExtLibrary.RubyEvent> queue = new ArrayBlockingQueue<>(10);
         queue.add(testEvent);
 
-        final JrubyMemoryReadClientExt client = JrubyMemoryReadClientExt.create(queue, 5, 50);
+        final JrubyMemoryReadClientExt client = JrubyMemoryReadClientExt.create(queue, 5, 50, BatchSizeSamplingType.FULL);
 
         AbstractNamespacedMetricExt metric = MockNamespacedMetric.create();
         client.setPipelineMetric(metric);
@@ -84,4 +85,24 @@ public void givenNonEmptyQueueAndEnabledBatchSizeMetricThenHistogramContainsData
         assertEquals(1.0, metricSnapshot.get75Percentile(), 0.0001);
         assertEquals(1.0, metricSnapshot.get90Percentile(), 0.0001);
     }
+
+    @Test
+    public void givenNonEmptyQueueAndDisabledBatchSizeMetricThenHistogramIsNotPopulated() throws InterruptedException {
+        final JrubyEventExtLibrary.RubyEvent testEvent = JrubyEventExtLibrary.RubyEvent.newRubyEvent(RubyUtil.RUBY, new Event());
+        final BlockingQueue<JrubyEventExtLibrary.RubyEvent> queue = new ArrayBlockingQueue<>(10);
+        queue.add(testEvent);
+
+        final JrubyMemoryReadClientExt client = JrubyMemoryReadClientExt.create(queue, 5, 50, BatchSizeSamplingType.NONE);
+
+        AbstractNamespacedMetricExt metric = MockNamespacedMetric.create();
+        client.setPipelineMetric(metric);
+
+        final QueueBatch batch = client.readBatch();
+        assertEquals(1, batch.filteredSize());
+
+        HistogramMetric histogram = HistogramMetric.fromRubyBase(metric, MetricKeys.BATCH_SIZE_KEY);
+        HistogramSnapshot metricSnapshot = histogram.getValue();
+        assertEquals(0.0, metricSnapshot.get75Percentile(), 0.0001);
+        assertEquals(0.0, metricSnapshot.get90Percentile(), 0.0001);
+    }
 }