335 lines
7.6 KiB
Go
335 lines
7.6 KiB
Go
package broker
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"go.unistack.org/micro/v4/codec"
|
|
"go.unistack.org/micro/v4/logger"
|
|
"go.unistack.org/micro/v4/metadata"
|
|
"go.unistack.org/micro/v4/options"
|
|
"go.unistack.org/micro/v4/semconv"
|
|
maddr "go.unistack.org/micro/v4/util/addr"
|
|
"go.unistack.org/micro/v4/util/id"
|
|
mnet "go.unistack.org/micro/v4/util/net"
|
|
"go.unistack.org/micro/v4/util/rand"
|
|
)
|
|
|
|
type MemoryBroker struct {
|
|
subscribers map[string][]*memorySubscriber
|
|
addr string
|
|
opts Options
|
|
sync.RWMutex
|
|
connected bool
|
|
}
|
|
|
|
func (m *MemoryBroker) Options() Options {
|
|
return m.opts
|
|
}
|
|
|
|
func (m *MemoryBroker) Address() string {
|
|
return m.addr
|
|
}
|
|
|
|
func (m *MemoryBroker) Connect(ctx context.Context) error {
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
|
|
if m.connected {
|
|
return nil
|
|
}
|
|
|
|
// use 127.0.0.1 to avoid scan of all network interfaces
|
|
addr, err := maddr.Extract("127.0.0.1")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var rng rand.Rand
|
|
i := rng.Intn(20000)
|
|
// set addr with port
|
|
addr = mnet.HostPort(addr, 10000+i)
|
|
|
|
m.addr = addr
|
|
m.connected = true
|
|
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryBroker) Disconnect(ctx context.Context) error {
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
default:
|
|
if m.connected {
|
|
m.connected = false
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryBroker) Init(opts ...options.Option) error {
|
|
var err error
|
|
for _, o := range opts {
|
|
if err = o(&m.opts); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryBroker) Publish(ctx context.Context, message interface{}, opts ...options.Option) error {
|
|
m.RLock()
|
|
if !m.connected {
|
|
m.RUnlock()
|
|
return ErrNotConnected
|
|
}
|
|
m.RUnlock()
|
|
|
|
var err error
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
default:
|
|
options := NewPublishOptions(opts...)
|
|
var msgs []Message
|
|
switch v := message.(type) {
|
|
case []Message:
|
|
msgs = v
|
|
case Message:
|
|
msgs = append(msgs, v)
|
|
default:
|
|
return ErrInvalidMessage
|
|
}
|
|
msgTopicMap := make(map[string][]*memoryMessage)
|
|
for _, msg := range msgs {
|
|
p := &memoryMessage{opts: options}
|
|
p.topic, _ = msg.Header().Get(metadata.HeaderTopic)
|
|
if v, ok := msg.Body().(*codec.Frame); ok {
|
|
p.body = msg.Body()
|
|
} else if len(m.opts.Codecs) == 0 {
|
|
p.body = msg.Body()
|
|
} else {
|
|
cf, ok := m.opts.Codecs[options.ContentType]
|
|
if !ok {
|
|
return fmt.Errorf("%s: %s", codec.ErrUnknownContentType, options.ContentType)
|
|
}
|
|
p.body, err = cf.Marshal(v)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
msgTopicMap[p.topic] = append(msgTopicMap[p.topic], p)
|
|
}
|
|
|
|
eh := m.opts.ErrorHandler
|
|
|
|
for t, ms := range msgTopicMap {
|
|
ts := time.Now()
|
|
|
|
m.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", t).Add(len(ms))
|
|
m.opts.Meter.Counter(semconv.SubscribeMessageInflight, "endpoint", t).Add(len(ms))
|
|
|
|
m.RLock()
|
|
subs, ok := m.subscribers[t]
|
|
m.RUnlock()
|
|
if !ok {
|
|
m.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", t, "status", "failure").Add(len(ms))
|
|
m.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", t).Add(-len(ms))
|
|
m.opts.Meter.Counter(semconv.SubscribeMessageInflight, "endpoint", t).Add(-len(ms))
|
|
continue
|
|
}
|
|
|
|
m.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", t, "status", "success").Add(len(ms))
|
|
for _, sub := range subs {
|
|
if sub.opts.ErrorHandler != nil {
|
|
eh = sub.opts.ErrorHandler
|
|
}
|
|
|
|
switch mh := sub.handler.(type) {
|
|
case MessagesHandler:
|
|
mhs := make([]Message, 0, len(ms))
|
|
for _, m := range ms {
|
|
mhs = append(mhs, m)
|
|
}
|
|
if err = mh(mhs); err != nil {
|
|
m.opts.Meter.Counter(semconv.SubscribeMessageTotal, "endpoint", t, "status", "failure").Add(len(ms))
|
|
if eh != nil {
|
|
switch meh := eh.(type) {
|
|
case MessagesHandler:
|
|
_ = meh(mhs)
|
|
case MessageHandler:
|
|
for _, me := range mhs {
|
|
_ = meh(me)
|
|
}
|
|
}
|
|
} else if m.opts.Logger.V(logger.ErrorLevel) {
|
|
m.opts.Logger.Error(m.opts.Context, err.Error())
|
|
}
|
|
}
|
|
case MessageHandler:
|
|
for _, p := range ms {
|
|
if err = mh(p); err != nil {
|
|
m.opts.Meter.Counter(semconv.SubscribeMessageTotal, "endpoint", t, "status", "failure").Inc()
|
|
if eh != nil {
|
|
switch meh := eh.(type) {
|
|
case MessageHandler:
|
|
_ = meh(p)
|
|
case MessagesHandler:
|
|
_ = meh([]Message{p})
|
|
}
|
|
} else if m.opts.Logger.V(logger.ErrorLevel) {
|
|
m.opts.Logger.Error(m.opts.Context, err.Error())
|
|
}
|
|
} else {
|
|
if sub.opts.AutoAck {
|
|
if err = p.Ack(); err != nil {
|
|
m.opts.Logger.Error(m.opts.Context, "ack failed: "+err.Error())
|
|
m.opts.Meter.Counter(semconv.SubscribeMessageTotal, "endpoint", t, "status", "failure").Inc()
|
|
} else {
|
|
m.opts.Meter.Counter(semconv.SubscribeMessageTotal, "endpoint", t, "status", "success").Inc()
|
|
}
|
|
} else {
|
|
m.opts.Meter.Counter(semconv.SubscribeMessageTotal, "endpoint", t, "status", "success").Inc()
|
|
}
|
|
}
|
|
m.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", t).Add(-1)
|
|
m.opts.Meter.Counter(semconv.SubscribeMessageInflight, "endpoint", t).Add(-1)
|
|
}
|
|
}
|
|
}
|
|
|
|
te := time.Since(ts)
|
|
m.opts.Meter.Summary(semconv.PublishMessageLatencyMicroseconds, "endpoint", t).Update(te.Seconds())
|
|
m.opts.Meter.Histogram(semconv.PublishMessageDurationSeconds, "endpoint", t).Update(te.Seconds())
|
|
m.opts.Meter.Summary(semconv.SubscribeMessageLatencyMicroseconds, "endpoint", t).Update(te.Seconds())
|
|
m.opts.Meter.Histogram(semconv.SubscribeMessageDurationSeconds, "endpoint", t).Update(te.Seconds())
|
|
}
|
|
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryBroker) Subscribe(ctx context.Context, topic string, handler interface{}, opts ...options.Option) (Subscriber, error) {
|
|
m.RLock()
|
|
if !m.connected {
|
|
m.RUnlock()
|
|
return nil, ErrNotConnected
|
|
}
|
|
m.RUnlock()
|
|
|
|
sid, err := id.New()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
options := NewSubscribeOptions(opts...)
|
|
|
|
sub := &memorySubscriber{
|
|
exit: make(chan bool, 1),
|
|
id: sid,
|
|
topic: topic,
|
|
handler: handler,
|
|
opts: options,
|
|
ctx: ctx,
|
|
}
|
|
|
|
m.Lock()
|
|
m.subscribers[topic] = append(m.subscribers[topic], sub)
|
|
m.Unlock()
|
|
|
|
go func() {
|
|
<-sub.exit
|
|
m.Lock()
|
|
newSubscribers := make([]*memorySubscriber, 0, len(m.subscribers)-1)
|
|
for _, sb := range m.subscribers[topic] {
|
|
if sb.id == sub.id {
|
|
continue
|
|
}
|
|
newSubscribers = append(newSubscribers, sb)
|
|
}
|
|
m.subscribers[topic] = newSubscribers
|
|
m.Unlock()
|
|
}()
|
|
|
|
return sub, nil
|
|
}
|
|
|
|
func (m *MemoryBroker) String() string {
|
|
return "memory"
|
|
}
|
|
|
|
func (m *MemoryBroker) Name() string {
|
|
return m.opts.Name
|
|
}
|
|
|
|
type memoryMessage struct {
|
|
err error
|
|
body interface{}
|
|
topic string
|
|
header metadata.Metadata
|
|
opts PublishOptions
|
|
ctx context.Context
|
|
}
|
|
|
|
func (m *memoryMessage) Topic() string {
|
|
return m.topic
|
|
}
|
|
|
|
func (m *memoryMessage) Header() metadata.Metadata {
|
|
return m.header
|
|
}
|
|
|
|
func (m *memoryMessage) Body() interface{} {
|
|
return m.body
|
|
}
|
|
|
|
func (m *memoryMessage) Ack() error {
|
|
return nil
|
|
}
|
|
|
|
func (m *memoryMessage) Error() error {
|
|
return m.err
|
|
}
|
|
|
|
func (m *memoryMessage) Context() context.Context {
|
|
return m.ctx
|
|
}
|
|
|
|
type memorySubscriber struct {
|
|
ctx context.Context
|
|
exit chan bool
|
|
handler interface{}
|
|
id string
|
|
topic string
|
|
opts SubscribeOptions
|
|
}
|
|
|
|
func (m *memorySubscriber) Options() SubscribeOptions {
|
|
return m.opts
|
|
}
|
|
|
|
func (m *memorySubscriber) Topic() string {
|
|
return m.topic
|
|
}
|
|
|
|
func (m *memorySubscriber) Unsubscribe(ctx context.Context) error {
|
|
m.exit <- true
|
|
return nil
|
|
}
|
|
|
|
// NewBroker return new memory broker
|
|
func NewBroker(opts ...options.Option) *MemoryBroker {
|
|
return &MemoryBroker{
|
|
opts: NewOptions(opts...),
|
|
subscribers: make(map[string][]*memorySubscriber),
|
|
}
|
|
}
|