chore: import clean Agent Desk snapshot (2/7)
This commit is contained in:
191
internal/pkg/eventbus/bus.go
Normal file
191
internal/pkg/eventbus/bus.go
Normal file
@@ -0,0 +1,191 @@
|
||||
package eventbus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"runtime/debug"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
type Handler[T any] func(ctx context.Context, event T) error
|
||||
|
||||
type ErrorHandler func(ctx context.Context, err error)
|
||||
|
||||
type handlerEntry[T any] struct {
|
||||
id uint64
|
||||
handler Handler[T]
|
||||
}
|
||||
|
||||
type Bus[T any] struct {
|
||||
mu sync.RWMutex
|
||||
handlers []handlerEntry[T]
|
||||
nextID uint64
|
||||
|
||||
onError ErrorHandler
|
||||
|
||||
// 用于限制 PublishAsync 的 goroutine 并发数;nil 表示不限制
|
||||
asyncSem chan struct{}
|
||||
}
|
||||
|
||||
func New[T any](opts ...Option[T]) *Bus[T] {
|
||||
b := &Bus[T]{}
|
||||
for _, opt := range opts {
|
||||
opt(b)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
type Option[T any] func(*Bus[T])
|
||||
|
||||
func WithErrorHandler[T any](fn ErrorHandler) Option[T] {
|
||||
return func(b *Bus[T]) {
|
||||
b.onError = fn
|
||||
}
|
||||
}
|
||||
|
||||
func WithAsyncConcurrency[T any](n int) Option[T] {
|
||||
return func(b *Bus[T]) {
|
||||
if n > 0 {
|
||||
b.asyncSem = make(chan struct{}, n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Subscribe 返回 handlerID 和取消订阅函数
|
||||
func (b *Bus[T]) Subscribe(h Handler[T]) (uint64, func()) {
|
||||
if h == nil {
|
||||
panic("eventbus: nil handler")
|
||||
}
|
||||
|
||||
id := atomic.AddUint64(&b.nextID, 1)
|
||||
|
||||
b.mu.Lock()
|
||||
b.handlers = append(b.handlers, handlerEntry[T]{
|
||||
id: id,
|
||||
handler: h,
|
||||
})
|
||||
b.mu.Unlock()
|
||||
|
||||
return id, func() {
|
||||
b.Unsubscribe(id)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bus[T]) SubscribeOnce(h Handler[T]) (uint64, func()) {
|
||||
if h == nil {
|
||||
panic("eventbus: nil handler")
|
||||
}
|
||||
|
||||
var id uint64
|
||||
var called atomic.Bool
|
||||
|
||||
wrapper := func(ctx context.Context, event T) error {
|
||||
if !called.CompareAndSwap(false, true) {
|
||||
return nil
|
||||
}
|
||||
b.Unsubscribe(id)
|
||||
return h(ctx, event)
|
||||
}
|
||||
|
||||
id = atomic.AddUint64(&b.nextID, 1)
|
||||
|
||||
b.mu.Lock()
|
||||
b.handlers = append(b.handlers, handlerEntry[T]{
|
||||
id: id,
|
||||
handler: wrapper,
|
||||
})
|
||||
b.mu.Unlock()
|
||||
|
||||
return id, func() {
|
||||
b.Unsubscribe(id)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bus[T]) Unsubscribe(id uint64) {
|
||||
b.mu.Lock()
|
||||
for i, entry := range b.handlers {
|
||||
if entry.id != id {
|
||||
continue
|
||||
}
|
||||
b.handlers = append(b.handlers[:i], b.handlers[i+1:]...)
|
||||
break
|
||||
}
|
||||
b.mu.Unlock()
|
||||
}
|
||||
|
||||
func (b *Bus[T]) Publish(ctx context.Context, event T) error {
|
||||
handlers := b.snapshotHandlers()
|
||||
errs := make([]error, 0, len(handlers))
|
||||
for _, h := range handlers {
|
||||
if err := b.callHandler(ctx, h, event); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
func (b *Bus[T]) PublishAsync(ctx context.Context, event T) {
|
||||
handlers := b.snapshotHandlers()
|
||||
for _, h := range handlers {
|
||||
go b.callHandlerAsync(ctx, h, event)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bus[T]) HandlerCount() int {
|
||||
b.mu.RLock()
|
||||
defer b.mu.RUnlock()
|
||||
return len(b.handlers)
|
||||
}
|
||||
|
||||
func (b *Bus[T]) snapshotHandlers() []Handler[T] {
|
||||
b.mu.RLock()
|
||||
defer b.mu.RUnlock()
|
||||
|
||||
handlers := make([]Handler[T], 0, len(b.handlers))
|
||||
for _, entry := range b.handlers {
|
||||
handlers = append(handlers, entry.handler)
|
||||
}
|
||||
return handlers
|
||||
}
|
||||
|
||||
func (b *Bus[T]) callHandlerAsync(ctx context.Context, h Handler[T], event T) {
|
||||
if b.asyncSem != nil {
|
||||
select {
|
||||
case b.asyncSem <- struct{}{}:
|
||||
defer func() { <-b.asyncSem }()
|
||||
case <-ctx.Done():
|
||||
b.handleError(ctx, ctx.Err())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err := b.callHandler(ctx, h, event); err != nil {
|
||||
b.handleError(ctx, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bus[T]) callHandler(ctx context.Context, h Handler[T], event T) (err error) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
err = fmt.Errorf("event handler panic: %v\n%s", r, debug.Stack())
|
||||
}
|
||||
}()
|
||||
if err = h(ctx, event); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *Bus[T]) handleError(ctx context.Context, err error) {
|
||||
if err == nil {
|
||||
return
|
||||
}
|
||||
if b.onError != nil {
|
||||
b.onError(ctx, err)
|
||||
return
|
||||
}
|
||||
slog.Error("eventbus handler failed", "error", err)
|
||||
}
|
||||
199
internal/pkg/eventbus/bus_test.go
Normal file
199
internal/pkg/eventbus/bus_test.go
Normal file
@@ -0,0 +1,199 @@
|
||||
package eventbus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"reflect"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type testEvent struct {
|
||||
ID int64
|
||||
}
|
||||
|
||||
func TestPackageSubscribeAndPublish(t *testing.T) {
|
||||
resetManagerForTest(t)
|
||||
|
||||
var got atomic.Int64
|
||||
_, unsubscribe := Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
got.Store(event.ID)
|
||||
return nil
|
||||
})
|
||||
defer unsubscribe()
|
||||
|
||||
if err := Publish(context.Background(), testEvent{ID: 42}); err != nil {
|
||||
t.Fatalf("publish failed: %v", err)
|
||||
}
|
||||
|
||||
if got.Load() != 42 {
|
||||
t.Fatalf("expected event ID 42, got %d", got.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublishReturnsJoinedHandlerErrors(t *testing.T) {
|
||||
firstErr := errors.New("first")
|
||||
secondErr := errors.New("second")
|
||||
bus := New[testEvent]()
|
||||
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
return firstErr
|
||||
})
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
return secondErr
|
||||
})
|
||||
|
||||
err := bus.Publish(context.Background(), testEvent{})
|
||||
if !errors.Is(err, firstErr) {
|
||||
t.Fatalf("expected joined error to include first error, got %v", err)
|
||||
}
|
||||
if !errors.Is(err, secondErr) {
|
||||
t.Fatalf("expected joined error to include second error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublishCallsHandlersInSubscribeOrder(t *testing.T) {
|
||||
bus := New[testEvent]()
|
||||
calls := make([]int, 0, 3)
|
||||
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
calls = append(calls, 1)
|
||||
return nil
|
||||
})
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
calls = append(calls, 2)
|
||||
return nil
|
||||
})
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
calls = append(calls, 3)
|
||||
return nil
|
||||
})
|
||||
|
||||
if err := bus.Publish(context.Background(), testEvent{}); err != nil {
|
||||
t.Fatalf("publish failed: %v", err)
|
||||
}
|
||||
|
||||
if !reflect.DeepEqual(calls, []int{1, 2, 3}) {
|
||||
t.Fatalf("expected handlers to run in subscribe order, got %#v", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnsubscribeKeepsRemainingHandlerOrder(t *testing.T) {
|
||||
bus := New[testEvent]()
|
||||
calls := make([]int, 0, 2)
|
||||
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
calls = append(calls, 1)
|
||||
return nil
|
||||
})
|
||||
_, unsubscribe := bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
calls = append(calls, 2)
|
||||
return nil
|
||||
})
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
calls = append(calls, 3)
|
||||
return nil
|
||||
})
|
||||
|
||||
unsubscribe()
|
||||
|
||||
if err := bus.Publish(context.Background(), testEvent{}); err != nil {
|
||||
t.Fatalf("publish failed: %v", err)
|
||||
}
|
||||
|
||||
if !reflect.DeepEqual(calls, []int{1, 3}) {
|
||||
t.Fatalf("expected remaining handlers to keep order, got %#v", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPanicReturnsError(t *testing.T) {
|
||||
bus := New[testEvent]()
|
||||
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
panic("boom")
|
||||
})
|
||||
|
||||
err := bus.Publish(context.Background(), testEvent{})
|
||||
if err == nil {
|
||||
t.Fatalf("expected panic to be returned as error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublishAsyncCallsErrorHandler(t *testing.T) {
|
||||
handlerErr := errors.New("async failed")
|
||||
var handled atomic.Int64
|
||||
done := make(chan struct{})
|
||||
bus := New[testEvent](WithErrorHandler[testEvent](func(ctx context.Context, err error) {
|
||||
if errors.Is(err, handlerErr) {
|
||||
handled.Add(1)
|
||||
}
|
||||
close(done)
|
||||
}))
|
||||
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
return handlerErr
|
||||
})
|
||||
|
||||
bus.PublishAsync(context.Background(), testEvent{})
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatalf("expected async error handler to be called")
|
||||
}
|
||||
|
||||
if handled.Load() != 1 {
|
||||
t.Fatalf("expected error handler to be called once, got %d", handled.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubscribeOnceConcurrentPublishOnlyCallsOnce(t *testing.T) {
|
||||
bus := New[testEvent]()
|
||||
var calls atomic.Int64
|
||||
|
||||
bus.SubscribeOnce(func(ctx context.Context, event testEvent) error {
|
||||
calls.Add(1)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
return nil
|
||||
})
|
||||
|
||||
const workers = 32
|
||||
var wg sync.WaitGroup
|
||||
for range workers {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = bus.Publish(context.Background(), testEvent{})
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if calls.Load() != 1 {
|
||||
t.Fatalf("expected once handler to be called once, got %d", calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublishAsyncAllowsReentrantPublishWithConcurrencyLimit(t *testing.T) {
|
||||
bus := New[testEvent](WithAsyncConcurrency[testEvent](1))
|
||||
done := make(chan struct{})
|
||||
var calls atomic.Int64
|
||||
|
||||
bus.Subscribe(func(ctx context.Context, event testEvent) error {
|
||||
if calls.Add(1) == 1 {
|
||||
bus.PublishAsync(ctx, testEvent{ID: 2})
|
||||
return nil
|
||||
}
|
||||
close(done)
|
||||
return nil
|
||||
})
|
||||
|
||||
bus.PublishAsync(context.Background(), testEvent{ID: 1})
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatalf("expected reentrant async publish to complete")
|
||||
}
|
||||
}
|
||||
80
internal/pkg/eventbus/manager.go
Normal file
80
internal/pkg/eventbus/manager.go
Normal file
@@ -0,0 +1,80 @@
|
||||
package eventbus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var (
|
||||
mu sync.RWMutex
|
||||
buss = make(map[reflect.Type]any)
|
||||
)
|
||||
|
||||
// Register initializes the global bus for T. If the bus already exists, the
|
||||
// existing instance is returned and later options are ignored.
|
||||
func Register[T any](opts ...Option[T]) *Bus[T] {
|
||||
key := eventTypeOf[T]()
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
if bus, ok := buss[key]; ok {
|
||||
return bus.(*Bus[T])
|
||||
}
|
||||
|
||||
created := New(opts...)
|
||||
buss[key] = created
|
||||
return created
|
||||
}
|
||||
|
||||
// Get returns the global bus for T, creating it with default options if needed.
|
||||
func Get[T any]() *Bus[T] {
|
||||
key := eventTypeOf[T]()
|
||||
|
||||
mu.RLock()
|
||||
bus, ok := buss[key]
|
||||
mu.RUnlock()
|
||||
if ok {
|
||||
return bus.(*Bus[T])
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
if bus, ok = buss[key]; ok {
|
||||
return bus.(*Bus[T])
|
||||
}
|
||||
|
||||
created := New[T]()
|
||||
buss[key] = created
|
||||
return created
|
||||
}
|
||||
|
||||
func eventTypeOf[T any]() reflect.Type {
|
||||
typ := reflect.TypeFor[T]()
|
||||
if typ == nil {
|
||||
panic("eventbus: nil event type")
|
||||
}
|
||||
return typ
|
||||
}
|
||||
|
||||
// Subscribe registers a handler on the global bus for T.
|
||||
func Subscribe[T any](h Handler[T]) (uint64, func()) {
|
||||
return Get[T]().Subscribe(h)
|
||||
}
|
||||
|
||||
// SubscribeOnce registers a handler that runs at most once on the global bus for T.
|
||||
func SubscribeOnce[T any](h Handler[T]) (uint64, func()) {
|
||||
return Get[T]().SubscribeOnce(h)
|
||||
}
|
||||
|
||||
// Publish synchronously publishes event to the global bus for T.
|
||||
func Publish[T any](ctx context.Context, event T) error {
|
||||
return Get[T]().Publish(ctx, event)
|
||||
}
|
||||
|
||||
// PublishAsync asynchronously publishes event to the global bus for T.
|
||||
func PublishAsync[T any](ctx context.Context, event T) {
|
||||
Get[T]().PublishAsync(ctx, event)
|
||||
}
|
||||
102
internal/pkg/eventbus/manager_test.go
Normal file
102
internal/pkg/eventbus/manager_test.go
Normal file
@@ -0,0 +1,102 @@
|
||||
package eventbus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type userCreatedEvent struct {
|
||||
UserID int64
|
||||
}
|
||||
|
||||
type orderCreatedEvent struct {
|
||||
OrderID int64
|
||||
}
|
||||
|
||||
func TestGetReturnsSameBusForSameType(t *testing.T) {
|
||||
resetManagerForTest(t)
|
||||
|
||||
first := Get[userCreatedEvent]()
|
||||
second := Get[userCreatedEvent]()
|
||||
|
||||
if first != second {
|
||||
t.Fatalf("expected same bus instance for same event type")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetReturnsDifferentBusForDifferentTypes(t *testing.T) {
|
||||
resetManagerForTest(t)
|
||||
|
||||
userBus := Get[userCreatedEvent]()
|
||||
orderBus := Get[orderCreatedEvent]()
|
||||
|
||||
if userBus == any(orderBus) {
|
||||
t.Fatalf("expected different bus instances for different event types")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetCreatesOnlyOneBusUnderConcurrency(t *testing.T) {
|
||||
resetManagerForTest(t)
|
||||
|
||||
const workers = 32
|
||||
|
||||
results := make(chan *Bus[userCreatedEvent], workers)
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for range workers {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
results <- Get[userCreatedEvent]()
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
close(results)
|
||||
|
||||
var first *Bus[userCreatedEvent]
|
||||
for bus := range results {
|
||||
if first == nil {
|
||||
first = bus
|
||||
continue
|
||||
}
|
||||
if bus != first {
|
||||
t.Fatalf("expected same bus instance for concurrent access")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterReturnsConfiguredGlobalBus(t *testing.T) {
|
||||
resetManagerForTest(t)
|
||||
|
||||
var handled bool
|
||||
registered := Register[testEvent](WithErrorHandler[testEvent](func(ctx context.Context, err error) {
|
||||
handled = true
|
||||
}))
|
||||
got := Get[testEvent]()
|
||||
|
||||
if registered != got {
|
||||
t.Fatalf("expected registered bus to be returned by Get")
|
||||
}
|
||||
|
||||
got.handleError(context.Background(), assertErr{})
|
||||
if !handled {
|
||||
t.Fatalf("expected registered error handler to be used")
|
||||
}
|
||||
}
|
||||
|
||||
type assertErr struct{}
|
||||
|
||||
func (assertErr) Error() string {
|
||||
return "assert"
|
||||
}
|
||||
|
||||
func resetManagerForTest(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
mu.Lock()
|
||||
buss = make(map[reflect.Type]any)
|
||||
mu.Unlock()
|
||||
}
|
||||
Reference in New Issue
Block a user