geo/server/internal/shared/auditlog/async.go

package auditlog

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"go.uber.org/zap"
)

const (
	defaultBatchSize     = 64
	defaultFlushInterval = time.Second
	defaultFlushTimeout  = 5 * time.Second
)

type Entry struct {
	OperatorID   int64
	TenantID     *int64
	Module       string
	Action       string
	ResourceType *string
	ResourceID   *int64
	RequestID    *string
	BeforeJSON   []byte
	AfterJSON    []byte
	Result       *string
	ErrorMessage *string
}

type AsyncWriter struct {
	pool          *pgxpool.Pool
	logger        *zap.Logger
	batchSize     int
	flushInterval time.Duration
	flushTimeout  time.Duration

	mu      sync.Mutex
	pending []Entry
	readPos int
	signal  chan struct{}
	closed  bool
	wg      sync.WaitGroup
}

func NewAsyncWriter(pool *pgxpool.Pool, logger *zap.Logger) *AsyncWriter {
	writer := &AsyncWriter{
		pool:          pool,
		logger:        logger,
		batchSize:     defaultBatchSize,
		flushInterval: defaultFlushInterval,
		flushTimeout:  defaultFlushTimeout,
		signal:        make(chan struct{}, 1),
	}
	writer.wg.Add(1)
	go writer.run()
	return writer
}

func (w *AsyncWriter) Log(entry Entry) bool {
	if w == nil {
		return false
	}

	w.mu.Lock()
	if w.closed {
		w.mu.Unlock()
		return false
	}
	w.pending = append(w.pending, entry)
	w.mu.Unlock()

	select {
	case w.signal <- struct{}{}:
	default:
	}

	return true
}

func (w *AsyncWriter) Close(ctx context.Context) error {
	if w == nil {
		return nil
	}

	w.mu.Lock()
	if !w.closed {
		w.closed = true
	}
	w.mu.Unlock()

	select {
	case w.signal <- struct{}{}:
	default:
	}

	done := make(chan struct{})
	go func() {
		defer close(done)
		w.wg.Wait()
	}()

	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (w *AsyncWriter) run() {
	defer w.wg.Done()

	ticker := time.NewTicker(w.flushInterval)
	defer ticker.Stop()

	for {
		select {
		case <-w.signal:
		case <-ticker.C:
		}

		for {
			batch, closed := w.nextBatch()
			if len(batch) == 0 {
				if closed {
					return
				}
				break
			}
			if err := w.flush(batch); err != nil {
				w.logger.Warn("audit log batch flush failed; keeping backlog for retry",
					zap.Int("count", len(batch)),
					zap.Error(err),
				)
				w.requeueFront(batch)
				time.Sleep(500 * time.Millisecond)
				break
			}
		}
	}
}

func (w *AsyncWriter) flush(entries []Entry) error {
	ctx, cancel := context.WithTimeout(context.Background(), w.flushTimeout)
	defer cancel()

	query, args := buildInsert(entries)
	if _, err := w.pool.Exec(ctx, query, args...); err != nil {
		return err
	}
	return nil
}

func (w *AsyncWriter) nextBatch() ([]Entry, bool) {
	w.mu.Lock()
	defer w.mu.Unlock()

	available := len(w.pending) - w.readPos
	if available == 0 {
		if w.readPos > 0 {
			w.pending = nil
			w.readPos = 0
		}
		return nil, w.closed
	}

	size := w.batchSize
	if available < size {
		size = available
	}

	batch := append([]Entry(nil), w.pending[w.readPos:w.readPos+size]...)
	w.readPos += size

	if w.readPos == len(w.pending) {
		w.pending = nil
		w.readPos = 0
	} else if w.readPos >= len(w.pending)/2 {
		remaining := append([]Entry(nil), w.pending[w.readPos:]...)
		w.pending = remaining
		w.readPos = 0
	}

	return batch, w.closed
}

func (w *AsyncWriter) requeueFront(entries []Entry) {
	w.mu.Lock()
	defer w.mu.Unlock()

	if len(entries) == 0 {
		return
	}

	if w.readPos > 0 {
		remaining := append([]Entry(nil), w.pending[w.readPos:]...)
		w.pending = remaining
		w.readPos = 0
	}

	w.pending = append(append([]Entry(nil), entries...), w.pending...)
}

func buildInsert(entries []Entry) (string, []interface{}) {
	var builder strings.Builder
	builder.WriteString(`INSERT INTO audit_logs (
operator_id,
tenant_id,
module,
action,
resource_type,
resource_id,
request_id,
before_json,
after_json,
result,
error_message
) VALUES `)

	args := make([]interface{}, 0, len(entries)*11)
	placeholder := 1

	for i, entry := range entries {
		if i > 0 {
			builder.WriteString(",")
		}
		builder.WriteString("(")
		for col := 0; col < 11; col++ {
			if col > 0 {
				builder.WriteString(",")
			}
			builder.WriteString(fmt.Sprintf("$%d", placeholder))
			placeholder++
		}
		builder.WriteString(")")

		args = append(args,
			entry.OperatorID,
			derefInt64(entry.TenantID),
			entry.Module,
			entry.Action,
			derefString(entry.ResourceType),
			derefInt64(entry.ResourceID),
			derefString(entry.RequestID),
			nilIfEmptyBytes(entry.BeforeJSON),
			nilIfEmptyBytes(entry.AfterJSON),
			derefString(entry.Result),
			derefString(entry.ErrorMessage),
		)
	}

	return builder.String(), args
}

func derefString(value *string) interface{} {
	if value == nil {
		return nil
	}
	return *value
}

func derefInt64(value *int64) interface{} {
	if value == nil {
		return nil
	}
	return *value
}

func nilIfEmptyBytes(value []byte) interface{} {
	if len(value) == 0 {
		return nil
	}
	return value
}