kvm/internal/audio/zero_copy.go

package audio

import (
	"sync"
	"sync/atomic"
	"unsafe"
)

// ZeroCopyAudioFrame represents an audio frame that can be passed between
// components without copying the underlying data
type ZeroCopyAudioFrame struct {
	data     []byte
	length   int
	capacity int
	refCount int32
	mutex    sync.RWMutex
	pooled   bool
}

// ZeroCopyFramePool manages reusable zero-copy audio frames
type ZeroCopyFramePool struct {
	// Atomic fields MUST be first for ARM32 alignment (int64 fields need 8-byte alignment)
	counter         int64 // Frame counter (atomic)
	hitCount        int64 // Pool hit counter (atomic)
	missCount       int64 // Pool miss counter (atomic)
	allocationCount int64 // Total allocations counter (atomic)

	// Other fields
	pool    sync.Pool
	maxSize int
	mutex   sync.RWMutex
	// Memory optimization fields
	preallocated []*ZeroCopyAudioFrame // Pre-allocated frames for immediate use
	preallocSize int                   // Number of pre-allocated frames
	maxPoolSize  int                   // Maximum pool size to prevent memory bloat
}

// NewZeroCopyFramePool creates a new zero-copy frame pool
func NewZeroCopyFramePool(maxFrameSize int) *ZeroCopyFramePool {
	// Pre-allocate frames for immediate availability
	preallocSizeBytes := GetConfig().PreallocSize
	maxPoolSize := GetConfig().MaxPoolSize // Limit total pool size

	// Calculate number of frames based on memory budget, not frame count
	preallocFrameCount := preallocSizeBytes / maxFrameSize
	if preallocFrameCount > maxPoolSize {
		preallocFrameCount = maxPoolSize
	}
	if preallocFrameCount < 1 {
		preallocFrameCount = 1 // Always preallocate at least one frame
	}

	preallocated := make([]*ZeroCopyAudioFrame, 0, preallocFrameCount)

	// Pre-allocate frames to reduce initial allocation overhead
	for i := 0; i < preallocFrameCount; i++ {
		frame := &ZeroCopyAudioFrame{
			data:     make([]byte, 0, maxFrameSize),
			capacity: maxFrameSize,
			pooled:   true,
		}
		preallocated = append(preallocated, frame)
	}

	return &ZeroCopyFramePool{
		maxSize:      maxFrameSize,
		preallocated: preallocated,
		preallocSize: preallocFrameCount,
		maxPoolSize:  maxPoolSize,
		pool: sync.Pool{
			New: func() interface{} {
				return &ZeroCopyAudioFrame{
					data:     make([]byte, 0, maxFrameSize),
					capacity: maxFrameSize,
					pooled:   true,
				}
			},
		},
	}
}

// Get retrieves a zero-copy frame from the pool
func (p *ZeroCopyFramePool) Get() *ZeroCopyAudioFrame {
	// Memory guard: Track allocation count to prevent excessive memory usage
	allocationCount := atomic.LoadInt64(&p.allocationCount)
	if allocationCount > int64(p.maxPoolSize*2) {
		// If we've allocated too many frames, force pool reuse
		atomic.AddInt64(&p.missCount, 1)
		frame := p.pool.Get().(*ZeroCopyAudioFrame)
		frame.mutex.Lock()
		frame.refCount = 1
		frame.length = 0
		frame.data = frame.data[:0]
		frame.mutex.Unlock()
		return frame
	}

	// First try pre-allocated frames for fastest access
	p.mutex.Lock()
	if len(p.preallocated) > 0 {
		frame := p.preallocated[len(p.preallocated)-1]
		p.preallocated = p.preallocated[:len(p.preallocated)-1]
		p.mutex.Unlock()

		frame.mutex.Lock()
		frame.refCount = 1
		frame.length = 0
		frame.data = frame.data[:0]
		frame.mutex.Unlock()

		atomic.AddInt64(&p.hitCount, 1)
		return frame
	}
	p.mutex.Unlock()

	// Try sync.Pool next and track allocation
	atomic.AddInt64(&p.allocationCount, 1)
	frame := p.pool.Get().(*ZeroCopyAudioFrame)
	frame.mutex.Lock()
	frame.refCount = 1
	frame.length = 0
	frame.data = frame.data[:0]
	frame.mutex.Unlock()

	atomic.AddInt64(&p.hitCount, 1)
	return frame
}

// Put returns a zero-copy frame to the pool
func (p *ZeroCopyFramePool) Put(frame *ZeroCopyAudioFrame) {
	if frame == nil || !frame.pooled {
		return
	}

	frame.mutex.Lock()
	frame.refCount--
	if frame.refCount <= 0 {
		frame.refCount = 0
		frame.length = 0
		frame.data = frame.data[:0]
		frame.mutex.Unlock()

		// First try to return to pre-allocated pool for fastest reuse
		p.mutex.Lock()
		if len(p.preallocated) < p.preallocSize {
			p.preallocated = append(p.preallocated, frame)
			p.mutex.Unlock()
			return
		}
		p.mutex.Unlock()

		// Check pool size limit to prevent excessive memory usage
		p.mutex.RLock()
		currentCount := atomic.LoadInt64(&p.counter)
		p.mutex.RUnlock()

		if currentCount >= int64(p.maxPoolSize) {
			return // Pool is full, let GC handle this frame
		}

		// Return to sync.Pool
		p.pool.Put(frame)
		atomic.AddInt64(&p.counter, 1)
	} else {
		frame.mutex.Unlock()
	}
}

// Data returns the frame data as a slice (zero-copy view)
func (f *ZeroCopyAudioFrame) Data() []byte {
	f.mutex.RLock()
	defer f.mutex.RUnlock()
	return f.data[:f.length]
}

// SetData sets the frame data (zero-copy if possible)
func (f *ZeroCopyAudioFrame) SetData(data []byte) error {
	f.mutex.Lock()
	defer f.mutex.Unlock()

	if len(data) > f.capacity {
		// Need to reallocate - not zero-copy but necessary
		f.data = make([]byte, len(data))
		f.capacity = len(data)
		f.pooled = false // Can't return to pool anymore
	}

	// Zero-copy assignment when data fits in existing buffer
	if cap(f.data) >= len(data) {
		f.data = f.data[:len(data)]
		copy(f.data, data)
	} else {
		f.data = append(f.data[:0], data...)
	}
	f.length = len(data)
	return nil
}

// SetDataDirect sets frame data using direct buffer assignment (true zero-copy)
// WARNING: The caller must ensure the buffer remains valid for the frame's lifetime
func (f *ZeroCopyAudioFrame) SetDataDirect(data []byte) {
	f.mutex.Lock()
	defer f.mutex.Unlock()
	f.data = data
	f.length = len(data)
	f.capacity = cap(data)
	f.pooled = false // Direct assignment means we can't pool this frame
}

// AddRef increments the reference count for shared access
func (f *ZeroCopyAudioFrame) AddRef() {
	f.mutex.Lock()
	f.refCount++
	f.mutex.Unlock()
}

// Release decrements the reference count
func (f *ZeroCopyAudioFrame) Release() {
	f.mutex.Lock()
	f.refCount--
	f.mutex.Unlock()
}

// Length returns the current data length
func (f *ZeroCopyAudioFrame) Length() int {
	f.mutex.RLock()
	defer f.mutex.RUnlock()
	return f.length
}

// Capacity returns the buffer capacity
func (f *ZeroCopyAudioFrame) Capacity() int {
	f.mutex.RLock()
	defer f.mutex.RUnlock()
	return f.capacity
}

// UnsafePointer returns an unsafe pointer to the data for CGO calls
// WARNING: Only use this for CGO interop, ensure frame lifetime
func (f *ZeroCopyAudioFrame) UnsafePointer() unsafe.Pointer {
	f.mutex.RLock()
	defer f.mutex.RUnlock()
	if len(f.data) == 0 {
		return nil
	}
	return unsafe.Pointer(&f.data[0])
}

// Global zero-copy frame pool
// GetZeroCopyPoolStats returns detailed statistics about the zero-copy frame pool
func (p *ZeroCopyFramePool) GetZeroCopyPoolStats() ZeroCopyFramePoolStats {
	p.mutex.RLock()
	preallocatedCount := len(p.preallocated)
	currentCount := atomic.LoadInt64(&p.counter)
	p.mutex.RUnlock()

	hitCount := atomic.LoadInt64(&p.hitCount)
	missCount := atomic.LoadInt64(&p.missCount)
	allocationCount := atomic.LoadInt64(&p.allocationCount)
	totalRequests := hitCount + missCount

	var hitRate float64
	if totalRequests > 0 {
		hitRate = float64(hitCount) / float64(totalRequests) * GetConfig().PercentageMultiplier
	}

	return ZeroCopyFramePoolStats{
		MaxFrameSize:      p.maxSize,
		MaxPoolSize:       p.maxPoolSize,
		CurrentPoolSize:   currentCount,
		PreallocatedCount: int64(preallocatedCount),
		PreallocatedMax:   int64(p.preallocSize),
		HitCount:          hitCount,
		MissCount:         missCount,
		AllocationCount:   allocationCount,
		HitRate:           hitRate,
	}
}

// ZeroCopyFramePoolStats provides detailed zero-copy pool statistics
type ZeroCopyFramePoolStats struct {
	MaxFrameSize      int
	MaxPoolSize       int
	CurrentPoolSize   int64
	PreallocatedCount int64
	PreallocatedMax   int64
	HitCount          int64
	MissCount         int64
	AllocationCount   int64
	HitRate           float64 // Percentage
}

var (
	globalZeroCopyPool = NewZeroCopyFramePool(GetMaxAudioFrameSize())
)

// GetZeroCopyFrame gets a frame from the global pool
func GetZeroCopyFrame() *ZeroCopyAudioFrame {
	return globalZeroCopyPool.Get()
}

// GetGlobalZeroCopyPoolStats returns statistics for the global zero-copy pool
func GetGlobalZeroCopyPoolStats() ZeroCopyFramePoolStats {
	return globalZeroCopyPool.GetZeroCopyPoolStats()
}

// PutZeroCopyFrame returns a frame to the global pool
func PutZeroCopyFrame(frame *ZeroCopyAudioFrame) {
	globalZeroCopyPool.Put(frame)
}

// ZeroCopyAudioReadEncode performs audio read and encode with zero-copy optimization
func ZeroCopyAudioReadEncode() (*ZeroCopyAudioFrame, error) {
	frame := GetZeroCopyFrame()

	maxFrameSize := GetMaxAudioFrameSize()
	// Ensure frame has enough capacity
	if frame.Capacity() < maxFrameSize {
		// Reallocate if needed
		frame.data = make([]byte, maxFrameSize)
		frame.capacity = maxFrameSize
		frame.pooled = false
	}

	// Use unsafe pointer for direct CGO call
	n, err := CGOAudioReadEncode(frame.data[:maxFrameSize])
	if err != nil {
		PutZeroCopyFrame(frame)
		return nil, err
	}

	if n == 0 {
		PutZeroCopyFrame(frame)
		return nil, nil
	}

	// Set the actual data length
	frame.mutex.Lock()
	frame.length = n
	frame.data = frame.data[:n]
	frame.mutex.Unlock()

	return frame, nil
}