Add more test cases

Signed-off-by: Andriy Redko <[email protected]>
opensearch-project · Aug 28, 2024 · 62b27d2 · 62b27d2
1 parent 5b3087a
commit 62b27d2
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 3 deletions.
diff --git a/...rt-reactor-netty4/src/javaRestTest/java/org/opensearch/rest/ReactorNetty4StreamingIT.java b/...rt-reactor-netty4/src/javaRestTest/java/org/opensearch/rest/ReactorNetty4StreamingIT.java
@@ -169,6 +169,49 @@ public void testStreamingRequestManyBatchesBySize() throws IOException {
         assertThat(count, equalTo(5));
     }
 
+    public void testStreamingRequestManyBatchesByInterval() throws IOException {
+        final Stream<String> stream = IntStream.range(1, 6)
+            .mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
+
+        final Duration delay = Duration.ofMillis(500);
+        final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
+            "POST",
+            "/_bulk/stream",
+            Flux.fromStream(stream).delayElements(delay).map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8)))
+        );
+        streamingRequest.addParameter("refresh", "true");
+        streamingRequest.addParameter("batch_interval", "5s");
+
+        final StreamingResponse<ByteBuffer> streamingResponse = client().streamRequest(streamingRequest);
+
+        // We don't check for a other documents here since those may appear in any of the chunks (it is very
+        // difficult to get the timing right). But at the end, the total number of the documents is being checked.
+        StepVerifier.create(Flux.from(streamingResponse.getBody()).map(b -> new String(b.array(), StandardCharsets.UTF_8)))
+            .expectNextMatches(
+                s -> s.contains("\"result\":\"created\"")
+                    && s.contains("\"_id\":\"1\"")
+                    && s.contains("\"result\":\"created\"")
+                    && s.contains("\"_id\":\"2\"")
+                    && s.contains("\"result\":\"created\"")
+                    && s.contains("\"_id\":\"3\"")
+                    && s.contains("\"result\":\"created\"")
+                    && s.contains("\"_id\":\"4\"")
+                    && s.contains("\"result\":\"created\"")
+                    && s.contains("\"_id\":\"5\"")
+            )
+            .expectComplete()
+            .verify();
+
+        assertThat(streamingResponse.getStatusLine().getStatusCode(), equalTo(200));
+        assertThat(streamingResponse.getWarnings(), empty());
+
+        final Request request = new Request("GET", "/test-streaming/_count");
+        final Response response = client().performRequest(request);
+        final ObjectPath objectPath = ObjectPath.createFromResponse(response);
+        final Integer count = objectPath.evaluate("count");
+        assertThat(count, equalTo(5));
+    }
+
     public void testStreamingRequestManyBatchesByIntervalAndSize() throws IOException {
         final Stream<String> stream = IntStream.range(1, 6)
             .mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");

diff --git a/server/src/main/java/org/opensearch/rest/action/document/RestBulkStreamingAction.java b/server/src/main/java/org/opensearch/rest/action/document/RestBulkStreamingAction.java
@@ -98,6 +98,7 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         final String refresh = request.param("refresh");
         final TimeValue batchInterval = request.paramAsTime("batch_interval", null);
         final int batchSize = request.paramAsInt("batch_size", 1); /* by default, batch size of 1 */
+        final boolean hasBatchSize = request.hasParam("batch_size"); /* is batch_size explicitly specified or default is used */
 
         if (batchInterval != null && batchInterval.duration() <= 0) {
             throw new IllegalArgumentException("The batch_interval value should be non-negative [" + batchInterval.millis() + "ms].");
@@ -127,7 +128,7 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
 
             // TODOs:
             // - eliminate serialization inefficiencies
-            createBufferedFlux(batchInterval, batchSize, channel).zipWith(Flux.fromStream(Stream.generate(() -> {
+            createBufferedFlux(batchInterval, batchSize, hasBatchSize, channel).zipWith(Flux.fromStream(Stream.generate(() -> {
                 BulkRequest bulkRequest = Requests.bulkRequest();
                 bulkRequest.waitForActiveShards(prepareBulkRequest.waitForActiveShards());
                 bulkRequest.timeout(prepareBulkRequest.timeout());
@@ -233,9 +234,19 @@ public boolean allowsUnsafeBuffers() {
         return true;
     }
 
-    private Flux<List<HttpChunk>> createBufferedFlux(final TimeValue batchInterval, final int batchSize, StreamingRestChannel channel) {
+    private Flux<List<HttpChunk>> createBufferedFlux(
+        final TimeValue batchInterval,
+        final int batchSize,
+        final boolean hasBatchSize,
+        StreamingRestChannel channel
+    ) {
         if (batchInterval != null) {
-            return Flux.from(channel).bufferTimeout(batchSize, Duration.ofMillis(batchInterval.millis()));
+            // If non-default batch size is specified, buffer by interval and batch
+            if (hasBatchSize) {
+                return Flux.from(channel).bufferTimeout(batchSize, Duration.ofMillis(batchInterval.millis()));
+            } else {
+                return Flux.from(channel).buffer(Duration.ofMillis(batchInterval.millis()));
+            }
         } else {
             return Flux.from(channel).buffer(batchSize);
         }