159 lines
3.6 KiB
Go
159 lines
3.6 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log"
|
|
"strings"
|
|
"time"
|
|
"unicode/utf8"
|
|
)
|
|
|
|
const (
|
|
maxLogFieldBytes = 512
|
|
maxLogErrorBytes = 4096
|
|
maxLogUsageBytes = 16384
|
|
logCleanupBatch = 1000
|
|
)
|
|
|
|
func (s *Server) SetLogRetention(days int) {
|
|
if days < 0 {
|
|
days = 0
|
|
}
|
|
s.logRetentionDays.Store(int64(days))
|
|
}
|
|
|
|
func truncateLogField(value string, limit int) string {
|
|
if len(value) <= limit {
|
|
return value
|
|
}
|
|
for limit > 0 && !utf8.RuneStart(value[limit]) {
|
|
limit--
|
|
}
|
|
return strings.Clone(value[:limit])
|
|
}
|
|
|
|
func boundedLog(l requestLog) requestLog {
|
|
for _, field := range []*string{&l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream} {
|
|
*field = truncateLogField(*field, maxLogFieldBytes)
|
|
}
|
|
l.Error = truncateLogField(l.Error, maxLogErrorBytes)
|
|
if len(l.Usage.Raw) > maxLogUsageBytes {
|
|
// Keep valid JSON without retaining an attacker-controlled oversized object.
|
|
l.Usage.Raw = []byte(`{"truncated":true}`)
|
|
}
|
|
l.RawUsage = nil
|
|
return l
|
|
}
|
|
|
|
func (s *Server) purgeExpiredLogs(ctx context.Context, days int) error {
|
|
if days <= 0 {
|
|
return nil
|
|
}
|
|
cutoff := time.Now().UTC().AddDate(0, 0, -days).Truncate(time.Second)
|
|
_, err := s.db.ExecContext(ctx, "DELETE FROM logs WHERE id IN (SELECT id FROM logs WHERE created_at<? AND julianday(created_at)<julianday(?) LIMIT ?)", cutoff.Format("2006-01-02T15:04:05"), cutoff.Format(time.RFC3339), logCleanupBatch)
|
|
return err
|
|
}
|
|
|
|
type logEntry struct {
|
|
value requestLog
|
|
barrier chan error
|
|
}
|
|
|
|
func (s *Server) record(l requestLog) {
|
|
s.logMu.RLock()
|
|
defer s.logMu.RUnlock()
|
|
if s.logsClosed {
|
|
return
|
|
}
|
|
l = boundedLog(l)
|
|
l.Usage.Raw = append([]byte(nil), l.Usage.Raw...)
|
|
s.logQueue <- logEntry{value: l}
|
|
}
|
|
|
|
func (s *Server) flushLogs() error {
|
|
s.logMu.RLock()
|
|
if s.logsClosed {
|
|
s.logMu.RUnlock()
|
|
<-s.logDone
|
|
return s.logErr
|
|
}
|
|
done := make(chan error, 1)
|
|
s.logQueue <- logEntry{barrier: done}
|
|
s.logMu.RUnlock()
|
|
return <-done
|
|
}
|
|
|
|
func (s *Server) writeLogs() {
|
|
defer close(s.logDone)
|
|
ticker := time.NewTicker(100 * time.Millisecond)
|
|
defer ticker.Stop()
|
|
cleanup := time.NewTicker(time.Minute)
|
|
defer cleanup.Stop()
|
|
batch := make([]requestLog, 0, 100)
|
|
flush := func() {
|
|
if len(batch) == 0 {
|
|
return
|
|
}
|
|
for attempt := 0; ; attempt++ {
|
|
if s.persistLogs(batch) == nil {
|
|
break
|
|
}
|
|
if attempt == 3 {
|
|
s.logErr = errors.New("request log persistence failed after 3 retries")
|
|
log.Printf("ai-router: request log persistence failed; batch_size=%d retries=3", len(batch))
|
|
break
|
|
}
|
|
log.Printf("ai-router: retrying request log persistence; batch_size=%d retry=%d", len(batch), attempt+1)
|
|
time.Sleep(time.Duration(attempt+1) * 10 * time.Millisecond)
|
|
}
|
|
clear(batch)
|
|
batch = batch[:0]
|
|
}
|
|
for {
|
|
select {
|
|
case entry, ok := <-s.logQueue:
|
|
if !ok {
|
|
flush()
|
|
return
|
|
}
|
|
if entry.barrier != nil {
|
|
flush()
|
|
entry.barrier <- s.logErr
|
|
continue
|
|
}
|
|
batch = append(batch, entry.value)
|
|
if len(batch) == cap(batch) {
|
|
flush()
|
|
}
|
|
case <-ticker.C:
|
|
flush()
|
|
case <-cleanup.C:
|
|
if days := s.logRetentionDays.Load(); days > 0 {
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
err := s.purgeExpiredLogs(ctx, int(days))
|
|
cancel()
|
|
if err != nil {
|
|
log.Print("ai-router: request log retention cleanup failed")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *Server) persistLogs(batch []requestLog) error {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
tx, err := s.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback()
|
|
for _, l := range batch {
|
|
if err := insertLog(ctx, tx, l); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return tx.Commit()
|
|
}
|