Agent-System/internal/events/bus.go

package events

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"io"
	"sync"
	"time"
)

const MaxAttempts = 12

// MemoryBus is a process-local outbox/delivery/inbox used for dispatcher
// characterization. PostgreSQL SKIP LOCKED is the production claim primitive.
type MemoryBus struct {
	mu              sync.Mutex
	envelopes       map[string]Envelope
	order           []string
	deliveries      map[string]map[string]*Delivery // consumer -> eventID
	heads           map[string]int64                // aggregate -> last applied version
	side            map[string]int                  // consumer:event visible side effects
	consumerProject map[string]int64
}

func NewMemoryBus() *MemoryBus {
	return &MemoryBus{
		envelopes:       map[string]Envelope{},
		deliveries:      map[string]map[string]*Delivery{},
		heads:           map[string]int64{},
		side:            map[string]int{},
		consumerProject: map[string]int64{},
	}
}

func (b *MemoryBus) BindConsumer(consumer string, projectID int64) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.consumerProject[consumer] = projectID
}

func (b *MemoryBus) Append(ctx context.Context, env Envelope, consumers []string) error {
	_ = ctx
	if err := ValidateEnvelope(env); err != nil {
		b.quarantine(env, consumers)
		return err
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	if _, ok := b.envelopes[env.EventID]; ok {
		return nil
	}
	b.envelopes[env.EventID] = env
	b.order = append(b.order, env.EventID)
	for _, c := range consumers {
		if b.deliveries[c] == nil {
			b.deliveries[c] = map[string]*Delivery{}
		}
		b.deliveries[c][env.EventID] = &Delivery{
			EventID: env.EventID, Consumer: c, Status: "pending", NextAttemptAt: time.Now(),
		}
	}
	return nil
}

func (b *MemoryBus) quarantine(env Envelope, consumers []string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.envelopes[env.EventID] = env
	for _, c := range consumers {
		if b.deliveries[c] == nil {
			b.deliveries[c] = map[string]*Delivery{}
		}
		b.deliveries[c][env.EventID] = &Delivery{EventID: env.EventID, Consumer: c, Status: "quarantine"}
	}
}

func (b *MemoryBus) Claim(consumer, worker string, n int, now time.Time) ([]*Delivery, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	rows := b.deliveries[consumer]
	if rows == nil {
		return nil, ErrUnknownConsumer
	}
	out := []*Delivery{}
	for _, id := range b.order {
		if len(out) >= n {
			break
		}
		d := rows[id]
		if d == nil {
			continue
		}
		if d.Status == "acked" || d.Status == "dead" || d.Status == "quarantine" {
			continue
		}
		if d.Status == "leased" && d.LeaseUntil.After(now) {
			continue
		}
		env := b.envelopes[id]
		head := b.heads[env.AggregateID]
		if env.AggregateVersion != 0 && env.AggregateVersion > head+1 {
			continue // gap: do not skip missing versions
		}
		token, err := randomToken()
		if err != nil {
			return nil, err
		}
		d.Status = "leased"
		d.LeaseOwner = worker
		d.FencingToken = token
		d.LeaseUntil = now.Add(2 * time.Second)
		d.Attempt++
		cp := *d
		out = append(out, &cp)
	}
	return out, nil
}

func (b *MemoryBus) Apply(consumer string, claimed *Delivery, fail error) error {
	b.mu.Lock()
	defer b.mu.Unlock()
	d := b.deliveries[consumer][claimed.EventID]
	if d == nil || d.FencingToken != claimed.FencingToken || d.LeaseOwner != claimed.LeaseOwner {
		return ErrStaleFencing
	}
	env := b.envelopes[claimed.EventID]
	if pid := b.consumerProject[consumer]; pid != 0 && env.ProjectID != 0 && env.ProjectID != pid {
		d.Status = "quarantine"
		d.LeaseOwner = ""
		d.FencingToken = ""
		return ErrQuarantine
	}
	if fail == ErrAuthFailure || fail == ErrIntegrityFailure {
		d.Status = "dead"
		return fail
	}
	if fail != nil {
		if d.Attempt >= MaxAttempts {
			d.Status = "dead"
			return fail
		}
		d.Status = "pending"
		d.LeaseOwner = ""
		d.FencingToken = ""
		d.NextAttemptAt = time.Now().Add(time.Millisecond)
		return fail
	}
	head := b.heads[env.AggregateID]
	if env.AggregateVersion != 0 && env.AggregateVersion > head+1 {
		d.Status = "pending"
		d.LeaseOwner = ""
		d.FencingToken = ""
		return ErrGapWait
	}
	key := consumer + ":" + claimed.EventID
	if !d.Inbox {
		b.side[key]++
		d.Inbox = true
		if env.AggregateVersion > head {
			b.heads[env.AggregateID] = env.AggregateVersion
		}
	}
	d.Status = "acked"
	d.LeaseOwner = ""
	return nil
}

func (b *MemoryBus) SideEffects(consumer, eventID string) int {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.side[consumer+":"+eventID]
}

func (b *MemoryBus) Status(consumer, eventID string) string {
	b.mu.Lock()
	defer b.mu.Unlock()
	d := b.deliveries[consumer][eventID]
	if d == nil {
		return ""
	}
	return d.Status
}

func (b *MemoryBus) Owners(consumer string) map[string]string {
	b.mu.Lock()
	defer b.mu.Unlock()
	out := map[string]string{}
	for id, d := range b.deliveries[consumer] {
		if d.Status == "leased" {
			out[id] = d.LeaseOwner
		}
	}
	return out
}

func randomToken() (string, error) {
	var b [8]byte
	if _, err := io.ReadFull(rand.Reader, b[:]); err != nil {
		return "", err
	}
	return hex.EncodeToString(b[:]), nil
}