rewrite ProxyStream with 64 KiB ring buffer to reduce memory pressure (previously 64*16*2 KiB at peak)

Author: pymq

proxy_stream.go

@@ -0,0 +1,171 @@
+package socks5
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"sync"
+)
+
+type closeWriter interface {
+	CloseWrite() error
+}
+
+type closeReader interface {
+	CloseRead() error
+}
+
+// Ring buffer size. Must be a power of two so that `x & mask` is equivalent
+// to `x % proxyBufSize` (see ProxyStream for the indexing scheme).
+const (
+	proxyBufSize uint64 = 64 << 10
+	proxyBufMask        = proxyBufSize - 1
+)
+
+// proxyBufPool recycles the per-connection 64 KiB ring buffers. The buffer
+// dominates ProxyStream's allocation footprint, and under sustained load
+// each connection would otherwise force a fresh 64 KiB heap allocation.
+// Storing a pointer to a fixed-size array (rather than a slice) avoids the
+// extra slice-header allocation that sync.Pool would otherwise introduce.
+var proxyBufPool = sync.Pool{
+	New: func() any {
+		return new([proxyBufSize]byte)
+	},
+}
+
+// ProxyStream forwards data from src to dst, similar to io.Copy, but with improved performance.
+// Unlike io.Copy’s sequential read/write model, it allows reads to continue while writes are in progress,
+// using a single 64 KiB ring buffer shared between the reader and writer.
+// ProxyStream closes both the read and write sides when the transfer completes.
+func ProxyStream(src io.Reader, dst io.Writer) error {
+	// head and tail are monotonically increasing byte counters, not positions
+	// inside buf. Indexing into buf is done via `& proxyBufMask`, which is
+	// equivalent to `% proxyBufSize` but cheaper — this requires proxyBufSize
+	// to be a power of two. With this scheme, `head - tail` directly yields
+	// the number of occupied bytes regardless of wrap, and we never need to
+	// reset the counters.
+	bufArr := proxyBufPool.Get().(*[proxyBufSize]byte)
+	defer proxyBufPool.Put(bufArr)
+	buf := bufArr[:]
+	var (
+		head     uint64 // total bytes written into buf by reader
+		tail     uint64 // total bytes consumed from buf by writer
+		mu       sync.Mutex
+		readDone bool
+		writeErr error
+	)
+	cond := sync.NewCond(&mu)
+
+	writerDone := make(chan struct{})
+
+	// Writer goroutine: drains [tail, head) from the ring buffer into dst.
+	go func() {
+		defer close(writerDone)
+		for {
+			mu.Lock()
+			for head == tail && !readDone {
+				cond.Wait()
+			}
+			if head == tail {
+				// Reader finished and buffer drained.
+				mu.Unlock()
+				return
+			}
+
+			// Contiguous readable region starting at rIdx. When the occupied
+			// range wraps past proxyBufSize, only the first segment
+			// [rIdx, proxyBufSize) is taken here; the remainder [0, hIdx) is
+			// handled on the next iteration once tail has advanced past the
+			// wrap point.
+			rIdx := tail & proxyBufMask
+			hIdx := head & proxyBufMask
+			var data []byte
+			if hIdx > rIdx {
+				data = buf[rIdx:hIdx]
+			} else {
+				data = buf[rIdx:proxyBufSize]
+			}
+			mu.Unlock()
+
+			n, err := dst.Write(data)
+
+			mu.Lock()
+			tail += uint64(n)
+			cond.Signal()
+			if err != nil {
+				writeErr = err
+				mu.Unlock()
+				// Unblock reader if it’s stuck on src.Read.
+				if cr, ok := src.(closeReader); ok {
+					_ = cr.CloseRead()
+				}
+				return
+			}
+			mu.Unlock()
+		}
+	}()
+
+	// Reader loop: fills [head, tail+proxyBufSize) in the ring buffer from src.
+	var readErr error
+	for {
+		mu.Lock()
+		for head-tail == proxyBufSize && writeErr == nil {
+			cond.Wait()
+		}
+		if writeErr != nil {
+			mu.Unlock()
+			break
+		}
+
+		// Contiguous writable region starting at wIdx. Mirrors the writer
+		// side: if the free range wraps, only [wIdx, proxyBufSize) is taken
+		// here and [0, tIdx) is picked up on the next iteration after head
+		// wraps.
+		wIdx := head & proxyBufMask
+		tIdx := tail & proxyBufMask
+		var space []byte
+		if tIdx > wIdx {
+			space = buf[wIdx:tIdx]
+		} else {
+			space = buf[wIdx:proxyBufSize]
+		}
+		mu.Unlock()
+
+		n, err := src.Read(space)
+
+		mu.Lock()
+		head += uint64(n)
+		if err != nil {
+			readDone = true
+			if !errors.Is(err, io.EOF) {
+				readErr = err
+			}
+			cond.Signal()
+			mu.Unlock()
+			break
+		}
+		if n > 0 {
+			cond.Signal()
+		}
+		mu.Unlock()
+	}
+
+	<-writerDone
+
+	// Close both sides.
+	if cr, ok := src.(closeReader); ok {
+		_ = cr.CloseRead()
+	}
+	if cw, ok := dst.(closeWriter); ok {
+		_ = cw.CloseWrite()
+	}
+
+	if readErr != nil {
+		readErr = fmt.Errorf("read error: %v", readErr)
+	}
+	if writeErr != nil {
+		writeErr = fmt.Errorf("write error: %v", writeErr)
+	}
+
+	return errors.Join(writeErr, readErr)
+}