[ADDED] Option to customize write buffer size (#2057)

Signed-off-by: Piotr Piotrowski <piotr@synadia.com>

Author: piotrpio

nats.go

@@ -61,6 +61,7 @@ const (
 	DefaultMaxPingOut         = 2
 	DefaultMaxChanLen         = 64 * 1024       // 64k
 	DefaultReconnectBufSize   = 8 * 1024 * 1024 // 8MB
+	DefaultWriteBufSize       = defaultBufSize
 	RequestChanLen            = 8
 	DefaultDrainTimeout       = 30 * time.Second
 	DefaultFlusherTimeout     = time.Minute
@@ -574,6 +575,14 @@ type Options struct {
 	// IgnoreDiscoveredServers will disable adding advertised server URLs
 	// from INFO messages to the server pool.
 	IgnoreDiscoveredServers bool
+
+	// WriteBufferSize is an advanced option that sets the flush threshold
+	// of the write buffer used to batch outgoing data before writing to
+	// the underlying connection. In most cases, the default value should
+	// not be changed. A smaller buffer reduces the amount of data that
+	// can be lost on blocked writes but may significantly reduce throughput.
+	// Defaults to 32768 bytes (32KB).
+	WriteBufferSize int
 }
 
 const (
@@ -1172,6 +1181,19 @@ func ReconnectBufSize(size int) Option {
 	}
 }
 
+// WriteBufferSize is an advanced option that sets the flush threshold
+// of the write buffer used to batch outgoing data before writing to
+// the underlying connection. In most cases, the default value should
+// not be changed. A smaller buffer reduces the amount of data that
+// can be lost on blocked writes but may significantly reduce throughput.
+// Defaults to 32768 bytes (32KB).
+func WriteBufferSize(size int) Option {
+	return func(o *Options) error {
+		o.WriteBufferSize = size
+		return nil
+	}
+}
+
 // Timeout is an Option to set the timeout for Dial on a connection.
 // Defaults to 2s.
 func Timeout(t time.Duration) Option {
@@ -1741,6 +1763,10 @@ func (o Options) Connect() (*Conn, error) {
 	if nc.Opts.ReconnectBufSize == 0 {
 		nc.Opts.ReconnectBufSize = DefaultReconnectBufSize
 	}
+	// Default WriteBufferSize
+	if nc.Opts.WriteBufferSize <= 0 {
+		nc.Opts.WriteBufferSize = DefaultWriteBufSize
+	}
 	// Ensure that Timeout is not 0
 	if nc.Opts.Timeout == 0 {
 		nc.Opts.Timeout = DefaultTimeout
@@ -2080,7 +2106,7 @@ func (nc *Conn) newReaderWriter() {
 		off: -1,
 	}
 	nc.bw = &natsWriter{
-		limit:  defaultBufSize,
+		limit:  nc.Opts.WriteBufferSize,
 		plimit: nc.Opts.ReconnectBufSize,
 	}
 }

nats_test.go

@@ -1,4 +1,4 @@
-// Copyright 2012-2023 The NATS Authors
+// Copyright 2012-2026 The NATS Authors
 // 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
@@ -1915,3 +1915,30 @@ func TestValidateSubject(t *testing.T) {
 		})
 	}
 }
+
+func TestWriteBufferSize(t *testing.T) {
+	opts := GetDefaultOptions()
+	opts.WriteBufferSize = 64 * 1024
+	opts.Servers = []string{"nats://127.0.0.1:4222"}
+	nc := &Conn{Opts: opts}
+	nc.newReaderWriter()
+
+	if nc.bw.limit != 64*1024 {
+		t.Fatalf("Expected write buffer limit of %d, got %d", 64*1024, nc.bw.limit)
+	}
+	if len(nc.br.buf) != defaultBufSize {
+		t.Fatalf("Expected read buffer size of %d, got %d", defaultBufSize, len(nc.br.buf))
+	}
+}
+
+func TestWriteBufferSizeDefault(t *testing.T) {
+	opts := GetDefaultOptions()
+	opts.WriteBufferSize = DefaultWriteBufSize
+	opts.Servers = []string{"nats://127.0.0.1:4222"}
+	nc := &Conn{Opts: opts}
+	nc.newReaderWriter()
+
+	if nc.bw.limit != defaultBufSize {
+		t.Fatalf("Expected default write buffer limit of %d, got %d", defaultBufSize, nc.bw.limit)
+	}
+}

test/bench_test.go

@@ -14,6 +14,7 @@
 package test
 
 import (
+	"fmt"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -255,3 +256,41 @@ func BenchmarkPublishValidation(b *testing.B) {
 		b.StopTimer()
 	})
 }
+
+func BenchmarkPublishWithWriteBufferSize(b *testing.B) {
+	payloads := []struct {
+		name string
+		size int
+	}{
+		{"16B", 16},
+		{"128B", 128},
+		{"512B", 512},
+	}
+	bufSizes := []int{512, 4096, 8192, 16384, 32768, 65536, 131072}
+
+	for _, p := range payloads {
+		msg := make([]byte, p.size)
+		for i := range msg {
+			msg[i] = byte(i)
+		}
+		for _, sz := range bufSizes {
+			b.Run(fmt.Sprintf("payload_%s/buf_%d", p.name, sz), func(b *testing.B) {
+				s := RunDefaultServer()
+				defer s.Shutdown()
+				nc, err := nats.Connect(s.ClientURL(), nats.WriteBufferSize(sz))
+				if err != nil {
+					b.Fatalf("Failed to connect: %v", err)
+				}
+				defer nc.Close()
+				b.ResetTimer()
+				for range b.N {
+					if err := nc.Publish("foo", msg); err != nil {
+						b.Fatalf("Error publishing: %v", err)
+					}
+				}
+				b.StopTimer()
+				nc.Flush()
+			})
+		}
+	}
+}

test/conn_test.go

@@ -3563,3 +3563,29 @@ func TestTLSEOFAfterHandshakeBrokenPipe(t *testing.T) {
 		t.Fatalf("Expected error to wrap nats.ErrTLS, got: %v", err)
 	}
 }
+
+func TestWriteBufferSizeOption(t *testing.T) {
+	s := RunDefaultServer()
+	defer s.Shutdown()
+
+	nc, err := nats.Connect(s.ClientURL(), nats.WriteBufferSize(64*1024))
+	if err != nil {
+		t.Fatalf("Expected to connect, got: %v", err)
+	}
+	defer nc.Close()
+
+	if nc.Opts.WriteBufferSize != 64*1024 {
+		t.Fatalf("Expected WriteBufferSize 64KB, got %d", nc.Opts.WriteBufferSize)
+	}
+
+	sub, err := nc.SubscribeSync("foo")
+	if err != nil {
+		t.Fatalf("Error subscribing: %v", err)
+	}
+	if err := nc.Publish("foo", []byte("hello")); err != nil {
+		t.Fatalf("Error publishing: %v", err)
+	}
+	if _, err := sub.NextMsg(2 * time.Second); err != nil {
+		t.Fatalf("Error receiving message: %v", err)
+	}
+}