From f8085007189668439b807c7f5a1019beb69b0687 Mon Sep 17 00:00:00 2001 From: KirillPamPam Date: Tue, 19 Dec 2023 18:19:43 +0400 Subject: [PATCH] Accumulate first bytes (#366) --- .../rpcclient/stream/JsonRpcStreamParser.kt | 16 ++++++++++++++-- .../rpcclient/stream/JsonRpcStreamParserTest.kt | 2 +- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/src/main/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParser.kt b/src/main/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParser.kt index 61b665cf..bc4cec37 100644 --- a/src/main/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParser.kt +++ b/src/main/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParser.kt @@ -15,7 +15,9 @@ import java.util.concurrent.atomic.AtomicBoolean import java.util.concurrent.atomic.AtomicInteger import java.util.concurrent.atomic.AtomicReference -class JsonRpcStreamParser { +class JsonRpcStreamParser( + private val firstChunkMaxSize: Int = 8192, +) { companion object { private val log = LoggerFactory.getLogger(JsonRpcStreamParser::class.java) @@ -31,7 +33,17 @@ class JsonRpcStreamParser { } fun streamParse(statusCode: Int, response: Flux): Mono { - return response.switchOnFirst({ first, responseStream -> + val firstPartSize = AtomicInteger() + return response.bufferUntil { + if (firstPartSize.get() > firstChunkMaxSize) { + true + } else { + firstPartSize.addAndGet(it.size) + firstPartSize.get() > firstChunkMaxSize // accumulate bytes until chunk is full + } + }.map { + it.reduce { acc, bytes -> acc.plus(bytes) } + }.switchOnFirst({ first, responseStream -> if (first.get() == null || statusCode != 200) { aggregateResponse(responseStream, statusCode) } else { diff --git a/src/test/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParserTest.kt b/src/test/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParserTest.kt index 4068cf66..4faa4d03 100644 --- a/src/test/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParserTest.kt +++ b/src/test/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParserTest.kt @@ -13,7 +13,7 @@ import reactor.test.StepVerifier import java.time.Duration class JsonRpcStreamParserTest { - private val streamParser = JsonRpcStreamParser() + private val streamParser = JsonRpcStreamParser(5) @Test fun `if first part couldn't be parsed then aggregate response`() {