Add a buffering functionality to the SseWorkerConnection (#721)

* Add a buffering functionality to the SsseWorkerConnection
Netflix · Nov 1, 2024 · f9d3a5c · f9d3a5c
1 parent 7f996f8
commit f9d3a5c
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 0 deletions.
diff --git a/...ver-worker-client/src/main/java/io/mantisrx/server/worker/client/SseWorkerConnection.java b/...ver-worker-client/src/main/java/io/mantisrx/server/worker/client/SseWorkerConnection.java
@@ -305,6 +305,7 @@ protected Observable<MantisServerSentEvent> streamContent(HttpClientResponse<Ser
         }
         return response.getContent()
                 .lift(new DropOperator<ServerSentEvent>(metricGroupId))
+                .rebatchRequests(this.bufferSize <= 0 ? 1 : this.bufferSize)
                 .flatMap((ServerSentEvent t1) -> {
                     lastDataReceived.set(System.currentTimeMillis());
                     if (isConnected.get() && isReceivingData.compareAndSet(false, true))

diff --git a/...worker-client/src/test/java/io/mantisrx/server/worker/client/SseWorkerConnectionTest.java b/...worker-client/src/test/java/io/mantisrx/server/worker/client/SseWorkerConnectionTest.java
@@ -148,4 +148,51 @@ public void testMantisHttpClientUsage() throws Exception {
         logger.info("Connection tracker size: {}", client.connectionTrackerSize());
         assertEquals(0, client.connectionTrackerSize());
     }
+
+    @Test
+    public void testStreamContentBuffersBeforeDrop() throws Exception {
+        int bufferSize = 20;
+        int totalEvents = 100;
+        SpectatorRegistryFactory.setRegistry(new DefaultRegistry());
+        String metricGroupString = "testmetric_buffer";
+        MetricGroupId metricGroupId = new MetricGroupId(metricGroupString);
+        SseWorkerConnection workerConnection = new SseWorkerConnection("connection_type",
+            "hostname",
+            80,
+            b -> {},
+            b -> {},
+            t -> {},
+            600,
+            false,
+            new CopyOnWriteArraySet<>(),
+            bufferSize,
+            null,
+            true,
+            metricGroupId);
+        HttpClientResponse<ServerSentEvent> response = mock(HttpClientResponse.class);
+        TestScheduler testScheduler = Schedulers.test();
+
+        // Events are just "0", "1", "2", ...
+        Observable<ServerSentEvent> contentObs = Observable.interval(1, TimeUnit.SECONDS, testScheduler)
+            .map(t -> new ServerSentEvent(Unpooled.copiedBuffer(Long.toString(t), Charset.defaultCharset())));
+
+        when(response.getContent()).thenReturn(contentObs);
+
+        TestSubscriber<MantisServerSentEvent> subscriber = new TestSubscriber<>(1);
+
+        workerConnection.streamContent(response, b -> {}, 600, "delimiter").subscribeOn(testScheduler).subscribe(subscriber);
+
+        testScheduler.advanceTimeBy(totalEvents, TimeUnit.SECONDS);
+        subscriber.assertValueCount(1);
+        List<MantisServerSentEvent> events = subscriber.getOnNextEvents();
+        assertEquals("0", events.get(0).getEventAsString());
+
+        Metrics metrics = MetricsRegistry.getInstance().getMetric(metricGroupId);
+        Counter onNextCounter = metrics.getCounter(DropOperator.Counters.onNext.toString());
+        Counter droppedCounter = metrics.getCounter(DropOperator.Counters.dropped.toString());
+        logger.info("next: {}", onNextCounter.value());
+        logger.info("drop: {}", droppedCounter.value());
+        assertTrue(onNextCounter.value() >= bufferSize); // Should request at least the buffer even though we requested 1.
+        assertTrue(droppedCounter.value() <= totalEvents - bufferSize ); // We should not drop any of the buffer.
+    }
 }