Cyrene_v2/internal/bus/impl.go

package bus

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
)

// Compile-time assertion that *busImpl satisfies Bus.
var _ Bus = (*busImpl)(nil)

// busImpl is the default in-process Bus implementation. It is safe for
// concurrent use.
type busImpl struct {
	queueSize int
	mu        sync.RWMutex
	subs      map[uint64]*subscription
	nextID    atomic.Uint64
	dropped   atomic.Uint64
	closed    atomic.Bool
	ctx       context.Context
	cancel    context.CancelFunc
	closeOnce sync.Once
}

// Dropped returns the total number of events dropped because a subscription's
// buffer was full. It is useful for observing backpressure on the bus.
func (b *busImpl) Dropped() uint64 {
	return b.dropped.Load()
}

// subscription is one registered handler plus its bounded queue and the single
// worker goroutine that drains it.
type subscription struct {
	id      uint64
	pattern string
	handler Handler
	events  chan Event
	done    chan struct{}
	ctx     context.Context
	cancel  context.CancelFunc
	once    sync.Once
}

// stop idempotently cancels the handler context and signals the worker to
// exit. It is safe to call more than once and from more than one goroutine.
func (s *subscription) stop() {
	s.once.Do(func() {
		s.cancel()
		close(s.done)
	})
}

// Publish fans an event out to every matching subscription without waiting for
// any handler. It holds the read lock across the fan-out so a concurrent
// Close cannot stop a worker mid-broadcast. Events that cannot be queued
// because a subscription buffer is full are dropped and counted.
func (b *busImpl) Publish(ctx context.Context, subject string, payload any) error {
	if err := validateSubject(subject); err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		return fmt.Errorf("bus: publish %q: %w", subject, err)
	}

	b.mu.RLock()
	defer b.mu.RUnlock()
	if b.closed.Load() {
		return fmt.Errorf("bus is closed: %w", ErrClosed)
	}

	ev := Event{Subject: subject, Payload: payload}
	for _, sub := range b.subs {
		if !match(sub.pattern, subject) {
			continue
		}
		if err := ctx.Err(); err != nil {
			return fmt.Errorf("bus: publish %q: %w", subject, err)
		}
		select {
		case sub.events <- ev:
		default:
			b.dropped.Add(1)
		}
	}
	return nil
}

// Subscribe registers a handler and returns an unsubscribe function. The
// registration is removed automatically when the caller's context is canceled.
func (b *busImpl) Subscribe(ctx context.Context, pattern string, h Handler) (func(), error) {
	if err := validatePattern(pattern); err != nil {
		return nil, err
	}
	if h == nil {
		return nil, fmt.Errorf("bus: handler must not be nil: %w", ErrInvalidArgument)
	}
	if err := ctx.Err(); err != nil {
		return nil, fmt.Errorf("bus: subscribe %q: %w", pattern, err)
	}

	subCtx, subCancel := context.WithCancel(b.ctx)
	sub := &subscription{
		id:      b.nextID.Add(1),
		pattern: pattern,
		handler: h,
		events:  make(chan Event, b.queueSize),
		done:    make(chan struct{}),
		ctx:     subCtx,
		cancel:  subCancel,
	}

	b.mu.Lock()
	if b.closed.Load() {
		b.mu.Unlock()
		subCancel()
		return nil, fmt.Errorf("bus is closed: %w", ErrClosed)
	}
	b.subs[sub.id] = sub
	b.mu.Unlock()

	b.startWorker(sub)

	go func() {
		select {
		case <-ctx.Done():
			b.removeSub(sub.id, sub)
		case <-sub.done:
			// Already removed; nothing to do.
		}
	}()

	return func() { b.removeSub(sub.id, sub) }, nil
}

// Close is idempotent and safe to call concurrently. It marks the bus closed,
// cancels every handler context, unsubscribes all registrations, and signals
// every worker to stop. It does not wait for in-flight handlers, which instead
// observe their canceled context.
func (b *busImpl) Close() error {
	b.closeOnce.Do(func() {
		b.mu.Lock()
		b.closed.Store(true)
		b.cancel()
		for _, sub := range b.subs {
			sub.stop()
		}
		b.subs = make(map[uint64]*subscription)
		b.mu.Unlock()
	})
	return nil
}

// startWorker launches the single goroutine that drains the subscription queue
// and runs the handler for each event.
func (b *busImpl) startWorker(sub *subscription) {
	go func() {
		for {
			select {
			case <-sub.done:
				return
			case ev := <-sub.events:
				b.deliver(sub, ev)
			}
		}
	}()
}

// deliver runs one handler invocation. It recovers panics so a misbehaving
// subscriber cannot crash the bus, and it ignores the handler's returned error
// (see the package comment).
func (b *busImpl) deliver(sub *subscription, ev Event) {
	defer func() {
		_ = recover()
	}()
	_ = sub.handler(sub.ctx, ev)
}

// removeSub unsubscribes a registration exactly once. The events channel is
// never closed by a sender, so a concurrent Publish cannot panic on a send;
// this method only closes the done channel after the registration is gone.
func (b *busImpl) removeSub(id uint64, sub *subscription) {
	b.mu.Lock()
	if _, ok := b.subs[id]; ok {
		delete(b.subs, id)
		b.mu.Unlock()
		sub.stop()
		return
	}
	b.mu.Unlock()
}