Agent-System/internal/events/worker.go

package events

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Handler runs outside the claim transaction but must apply side effects on the
// provided tx so fencing/ack/inbox share one unit of work.
type Handler func(ctx context.Context, lease Lease, tx pgx.Tx) error

type Dispatcher struct {
	Pool        *pgxpool.Pool
	Environment string
	Owner       string
	Batch       int
	Lease       time.Duration
	Heartbeat   time.Duration
	Handlers    map[string]Handler
	Metrics     *Metrics

	mu        sync.Mutex
	claiming  bool
	inflight  map[string]struct{}
	cancelFns []context.CancelFunc
	lastBeat  atomic.Int64
}

type Metrics struct {
	Claimed    atomic.Int64
	Acked      atomic.Int64
	Failed     atomic.Int64
	Fenced     atomic.Int64
	Dead       atomic.Int64
	Quarantine atomic.Int64
	Heartbeat  atomic.Int64
}

func NewDispatcher(pool *pgxpool.Pool, environment, owner string) *Dispatcher {
	d := &Dispatcher{
		Pool:        pool,
		Environment: environment,
		Owner:       owner,
		Batch:       DefaultBatch,
		Lease:       DefaultLease,
		Heartbeat:   DefaultHeartbeat,
		Handlers:    map[string]Handler{},
		Metrics:     &Metrics{},
		claiming:    true,
		inflight:    map[string]struct{}{},
	}
	d.Touch()
	return d
}

func (d *Dispatcher) Touch() { d.lastBeat.Store(time.Now().UnixNano()) }

func (d *Dispatcher) Alive(maxAge time.Duration) bool {
	if maxAge <= 0 {
		maxAge = 30 * time.Second
	}
	return time.Since(time.Unix(0, d.lastBeat.Load())) < maxAge
}

func (d *Dispatcher) StopClaim() {
	d.mu.Lock()
	d.claiming = false
	cancels := append([]context.CancelFunc(nil), d.cancelFns...)
	d.mu.Unlock()
	for _, c := range cancels {
		c()
	}
}

func (d *Dispatcher) Register(consumer string, h Handler) error {
	if err := NewRegistry().Register(consumer); err != nil {
		return err
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	d.Handlers[consumer] = h
	return nil
}

func (d *Dispatcher) claimingEnabled() bool {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.claiming
}

func (d *Dispatcher) Serve(ctx context.Context, consumers []string) error {
	tick := time.NewTicker(time.Second)
	defer tick.Stop()
	d.Touch()
	for {
		select {
		case <-ctx.Done():
			d.StopClaim()
			return ctx.Err()
		case <-tick.C:
			d.Touch()
			for _, c := range consumers {
				_, _ = d.RunOnce(ctx, c)
			}
		}
	}
}

func (d *Dispatcher) RunOnce(ctx context.Context, consumer string) (int, error) {
	d.Touch()
	if !d.claimingEnabled() {
		return 0, nil
	}
	mode, _, err := LoadMode(ctx, d.Pool, d.Environment, consumer)
	if err != nil {
		return 0, err
	}
	if mode == ModeOff || consumer == ConsumerFAV {
		return 0, nil
	}
	if mode == ModeShadow {
		return d.shadowScan(ctx, consumer)
	}
	tx, err := d.Pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	leases, err := ClaimLeases(ctx, tx, consumer, d.Owner, d.Batch, d.Lease)
	if err != nil {
		return 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	d.Metrics.Claimed.Add(int64(len(leases)))
	for _, lease := range leases {
		d.handle(ctx, lease)
	}
	return len(leases), nil
}

func (d *Dispatcher) shadowScan(ctx context.Context, consumer string) (int, error) {
	leases, err := PeekPendingLeases(ctx, d.Pool, consumer, d.Batch)
	if err != nil {
		return 0, err
	}
	h := d.Handlers[consumer]
	var n int
	for _, lease := range leases {
		digest := Digest([]byte(lease.EventID + ":" + consumer + ":" + ModeShadow))
		if h != nil {
			tx, err := d.Pool.Begin(ctx)
			if err != nil {
				return n, err
			}
			herr := h(ctx, lease, tx)
			if herr != nil {
				digest = Digest([]byte("fail:" + herr.Error()))
			} else if dgst, herr := HashVisibleNotifications(ctx, tx, lease.EventID); herr == nil {
				digest = dgst
			}
			_ = tx.Rollback(ctx)
		}
		wtx, err := d.Pool.Begin(ctx)
		if err != nil {
			return n, err
		}
		if err := WriteReconciliation(ctx, wtx, lease.EventID, consumer, ModeShadow, digest); err != nil {
			_ = wtx.Rollback(ctx)
			return n, err
		}
		if err := wtx.Commit(ctx); err != nil {
			return n, err
		}
		n++
	}
	d.Touch()
	return n, nil
}

func (d *Dispatcher) handle(ctx context.Context, lease Lease) {
	key := lease.Consumer + ":" + lease.AggregateType + ":" + lease.AggregateID
	d.mu.Lock()
	d.inflight[key] = struct{}{}
	hctx, cancel := context.WithCancel(ctx)
	d.cancelFns = append(d.cancelFns, cancel)
	d.mu.Unlock()
	defer func() {
		cancel()
		d.mu.Lock()
		delete(d.inflight, key)
		d.mu.Unlock()
	}()

	done := make(chan struct{})
	go func() {
		ticker := time.NewTicker(d.Heartbeat)
		defer ticker.Stop()
		for {
			select {
			case <-done:
				return
			case <-ticker.C:
				tx, err := d.Pool.Begin(hctx)
				if err != nil {
					cancel()
					return
				}
				err = HeartbeatLease(hctx, tx, lease, d.Lease)
				if err != nil {
					_ = tx.Rollback(hctx)
					d.Metrics.Fenced.Add(1)
					cancel()
					return
				}
				if err := tx.Commit(hctx); err != nil {
					cancel()
					return
				}
				d.Metrics.Heartbeat.Add(1)
				d.Touch()
			}
		}
	}()
	defer close(done)

	h := d.Handlers[lease.Consumer]
	if h == nil {
		_ = d.fail(ctx, lease, "NO_HANDLER", true)
		return
	}
	tx, err := d.Pool.Begin(hctx)
	if err != nil {
		_ = d.fail(ctx, lease, "TRANSIENT", true)
		return
	}
	if err := h(hctx, lease, tx); err != nil {
		_ = tx.Rollback(hctx)
		if hctx.Err() != nil {
			return
		}
		code, retry := classifyHandlerErr(err)
		_ = d.fail(ctx, lease, code, retry)
		return
	}
	if err := AckLease(hctx, tx, lease); err != nil {
		_ = tx.Rollback(hctx)
		if errors.Is(err, ErrStaleFencing) {
			d.Metrics.Fenced.Add(1)
			return
		}
		_ = d.fail(ctx, lease, "TRANSIENT", true)
		return
	}
	if err := tx.Commit(hctx); err != nil {
		_ = d.fail(ctx, lease, "TRANSIENT", true)
		return
	}
	d.Metrics.Acked.Add(1)
	d.Touch()
}

func (d *Dispatcher) fail(ctx context.Context, lease Lease, code string, retryable bool) error {
	tx, err := d.Pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if err := FailLease(ctx, tx, lease, code, retryable); err != nil {
		if errors.Is(err, ErrStaleFencing) {
			d.Metrics.Fenced.Add(1)
			return err
		}
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return err
	}
	d.Metrics.Failed.Add(1)
	if code == "QUARANTINE" || code == "AUTHZ" || code == "SCOPE" {
		d.Metrics.Quarantine.Add(1)
	}
	if !retryable || lease.Attempt >= MaxAttempts {
		d.Metrics.Dead.Add(1)
	}
	d.Touch()
	return nil
}

func classifyHandlerErr(err error) (string, bool) {
	switch {
	case errors.Is(err, ErrQuarantine), errors.Is(err, ErrAuthFailure):
		return "QUARANTINE", false
	case errors.Is(err, ErrIntegrityFailure):
		return "INTEGRITY", false
	case errors.Is(err, ErrInvalidScope):
		return "SCOPE", false
	case errors.Is(err, ErrStaleFencing):
		return "FENCED", false
	default:
		return "TRANSIENT", true
	}
}