team-api/internal/logic/common/batch_writer.go

package common

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gctx"
)

// BatchWriter provides buffered batch writing for high-throughput log tables.
// It accumulates records in memory and flushes them to the database in batches.
//
// Three flush strategies:
//   - Synchronous: Flush immediately (for single writes)
//   - Async batch: Flush when buffer reaches capacity (default 64)
//   - Best-effort: Flush when buffer reaches max capacity (default 256), drop on failure
//
// Usage:
//
//	writer := common.NewBatchWriter("bil_usage_logs", 64, 256)
//	writer.Write(ctx, record)
//	writer.Flush(ctx) // manual flush
//	writer.Close()     // auto flush on shutdown
type BatchWriter struct {
	table      string
	buffer     []any
	bufferMu   sync.Mutex
	flushSize  int
	maxSize    int
	flushTimer *time.Timer
	closed     bool
	closeMu    sync.Mutex
	dropped    int64 // 缓冲区溢出累计丢弃的记录数（原子访问）
}

// NewBatchWriter creates a new BatchWriter.
//   - table: target database table name
//   - flushSize: flush when buffer reaches this size (0 = no auto-flush)
//   - maxSize: maximum buffer size, drop oldest on overflow
func NewBatchWriter(table string, flushSize, maxSize int) *BatchWriter {
	if maxSize == 0 {
		maxSize = 256
	}
	if flushSize == 0 {
		flushSize = maxSize
	}

	bw := &BatchWriter{
		table:     table,
		buffer:    make([]any, 0, maxSize),
		flushSize: flushSize,
		maxSize:   maxSize,
	}

	// Auto-flush every 5 seconds if there are pending records
	bw.flushTimer = time.AfterFunc(5*time.Second, func() {
		bw.autoFlush()
	})

	return bw
}

// Write adds a record to the buffer.
// If flushSize is reached, triggers an asynchronous flush.
func (bw *BatchWriter) Write(ctx context.Context, record any) {
	bw.bufferMu.Lock()
	defer bw.bufferMu.Unlock()

	if bw.closed {
		return
	}

	if len(bw.buffer) >= bw.maxSize {
		// 缓冲区已满（下游 flush 跟不上写入速度），丢弃最早记录以让出空间。
		// 对 bil_usage_logs 等计费高频表意味着数据丢失，必须计数并告警，避免静默丢弃。
		bw.buffer = bw.buffer[1:]
		total := atomic.AddInt64(&bw.dropped, 1)
		g.Log().Warningf(ctx, "batch writer buffer overflow on %s: dropped oldest record (total dropped=%d, maxSize=%d)", bw.table, total, bw.maxSize)
	}

	bw.buffer = append(bw.buffer, record)

	// Trigger async flush if threshold reached
	if len(bw.buffer) >= bw.flushSize {
		go bw.flush(gctx.New())
	}
}

// Flush synchronously flushes all buffered records to the database.
func (bw *BatchWriter) Flush(ctx context.Context) error {
	bw.bufferMu.Lock()
	if len(bw.buffer) == 0 {
		bw.bufferMu.Unlock()
		return nil
	}
	records := make([]any, len(bw.buffer))
	copy(records, bw.buffer)
	bw.buffer = bw.buffer[:0]
	bw.bufferMu.Unlock()

	if len(records) == 0 {
		return nil
	}

	_, err := g.DB().Model(bw.table).Ctx(ctx).Data(records).Batch(len(records)).Insert()
	if err != nil {
		g.Log().Errorf(ctx, "batch write to %s failed (%d records): %v", bw.table, len(records), err)
		return err
	}
	return nil
}

// flush is the internal flush without locking the buffer.
func (bw *BatchWriter) flush(ctx context.Context) {
	_ = bw.Flush(ctx)
}

// autoFlush is called periodically by the timer.
func (bw *BatchWriter) autoFlush() {
	bw.closeMu.Lock()
	defer bw.closeMu.Unlock()

	if bw.closed {
		return
	}

	_ = bw.Flush(gctx.New())

	if bw.closed {
		return
	}

	// Stop before Reset per Go docs recommendation
	bw.flushTimer.Stop()
	bw.flushTimer.Reset(5 * time.Second)
}

// Len returns the current buffer length.
func (bw *BatchWriter) Len() int {
	bw.bufferMu.Lock()
	defer bw.bufferMu.Unlock()
	return len(bw.buffer)
}

// DroppedCount 返回因缓冲区溢出累计丢弃的记录数，供监控/告警使用。
func (bw *BatchWriter) DroppedCount() int64 {
	return atomic.LoadInt64(&bw.dropped)
}

// Close flushes all remaining records and stops the auto-flush timer.
func (bw *BatchWriter) Close() {
	bw.closeMu.Lock()
	defer bw.closeMu.Unlock()

	if bw.closed {
		return
	}
	bw.closed = true

	if bw.flushTimer != nil {
		bw.flushTimer.Stop()
	}

	// Final flush
	_ = bw.Flush(gctx.New())
}

// SyncWrite writes a single record synchronously (bypasses buffer).
// Use this for critical records that must be persisted immediately.
func SyncWrite(ctx context.Context, table string, record any) error {
	_, err := g.DB().Model(table).Ctx(ctx).Data(record).Insert()
	return err
}