tweak: consolidate bytes pool management and enhance CopyClose functionality for improved performance

internal/utils/io.go

@@ -72,7 +72,7 @@ func NewPipe(ctx context.Context, r io.ReadCloser, w io.WriteCloser) *Pipe {
 }
 
 func (p *Pipe) Start() (err error) {
-	err = CopyClose(&p.w, &p.r)
+	err = CopyClose(&p.w, &p.r, 0)
 	switch {
 	case
 		// NOTE: ignoring broken pipe and connection reset by peer
@@ -117,7 +117,7 @@ func getHTTPFlusher(dst io.Writer) httpFlusher {
 	return nil
 }
 
-const copyBufSize = 32 * 1024
+const copyBufSize = synk.SizedPoolThreshold
 
 var bytesPool = synk.NewBytesPool()
 
@@ -125,7 +125,7 @@ var bytesPool = synk.NewBytesPool()
 // Use of this source code is governed by a BSD-style
 // This is a copy of io.Copy with context and HTTP flusher handling
 // Author: yusing <yusing@6uo.me>.
-func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
+func CopyClose(dst *ContextWriter, src *ContextReader, sizeHint int) (err error) {
 	size := copyBufSize
 	if l, ok := src.Reader.(*io.LimitedReader); ok {
 		if int64(size) > l.N {
@@ -135,6 +135,8 @@ func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
 				size = int(l.N)
 			}
 		}
+	} else if sizeHint > 0 {
+		size = min(size, sizeHint)
 	}
 	buf := bytesPool.GetSized(size)
 	defer bytesPool.Put(buf)
@@ -196,6 +198,6 @@ func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
 	}
 }
 
-func CopyCloseWithContext(ctx context.Context, dst io.Writer, src io.Reader) (err error) {
-	return CopyClose(NewContextWriter(ctx, dst), NewContextReader(ctx, src))
+func CopyCloseWithContext(ctx context.Context, dst io.Writer, src io.Reader, sizeHint int) (err error) {
+	return CopyClose(NewContextWriter(ctx, dst), NewContextReader(ctx, src), sizeHint)
 }

internal/utils/synk/pool.go

@@ -9,8 +9,8 @@ type weakBuf = unsafe.Pointer
 
 func makeWeak(b *[]byte) weakBuf {
 	ptr := runtime_registerWeakPointer(unsafe.Pointer(b))
-	runtime.KeepAlive(ptr)
 	addCleanup(b, addGCed, cap(*b))
+	runtime.KeepAlive(ptr)
 	return weakBuf(ptr)
 }
 
@@ -29,8 +29,9 @@ func runtime_registerWeakPointer(unsafe.Pointer) unsafe.Pointer
 func runtime_makeStrongFromWeak(unsafe.Pointer) unsafe.Pointer
 
 type BytesPool struct {
-	pool     chan weakBuf
-	initSize int
+	sizedPool   chan weakBuf
+	unsizedPool chan weakBuf
+	initSize    int
 }
 
 const (
@@ -41,16 +42,18 @@ const (
 const (
 	InPoolLimit = 32 * mb
 
-	DefaultInitBytes  = 4 * kb
-	PoolThreshold     = 256 * kb
-	DropThresholdHigh = 4 * mb
+	UnsizedAvg         = 4 * kb
+	SizedPoolThreshold = 256 * kb
+	DropThreshold      = 4 * mb
 
-	PoolSize = InPoolLimit / PoolThreshold
+	SizedPoolSize   = InPoolLimit * 8 / 10 / SizedPoolThreshold
+	UnsizedPoolSize = InPoolLimit * 2 / 10 / UnsizedAvg
 )
 
 var bytesPool = &BytesPool{
-	pool:     make(chan weakBuf, PoolSize),
-	initSize: DefaultInitBytes,
+	sizedPool:   make(chan weakBuf, SizedPoolSize),
+	unsizedPool: make(chan weakBuf, UnsizedPoolSize),
+	initSize:    UnsizedAvg,
 }
 
 func NewBytesPool() *BytesPool {
@@ -60,7 +63,7 @@ func NewBytesPool() *BytesPool {
 func (p *BytesPool) Get() []byte {
 	for {
 		select {
-		case bWeak := <-p.pool:
+		case bWeak := <-p.unsizedPool:
 			bPtr := getBufFromWeak(bWeak)
 			if bPtr == nil {
 				continue
@@ -68,18 +71,19 @@ func (p *BytesPool) Get() []byte {
 			addReused(cap(bPtr))
 			return bPtr
 		default:
-			return make([]byte, 0, p.initSize)
+			return make([]byte, 0)
 		}
 	}
 }
 
 func (p *BytesPool) GetSized(size int) []byte {
-	if size <= PoolThreshold {
+	if size <= SizedPoolThreshold {
+		addNonPooled(size)
 		return make([]byte, size)
 	}
 	for {
 		select {
-		case bWeak := <-p.pool:
+		case bWeak := <-p.sizedPool:
 			bPtr := getBufFromWeak(bWeak)
 			if bPtr == nil {
 				continue
@@ -90,25 +94,35 @@ func (p *BytesPool) GetSized(size int) []byte {
 				return (bPtr)[:size]
 			}
 			select {
-			case p.pool <- bWeak:
+			case p.sizedPool <- bWeak:
 			default:
 				// just drop it
 			}
 		default:
 		}
+		addNonPooled(size)
 		return make([]byte, size)
 	}
 }
 
 func (p *BytesPool) Put(b []byte) {
 	size := cap(b)
-	if size <= PoolThreshold || size > DropThresholdHigh {
+	if size > DropThreshold {
 		return
 	}
 	b = b[:0]
 	w := makeWeak(&b)
+	if size <= SizedPoolThreshold {
+		p.put(w, p.unsizedPool)
+	} else {
+		p.put(w, p.sizedPool)
+	}
+}
+
+//go:inline
+func (p *BytesPool) put(w weakBuf, pool chan weakBuf) {
 	select {
-	case p.pool <- w:
+	case pool <- w:
 	default:
 		// just drop it
 	}

internal/utils/synk/pool_debug.go

@@ -14,10 +14,16 @@ import (
 )
 
 var (
-	numReused, sizeReused uint64
-	numGCed, sizeGCed     uint64
+	numNonPooled, sizeNonPooled uint64
+	numReused, sizeReused       uint64
+	numGCed, sizeGCed           uint64
 )
 
+func addNonPooled(size int) {
+	atomic.AddUint64(&numNonPooled, 1)
+	atomic.AddUint64(&sizeNonPooled, uint64(size))
+}
+
 func addReused(size int) {
 	atomic.AddUint64(&numReused, 1)
 	atomic.AddUint64(&sizeReused, uint64(size))
@@ -48,6 +54,8 @@ func initPoolStats() {
 					Str("sizeReused", strutils.FormatByteSize(atomic.LoadUint64(&sizeReused))).
 					Uint64("numGCed", atomic.LoadUint64(&numGCed)).
 					Str("sizeGCed", strutils.FormatByteSize(atomic.LoadUint64(&sizeGCed))).
+					Uint64("numNonPooled", atomic.LoadUint64(&numNonPooled)).
+					Str("sizeNonPooled", strutils.FormatByteSize(atomic.LoadUint64(&sizeNonPooled))).
 					Msg("bytes pool stats")
 			}
 		}

internal/utils/synk/pool_prod.go

@@ -2,6 +2,7 @@
 
 package synk
 
+func addNonPooled(size int)                              {}
 func addReused(size int)                                 {}
 func addGCed(size int)                                   {}
 func initPoolStats()                                     {}