micro/broker/inprocess.go

package broker

import (
	"context"
	"fmt"
	"runtime"
	"sync"

	"git.nspix.com/golang/micro/log"
)

type InProcBus struct {
	ctx         context.Context
	once        sync.Once
	eventChan   chan *Event
	subscribers sync.Map
}

func (bus *InProcBus) worker() {
	for {
		select {
		case e, ok := <-bus.eventChan:
			if ok {
				_ = bus.Dispatch(e)
			}
		case <-bus.ctx.Done():
			return
		}
	}
}

//Subscribers 获取所有订阅者的快照
func (bus *InProcBus) Subscribers() []Subscriber {
	vs := make([]Subscriber, 0)
	bus.subscribers.Range(func(key, value interface{}) bool {
		vs = append(vs, value.(Subscriber))
		return true
	})
	return vs
}

// Publish 发布一个事件
func (bus *InProcBus) Publish(e *Event) {
	select {
	case bus.eventChan <- e:
	default:
	}
	return
}

// Dispatch 分配事件，关心返回值
func (bus *InProcBus) Dispatch(e *Event) (err error) {
	return bus.DispatchCtx(bus.ctx, e)
}

//DispatchCtx 分配一个事件
func (bus *InProcBus) DispatchCtx(ctx context.Context, e *Event) (err error) {
	bus.subscribers.Range(func(key, value interface{}) bool {
		sub := value.(Subscriber)
		if sub.HasTopic(e.Name) {
			if err = sub.Process(ctx, e); err != nil {
				log.Warnf("subscriber %s process event %s error: %s", sub.ID(), e.Name, err.Error())
			}
		}
		return true
	})
	releaseEvent(e)
	return
}

//Subscribe 订阅一个事件
func (bus *InProcBus) Subscribe(sub Subscriber) (err error) {
	if _, ok := bus.subscribers.Load(sub.ID()); ok {
		err = fmt.Errorf("subscriber %s already exists", sub.ID())
		return
	}
	if err = sub.OnAttach(); err != nil {
		return
	}
	bus.subscribers.Store(sub.ID(), sub)
	return
}

//UnSubscribe 取消一个订阅
func (bus *InProcBus) UnSubscribe(sub Subscriber) (err error) {
	bus.subscribers.Delete(sub.ID())
	return
}

func (bus *InProcBus) WithContext(ctx context.Context) {
	bus.ctx = ctx
}

func NewInPrcBus(ctx context.Context) *InProcBus {
	bus := &InProcBus{
		ctx:       ctx,
		eventChan: make(chan *Event, 1024),
	}
	bus.once.Do(func() {
		for i := 0; i < runtime.NumCPU(); i++ {
			go bus.worker()
		}
	})
	return bus
}