From cc3d1195a220cf6413a630e9ae99040e7b12cf49 Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Thu, 15 Oct 2026 04:37:53 +0000 Subject: [PATCH 1/3] perf(storage): move grouped String payloads into pages --- src/storage/engine/grouped/string.cpp | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/storage/engine/grouped/string.cpp b/src/storage/engine/grouped/string.cpp index fd95d9a49..9335acf32 100644 --- a/src/storage/engine/grouped/string.cpp +++ b/src/storage/engine/grouped/string.cpp @@ -32,12 +32,16 @@ OrderedCollectionRoot StringRoot(std::size_t bytes, std::uint64_t incarnation) { OrderedGroupSnapshot StringPage(const OrderedCollectionRoot& root, std::uint64_t id, std::string bytes) { - return {.kind_ = OrderedCollectionKind::kString, - .incarnation_ = root.incarnation_, - .id_ = id, - .previous_ = id - 1, - .next_ = id == root.last_group_ ? 0 : id + 1, - .entries_ = {{.value_ = std::move(bytes)}}}; + OrderedGroupSnapshot page{.kind_ = OrderedCollectionKind::kString, + .incarnation_ = root.incarnation_, + .id_ = id, + .previous_ = id - 1, + .next_ = id == root.last_group_ ? 0 : id + 1, + .entries_ = {}}; + // An initializer_list exposes const entries, forcing the vector to copy + // the segment even when its initializer moves the string. + page.entries_.emplace_back().value_ = std::move(bytes); + return page; } } // namespace From 7c4c2b1ffcade403a7daccbf197c43580e6805bd Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Thu, 15 Oct 2026 04:37:53 +0000 Subject: [PATCH 2/3] perf(redis): receive large bulk arguments into final storage --- CMakeLists.txt | 8 + bycorf | 2 +- docs/architecture/02-request-serving.md | 28 +- docs/operations/metrics.md | 12 +- include/lavik/memory.h | 15 +- include/lavik/resp.h | 25 ++ src/metrics.cpp | 2 +- src/redis/resp.cpp | 69 ++++- src/redis/server.cpp | 51 ++++ tests/large_bulk_e2e.py | 360 ++++++++++++++++++++++++ tests/resp_test.cpp | 189 +++++++++++++ 11 files changed, 733 insertions(+), 28 deletions(-) create mode 100644 tests/large_bulk_e2e.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 1202efc05..6fb45fc22 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2189,6 +2189,14 @@ endforeach() ) set_tests_properties(lavik_reply_fragments_e2e PROPERTIES TIMEOUT 120) + add_test( + NAME lavik_large_bulk_e2e + COMMAND ${Python3_EXECUTABLE} + ${CMAKE_CURRENT_SOURCE_DIR}/tests/large_bulk_e2e.py + $ + ) + set_tests_properties(lavik_large_bulk_e2e PROPERTIES TIMEOUT 180) + add_executable(lavik_grouped_recovery_e2e_test tests/grouped/recovery_e2e_test.cpp ) diff --git a/bycorf b/bycorf index 8caaf75cb..79a0f2c75 160000 --- a/bycorf +++ b/bycorf @@ -1 +1 @@ -Subproject commit 8caaf75cb131b591922af51a19861f92e2ffc893 +Subproject commit 79a0f2c754330280f0a1487093115b7648fd822c diff --git a/docs/architecture/02-request-serving.md b/docs/architecture/02-request-serving.md index 5af85d956..d1a19a609 100644 --- a/docs/architecture/02-request-serving.md +++ b/docs/architecture/02-request-serving.md @@ -116,16 +116,23 @@ racing closure is either rejected or remains visible to the drain. ## Request pipeline -1. `ReadCommandBatch` reads into a connection-local buffer and feeds - `RespCommandParser`. Up to 128 parsed commands are retained in wire order. +1. `ReadCommandBatch` feeds `RespCommandParser` through a small connection-local + buffer. Large bulk bodies can instead receive directly into parser-owned + argument storage after their first payload bytes arrive. The opt-in Bycorf + read drains already-received buffers without waiting to fill the destination; + TLS continues through its decrypted read path. Up to 128 parsed commands are + retained in wire order. Ordinary connections collectively hold at most `maxmemory-clients`, which accepts an absolute byte size or a percentage of effective maxmemory and defaults to 5%. The effective limit is divided into fixed worker shares. This abuse-control quota is independent of retained-memory admission; the server does not reserve temporary request bytes against `maxmemory` a second - time. Admission runs once after each socket read; - parsed wire bytes follow their command until execution finishes, while blank - input is retired as soon as the parser returns to its idle state. The + time. Ordinary reads admit received bytes; direct bulk reads admit bounded + destination windows before allocation. Prepaid bytes become received-byte + charges without a second admission, and unread reservations are released + on connection teardown. Parsed wire bytes follow their command until + execution finishes, while blank input is retired as soon as the parser + returns to its idle state. The offending connection receives an error and closes when its worker's share is exhausted. Zero disables this client-specific limit. As in Valkey, every nonzero process-wide client allowance has a 128 KiB @@ -136,14 +143,13 @@ racing closure is either rejected or remains visible to the drain. Independently, `client-query-buffer-limit` bounds the wire bytes retained while one connection incrementally assembles a command. It defaults to 1 GiB and, like Redis, accepts an absolute value from 1 MiB through - `LONG_MAX`. Lavik consumes the small socket input window directly into - argument strings, so this is a parser-retained-byte limit rather than a - requirement for a second contiguous query buffer. Completed pipelined and + `LONG_MAX`. Both receive paths assemble argument strings without requiring + a second contiguous command buffer. Completed pipelined and `MULTI` commands remain governed by `maxmemory-clients` after parser ownership ends. `CONFIG SET client-query-buffer-limit` updates a process-wide - atomic. Existing connections refresh it once per socket-read parsing round, - avoiding an atomic load for each RESP token while ensuring a partially read - command is checked against the new limit when more bytes arrive. + atomic. Existing connections refresh it once per socket-read parsing round + and before committing a direct read, so a partially read command is checked + against the new limit when more bytes arrive. 2. `BuildCommandRequest` resolves case-insensitive static metadata from the command table and copies the connection's current database ID into the request. diff --git a/docs/operations/metrics.md b/docs/operations/metrics.md index 2955b426a..0be71428d 100644 --- a/docs/operations/metrics.md +++ b/docs/operations/metrics.md @@ -62,7 +62,11 @@ It accepts either a percentage such as `5%` or an absolute size such as `512mb`; the default is `5%`, and `0` disables this client-specific limit. Every nonzero value has a 128 KiB effective minimum so an accidentally tiny setting does not make administrative commands unusable. The process-wide value -is divided into fixed worker shares and checked once after each socket read. +is divided into fixed worker shares. Ordinary reads charge received bytes; +large bulk reads charge bounded destination windows before allocation, then +consume that reservation as data arrives. Unread reservations are released +when the connection closes. The quota tracks received and reserved request +bytes, not allocator capacity or exact heap usage. `--client-query-buffer-limit` is the independent per-connection hard limit for an incomplete command. It accepts an absolute Redis-style memory size, defaults @@ -72,7 +76,8 @@ configuration files. The connection closes if incremental parsing crosses the limit; completed commands continue to count against `maxmemory-clients` until execution or transaction teardown releases them. `CONFIG GET/SET client-query-buffer-limit` reads or changes the live process-wide value; -existing connections apply a change before their next parsing round. +existing connections apply a change before their next parsing round and before +committing bytes received directly into a large argument. The process budget is split evenly across workers and unused capacity is not borrowed across workers. SET/MSET/INCR performs one worker-local retained-growth @@ -218,7 +223,8 @@ sum by (result) (rate(lavik_storage_defrag_runs_total[5m])) - `lavik_client_request_buffer_limit_bytes`: effective client request-buffer limit after resolving percentages and the nonzero 128 KiB floor. - `lavik_client_buffered_request_bytes`: wire bytes currently retained by - ordinary parsers, ready command batches, and queued transactions. + ordinary parsers, ready command batches, and queued transactions, plus + prepaid destination bytes for incomplete direct bulk reads. - `lavik_fullsync_reserved_memory_bytes`: reusable retained-memory headroom promised to active full-sync coverage maps. - `lavik_memory_admission_pending_bytes`: short-lived worker-local permits diff --git a/include/lavik/memory.h b/include/lavik/memory.h index 1e1d13b10..a82cfffe4 100644 --- a/include/lavik/memory.h +++ b/include/lavik/memory.h @@ -133,10 +133,10 @@ class RetainedMemoryCharge { std::size_t bytes_ = 0; }; -// Accounts request bytes read from ordinary client connections until their -// command finishes. One instance belongs to one connection and must remain on -// the worker where it was created; this keeps acquire/release operations local -// to that worker's cache line. +// Accounts received request bytes and prepaid direct-bulk destinations from +// ordinary client connections until their command finishes or the connection +// closes. One instance belongs to one connection and must remain on the worker +// where it was created; acquire/release touches only that worker's cache line. class ClientBufferReservation { public: ClientBufferReservation() noexcept; @@ -144,9 +144,10 @@ class ClientBufferReservation { ClientBufferReservation& operator=(const ClientBufferReservation&) = delete; ~ClientBufferReservation(); - // Acquires bytes from the calling worker's share of the 5% client-buffer - // budget. The check happens once after a socket read, not for every parser - // append. Failure leaves the reservation unchanged. + // Acquires bytes from the calling worker's configured client-buffer share. + // Ordinary reads acquire received bytes; direct bulk reads acquire a bounded + // destination before allocation and do not recharge it as data arrives. + // Failure leaves the reservation unchanged. [[nodiscard]] bool TryAcquire(std::size_t bytes) noexcept; // Releases bytes after parser-discarded input or a completed command. The // caller must not release more than it has acquired. diff --git a/include/lavik/resp.h b/include/lavik/resp.h index 7c721d058..000693b20 100644 --- a/include/lavik/resp.h +++ b/include/lavik/resp.h @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -64,6 +65,27 @@ class RespCommandParser { RespParseResult Parse(std::string_view input); + // Suggests bounded additional body storage for a partially received large + // bulk argument. Returns zero until payload arrives, while prepared storage + // remains, or when the current query limit leaves no allowance. The caller + // must admit this many request-buffer bytes before allocating them. + [[nodiscard]] std::size_t BulkReadReservationSize() const noexcept; + // Preallocates 1..BulkReadReservationSize() bytes in the final argument. + // This does not consume input or account wire bytes; unused space stays + // owned by the parser until filled or Reset. Invalid sizes leave it intact. + absl::Status PrepareBulkRead(std::size_t bytes); + // Returns the next bounded writable body window, excluding the terminator. + // Commit bytes written here instead of passing them to Parse. The span is + // invalidated by Parse, Reset, or another preparation; none may run while a + // transport read owns it. The caller retains its admission across the read. + // A lower query limit does not revoke this admitted span; Commit rechecks it. + [[nodiscard]] std::span BulkReadBuffer() noexcept; + // Commits at most the exposed window's bytes and rechecks the latest query + // limit. Partial/zero commits are allowed while a prepared window remains; + // failure leaves logical parser state intact. Parse still consumes CRLF and + // returns completed commands. It can also fill a partially used window. + absl::Status CommitBulkRead(std::size_t bytes); + void Reset(); // Changes the guard without disturbing partially parsed state. If the new // limit is already below the retained command size, the next input byte is @@ -97,6 +119,9 @@ class RespCommandParser { std::string current_argument_; std::size_t arguments_remaining_ = 0; std::size_t bulk_remaining_ = 0; + // String size includes this prepared but unwritten tail. Received payload + // ends before it; bulk_remaining_ includes it. Parse-only callers keep zero. + std::size_t bulk_prepared_ = 0; std::size_t terminator_bytes_ = 0; std::size_t command_bytes_ = 0; std::size_t query_buffer_limit_ = kDefaultClientQueryBufferLimit; diff --git a/src/metrics.cpp b/src/metrics.cpp index 879a49792..33920b027 100644 --- a/src/metrics.cpp +++ b/src/metrics.cpp @@ -667,7 +667,7 @@ bycorf::Task RenderPrometheusMetrics( "lavik_client_request_buffer_limit_bytes ", memory_metrics.client_buffer_limit_bytes_, "\n", "# HELP lavik_client_buffered_request_bytes Ordinary client request " - "bytes read but not yet retired.\n" + "bytes received or reserved for direct reads but not yet retired.\n" "# TYPE lavik_client_buffered_request_bytes gauge\n" "lavik_client_buffered_request_bytes ", memory_metrics.client_buffered_bytes_, "\n", diff --git a/src/redis/resp.cpp b/src/redis/resp.cpp index b0e2e5cfe..6f23db19d 100644 --- a/src/redis/resp.cpp +++ b/src/redis/resp.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -36,10 +37,9 @@ namespace { constexpr long long kMaxArrayLen = std::numeric_limits::max(); constexpr std::size_t kInitialArgCapacity = 1024; constexpr std::size_t kMaxBulkLen = 512ULL * 1024 * 1024; -// TODO(memory-control): Once retained-memory admission is settled, let bulk -// arguments >= 32 KiB expose their final string storage as the ReadSome() -// destination. This removes the RequestInputBuffer-to-string copy without -// pinning a shared provided buffer across an await or a stalled client. +constexpr std::size_t kDirectBulkThreshold = 32 * 1024; +constexpr std::size_t kDirectBulkReservation = 1024 * 1024; +constexpr std::size_t kDirectBulkReadBytes = 256 * 1024; constexpr std::size_t kMaxLengthTextBytes = 32; absl::Status AppendArgument(std::string* output, std::string_view value) { if (value.size() > std::numeric_limits::max() - output->size()) { @@ -168,12 +168,59 @@ void RespCommandParser::ResetCommand() { command_ = RespCommand{}; arguments_remaining_ = 0; bulk_remaining_ = 0; + bulk_prepared_ = 0; terminator_bytes_ = 0; command_bytes_ = 0; } void RespCommandParser::Reset() { ResetCommand(); } +std::size_t RespCommandParser::BulkReadReservationSize() const noexcept { + if (state_ != State::kBulkData || current_argument_.empty() || + bulk_prepared_ != 0 || + current_argument_.size() + bulk_remaining_ < kDirectBulkThreshold || + command_bytes_ >= query_buffer_limit_) + return 0; + return std::min({bulk_remaining_, kDirectBulkReservation, + query_buffer_limit_ - command_bytes_}); +} + +absl::Status RespCommandParser::PrepareBulkRead(std::size_t bytes) { + if (bytes == 0 || bytes > BulkReadReservationSize()) + return absl::InvalidArgumentError("invalid RESP bulk read reservation"); + // Resize makes the whole writable range valid string storage. Tracking its + // unread tail prevents initialized but unreceived bytes becoming payload. + // Allocation failure follows Parse's process-fatal allocation policy. + current_argument_.resize(current_argument_.size() + bytes); + bulk_prepared_ = bytes; + return absl::OkStatus(); +} + +std::span RespCommandParser::BulkReadBuffer() noexcept { + if (state_ != State::kBulkData) return {}; + const auto bytes = std::min(bulk_prepared_, kDirectBulkReadBytes); + return {current_argument_.data() + current_argument_.size() - bulk_prepared_, + bytes}; +} + +absl::Status RespCommandParser::CommitBulkRead(std::size_t bytes) { + if (state_ != State::kBulkData || bulk_prepared_ == 0 || + bytes > std::min(bulk_prepared_, kDirectBulkReadBytes)) + return absl::InvalidArgumentError("invalid RESP bulk read length"); + // The query limit can change while a socket read is suspended. Recheck + // accounting before advancing any logical state, even for an earlier span. + if (!Account(bytes)) + return absl::ResourceExhaustedError( + "client request exceeds the query buffer limit"); + bulk_prepared_ -= bytes; + bulk_remaining_ -= bytes; + if (bulk_remaining_ == 0) { + terminator_bytes_ = 0; + state_ = State::kBulkTerminator; + } + return absl::OkStatus(); +} + bool RespCommandParser::Account(std::size_t bytes) { if (bytes > query_buffer_limit_ || command_bytes_ > query_buffer_limit_ - bytes) { @@ -312,6 +359,7 @@ RespParseResult RespCommandParser::Parse(std::string_view input) { absl::OutOfRangeError("RESP bulk string too large"), pos); } bulk_remaining_ = static_cast(parsed); + assert(bulk_prepared_ == 0); current_argument_.clear(); state_ = State::kBulkData; } @@ -409,7 +457,18 @@ RespParseResult RespCommandParser::Parse(std::string_view input) { // length_error handler on every ordinary RESP argument is redundant. // Allocation failure remains process-fatal through the coroutine's // unhandled-exception policy. - current_argument_.append(input.substr(pos, take)); + if (bulk_prepared_ == 0) [[likely]] { + current_argument_.append(input.substr(pos, take)); + } else { + const auto prepared = std::min(take, bulk_prepared_); + std::memcpy(current_argument_.data() + current_argument_.size() - + bulk_prepared_, + input.data() + pos, prepared); + bulk_prepared_ -= prepared; + if (take > prepared) + current_argument_.append( + input.substr(pos + prepared, take - prepared)); + } bulk_remaining_ -= take; if (!consume(take)) { return Error(absl::ResourceExhaustedError( diff --git a/src/redis/server.cpp b/src/redis/server.cpp index 68d09a765..a3b0cc2bf 100644 --- a/src/redis/server.cpp +++ b/src/redis/server.cpp @@ -1044,6 +1044,8 @@ Task ReadCommandBatch( const std::atomic* query_buffer_limit, std::size_t* unassigned_input_bytes, std::optional* deferred_error) { + constexpr std::size_t kBulkBytesPerYield = 256 * 1024; + std::size_t bulk_bytes_since_yield = 0; while (ready->empty()) { // CONFIG SET may run on another worker while this coroutine is suspended // in ReadSome(). Refresh once per parsing round; token-level loads would @@ -1082,6 +1084,55 @@ Task ReadCommandBatch( } if (!ready->empty()) co_return absl::OkStatus(); + const std::size_t bulk_reservation = parser->BulkReadReservationSize(); + if (bulk_reservation != 0) { + // A direct destination includes bytes the peer has not sent yet. Charge + // that bounded window before allocation, then let received bytes inherit + // its allowance instead of charging them again. The connection owns any + // unread remainder until teardown; completed commands retain the usual + // wire-byte charge through execution (or MULTI). + if (!client_buffers->TryAcquire(bulk_reservation)) [[unlikely]] { + RecordMemoryRejection(); + co_return absl::ResourceExhaustedError( + "client request buffers exceed the memory limit"); + } + const auto prepared = parser->PrepareBulkRead(bulk_reservation); + if (!prepared.ok()) [[unlikely]] { + client_buffers->Release(bulk_reservation); + co_return prepared; + } + } + if (auto body = parser->BulkReadBuffer(); !body.empty()) { + body = body.first( + std::min(body.size(), kBulkBytesPerYield - bulk_bytes_since_yield)); + auto read_result = + co_await stream.ReadSomeBuffered(std::as_writable_bytes(body)); + if (!read_result.ok()) [[unlikely]] + co_return read_result.status(); + if (*read_result == 0) [[unlikely]] { + co_return absl::UnavailableError("peer closed connection"); + } + // The limit may have shrunk while the direct read was suspended. Never + // dispatch a command using a stale guard merely because its destination + // was admitted before CONFIG SET ran. + parser->SetQueryBufferLimit( + query_buffer_limit->load(std::memory_order_acquire)); + const auto committed = parser->CommitBulkRead(*read_result); + if (!committed.ok()) [[unlikely]] + co_return committed; + assert(*read_result <= + std::numeric_limits::max() - *unassigned_input_bytes); + *unassigned_input_bytes += *read_result; + // Buffered reads can finish without suspending. Bound uninterrupted + // receive work so a large body cannot monopolize the connection worker. + bulk_bytes_since_yield += *read_result; + if (bulk_bytes_since_yield == kBulkBytesPerYield) { + bulk_bytes_since_yield = 0; + co_await Yield(*ThisWorker().self_); + } + continue; + } + std::span buffer = input->AppendBuffer(); auto read_result = co_await stream.ReadSome(buffer); if (!read_result.ok()) [[unlikely]] { diff --git a/tests/large_bulk_e2e.py b/tests/large_bulk_e2e.py new file mode 100644 index 000000000..2247901d1 --- /dev/null +++ b/tests/large_bulk_e2e.py @@ -0,0 +1,360 @@ +#!/usr/bin/env python3 +# Copyright (C) 2026 EloqData Inc. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Large request ownership, incremental admission, and cancellation over TCP/TLS.""" + +from contextlib import contextmanager +import os +from pathlib import Path +import socket +import ssl +import sys +import tempfile +import time + +import redis_follower_smoke as S + + +MIB = 1024 * 1024 + + +def bulk(value): + return b"$%d\r\n" % len(value) + value + b"\r\n" + + +def wire(*args): + return b"*%d\r\n" % len(args) + b"".join( + bulk(value if isinstance(value, bytes) else str(value).encode()) + for value in args + ) + + +def set_prefix(key, size): + return b"*3\r\n" + bulk(b"SET") + bulk(key.encode()) + b"$%d\r\n" % size + + +def payload(size): + # Embedded NUL, CRLF, non-UTF8 bytes, and RESP punctuation must survive + # direct reads just as they survive the ordinary incremental parser. + return (bytes(range(256)) * ((size + 255) // 256))[:size] + + +class ErrorReply(Exception): + pass + + +def response(reader): + line = reader.readline() + assert line.endswith(b"\r\n"), ("incomplete RESP response", line) + tag, body = line[:1], line[1:-2] + if tag == b"+": + return body + if tag == b"-": + raise ErrorReply(body) + if tag == b":": + return int(body) + if tag == b"$": + length = int(body) + if length == -1: + return None + value = reader.read(length) + assert len(value) == length, ("short bulk reply", len(value), length) + assert reader.read(2) == b"\r\n" + return value + if tag == b"*": + return [response(reader) for _ in range(int(body))] + raise AssertionError(("unexpected RESP response", line)) + + +class Client: + def __init__(self, sock, reader): + self.sock = sock + self.reader = reader + + def call(self, *args): + self.sock.sendall(wire(*args)) + return response(self.reader) + + +@contextmanager +def connection(port, tls=False): + sock = socket.create_connection(("127.0.0.1", port), timeout=10) + if tls: + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + sock = context.wrap_socket(sock, server_hostname="localhost") + with sock, sock.makefile("rb") as reader: + yield Client(sock, reader) + + +def buffered_bytes(control): + info = control.call("INFO", "memory") + for line in info.split(b"\r\n"): + if line.startswith(b"client_buffered_request_bytes:"): + return int(line.split(b":", 1)[1]) + raise AssertionError("INFO memory omitted client_buffered_request_bytes") + + +def wait_buffered(control, predicate, description): + deadline = time.monotonic() + 5 + while True: + actual = buffered_bytes(control) + if predicate(actual): + return actual + if time.monotonic() >= deadline: + raise AssertionError((description, "last buffered bytes", actual)) + time.sleep(0.01) + + +def wait_released(control, baseline): + # INFO's own request is charged while its snapshot is generated. Reusing + # exactly the same request gives a stable nonzero baseline, including when + # the tested connection still has a prepared but unfilled body window. + wait_buffered(control, lambda actual: actual == baseline, "request release") + + +def expect_error_and_close(client, message): + try: + result = response(client.reader) + except ErrorReply as error: + assert message in error.args[0], error.args[0] + else: + raise AssertionError(("expected error", message, result)) + try: + assert client.reader.read(1) == b"", "failed request connection stayed open" + except (ConnectionResetError, ssl.SSLEOFError): + pass + + +def check_roundtrips(port, tls, control, baseline, prefix): + with connection(port, tls) as client: + for index, size in enumerate((MIB, 2 * MIB + 137)): + key = f"{prefix}:fragmented:{index}" + value = payload(size) + header = set_prefix(key, size) + sent = 0 + # Observe each partial header before sending the next fragment so + # TCP coalescing cannot hide the split length/CRLF parser states. + for end in (3, 4, len(header) - 1, len(header)): + client.sock.sendall(header[sent:end]) + sent = end + wait_buffered( + control, + lambda actual: actual == baseline + sent, + "fragmented header", + ) + client.sock.sendall(value[:1024]) + wait_buffered( + control, + lambda actual: actual >= baseline + len(header) + min(size, MIB), + "large body admission", + ) + for offset in range(1024, size, 65537): + client.sock.sendall(value[offset : offset + 65537]) + client.sock.sendall(b"\r") + wait_buffered( + control, + lambda actual: actual == baseline + len(header) + size + 1, + "fragmented final CRLF", + ) + small_key = f"{prefix}:small:{index}" + client.sock.sendall( + b"\n" + + wire("SET", small_key, b"small\x00value") + + wire("GET", small_key) + + wire("GET", key) + + wire("PING") + ) + for expected in (b"OK", b"OK", b"small\x00value", value, b"PONG"): + assert response(client.reader) == expected + wait_released(control, baseline) + + first, second = payload(MIB + 17), payload(MIB + 8193) + keys = [f"{prefix}:mset:{index}" for index in range(2)] + client.sock.sendall( + wire("MSET", keys[0], first, keys[1], second) + + wire("GET", keys[1]) + + wire("SET", f"{prefix}:after-mset", "small") + + wire("GET", keys[0]) + ) + for expected in (b"OK", second, b"OK", first): + assert response(client.reader) == expected + wait_released(control, baseline) + + +def check_transactions(port, tls, control, baseline, prefix): + value = payload(MIB + 19) + for ending in ("DISCARD", "EXEC", "disconnect"): + keys = [f"{prefix}:transaction:{ending}:{index}" for index in range(2)] + frames = [wire("SET", key, value) for key in keys] + with connection(port, tls) as client: + assert client.call("MULTI") == b"OK" + client.sock.sendall(b"".join(frames)) + assert response(client.reader) == b"QUEUED" + assert response(client.reader) == b"QUEUED" + wait_buffered( + control, + lambda actual: actual == baseline + sum(map(len, frames)), + "MULTI retains exactly the received command bytes", + ) + assert control.call("GET", keys[0]) is None + if ending == "EXEC": + assert client.call(ending) == [b"OK", b"OK"] + elif ending == "DISCARD": + assert client.call(ending) == b"OK" + wait_released(control, baseline) + for key in keys: + assert control.call("GET", key) == (value if ending == "EXEC" else None) + + +def check_incomplete_and_invalid(port, tls, control, baseline, prefix): + key = f"{prefix}:incomplete" + header = set_prefix(key, 64 * MIB) + with connection(port, tls) as client: + client.sock.sendall(header) + # A declared length alone must not allocate/admit the entire body or + # even the first direct-read window. Progress is observed, not timed. + wait_buffered( + control, + lambda actual: actual == baseline + len(header), + "header-only request retains only header bytes", + ) + client.sock.sendall(b"x" * 1024) + admitted = wait_buffered( + control, + lambda actual: actual >= baseline + MIB, + "partial body prepays one bounded window", + ) + assert admitted < baseline + 2 * MIB, admitted + # Close with an outstanding read into the parser's reserved string. + wait_released(control, baseline) + assert control.call("GET", key) is None + + key = f"{prefix}:malformed" + with connection(port, tls) as client: + client.sock.sendall(wire("PING") + set_prefix(key, MIB) + payload(MIB) + b"\rX") + assert response(client.reader) == b"PONG" + expect_error_and_close(client, b"terminator") + wait_released(control, baseline) + assert control.call("GET", key) is None + assert control.call("PING") == b"PONG" + + +def check_query_limit_shrink(port, tls, control, baseline, prefix): + config = "client-query-buffer-limit" + previous = control.call("CONFIG", "GET", config)[1] + key = f"{prefix}:query-limit" + try: + with connection(port, tls) as client: + client.sock.sendall(set_prefix(key, 3 * MIB) + payload(MIB + 65536)) + # A second prepaid window proves the first 1 MiB window was + # consumed. The smallest valid runtime limit is 1 MiB, so the + # retained command now exceeds it before the suspended read ends. + wait_buffered( + control, + lambda actual: actual > baseline + 2 * MIB, + "second receive window before runtime limit shrink", + ) + assert control.call("CONFIG", "SET", config, "1mb") == b"OK" + client.sock.sendall(b"x") + expect_error_and_close(client, b"query buffer limit") + wait_released(control, baseline) + assert control.call("GET", key) is None + finally: + assert control.call("CONFIG", "SET", config, previous) == b"OK" + + +def check(port, tls, control, baseline): + prefix = "large-bulk:tls" if tls else "large-bulk:tcp" + check_roundtrips(port, tls, control, baseline, prefix) + check_transactions(port, tls, control, baseline, prefix) + check_incomplete_and_invalid(port, tls, control, baseline, prefix) + check_query_limit_shrink(port, tls, control, baseline, prefix) + wait_released(control, baseline) + + +def check_quota(port, tls, control, baseline): + prefix = "large-bulk:quota:tls" if tls else "large-bulk:quota:tcp" + queued_key, partial_key = f"{prefix}:queued", f"{prefix}:partial" + value = payload(MIB) + with connection(port, tls) as client: + assert client.call("MULTI") == b"OK" + assert client.call("SET", queued_key, value) == b"QUEUED" + wait_buffered( + control, + lambda actual: actual == baseline + len(wire("SET", queued_key, value)), + "queued request before quota exhaustion", + ) + # The full declared body is never sent. Its next bounded window alone + # exceeds the remaining quota while MULTI owns the earlier command. + client.sock.sendall(set_prefix(partial_key, 2 * MIB) + b"x" * 1024) + expect_error_and_close(client, b"memory limit") + wait_released(control, baseline) + assert control.call("GET", queued_key) is None + assert control.call("GET", partial_key) is None + with connection(port, tls) as healthy: + assert healthy.call("PING") == b"PONG" + assert healthy.call("ECHO", value) == value + wait_released(control, baseline) + + +def tls_options(port): + certificates = Path(__file__).resolve().parent / "tls" + return [ + "--tls-port", + str(port), + "--tls-cert-file", + str(certificates / "server.crt"), + "--tls-key-file", + str(certificates / "server.key"), + "--tls-auth-clients", + "no", + ] + + +def run(root, binary): + tls_port = S.H.free_port() + with S.process( + binary, + root / "roundtrips", + "large-bulk", + workers=1, + extra=tls_options(tls_port), + ) as (_, port, _): + with connection(port) as control: + baseline = buffered_bytes(control) + check(port, False, control, baseline) + check(tls_port, True, control, baseline) + + tls_port = S.H.free_port() + with S.process( + binary, + root / "quota", + "large-bulk-quota", + workers=1, + extra=tls_options(tls_port) + ["--maxmemory-clients", "2mb"], + ) as (_, port, _): + with connection(port) as control: + baseline = buffered_bytes(control) + check_quota(port, False, control, baseline) + check_quota(tls_port, True, control, baseline) + + +if __name__ == "__main__": + assert len(sys.argv) == 2, "usage: large_bulk_e2e.py /path/to/lavik" + binary = str(Path(sys.argv[1]).resolve()) + with tempfile.TemporaryDirectory( + prefix="lavik-large-bulk-", dir=os.environ.get("LAVIK_TEST_DATA_DIR") + ) as directory: + run(Path(directory), binary) + print("large bulk receive and admission checks passed over TCP and TLS") diff --git a/tests/resp_test.cpp b/tests/resp_test.cpp index 9d54ffb11..fd14382ea 100644 --- a/tests/resp_test.cpp +++ b/tests/resp_test.cpp @@ -18,6 +18,7 @@ #include +#include #include #include #include @@ -119,6 +120,194 @@ TEST(RespParserTest, AppliesLowerRuntimeLimitToIncompleteCommand) { EXPECT_EQ(rejected.status_.code(), absl::StatusCode::kResourceExhausted); } +TEST(RespParserTest, DirectBulkReadsRequirePayloadAndValidateWindows) { + RespCommandParser parser; + EXPECT_EQ(parser.BulkReadReservationSize(), 0); + EXPECT_TRUE(parser.BulkReadBuffer().empty()); + EXPECT_EQ(parser.CommitBulkRead(0).code(), + absl::StatusCode::kInvalidArgument); + ASSERT_EQ(parser.Parse("*1\r\n$536870912\r\n").state_, + RespParseState::kNeedMoreData); + EXPECT_EQ(parser.BulkReadReservationSize(), 0); + EXPECT_EQ(parser.PrepareBulkRead(1).code(), + absl::StatusCode::kInvalidArgument); + ASSERT_EQ(parser.Parse("x").state_, RespParseState::kNeedMoreData); + const auto reservation = parser.BulkReadReservationSize(); + EXPECT_EQ(reservation, 1024 * 1024); + EXPECT_EQ(parser.PrepareBulkRead(0).code(), + absl::StatusCode::kInvalidArgument); + EXPECT_EQ(parser.PrepareBulkRead(reservation + 1).code(), + absl::StatusCode::kInvalidArgument); + ASSERT_TRUE(parser.PrepareBulkRead(reservation).ok()); + EXPECT_EQ(parser.BulkReadReservationSize(), 0); + EXPECT_EQ(parser.PrepareBulkRead(1).code(), + absl::StatusCode::kInvalidArgument); + auto window = parser.BulkReadBuffer(); + EXPECT_EQ(window.size(), 256 * 1024); + EXPECT_EQ(parser.CommitBulkRead(window.size() + 1).code(), + absl::StatusCode::kInvalidArgument); + ASSERT_TRUE(parser.CommitBulkRead(0).ok()); + EXPECT_EQ(parser.BulkReadBuffer().data(), window.data()); + window[0] = 'y'; + ASSERT_TRUE(parser.CommitBulkRead(1).ok()); + EXPECT_EQ(parser.BulkReadBuffer().data(), window.data() + 1); + parser.Reset(); + ASSERT_EQ(parser.Parse("*1\r\n$32767\r\nx").state_, + RespParseState::kNeedMoreData); + EXPECT_EQ(parser.BulkReadReservationSize(), 0); +} + +TEST(RespParserTest, DirectBulkReadsUseBoundedWindowsForMultiMegabytePayloads) { + std::string body(2 * 1024 * 1024 + 13, '\0'); + for (std::size_t i = 0; i < body.size(); ++i) body[i] = i % 251; + const std::string header = "*1\r\n$" + std::to_string(body.size()) + "\r\n"; + RespCommandParser parser(header.size() + body.size() + 2); + ASSERT_EQ(parser.Parse(header + body.substr(0, 1)).state_, + RespParseState::kNeedMoreData); + std::size_t position = 1, reservations = 0; + while (position < body.size()) { + if (parser.BulkReadBuffer().empty()) { + const auto bytes = parser.BulkReadReservationSize(); + ASSERT_GT(bytes, 0); + ASSERT_LE(bytes, 1024 * 1024); + ASSERT_TRUE(parser.PrepareBulkRead(bytes).ok()); + ++reservations; + } + const auto window = parser.BulkReadBuffer(); + ASSERT_FALSE(window.empty()); + ASSERT_LE(window.size(), 256 * 1024); + ASSERT_LE(window.size(), body.size() - position); + std::memcpy(window.data(), body.data() + position, window.size()); + ASSERT_TRUE(parser.CommitBulkRead(window.size()).ok()); + position += window.size(); + } + EXPECT_EQ(reservations, 3); + EXPECT_TRUE(parser.BulkReadBuffer().empty()); + EXPECT_EQ(parser.BulkReadReservationSize(), 0); + EXPECT_EQ(parser.CommitBulkRead(0).code(), + absl::StatusCode::kInvalidArgument); + auto parsed = parser.Parse("\r\n"); + ASSERT_EQ(parsed.state_, RespParseState::kOk) << parsed.status_; + ASSERT_EQ(parsed.command_.args_.size(), 1); + EXPECT_EQ(parsed.command_.args_[0], body); + EXPECT_TRUE(parser.idle()); +} + +TEST(RespParserTest, DirectBulkReadsAllowPartialCommitsAndParseFallback) { + std::string body(48000, '\0'); + for (std::size_t i = 0; i < body.size(); ++i) body[i] = i % 251; + RespCommandParser parser; + ASSERT_EQ(parser.Parse("*1\r\n$48000\r\n" + body.substr(0, 19)).state_, + RespParseState::kNeedMoreData); + ASSERT_TRUE(parser.PrepareBulkRead(512).ok()); + auto window = parser.BulkReadBuffer(); + std::memcpy(window.data(), body.data() + 19, 7); + ASSERT_TRUE(parser.CommitBulkRead(7).ok()); + ASSERT_EQ(parser.BulkReadBuffer().size(), 505); + auto partial = parser.Parse(std::string_view(body).substr(26, 100)); + ASSERT_EQ(partial.state_, RespParseState::kNeedMoreData); + EXPECT_TRUE(partial.command_.args_.empty()); + EXPECT_EQ(parser.BulkReadBuffer().data(), window.data() + 107); + EXPECT_EQ(parser.BulkReadBuffer().size(), 405); + // The same call fills the remaining prepared tail, then appends beyond it. + auto parsed = parser.Parse(body.substr(126) + "\r\n"); + ASSERT_EQ(parsed.state_, RespParseState::kOk) << parsed.status_; + ASSERT_EQ(parsed.command_.args_.size(), 1); + EXPECT_EQ(parsed.command_.args_[0], body); + EXPECT_TRUE(parser.BulkReadBuffer().empty()); +} + +TEST(RespParserTest, DirectBulkReadsPreserveMultipleArgumentsAndPipelining) { + RespCommandParser parser; + const std::string first(32768, 'a'), + second = "txt:" + std::string(48000, 'b'); + ASSERT_EQ(parser.Parse("*2\r\n$32768\r\na").state_, + RespParseState::kNeedMoreData); + for (const auto* body : {&first, &second}) { + ASSERT_EQ(parser.BulkReadReservationSize(), body->size() - 1); + ASSERT_TRUE(parser.PrepareBulkRead(body->size() - 1).ok()); + auto window = parser.BulkReadBuffer(); + ASSERT_EQ(window.size(), body->size() - 1); + std::memcpy(window.data(), body->data() + 1, window.size()); + ASSERT_TRUE(parser.CommitBulkRead(window.size()).ok()); + ASSERT_EQ(parser.Parse("\r").state_, RespParseState::kNeedMoreData); + if (body == &first) + ASSERT_EQ(parser.Parse("\n=48004\r\nt").state_, + RespParseState::kNeedMoreData); + } + const std::string pipeline = "\n*1\r\n$4\r\nPING\r\n"; + auto parsed = parser.Parse(pipeline); + ASSERT_EQ(parsed.state_, RespParseState::kOk) << parsed.status_; + EXPECT_EQ(parsed.consumed_, 1); + EXPECT_EQ(parsed.command_.args_, + (std::vector{first, second.substr(4)})); + auto next = parser.Parse(std::string_view(pipeline).substr(parsed.consumed_)); + ASSERT_EQ(next.state_, RespParseState::kOk); + EXPECT_EQ(next.command_.args_, (std::vector{"PING"})); +} + +TEST(RespParserTest, DirectBulkReadsHonorQueryLimitsAtPreparationAndCommit) { + const std::string prefix = "*1\r\n$48000\r\nx"; + RespCommandParser parser(prefix.size() + 10); + ASSERT_EQ(parser.Parse(prefix).state_, RespParseState::kNeedMoreData); + EXPECT_EQ(parser.BulkReadReservationSize(), 10); + EXPECT_EQ(parser.PrepareBulkRead(11).code(), + absl::StatusCode::kInvalidArgument); + ASSERT_TRUE(parser.PrepareBulkRead(10).ok()); + auto window = parser.BulkReadBuffer(); + std::memset(window.data(), 'y', window.size()); + ASSERT_TRUE(parser.CommitBulkRead(3).ok()); + parser.SetQueryBufferLimit(prefix.size() + 5); + // A runtime change preserves the prepaid storage across a suspended read. + EXPECT_EQ(parser.BulkReadBuffer().size(), 7); + EXPECT_EQ(parser.CommitBulkRead(3).code(), + absl::StatusCode::kResourceExhausted); + EXPECT_EQ(parser.BulkReadBuffer().data(), window.data() + 3); + ASSERT_TRUE(parser.CommitBulkRead(2).ok()); + parser.SetQueryBufferLimit(prefix.size() + 4); + EXPECT_EQ(parser.CommitBulkRead(1).code(), + absl::StatusCode::kResourceExhausted); + EXPECT_EQ(parser.BulkReadBuffer().data(), window.data() + 5); + auto fallback = parser.Parse("z"); + EXPECT_EQ(fallback.state_, RespParseState::kError); + EXPECT_EQ(fallback.status_.code(), absl::StatusCode::kResourceExhausted); +} + +TEST(RespParserTest, DirectBulkResetDiscardsUnwrittenBody) { + RespCommandParser parser; + ASSERT_EQ(parser.Parse("*1\r\n$48000\r\nx").state_, + RespParseState::kNeedMoreData); + ASSERT_TRUE(parser.PrepareBulkRead(47999).ok()); + parser.BulkReadBuffer()[0] = 'y'; + ASSERT_TRUE(parser.CommitBulkRead(1).ok()); + parser.Reset(); + EXPECT_TRUE(parser.idle()); + EXPECT_EQ(parser.BulkReadReservationSize(), 0); + EXPECT_TRUE(parser.BulkReadBuffer().empty()); + EXPECT_EQ(parser.CommitBulkRead(1).code(), + absl::StatusCode::kInvalidArgument); + auto parsed = parser.Parse("*1\r\n$4\r\nPING\r\n"); + ASSERT_EQ(parsed.state_, RespParseState::kOk); + EXPECT_EQ(parsed.command_.args_, (std::vector{"PING"})); +} + +TEST(RespParserTest, DirectBulkReadsRetainTerminatorAndVerbatimValidation) { + for (const char type : {'$', '='}) { + RespCommandParser parser; + ASSERT_EQ( + parser.Parse("*1\r\n" + std::string(1, type) + "48000\r\nx").state_, + RespParseState::kNeedMoreData); + ASSERT_TRUE(parser.PrepareBulkRead(47999).ok()); + auto window = parser.BulkReadBuffer(); + std::memset(window.data(), 'x', window.size()); + ASSERT_TRUE(parser.CommitBulkRead(window.size()).ok()); + const auto parsed = parser.Parse(type == '$' ? "\rx" : "\r\n"); + EXPECT_EQ(parsed.state_, RespParseState::kError); + EXPECT_EQ(parsed.status_.code(), absl::StatusCode::kInvalidArgument); + EXPECT_TRUE(parsed.command_.args_.empty()); + } +} + TEST(RespParserTest, RejectsArrayLengthsAboveValkeyLimit) { const auto too_many = static_cast(std::numeric_limits::max()) + 1; From 2665acc502455ec1977f4c87d9b223cbac41db9d Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Thu, 15 Oct 2026 04:46:58 +0000 Subject: [PATCH 3/3] test: use buffered receive fixture compatible with older kernels --- bycorf | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bycorf b/bycorf index 79a0f2c75..c4ec8208a 160000 --- a/bycorf +++ b/bycorf @@ -1 +1 @@ -Subproject commit 79a0f2c754330280f0a1487093115b7648fd822c +Subproject commit c4ec8208a15279e9484a07f74bdec51cfe8e6b5f