package connectionlog

import (

"context"

"io"

"sync"

"time"

"github.com/google/uuid"

"github.com/hashicorp/go-multierror"

"github.com/sqlc-dev/pqtype"

"cdr.dev/slog/v3"

"github.com/coder/coder/v2/coderd/database"

"github.com/coder/coder/v2/coderd/database/dbauthz"

auditbackends "github.com/coder/coder/v2/enterprise/audit/backends"

"github.com/coder/quartz"

)

const (

// defaultBatchSize is the maximum number of connection log entries

// to batch before forcing a flush.

defaultBatchSize = 1000

// defaultFlushInterval is how frequently to flush batched connection

// log entries to the database. Five seconds balances near-real-time

// audit visibility with write efficiency.

defaultFlushInterval = 5 * time.Second

// retryQueueSize is the capacity of the bounded retry channel.

// Failed batches beyond this limit are dropped.

retryQueueSize = 10

// shutdownWriteTimeout bounds how long a final write attempt

// can take during shutdown when the batcher context is already

// canceled.

shutdownWriteTimeout = 10 * time.Second

// maxRetries is the number of times to retry a failed batch

// write before dropping it and moving on.

maxRetries = 3

// retryInterval is the fixed delay between retry attempts.

retryInterval = time.Second

)

// Backend is a destination for connection log events. Backends that

// also implement io.Closer will be closed when the ConnectionLogger

// is closed.

type Backend interface {

Upsert(ctx context.Context, clog database.UpsertConnectionLogParams) error

}

// ConnectionLogger fans out each connection log event to every

// registered backend.

type ConnectionLogger struct {

backends []Backend

}

// New creates a ConnectionLogger that dispatches to the given

// backends.

func New(backends ...Backend) *ConnectionLogger {

return &ConnectionLogger{

backends: backends,

}

}

func (c *ConnectionLogger) Upsert(ctx context.Context, clog database.UpsertConnectionLogParams) error {

var errs error

for _, backend := range c.backends {

err := backend.Upsert(ctx, clog)

if err != nil {

errs = multierror.Append(errs, err)

}

}

return errs

}

// Close closes all backends that implement io.Closer.

func (c *ConnectionLogger) Close() error {

var errs error

for _, backend := range c.backends {

if closer, ok := backend.(io.Closer); ok {

if err := closer.Close(); err != nil {

errs = multierror.Append(errs, err)

}

}

}

return errs

}

// DBBatcherOption is a functional option for configuring a DBBatcher.

type DBBatcherOption func(b *DBBatcher)

// WithBatchSize sets the maximum number of entries to accumulate

// before forcing a flush.

func WithBatchSize(size int) DBBatcherOption {

return func(b *DBBatcher) {

b.maxBatchSize = size

}

}

// WithFlushInterval sets how frequently the batcher flushes to the

// database.

func WithFlushInterval(d time.Duration) DBBatcherOption {

return func(b *DBBatcher) {

b.interval = d

}

}

// WithClock sets the clock, useful for testing.

func WithClock(clock quartz.Clock) DBBatcherOption {

return func(b *DBBatcher) {

b.clock = clock

}

}

// DBBatcher batches connection log upserts and periodically flushes

// them to the database to reduce per-event write pressure.

type DBBatcher struct {

store database.Store

log slog.Logger

itemCh chan database.UpsertConnectionLogParams

// dedupedBatch holds entries keyed by connection ID so that

// PostgreSQL never sees the same row twice in one INSERT …

// ON CONFLICT DO UPDATE. Connection IDs are globally unique

// (each new session gets a fresh UUID). Entries with a NULL

// connection_id (web events) go into nullConnIDBatch instead

// because NULL != NULL in SQL unique constraints.

dedupedBatch map[uuid.UUID]batchEntry

nullConnIDBatch []batchEntry

maxBatchSize int

// retryCh is a bounded channel of failed batches awaiting

// retry. A single retry worker goroutine processes this

// channel, retrying each batch up to maxRetries times before

// dropping it. If the channel is full, new failures are

// dropped immediately.

retryCh chan database.BatchUpsertConnectionLogsParams

clock quartz.Clock

timer *quartz.Timer

interval time.Duration

ctx context.Context

cancel context.CancelFunc

wg sync.WaitGroup

}

// NewDBBatcher creates a DBBatcher that batches writes to the database

// and starts its background processing loop. Close must be called to

// flush remaining entries on shutdown.

func NewDBBatcher(ctx context.Context, store database.Store, log slog.Logger, opts ...DBBatcherOption) *DBBatcher {

b := &DBBatcher{

store: store,

log: log,

clock: quartz.NewReal(),

}

for _, opt := range opts {

opt(b)

}

if b.interval == 0 {

b.interval = defaultFlushInterval

}

if b.maxBatchSize == 0 {

b.maxBatchSize = defaultBatchSize

}

b.timer = b.clock.NewTimer(b.interval)

b.itemCh = make(chan database.UpsertConnectionLogParams, b.maxBatchSize)

b.dedupedBatch = make(map[uuid.UUID]batchEntry, b.maxBatchSize)

b.retryCh = make(chan database.BatchUpsertConnectionLogsParams, retryQueueSize)

b.ctx, b.cancel = context.WithCancel(ctx)

b.wg.Add(2)

go func() {

defer b.wg.Done()

b.run(b.ctx)

}()

go func() {

defer b.wg.Done()

b.retryLoop()

}()

return b

}

// Upsert enqueues a connection log entry for batched writing. It

// blocks if the internal buffer is full, ensuring no logs are dropped.

// It returns an error if the batcher or caller context is canceled.

func (b *DBBatcher) Upsert(ctx context.Context, clog database.UpsertConnectionLogParams) error {

if b.ctx.Err() != nil {

return b.ctx.Err()

}

select {

case b.itemCh <- clog:

return nil

case <-b.ctx.Done():

return b.ctx.Err()

case <-ctx.Done():

return ctx.Err()

}

}

// Close cancels the batcher context, waits for the run loop and

// retry worker to exit.

func (b *DBBatcher) Close() error {

b.cancel()

if b.timer != nil {

b.timer.Stop()

}

b.wg.Wait()

return nil

}

// addToBatch inserts an item into the batch, deduplicating by conflict

// key on the fly. For entries with the same key, disconnect events are

// preferred over connect events, and later events are preferred over

// earlier ones.

//

// This is safe because each new connection gets a fresh UUID (see

// agent/agent.go and agent/agentssh), so the only duplicate for the

// same (connection_id, workspace_id, agent_name) is a connect/disconnect

// pair for the same session. A "reconnect" always uses a new ID.

func (b *DBBatcher) addToBatch(item database.UpsertConnectionLogParams) {

entry := batchEntry{

UpsertConnectionLogParams: item,

}

if item.ConnectionStatus == database.ConnectionStatusDisconnected {

// For standalone disconnect events, use the disconnect

// time as both connect and disconnect time. This matches

// the single-row UpsertConnectionLog behavior which uses

// @time for connect_time regardless of status. The SQL

// LEAST logic will correct connect_time if the real

// connect event arrives in a later batch.

entry.connectTime = item.Time

entry.disconnectTime = item.Time

} else {

entry.connectTime = item.Time

}

if !item.ConnectionID.Valid {

b.nullConnIDBatch = append(b.nullConnIDBatch, entry)

return

}

connID := item.ConnectionID.UUID

existing, ok := b.dedupedBatch[connID]

if !ok {

b.dedupedBatch[connID] = entry

return

}

// When merging entries for the same connection, always preserve

// the earliest non-zero connect_time and latest disconnect_time

// so the row records the full session span.

if !existing.connectTime.IsZero() && existing.connectTime.Before(entry.connectTime) {

entry.connectTime = existing.connectTime

}

if existing.disconnectTime.After(entry.disconnectTime) {

entry.disconnectTime = existing.disconnectTime

}

// Prefer disconnect over connect (superset of info).

// If same status, prefer the later event.

if item.ConnectionStatus == database.ConnectionStatusDisconnected &&

existing.ConnectionStatus != database.ConnectionStatusDisconnected {

b.dedupedBatch[connID] = entry

} else if item.Time.After(existing.Time) {

b.dedupedBatch[connID] = entry

}

}

// batchLen returns the total number of entries currently buffered.

func (b *DBBatcher) batchLen() int {

return len(b.dedupedBatch) + len(b.nullConnIDBatch)

}

func (b *DBBatcher) run(ctx context.Context) {

//nolint:gocritic // System-level batch operation for connection logs.

authCtx := dbauthz.AsConnectionLogger(ctx)

for ctx.Err() == nil {

select {

case item := <-b.itemCh:

b.addToBatch(item)

if b.batchLen() >= b.maxBatchSize {

b.flush(authCtx)

b.timer.Reset(b.interval, "connectionLogBatcher", "capacityFlush")

}

case <-b.timer.C:

b.flush(authCtx)

b.timer.Reset(b.interval, "connectionLogBatcher", "scheduledFlush")

case <-ctx.Done():

}

}

b.log.Debug(ctx, "context done, flushing before exit")

// Drain any remaining items from the channel.

for {

select {

case item := <-b.itemCh:

b.addToBatch(item)

default:

if b.batchLen() > 0 {

b.shutdownBatch(b.buildParams())

}

// Signal the retry worker to skip delays and close

// the channel so it exits after processing any

// remaining items.

// Mark the batcher as closed so that any subsequent

// Upsert calls fail immediately instead of sending

// into itemCh after the run loop has exited.

close(b.retryCh)

return

}

}

}

// batchEntry wraps a connection log event with explicit connect and

// disconnect times. When a connect and disconnect for the same session

// are merged into one entry, connectTime preserves the original

// session start while disconnectTime records when it ended.

type batchEntry struct {

database.UpsertConnectionLogParams

connectTime time.Time

disconnectTime time.Time

}

// flush builds the batch params, clears the in-memory batch, and

// writes to the database. On failure, the batch is queued for retry

// by the single retry worker goroutine. If the retry queue is full,

// the batch is dropped.

func (b *DBBatcher) flush(ctx context.Context) {

count := b.batchLen()

if count == 0 {

return

}

params := b.buildParams()

// Clear the batch before writing so the run loop can start

// accumulating new entries.

b.dedupedBatch = make(map[uuid.UUID]batchEntry, b.maxBatchSize)

b.nullConnIDBatch = nil

// Use the batcher's context for normal operation so Close()

// can cancel hung writes. During shutdown (ctx already canceled),

// fall back to a bounded timeout.

writeCtx := b.ctx

if writeCtx.Err() != nil {

var cancel context.CancelFunc

writeCtx, cancel = context.WithTimeout(context.Background(), shutdownWriteTimeout)

defer cancel()

}

//nolint:gocritic // System-level batch operation for connection logs.

err := b.store.BatchUpsertConnectionLogs(dbauthz.AsConnectionLogger(writeCtx), params)

if err == nil {

return

}

b.log.Error(ctx, "batch upsert failed, queueing for retry",

slog.Error(err), slog.F("count", count))

// Don't retry on shutdown.

if ctx.Err() != nil {

return

}

select {

case b.retryCh <- params:

default:

b.log.Error(ctx, "retry queue full, dropping batch",

slog.F("dropped", count))

}

}

func (b *DBBatcher) buildParams() database.BatchUpsertConnectionLogsParams {

count := b.batchLen()

var (

ids = make([]uuid.UUID, 0, count)

connectTime = make([]time.Time, 0, count)

organizationID = make([]uuid.UUID, 0, count)

workspaceOwnerID = make([]uuid.UUID, 0, count)

workspaceID = make([]uuid.UUID, 0, count)

workspaceName = make([]string, 0, count)

agentName = make([]string, 0, count)

connType = make([]database.ConnectionType, 0, count)

code = make([]int32, 0, count)

codeValid = make([]bool, 0, count)

ip = make([]pqtype.Inet, 0, count)

userAgent = make([]string, 0, count)

userID = make([]uuid.UUID, 0, count)

slugOrPort = make([]string, 0, count)

connectionID = make([]uuid.UUID, 0, count)

disconnectReason = make([]string, 0, count)

disconnectTime = make([]time.Time, 0, count)

)

appendEntry := func(e batchEntry) {

ids = append(ids, e.ID)

connectTime = append(connectTime, e.connectTime)

organizationID = append(organizationID, e.OrganizationID)

workspaceOwnerID = append(workspaceOwnerID, e.WorkspaceOwnerID)

workspaceID = append(workspaceID, e.WorkspaceID)

workspaceName = append(workspaceName, e.WorkspaceName)

agentName = append(agentName, e.AgentName)

connType = append(connType, e.Type)

code = append(code, e.Code.Int32)

codeValid = append(codeValid, e.Code.Valid)

ip = append(ip, e.IP)

userAgent = append(userAgent, e.UserAgent.String)

userID = append(userID, e.UserID.UUID)

slugOrPort = append(slugOrPort, e.SlugOrPort.String)

connectionID = append(connectionID, e.ConnectionID.UUID)

disconnectReason = append(disconnectReason, e.DisconnectReason.String)

disconnectTime = append(disconnectTime, e.disconnectTime)

}

for _, entry := range b.dedupedBatch {

appendEntry(entry)

}

for _, entry := range b.nullConnIDBatch {

appendEntry(entry)

}

return database.BatchUpsertConnectionLogsParams{

ID: ids,

ConnectTime: connectTime,

OrganizationID: organizationID,

WorkspaceOwnerID: workspaceOwnerID,

WorkspaceID: workspaceID,

WorkspaceName: workspaceName,

AgentName: agentName,

Type: connType,

Code: code,

CodeValid: codeValid,

Ip: ip,

UserAgent: userAgent,

UserID: userID,

SlugOrPort: slugOrPort,

ConnectionID: connectionID,

DisconnectReason: disconnectReason,

DisconnectTime: disconnectTime,

}

}

// retryLoop is a single background goroutine that processes failed

// batches from retryCh. Each batch is retried up to maxRetries times

// with a fixed delay between attempts. When draining is set (shutdown),

// batches get a single immediate write attempt instead. The loop exits

// when retryCh is closed by the run goroutine.

func (b *DBBatcher) retryLoop() {

for params := range b.retryCh {

b.retryBatch(params)

}

}

// retryBatch retries writing a batch up to maxRetries times with a

// fixed delay between attempts. If the batcher context is canceled

// during a wait, one final attempt is made before returning.

func (b *DBBatcher) retryBatch(params database.BatchUpsertConnectionLogsParams) {

count := len(params.ID)

for attempt := range maxRetries {

t := b.clock.NewTimer(retryInterval, "connectionLogBatcher", "retryBackoff")

select {

case <-b.ctx.Done():

t.Stop()

b.shutdownBatch(params)

return

case <-t.C:

}

//nolint:gocritic // System-level batch operation for connection logs.

err := b.store.BatchUpsertConnectionLogs(dbauthz.AsConnectionLogger(b.ctx), params)

if err == nil {

return

}

b.log.Warn(b.ctx, "batch retry failed",

slog.Error(err),

slog.F("count", count),

slog.F("attempt", attempt+1),

slog.F("max_attempts", maxRetries),

)

}

b.log.Error(b.ctx, "batch retries exhausted, dropping batch",

slog.F("dropped", count))

}

// shutdownBatch makes a single write attempt during shutdown with a

// bounded timeout so it can't hang indefinitely.

func (b *DBBatcher) shutdownBatch(params database.BatchUpsertConnectionLogsParams) {

ctx, cancel := context.WithTimeout(context.Background(), shutdownWriteTimeout)

defer cancel()

//nolint:gocritic // System-level batch operation for connection logs.

err := b.store.BatchUpsertConnectionLogs(dbauthz.AsConnectionLogger(ctx), params)

if err != nil {

b.log.Error(b.ctx, "batch write failed on shutdown, dropping batch",

slog.Error(err), slog.F("dropped", len(params.ID)))

}

}

type connectionSlogBackend struct {

exporter *auditbackends.SlogExporter

}

// NewSlogBackend returns a Backend that logs connection events via

// the structured logger.

func NewSlogBackend(logger slog.Logger) Backend {

return &connectionSlogBackend{

exporter: auditbackends.NewSlogExporter(logger),

}

}

func (b *connectionSlogBackend) Upsert(ctx context.Context, clog database.UpsertConnectionLogParams) error {

return b.exporter.ExportStruct(ctx, clog, "connection_log")

}