vibedns/internal/querylog/querylog.go

// Package querylog buffers DNS query records in memory and flushes them to
// SQLite in batches.
//
// Logging must never slow a query down or block on disk, so Log() is a
// non-blocking send onto a bounded channel: if the writer falls behind, records
// are dropped and counted rather than backing up into the resolver.
package querylog

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

	"github.com/owen/vibedns/internal/blacklist"
	"github.com/owen/vibedns/internal/database"
	"github.com/owen/vibedns/internal/models"
	"github.com/owen/vibedns/internal/netutil"
)

// Config controls query logging.
type Config struct {
	Enabled        bool
	RetentionDays  int
	MaxRows        int
	CleanupMinutes int
	IgnoreNetworks []string
	IgnoreDomains  []string
}

const (
	bufferSize    = 8192
	batchSize     = 512
	flushInterval = time.Second
)

// Logger writes query log entries to the database.
type Logger struct {
	db  *database.DB
	log *slog.Logger
	ch  chan models.QueryLogEntry

	mu             sync.RWMutex
	enabled        bool
	ignoreNetworks *netutil.PrefixSet
	ignoreDomains  *blacklist.Set
	retentionDays  int
	maxRows        int
	cleanupMinutes int

	written atomic.Int64
	dropped atomic.Int64
	pruned  atomic.Int64

	wg   sync.WaitGroup
	once sync.Once
}

// New creates a query logger. Call Start to begin draining the buffer.
func New(db *database.DB, log *slog.Logger, cfg Config) *Logger {
	l := &Logger{
		db:  db,
		log: log,
		ch:  make(chan models.QueryLogEntry, bufferSize),
	}
	l.SetConfig(cfg)
	return l
}

// SetConfig replaces the logging configuration.
func (l *Logger) SetConfig(cfg Config) {
	b := blacklist.NewBuilder(0, "querylog-ignore", models.KindBlacklist, len(cfg.IgnoreDomains))
	for _, d := range cfg.IgnoreDomains {
		b.Add(d, true) // ignoring a domain ignores its subdomains too
	}
	l.mu.Lock()
	l.enabled = cfg.Enabled
	l.ignoreNetworks = netutil.NewPrefixSet(cfg.IgnoreNetworks)
	l.ignoreDomains = b.Build()
	l.retentionDays = cfg.RetentionDays
	l.maxRows = cfg.MaxRows
	l.cleanupMinutes = cfg.CleanupMinutes
	l.mu.Unlock()
}

// Enabled reports whether logging is currently on.
func (l *Logger) Enabled() bool {
	l.mu.RLock()
	defer l.mu.RUnlock()
	return l.enabled
}

// Log queues one entry. It never blocks: a full buffer means the writer cannot
// keep up, and dropping is preferable to delaying DNS responses.
func (l *Logger) Log(e models.QueryLogEntry) {
	l.mu.RLock()
	enabled := l.enabled
	ignoreNets := l.ignoreNetworks
	ignoreDoms := l.ignoreDomains
	l.mu.RUnlock()

	if !enabled {
		return
	}
	if !ignoreNets.Empty() {
		if addr, ok := netutil.AddrFromHostPort(e.ClientIP); ok && ignoreNets.Contains(addr) {
			return
		}
	}
	if ignoreDoms.Len() > 0 {
		if _, matched := ignoreDoms.Match(e.QName); matched {
			return
		}
	}

	select {
	case l.ch <- e:
	default:
		l.dropped.Add(1)
	}
}

// Start launches the writer and the retention cleaner.
func (l *Logger) Start(ctx context.Context) {
	l.once.Do(func() {
		l.wg.Add(2)
		go l.writeLoop(ctx)
		go l.cleanupLoop(ctx)
	})
}

// Stop waits for the writer to drain and exit.
func (l *Logger) Stop() { l.wg.Wait() }

func (l *Logger) writeLoop(ctx context.Context) {
	defer l.wg.Done()
	batch := make([]models.QueryLogEntry, 0, batchSize)
	t := time.NewTicker(flushInterval)
	defer t.Stop()

	flush := func() {
		if len(batch) == 0 {
			return
		}
		// Use a detached context so a shutdown does not discard buffered rows.
		wctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
		if err := l.db.InsertQueryLogs(wctx, batch); err != nil {
			l.log.Error("could not write query log batch", "error", err, "rows", len(batch))
		} else {
			l.written.Add(int64(len(batch)))
		}
		cancel()
		batch = batch[:0]
	}

	for {
		select {
		case <-ctx.Done():
			// Drain whatever is still queued before exiting.
			for {
				select {
				case e := <-l.ch:
					batch = append(batch, e)
					if len(batch) >= batchSize {
						flush()
					}
				default:
					flush()
					return
				}
			}
		case e := <-l.ch:
			batch = append(batch, e)
			if len(batch) >= batchSize {
				flush()
			}
		case <-t.C:
			flush()
		}
	}
}

func (l *Logger) cleanupLoop(ctx context.Context) {
	defer l.wg.Done()
	for {
		l.mu.RLock()
		every := time.Duration(l.cleanupMinutes) * time.Minute
		l.mu.RUnlock()
		if every <= 0 {
			every = 30 * time.Minute
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(every):
			l.Prune(ctx)
		}
	}
}

// Prune enforces the retention policy immediately.
func (l *Logger) Prune(ctx context.Context) int64 {
	l.mu.RLock()
	days, rows := l.retentionDays, l.maxRows
	l.mu.RUnlock()
	if days <= 0 && rows <= 0 {
		return 0
	}
	n, err := l.db.PruneQueryLogs(ctx, days, rows)
	if err != nil {
		l.log.Error("could not prune query log", "error", err)
		return 0
	}
	if n > 0 {
		l.pruned.Add(n)
		l.log.Debug("pruned query log", "rows", n)
	}
	return n
}

// Stats reports logger activity.
type Stats struct {
	Enabled  bool  `json:"enabled"`
	Written  int64 `json:"written"`
	Dropped  int64 `json:"dropped"`
	Pruned   int64 `json:"pruned"`
	Buffered int   `json:"buffered"`
}

// Stats returns the logger counters.
func (l *Logger) Stats() Stats {
	return Stats{
		Enabled:  l.Enabled(),
		Written:  l.written.Load(),
		Dropped:  l.dropped.Load(),
		Pruned:   l.pruned.Load(),
		Buffered: len(l.ch),
	}
}