Optimize memory allocation with buffer pools and slice reuse

pkg/proxy/parsercmd.go

@@ -101,7 +101,7 @@
 	s.commands = nil
 }
 
 func (s *TerminalParser) GetCursorRow() string {
 	switch s.screenType {
 	case LinuxScreen:
 		row := s.Screen.GetCursorRow()
@@ -368,25 +368,34 @@
 func (s *TerminalParser) TryTmuxInput() string {
 	lastLine := s.tmuxParser.TmuxScreen.GetCursorRow()
 	cmd := strings.TrimPrefix(lastLine.String(), s.Ps1sStr)
-	s.InputBuf.Reset()
+	s.resetInputBuf()
 	return strings.TrimSpace(cmd)
 }
 
 func (s *TerminalParser) TryInput() string {
 	lastLine := s.Screen.GetCursorRow()
 	cmd := strings.TrimPrefix(lastLine.String(), s.Ps1sStr)
-	s.InputBuf.Reset()
+	s.resetInputBuf()
 	return strings.TrimSpace(cmd)
 }
 
 func (s *TerminalParser) TryLastRowInput() string {
 	rowStr := s.GetCursorRow()
 	cmd := strings.TrimPrefix(rowStr, s.Ps1sStr)
-	s.InputBuf.Reset()
+	s.resetInputBuf()
 	return strings.TrimSpace(cmd)
 }
 
+// resetInputBuf resets the input buffer and prevents memory bloat
+func (s *TerminalParser) resetInputBuf() {
+	s.InputBuf.Reset()
+	// If capacity is too large, create a new buffer to release memory
+	if s.InputBuf.Cap() > maxBufSize {
+		s.InputBuf = bytes.Buffer{}
+	}
+}
+
 func (s *TerminalParser) GetPs1() string {
 	rowStr := s.GetCursorRow()
 	return strings.TrimSuffix(rowStr, s.InputBuf.String())
 }

pkg/proxy/recorder.go

@@ -18,6 +18,14 @@
 	"github.com/jumpserver/koko/pkg/logger"
 )
 
+// ftpReadBufferPool provides reusable 32KB buffers for FTP file operations
+var ftpReadBufferPool = sync.Pool{
+	New: func() interface{} {
+		buf := make([]byte, 32*1024)
+		return &buf
+	},
+}
+
 type CommandRecorder struct {
 	sessionID string
 	storage   CommandStorage
@@ -317,7 +325,7 @@
 	delete(r.ftpLogMap, id)
 }
 
 func (r *FTPFileRecorder) getFTPFile(id string) *FTPFileInfo {
 	r.lock.RLock()
 	defer r.lock.RUnlock()
 	return r.ftpLogMap[id]
@@ -482,7 +490,9 @@
 }
 
 func (f *FTPFileInfo) WriteFromReader(r io.Reader) error {
-	buf := make([]byte, 32*1024)
+	bufPtr := ftpReadBufferPool.Get().(*[]byte)
+	buf := *bufPtr
+	defer ftpReadBufferPool.Put(bufPtr)
 	var err error
 	for {
 		nr, er := r.Read(buf)
@@ -528,8 +538,8 @@
 	return nil
 }
 
 func GetFTPFileRecorder(jmsService *service.JMService) *FTPFileRecorder {
 	terminalConfig, _ := jmsService.GetTerminalConfig()
 	maxSize := int64(terminalConfig.MaxStoreFTPFileSize) * 1024 * 1024
 	recorder := NewFTPFileRecord(jmsService, NewFTPFileStorage(jmsService, &terminalConfig), maxSize)
 	return recorder

pkg/utils/buffer.go

@@ -5,6 +5,85 @@
 	"sync"
 )
 
+// Buffer pool sizes for different use cases
+const (
+	SmallBufferSize  = 1024       // 1KB for small reads
+	MediumBufferSize = 32 * 1024  // 32KB for medium reads
+	LargeBufferSize  = 64 * 1024  // 64KB for large reads
+)
+
+// BufferPool provides a sync.Pool based buffer pool with configurable size
+type BufferPool struct {
+	pool sync.Pool
+	size int
+}
+
+// NewBufferPool creates a new buffer pool with the specified buffer size
+func NewBufferPool(size int) *BufferPool {
+	return &BufferPool{
+		size: size,
+		pool: sync.Pool{
+			New: func() interface{} {
+				buf := make([]byte, size)
+				return &buf
+			},
+		},
+	}
+}
+
+// Get returns a buffer from the pool
+func (p *BufferPool) Get() *[]byte {
+	return p.pool.Get().(*[]byte)
+}
+
+// Put returns a buffer to the pool
+func (p *BufferPool) Put(buf *[]byte) {
+	if buf == nil {
+		return
+	}
+	// Only return buffers of the expected size to the pool
+	if cap(*buf) == p.size {
+		*buf = (*buf)[:p.size]
+		p.pool.Put(buf)
+	}
+}
+
+// Global buffer pools for common use cases
+var (
+	// SmallBufferPool for small reads (1KB)
+	SmallBufferPool = NewBufferPool(SmallBufferSize)
+	// MediumBufferPool for medium reads (32KB)
+	MediumBufferPool = NewBufferPool(MediumBufferSize)
+	// LargeBufferPool for large reads (64KB)
+	LargeBufferPool = NewBufferPool(LargeBufferSize)
+)
+
+// BytesBufferPool provides a sync.Pool for bytes.Buffer instances.
+// It helps reduce memory allocations by reusing bytes.Buffer objects.
+// Usage pattern:
+//   buf := GetBytesBuffer()
+//   defer PutBytesBuffer(buf)
+//   // use buf...
+var BytesBufferPool = sync.Pool{
+	New: func() interface{} {
+		return new(bytes.Buffer)
+	},
+}
+
+// GetBytesBuffer returns a bytes.Buffer from the pool
+func GetBytesBuffer() *bytes.Buffer {
+	return BytesBufferPool.Get().(*bytes.Buffer)
+}
+
+// PutBytesBuffer returns a bytes.Buffer to the pool after resetting it
+func PutBytesBuffer(buf *bytes.Buffer) {
+	if buf == nil {
+		return
+	}
+	buf.Reset()
+	BytesBufferPool.Put(buf)
+}
+
 type SyncBuffer struct {
 	maxSize int
 	mu      sync.Mutex

pkg/utils/buffer_test.go

@@ -0,0 +1,170 @@
+package utils
+
+import (
+	"bytes"
+	"sync"
+	"testing"
+)
+
+func TestBufferPool(t *testing.T) {
+	pool := NewBufferPool(1024)
+
+	// Get a buffer
+	buf := pool.Get()
+	if buf == nil {
+		t.Fatal("Expected non-nil buffer")
+	}
+	if len(*buf) != 1024 {
+		t.Fatalf("Expected buffer length 1024, got %d", len(*buf))
+	}
+
+	// Use the buffer
+	copy(*buf, []byte("test data"))
+
+	// Return buffer to pool
+	pool.Put(buf)
+
+	// Get another buffer - should reuse the pooled one
+	buf2 := pool.Get()
+	if buf2 == nil {
+		t.Fatal("Expected non-nil buffer")
+	}
+	if len(*buf2) != 1024 {
+		t.Fatalf("Expected buffer length 1024, got %d", len(*buf2))
+	}
+	pool.Put(buf2)
+}
+
+func TestBufferPoolConcurrent(t *testing.T) {
+	pool := NewBufferPool(512)
+	var wg sync.WaitGroup
+
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			buf := pool.Get()
+			if buf == nil {
+				t.Error("Expected non-nil buffer")
+				return
+			}
+			// Simulate some work
+			copy(*buf, []byte("concurrent test"))
+			pool.Put(buf)
+		}()
+	}
+	wg.Wait()
+}
+
+func TestGlobalBufferPools(t *testing.T) {
+	// Test SmallBufferPool
+	smallBuf := SmallBufferPool.Get()
+	if smallBuf == nil || len(*smallBuf) != SmallBufferSize {
+		t.Errorf("SmallBufferPool returned invalid buffer")
+	}
+	SmallBufferPool.Put(smallBuf)
+
+	// Test MediumBufferPool
+	mediumBuf := MediumBufferPool.Get()
+	if mediumBuf == nil || len(*mediumBuf) != MediumBufferSize {
+		t.Errorf("MediumBufferPool returned invalid buffer")
+	}
+	MediumBufferPool.Put(mediumBuf)
+
+	// Test LargeBufferPool
+	largeBuf := LargeBufferPool.Get()
+	if largeBuf == nil || len(*largeBuf) != LargeBufferSize {
+		t.Errorf("LargeBufferPool returned invalid buffer")
+	}
+	LargeBufferPool.Put(largeBuf)
+}
+
+func TestBytesBufferPool(t *testing.T) {
+	buf := GetBytesBuffer()
+	if buf == nil {
+		t.Fatal("Expected non-nil bytes.Buffer")
+	}
+
+	buf.WriteString("test data")
+	if buf.String() != "test data" {
+		t.Errorf("Expected 'test data', got '%s'", buf.String())
+	}
+
+	PutBytesBuffer(buf)
+
+	// Get another buffer
+	buf2 := GetBytesBuffer()
+	if buf2 == nil {
+		t.Fatal("Expected non-nil bytes.Buffer")
+	}
+	// Buffer should be reset
+	if buf2.Len() != 0 {
+		t.Errorf("Expected empty buffer, got length %d", buf2.Len())
+	}
+	PutBytesBuffer(buf2)
+}
+
+func TestSyncBuffer(t *testing.T) {
+	buf := NewMaxSizeBuffer(100)
+
+	// Test write within limit
+	data := []byte("hello world")
+	n, err := buf.Write(data)
+	if err != nil {
+		t.Fatalf("Write failed: %v", err)
+	}
+	if n != len(data) {
+		t.Errorf("Expected write of %d bytes, got %d", len(data), n)
+	}
+
+	if buf.String() != "hello world" {
+		t.Errorf("Expected 'hello world', got '%s'", buf.String())
+	}
+
+	// Test write exceeding limit
+	largeData := make([]byte, 100)
+	n, err = buf.Write(largeData)
+	if err != nil {
+		t.Fatalf("Write failed: %v", err)
+	}
+	if n != len(largeData) {
+		t.Errorf("Expected write of %d bytes, got %d", len(largeData), n)
+	}
+	// Should not have written as it would exceed max size
+	if buf.String() != "hello world" {
+		t.Errorf("Buffer should not have changed, got '%s'", buf.String())
+	}
+}
+
+func BenchmarkBufferPoolGet(b *testing.B) {
+	pool := NewBufferPool(1024)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		buf := pool.Get()
+		pool.Put(buf)
+	}
+}
+
+func BenchmarkRawBufferAlloc(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		buf := make([]byte, 1024)
+		_ = buf
+	}
+}
+
+func BenchmarkBytesBufferPoolGet(b *testing.B) {
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		buf := GetBytesBuffer()
+		buf.WriteString("test data for benchmarking")
+		PutBytesBuffer(buf)
+	}
+}
+
+func BenchmarkRawBytesBufferAlloc(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		buf := new(bytes.Buffer)
+		buf.WriteString("test data for benchmarking")
+		_ = buf
+	}
+}

pkg/zmodem/zsession.go

@@ -26,6 +26,10 @@
 const (
 	bin16HeaderLen = 7
 	bin32HeaderLen = 9
+
+	// maxSubPacketBufCap is the maximum capacity for subPacketBuf before resetting
+	// to release memory. This helps prevent memory bloat from large file transfers.
+	maxSubPacketBufCap = 64 * 1024
 )
 
 func DecodeHexFrameHeader(p []byte) (h ZmodemHeader, offset int, ok bool) {
@@ -252,6 +256,8 @@
 		gotZDLE      bool
 		endSubPacket bool
 	)
+	// Reuse parsedSubPacket slice capacity instead of allocating new
+	s.parsedSubPacket = s.parsedSubPacket[:0]
 	for i := range buf {
 		switch buf[i] {
 		case ZDLE:
@@ -280,8 +286,11 @@
 		}
 	}
 	s.subPacketBuf.Reset()
+	// Reset subPacketBuf capacity if it's too large to avoid memory bloat
+	if s.subPacketBuf.Cap() > maxSubPacketBufCap {
+		s.subPacketBuf = bytes.Buffer{}
+	}
 	s.onSubPacket(s.parsedSubPacket)
-	s.parsedSubPacket = nil
 	s.consume(buf[offset+1:])
 }
 
@@ -315,7 +324,7 @@
 	s.currentHd = nil
 }
 
 func (s *ZSession) getHexHeader(p []byte) {
 	if hd, offset, ok := DecodeHexFrameHeader(p); ok {
 		s.onHeader(&hd)
 		if s.IsNeedSubPacket() {
@@ -324,7 +333,7 @@
 	}
 }
 
 func (s *ZSession) getB16Header(p []byte) {
 	if hd, offset, ok := DecodeB16FrameHeader(p); ok {
 		s.onHeader(&hd)
 		if s.IsNeedSubPacket() {
@@ -333,7 +342,7 @@
 	}
 }
 
 func (s *ZSession) getB32Header(p []byte) {
 	if hd, offset, ok := DecodeB32FrameHeader(p); ok {
 		s.onHeader(&hd)
 		if s.IsNeedSubPacket() {