598 lines
14 KiB
Go
598 lines
14 KiB
Go
package logger
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log"
|
|
"math/rand"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
"unsafe"
|
|
|
|
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
|
|
|
|
"git.misaka.ren/M1saka/token_thief/db"
|
|
)
|
|
|
|
const (
|
|
maxAttempts = 3
|
|
BaseBackoff = 200 * time.Millisecond
|
|
)
|
|
|
|
type Batch interface {
|
|
Append(v ...any) error
|
|
Send() error
|
|
Abort() error
|
|
}
|
|
|
|
type BatchPreparer interface {
|
|
PrepareBatch(ctx context.Context, query string) (Batch, error)
|
|
}
|
|
|
|
// Backend is the database contract required by the asynchronous queue.
|
|
type Backend interface {
|
|
BatchPreparer
|
|
Healthy() bool
|
|
MarkUnhealthy()
|
|
}
|
|
|
|
type clickHouseBatchPreparer struct {
|
|
conn driver.Conn
|
|
}
|
|
|
|
// The explicit adapter keeps the local Batch contract aligned with the real driver.
|
|
func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string) (Batch, error) {
|
|
return p.conn.PrepareBatch(ctx, query)
|
|
}
|
|
|
|
type poolBackend struct {
|
|
pool *db.Pool
|
|
conn driver.Conn
|
|
generation uint64
|
|
release func()
|
|
}
|
|
|
|
func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, error) {
|
|
if b.conn == nil {
|
|
return nil, errors.New("pool nil")
|
|
}
|
|
return clickHouseBatchPreparer{conn: b.conn}.PrepareBatch(ctx, query)
|
|
}
|
|
|
|
func (b poolBackend) Healthy() bool { return b.pool.Healthy() }
|
|
func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthyGeneration(b.generation) }
|
|
func (b poolBackend) Snapshot() (Backend, uint64) {
|
|
conn, generation, release := b.pool.Acquire()
|
|
return poolBackend{pool: b.pool, conn: conn, generation: generation, release: release}, generation
|
|
}
|
|
func (b poolBackend) MarkUnhealthyGeneration(generation uint64) {
|
|
b.pool.MarkUnhealthyGeneration(generation)
|
|
}
|
|
|
|
type Stats struct {
|
|
Enqueued uint64
|
|
Dropped uint64
|
|
Failed uint64
|
|
Ambiguous uint64
|
|
Bytes int64
|
|
}
|
|
|
|
// FlushResult separates retry-safe Prepare failures from final or ambiguous failures.
|
|
type FlushResult struct {
|
|
Retry []*LogEntry
|
|
Failed int
|
|
Ambiguous int
|
|
Err error
|
|
}
|
|
|
|
type queuedEntry struct {
|
|
entry *LogEntry
|
|
size int64
|
|
}
|
|
|
|
// Queue is an asynchronous, bounded logger. Submit never waits for database work.
|
|
type Queue struct {
|
|
ch chan queuedEntry
|
|
backend Backend
|
|
batchSize int
|
|
batchInterval time.Duration
|
|
workers int
|
|
entryBudget int64
|
|
byteBudget int64
|
|
|
|
dropped atomic.Uint64
|
|
failed atomic.Uint64
|
|
ambiguous atomic.Uint64
|
|
enq atomic.Uint64
|
|
entries atomic.Int64
|
|
bytes atomic.Int64
|
|
|
|
mu sync.RWMutex
|
|
started bool
|
|
stopped bool
|
|
shutdownCtx context.Context
|
|
reporterCancel context.CancelFunc
|
|
workCancel context.CancelFunc
|
|
wg sync.WaitGroup
|
|
done chan struct{}
|
|
doneOnce sync.Once
|
|
}
|
|
|
|
// NewQueue accepts an optional byte budget. A non-positive or omitted budget disables byte limiting.
|
|
func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval time.Duration, byteBudget ...int64) *Queue {
|
|
var backend Backend
|
|
if pool != nil {
|
|
backend = poolBackend{pool: pool}
|
|
}
|
|
return NewQueueWithBackend(backend, queueSize, batchSize, workers, batchInterval, byteBudget...)
|
|
}
|
|
|
|
// NewQueueWithBackend builds a queue against the database operations used by workers.
|
|
func NewQueueWithBackend(backend Backend, queueSize, batchSize, workers int, batchInterval time.Duration, byteBudget ...int64) *Queue {
|
|
var budget int64
|
|
if len(byteBudget) > 0 {
|
|
budget = byteBudget[0]
|
|
}
|
|
if queueSize < 0 {
|
|
queueSize = 0
|
|
}
|
|
if batchSize < 1 {
|
|
batchSize = 1
|
|
}
|
|
if workers < 1 {
|
|
workers = 1
|
|
}
|
|
if batchInterval <= 0 {
|
|
batchInterval = time.Second
|
|
}
|
|
return &Queue{
|
|
ch: make(chan queuedEntry, queueSize),
|
|
backend: backend,
|
|
batchSize: batchSize,
|
|
batchInterval: batchInterval,
|
|
workers: workers,
|
|
entryBudget: int64(queueSize),
|
|
byteBudget: budget,
|
|
done: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
// EstimatedBytes covers all variable-size strings and byte slices retained by an entry.
|
|
func EstimatedBytes(e *LogEntry) int64 {
|
|
if e == nil {
|
|
return 0
|
|
}
|
|
return int64(unsafe.Sizeof(*e)+unsafe.Sizeof(queuedEntry{})) + int64(len(e.RequestID)+len(e.Method)+len(e.Path)+len(e.Query)+len(e.ClientIP)+len(e.Error)+
|
|
len(e.RequestHeaders)+len(e.RequestBody)+len(e.ResponseHeaders)+len(e.ResponseBody))
|
|
}
|
|
|
|
func (q *Queue) Stats() Stats {
|
|
return Stats{
|
|
Enqueued: q.enq.Load(),
|
|
Dropped: q.dropped.Load(),
|
|
Failed: q.failed.Load(),
|
|
Ambiguous: q.ambiguous.Load(),
|
|
Bytes: q.bytes.Load(),
|
|
}
|
|
}
|
|
|
|
// Start starts workers once. The root context controls reporting only; Stop owns worker shutdown.
|
|
func (q *Queue) Start(root context.Context) {
|
|
q.mu.Lock()
|
|
if q.started || q.stopped {
|
|
q.mu.Unlock()
|
|
return
|
|
}
|
|
q.started = true
|
|
workCtx, workCancel := context.WithCancel(context.Background())
|
|
reportCtx, reporterCancel := context.WithCancel(root)
|
|
q.workCancel = workCancel
|
|
q.reporterCancel = reporterCancel
|
|
for i := 0; i < q.workers; i++ {
|
|
q.wg.Add(1)
|
|
go q.run(workCtx)
|
|
}
|
|
q.wg.Add(1)
|
|
go q.reportLoop(reportCtx)
|
|
q.mu.Unlock()
|
|
}
|
|
|
|
// Stop closes submissions once and waits under the caller's single total deadline.
|
|
// The variadic form permits legacy Stop() calls while new callers should pass a context.
|
|
func (q *Queue) Stop(contexts ...context.Context) error {
|
|
ctx := context.Background()
|
|
if len(contexts) > 0 && contexts[0] != nil {
|
|
ctx = contexts[0]
|
|
}
|
|
q.mu.Lock()
|
|
shutdownOwner := false
|
|
if !q.stopped {
|
|
shutdownOwner = true
|
|
q.stopped = true
|
|
q.shutdownCtx = ctx
|
|
close(q.ch)
|
|
if q.reporterCancel != nil {
|
|
q.reporterCancel()
|
|
}
|
|
if q.started {
|
|
go func() {
|
|
q.wg.Wait()
|
|
q.doneOnce.Do(func() { close(q.done) })
|
|
}()
|
|
} else {
|
|
q.discardQueued()
|
|
q.doneOnce.Do(func() { close(q.done) })
|
|
}
|
|
}
|
|
done := q.done
|
|
workCancel := q.workCancel
|
|
q.mu.Unlock()
|
|
|
|
select {
|
|
case <-done:
|
|
return nil
|
|
case <-ctx.Done():
|
|
if shutdownOwner && workCancel != nil {
|
|
workCancel()
|
|
timer := time.NewTimer(100 * time.Millisecond)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-done:
|
|
case <-timer.C:
|
|
}
|
|
}
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
|
|
// Submit reserves memory and enqueues without waiting; full, over-budget, and stopped queues drop.
|
|
func (q *Queue) Submit(e *LogEntry) {
|
|
if e == nil {
|
|
q.dropped.Add(1)
|
|
return
|
|
}
|
|
size := EstimatedBytes(e)
|
|
if !q.mu.TryRLock() {
|
|
q.dropped.Add(1)
|
|
return
|
|
}
|
|
if q.stopped || !q.reserve(size) {
|
|
q.mu.RUnlock()
|
|
q.dropped.Add(1)
|
|
return
|
|
}
|
|
entry := cloneLogEntry(e)
|
|
select {
|
|
case q.ch <- queuedEntry{entry: entry, size: size}:
|
|
q.enq.Add(1)
|
|
default:
|
|
q.release(size)
|
|
q.dropped.Add(1)
|
|
}
|
|
q.mu.RUnlock()
|
|
}
|
|
|
|
func cloneLogEntry(entry *LogEntry) *LogEntry {
|
|
clone := *entry
|
|
clone.RequestHeaders = append([]byte(nil), entry.RequestHeaders...)
|
|
clone.RequestBody = append([]byte(nil), entry.RequestBody...)
|
|
clone.ResponseHeaders = append([]byte(nil), entry.ResponseHeaders...)
|
|
clone.ResponseBody = append([]byte(nil), entry.ResponseBody...)
|
|
return &clone
|
|
}
|
|
|
|
func (q *Queue) reserve(size int64) bool {
|
|
for {
|
|
used := q.entries.Load()
|
|
if used >= q.entryBudget {
|
|
return false
|
|
}
|
|
if q.entries.CompareAndSwap(used, used+1) {
|
|
break
|
|
}
|
|
}
|
|
|
|
if q.byteBudget <= 0 {
|
|
q.bytes.Add(size)
|
|
return true
|
|
}
|
|
for {
|
|
used := q.bytes.Load()
|
|
if size > q.byteBudget-used {
|
|
q.entries.Add(-1)
|
|
return false
|
|
}
|
|
if q.bytes.CompareAndSwap(used, used+size) {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
|
|
func (q *Queue) release(size int64) {
|
|
q.bytes.Add(-size)
|
|
q.entries.Add(-1)
|
|
}
|
|
|
|
func (q *Queue) discardQueued() {
|
|
for item := range q.ch {
|
|
q.release(item.size)
|
|
q.failed.Add(1)
|
|
}
|
|
}
|
|
|
|
func (q *Queue) run(workCtx context.Context) {
|
|
defer q.wg.Done()
|
|
batch := make([]queuedEntry, 0, q.batchSize)
|
|
ticker := time.NewTicker(q.batchInterval)
|
|
defer ticker.Stop()
|
|
|
|
flush := func(final bool) bool {
|
|
if len(batch) == 0 {
|
|
return true
|
|
}
|
|
if q.backend == nil || !q.backend.Healthy() {
|
|
if !final {
|
|
return false
|
|
}
|
|
q.failed.Add(uint64(len(batch)))
|
|
} else {
|
|
q.flush(q.flushContext(workCtx), entriesFromBatch(batch))
|
|
}
|
|
for _, item := range batch {
|
|
q.release(item.size)
|
|
}
|
|
clear(batch)
|
|
batch = batch[:0]
|
|
return true
|
|
}
|
|
|
|
for {
|
|
if len(batch) >= q.batchSize && (q.backend == nil || !q.backend.Healthy()) {
|
|
if q.isStopped() {
|
|
flush(true)
|
|
q.discardQueued()
|
|
return
|
|
}
|
|
select {
|
|
case <-ticker.C:
|
|
flush(false)
|
|
case <-workCtx.Done():
|
|
flush(true)
|
|
q.discardQueued()
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
select {
|
|
case item, ok := <-q.ch:
|
|
if !ok {
|
|
flush(true)
|
|
return
|
|
}
|
|
batch = append(batch, item)
|
|
if len(batch) >= q.batchSize {
|
|
flush(false)
|
|
}
|
|
case <-ticker.C:
|
|
flush(false)
|
|
case <-workCtx.Done():
|
|
flush(true)
|
|
q.discardQueued()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (q *Queue) isStopped() bool {
|
|
q.mu.RLock()
|
|
defer q.mu.RUnlock()
|
|
return q.stopped
|
|
}
|
|
|
|
func entriesFromBatch(batch []queuedEntry) []*LogEntry {
|
|
entries := make([]*LogEntry, len(batch))
|
|
for i, item := range batch {
|
|
entries[i] = item.entry
|
|
}
|
|
return entries
|
|
}
|
|
|
|
func (q *Queue) flushContext(workCtx context.Context) context.Context {
|
|
q.mu.RLock()
|
|
defer q.mu.RUnlock()
|
|
if q.shutdownCtx != nil {
|
|
return q.shutdownCtx
|
|
}
|
|
return workCtx
|
|
}
|
|
|
|
func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
|
|
retry := entries
|
|
var lastErr error
|
|
var failedGeneration uint64
|
|
for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ {
|
|
if err := ctx.Err(); err != nil {
|
|
lastErr = err
|
|
break
|
|
}
|
|
attemptBackend, generation, release := backendSnapshot(q.backend)
|
|
failedGeneration = generation
|
|
result := Flush(ctx, attemptBackend, retry)
|
|
if release != nil {
|
|
release()
|
|
}
|
|
q.failed.Add(uint64(result.Failed))
|
|
q.ambiguous.Add(uint64(result.Ambiguous))
|
|
if result.Ambiguous > 0 {
|
|
markBackendUnhealthy(q.backend, failedGeneration)
|
|
}
|
|
lastErr = result.Err
|
|
retry = result.Retry
|
|
if len(retry) == 0 {
|
|
return
|
|
}
|
|
if attempt < maxAttempts && !waitBackoff(ctx, attempt) {
|
|
lastErr = ctx.Err()
|
|
break
|
|
}
|
|
}
|
|
|
|
if len(retry) > 0 {
|
|
q.failed.Add(uint64(len(retry)))
|
|
markBackendUnhealthy(q.backend, failedGeneration)
|
|
log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr)
|
|
}
|
|
}
|
|
|
|
type generationBackend interface {
|
|
Snapshot() (Backend, uint64)
|
|
MarkUnhealthyGeneration(uint64)
|
|
}
|
|
|
|
func backendSnapshot(backend Backend) (Backend, uint64, func()) {
|
|
if versioned, ok := backend.(generationBackend); ok {
|
|
snapshot, generation := versioned.Snapshot()
|
|
if leased, ok := snapshot.(interface{ Release() }); ok {
|
|
return snapshot, generation, leased.Release
|
|
}
|
|
return snapshot, generation, nil
|
|
}
|
|
return backend, 0, nil
|
|
}
|
|
|
|
func (b poolBackend) Release() {
|
|
if b.release != nil {
|
|
b.release()
|
|
}
|
|
}
|
|
|
|
func markBackendUnhealthy(backend Backend, generation uint64) {
|
|
if versioned, ok := backend.(generationBackend); ok {
|
|
versioned.MarkUnhealthyGeneration(generation)
|
|
return
|
|
}
|
|
backend.MarkUnhealthy()
|
|
}
|
|
|
|
func waitBackoff(ctx context.Context, attempt int) bool {
|
|
wait := BaseBackoff
|
|
for i := 1; i < attempt; i++ {
|
|
wait *= 3
|
|
}
|
|
jitter := time.Duration((rand.Float64()*0.4 - 0.2) * float64(wait))
|
|
timer := time.NewTimer(wait + jitter)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-timer.C:
|
|
return true
|
|
case <-ctx.Done():
|
|
return false
|
|
}
|
|
}
|
|
|
|
const insertStatement = `INSERT INTO proxy_logs (
|
|
request_id, method, path, query, client_ip,
|
|
request_headers, request_body, request_truncated,
|
|
status_code, response_headers, response_body, response_truncated,
|
|
is_stream, latency_ms, started_at, finished_at, error
|
|
)`
|
|
|
|
// Flush executes one batch. Only Prepare failures are retryable. Append failures are isolated by row;
|
|
// Send failures are ambiguous and therefore never replayed.
|
|
func Flush(ctx context.Context, conn BatchPreparer, entries []*LogEntry) FlushResult {
|
|
if len(entries) == 0 {
|
|
return FlushResult{}
|
|
}
|
|
result := flushOnce(ctx, conn, entries)
|
|
if result.stage != flushAppend || len(entries) == 1 {
|
|
return result.public(entries)
|
|
}
|
|
|
|
combined := FlushResult{}
|
|
for i, entry := range entries {
|
|
single := flushOnce(ctx, conn, []*LogEntry{entry}).public([]*LogEntry{entry})
|
|
combined.Retry = append(combined.Retry, single.Retry...)
|
|
combined.Failed += single.Failed
|
|
combined.Ambiguous += single.Ambiguous
|
|
combined.Err = errors.Join(combined.Err, single.Err)
|
|
if single.Ambiguous > 0 {
|
|
if backend, ok := conn.(interface{ MarkUnhealthy() }); ok {
|
|
backend.MarkUnhealthy()
|
|
}
|
|
combined.Failed += len(combined.Retry) + len(entries) - i - 1
|
|
combined.Retry = nil
|
|
break
|
|
}
|
|
}
|
|
return combined
|
|
}
|
|
|
|
type flushStage uint8
|
|
|
|
const (
|
|
flushSuccess flushStage = iota
|
|
flushPrepare
|
|
flushAppend
|
|
flushSend
|
|
)
|
|
|
|
type flushAttempt struct {
|
|
stage flushStage
|
|
err error
|
|
}
|
|
|
|
func (r flushAttempt) public(entries []*LogEntry) FlushResult {
|
|
switch r.stage {
|
|
case flushSuccess:
|
|
return FlushResult{}
|
|
case flushPrepare:
|
|
return FlushResult{Retry: entries, Err: r.err}
|
|
case flushSend:
|
|
return FlushResult{Failed: len(entries), Ambiguous: len(entries), Err: r.err}
|
|
default:
|
|
return FlushResult{Failed: len(entries), Err: r.err}
|
|
}
|
|
}
|
|
|
|
func flushOnce(ctx context.Context, conn BatchPreparer, entries []*LogEntry) flushAttempt {
|
|
batch, err := conn.PrepareBatch(ctx, insertStatement)
|
|
if err != nil {
|
|
return flushAttempt{stage: flushPrepare, err: err}
|
|
}
|
|
abort := func(cause error) error {
|
|
return errors.Join(cause, batch.Abort())
|
|
}
|
|
for _, e := range entries {
|
|
if err := batch.Append(
|
|
e.RequestID, e.Method, e.Path, e.Query, e.ClientIP,
|
|
string(e.RequestHeaders), string(e.RequestBody), e.RequestTruncated,
|
|
int32(e.StatusCode), string(e.ResponseHeaders), string(e.ResponseBody), e.ResponseTruncated,
|
|
e.IsStream, e.LatencyMS, e.StartedAt, e.FinishedAt, e.Error,
|
|
); err != nil {
|
|
return flushAttempt{stage: flushAppend, err: abort(err)}
|
|
}
|
|
}
|
|
if err := batch.Send(); err != nil {
|
|
return flushAttempt{stage: flushSend, err: abort(err)}
|
|
}
|
|
return flushAttempt{stage: flushSuccess}
|
|
}
|
|
|
|
func (q *Queue) reportLoop(ctx context.Context) {
|
|
defer q.wg.Done()
|
|
ticker := time.NewTicker(30 * time.Second)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
stats := q.Stats()
|
|
healthy := q.backend != nil && q.backend.Healthy()
|
|
log.Printf("[logger] metrics: enq=%d dropped_queue_full=%d dropped_db_fail=%d ambiguous_send=%d queue_len=%d queue_bytes=%d db_healthy=%v",
|
|
stats.Enqueued, stats.Dropped, stats.Failed, stats.Ambiguous, len(q.ch), stats.Bytes, healthy)
|
|
}
|
|
}
|
|
}
|