192 lines
3.6 KiB
Go
192 lines
3.6 KiB
Go
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)
|
||
}
|