Optimize publish to avoid copying message data - #129
Merged
Conversation
Added appendMany() to ConcurrentWriteBuffer that accepts multiple slices and appends them in a single operation. Modified publishMsgInternal() to build control line and headers separately, then use appendMany() to send [control+headers, msg.data, "\r\n"] without copying the message data. This reduces memory allocation and eliminates one copy of the message data. For large messages, we now allocate ~4KB instead of allocating space for the entire message body.
Contributor
📝 WalkthroughWalkthroughThe PR introduces a new Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~15 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/queue.zig`:
- Around line 653-706: The appendMany implementation accumulates total_size and
later does self.queue.items_available += total_size without overflow checks;
replace the unchecked additions with overflow-checked arithmetic (use
`@addWithOverflow`) when summing total_size from slices and again when computing
new_items_available = self.queue.items_available + total_size, and if either
addition overflows return a suitable error (e.g., PushError.OutOfMemory or a new
overflow error); update the max_size checks to use the overflow-safe result so
the max_size guard cannot be bypassed by wrapping.
Comment on lines
+653
to
+706
| /// Append multiple slices to the buffer in a single operation | ||
| /// More efficient than calling append() multiple times as it only takes the lock once | ||
| pub fn appendMany(self: *Self, rt: *zio.Runtime, slices: []const []const u8) (zio.Cancelable || PushError)!void { | ||
| try self.queue.mutex.lock(rt); | ||
| defer self.queue.mutex.unlock(rt); | ||
|
|
||
| if (self.queue.is_closed) { | ||
| return PushError.QueueClosed; | ||
| } | ||
| if (self.queue.is_frozen) { | ||
| return PushError.BufferFrozen; | ||
| } | ||
|
|
||
| // Calculate total size needed | ||
| var total_size: usize = 0; | ||
| for (slices) |slice| { | ||
| total_size += slice.len; | ||
| } | ||
|
|
||
| // Check size limit before adding (overflow-safe) | ||
| if (self.queue.max_size > 0) { | ||
| const max_items = self.queue.max_size; | ||
| if (self.queue.items_available >= max_items) { | ||
| return PushError.OutOfMemory; | ||
| } | ||
| if (total_size > max_items - self.queue.items_available) { | ||
| return PushError.OutOfMemory; | ||
| } | ||
| } | ||
|
|
||
| // Append each slice | ||
| for (slices) |slice| { | ||
| var remaining = slice; | ||
|
|
||
| while (remaining.len > 0) { | ||
| const chunk = try self.queue.ensureWritableChunk(); | ||
|
|
||
| const available = chunk.availableToWrite(); | ||
| const to_write = @min(available, remaining.len); | ||
|
|
||
| @memcpy(chunk.getWriteSlice()[0..to_write], remaining[0..to_write]); | ||
|
|
||
| chunk.write_pos += to_write; | ||
| remaining = remaining[to_write..]; | ||
|
|
||
| if (chunk.availableToWrite() == 0) { | ||
| chunk.is_sealed = true; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| self.queue.items_available += total_size; | ||
| self.queue.data_cond.signal(rt); | ||
| } |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's check if the file exists and examine the relevant lines
if [ -f "src/queue.zig" ]; then
wc -l src/queue.zig
echo "---"
sed -n '653,706p' src/queue.zig
else
echo "File not found, searching for queue.zig"
find . -name "queue.zig" -type f
fiRepository: lalinsky/nats.zig
Length of output: 2109
🏁 Script executed:
# Also check for any overflow handling utilities or patterns in the codebase
rg "@addWithOverflow|@mulWithOverflow|std.math.add" --type zig -B 2 -A 2Repository: lalinsky/nats.zig
Length of output: 43
🏁 Script executed:
# Check if there's any existing overflow handling in queue operations
rg "appendMany|overflow|OutOfMemory" src/queue.zig -B 3 -A 3Repository: lalinsky/nats.zig
Length of output: 3224
Guard against usize overflow in appendMany size accounting.
The code accumulates total_size and updates items_available with wrapping arithmetic (+=), not overflow-checked operations. If either wraps, the max_size checks are bypassed and the counter becomes corrupted. The misleading "overflow-safe" comment doesn't match the implementation. Use @addWithOverflow with error handling and fail if either accumulation overflows.
Proposed fix
- var total_size: usize = 0;
- for (slices) |slice| {
- total_size += slice.len;
- }
+ var total_size: usize = 0;
+ for (slices) |slice| {
+ const sum = `@addWithOverflow`(total_size, slice.len);
+ if (sum[1] != 0) return PushError.OutOfMemory;
+ total_size = sum[0];
+ }
@@
- self.queue.items_available += total_size;
+ const new_total = `@addWithOverflow`(self.queue.items_available, total_size);
+ if (new_total[1] != 0) return PushError.OutOfMemory;
+ self.queue.items_available = new_total[0];
self.queue.data_cond.signal(rt);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| /// Append multiple slices to the buffer in a single operation | |
| /// More efficient than calling append() multiple times as it only takes the lock once | |
| pub fn appendMany(self: *Self, rt: *zio.Runtime, slices: []const []const u8) (zio.Cancelable || PushError)!void { | |
| try self.queue.mutex.lock(rt); | |
| defer self.queue.mutex.unlock(rt); | |
| if (self.queue.is_closed) { | |
| return PushError.QueueClosed; | |
| } | |
| if (self.queue.is_frozen) { | |
| return PushError.BufferFrozen; | |
| } | |
| // Calculate total size needed | |
| var total_size: usize = 0; | |
| for (slices) |slice| { | |
| total_size += slice.len; | |
| } | |
| // Check size limit before adding (overflow-safe) | |
| if (self.queue.max_size > 0) { | |
| const max_items = self.queue.max_size; | |
| if (self.queue.items_available >= max_items) { | |
| return PushError.OutOfMemory; | |
| } | |
| if (total_size > max_items - self.queue.items_available) { | |
| return PushError.OutOfMemory; | |
| } | |
| } | |
| // Append each slice | |
| for (slices) |slice| { | |
| var remaining = slice; | |
| while (remaining.len > 0) { | |
| const chunk = try self.queue.ensureWritableChunk(); | |
| const available = chunk.availableToWrite(); | |
| const to_write = @min(available, remaining.len); | |
| @memcpy(chunk.getWriteSlice()[0..to_write], remaining[0..to_write]); | |
| chunk.write_pos += to_write; | |
| remaining = remaining[to_write..]; | |
| if (chunk.availableToWrite() == 0) { | |
| chunk.is_sealed = true; | |
| } | |
| } | |
| } | |
| self.queue.items_available += total_size; | |
| self.queue.data_cond.signal(rt); | |
| } | |
| /// Append multiple slices to the buffer in a single operation | |
| /// More efficient than calling append() multiple times as it only takes the lock once | |
| pub fn appendMany(self: *Self, rt: *zio.Runtime, slices: []const []const u8) (zio.Cancelable || PushError)!void { | |
| try self.queue.mutex.lock(rt); | |
| defer self.queue.mutex.unlock(rt); | |
| if (self.queue.is_closed) { | |
| return PushError.QueueClosed; | |
| } | |
| if (self.queue.is_frozen) { | |
| return PushError.BufferFrozen; | |
| } | |
| // Calculate total size needed | |
| var total_size: usize = 0; | |
| for (slices) |slice| { | |
| const sum = `@addWithOverflow`(total_size, slice.len); | |
| if (sum[1] != 0) return PushError.OutOfMemory; | |
| total_size = sum[0]; | |
| } | |
| // Check size limit before adding (overflow-safe) | |
| if (self.queue.max_size > 0) { | |
| const max_items = self.queue.max_size; | |
| if (self.queue.items_available >= max_items) { | |
| return PushError.OutOfMemory; | |
| } | |
| if (total_size > max_items - self.queue.items_available) { | |
| return PushError.OutOfMemory; | |
| } | |
| } | |
| // Append each slice | |
| for (slices) |slice| { | |
| var remaining = slice; | |
| while (remaining.len > 0) { | |
| const chunk = try self.queue.ensureWritableChunk(); | |
| const available = chunk.availableToWrite(); | |
| const to_write = `@min`(available, remaining.len); | |
| `@memcpy`(chunk.getWriteSlice()[0..to_write], remaining[0..to_write]); | |
| chunk.write_pos += to_write; | |
| remaining = remaining[to_write..]; | |
| if (chunk.availableToWrite() == 0) { | |
| chunk.is_sealed = true; | |
| } | |
| } | |
| } | |
| const new_total = `@addWithOverflow`(self.queue.items_available, total_size); | |
| if (new_total[1] != 0) return PushError.OutOfMemory; | |
| self.queue.items_available = new_total[0]; | |
| self.queue.data_cond.signal(rt); | |
| } |
🤖 Prompt for AI Agents
In `@src/queue.zig` around lines 653 - 706, The appendMany implementation
accumulates total_size and later does self.queue.items_available += total_size
without overflow checks; replace the unchecked additions with overflow-checked
arithmetic (use `@addWithOverflow`) when summing total_size from slices and again
when computing new_items_available = self.queue.items_available + total_size,
and if either addition overflows return a suitable error (e.g.,
PushError.OutOfMemory or a new overflow error); update the max_size checks to
use the overflow-safe result so the max_size guard cannot be bypassed by
wrapping.
Owner
Author
There was a problem hiding this comment.
I'm really not concened about the total size of each message
Contributor
There was a problem hiding this comment.
@lalinsky, understood. Zig's overflow detection in Debug and ReleaseSafe modes will catch any issues during development, and in practice message sizes won't approach usize limits.
✏️ Learnings added
Learnt from: lalinsky
Repo: lalinsky/nats.zig PR: 129
File: src/queue.zig:653-706
Timestamp: 2026-01-28T05:45:40.445Z
Learning: In the nats.zig project's ConcurrentWriteBuffer.appendMany() method in src/queue.zig, the maintainer is not concerned about explicit overflow checking for total message size calculations, preferring to rely on Zig's built-in overflow detection in Debug and ReleaseSafe modes rather than using addWithOverflow for size accounting.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
🧠 Learnings used
Learnt from: lalinsky
Repo: lalinsky/nats.zig PR: 108
File: src/dispatcher.zig:60-65
Timestamp: 2025-09-09T19:51:02.201Z
Learning: In the NATS Zig codebase dispatcher.zig, when decrementing pending_bytes (atomic u64) with message_data_len (usize) using fetchSub, explicit casting is not required as Zig handles the implicit conversion automatically and the code builds successfully.
Learnt from: lalinsky
Repo: lalinsky/nats.zig PR: 60
File: src/queue.zig:390-393
Timestamp: 2025-08-30T07:40:28.186Z
Learning: In the nats.zig project, the maintainer prefers to rely on Zig's built-in integer overflow detection rather than implementing manual saturated conversions for timeout calculations (timeout_ms * std.time.ns_per_ms). Zig automatically fails on integer overflow in Debug and ReleaseSafe modes.
Learnt from: lalinsky
Repo: lalinsky/nats.zig PR: 62
File: src/queue.zig:716-718
Timestamp: 2025-08-30T09:35:56.288Z
Learning: In ConcurrentWriteBuffer, gatherReadVectors should be a blocking operation with a timeout parameter (like getSlice and pop), blocking when no data is available OR when the buffer is frozen. It should return a VectorGather structure containing reset ID, first chunk reference, iovecs array, total bytes, and a consume method that validates the reset ID and first chunk to ensure thread safety and prevent concurrent consumers. The consume method should not check freeze state since it's consuming already-gathered data.
Learnt from: lalinsky
Repo: lalinsky/nats.zig PR: 63
File: src/socket.zig:79-83
Timestamp: 2025-08-30T12:18:09.213Z
Learning: In the nats.zig codebase, the maintainer (lalinsky) states that Socket.write() is atomic - either it writes some data or returns an error, never partial data with error. This may differ from standard std.net.Stream.write() semantics. writeAll() is described as NOT atomic because it can write partial data and then fail.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Added
appendMany()toConcurrentWriteBufferthat accepts multiple slices and appends them in a single operation. ModifiedpublishMsgInternal()to build control line and headers separately, then useappendMany()to send[control+headers, msg.data, "\r\n"]without copying the message data.Changes
ConcurrentWriteBuffer.appendMany()method that takes multiple slicespublishMsgInternal()to build control line and headers separatelyappendMany()to avoid copyingmsg.dataBenefits
While socket I/O is still the main bottleneck, this optimization helps avoid unnecessary memory allocation and copying, especially for large messages.