2021-02-12 16:33:16 +03:00
|
|
|
package broker
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"sync"
|
|
|
|
|
|
|
|
"github.com/google/uuid"
|
|
|
|
"github.com/unistack-org/micro/v3/logger"
|
2021-07-23 12:03:18 +03:00
|
|
|
"github.com/unistack-org/micro/v3/metadata"
|
2021-02-12 16:33:16 +03:00
|
|
|
maddr "github.com/unistack-org/micro/v3/util/addr"
|
|
|
|
mnet "github.com/unistack-org/micro/v3/util/net"
|
2021-02-17 22:58:55 +03:00
|
|
|
"github.com/unistack-org/micro/v3/util/rand"
|
2021-02-12 16:33:16 +03:00
|
|
|
)
|
|
|
|
|
|
|
|
type memoryBroker struct {
|
2021-07-22 22:53:44 +03:00
|
|
|
subscribers map[string][]*memorySubscriber
|
|
|
|
batchsubscribers map[string][]*memoryBatchSubscriber
|
|
|
|
addr string
|
|
|
|
opts Options
|
2021-03-06 19:45:13 +03:00
|
|
|
sync.RWMutex
|
|
|
|
connected bool
|
2021-02-12 16:33:16 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
type memoryEvent struct {
|
|
|
|
err error
|
|
|
|
message interface{}
|
2021-03-06 19:45:13 +03:00
|
|
|
topic string
|
2021-04-25 11:58:24 +03:00
|
|
|
opts Options
|
2021-02-12 16:33:16 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
type memorySubscriber struct {
|
|
|
|
ctx context.Context
|
2021-03-06 19:45:13 +03:00
|
|
|
exit chan bool
|
|
|
|
handler Handler
|
|
|
|
id string
|
|
|
|
topic string
|
2021-04-25 11:58:24 +03:00
|
|
|
opts SubscribeOptions
|
2021-02-12 16:33:16 +03:00
|
|
|
}
|
|
|
|
|
2021-07-22 22:53:44 +03:00
|
|
|
type memoryBatchSubscriber struct {
|
|
|
|
ctx context.Context
|
|
|
|
exit chan bool
|
|
|
|
handler BatchHandler
|
|
|
|
id string
|
|
|
|
topic string
|
|
|
|
opts SubscribeOptions
|
|
|
|
}
|
|
|
|
|
2021-02-12 16:33:16 +03:00
|
|
|
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
|
|
|
|
}
|
2021-02-17 22:58:55 +03:00
|
|
|
var rng rand.Rand
|
|
|
|
i := rng.Intn(20000)
|
2021-02-12 16:33:16 +03:00
|
|
|
// 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()
|
|
|
|
|
|
|
|
if !m.connected {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
m.connected = false
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryBroker) Init(opts ...Option) error {
|
|
|
|
for _, o := range opts {
|
|
|
|
o(&m.opts)
|
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-07-22 22:53:44 +03:00
|
|
|
func (m *memoryBroker) BatchPublish(ctx context.Context, msgs []*Message, opts ...PublishOption) error {
|
|
|
|
m.RLock()
|
|
|
|
if !m.connected {
|
|
|
|
m.RUnlock()
|
|
|
|
return ErrNotConnected
|
|
|
|
}
|
|
|
|
m.RUnlock()
|
|
|
|
|
|
|
|
type msgWrapper struct {
|
|
|
|
topic string
|
|
|
|
body interface{}
|
|
|
|
}
|
|
|
|
|
|
|
|
vs := make([]msgWrapper, 0, len(msgs))
|
|
|
|
if m.opts.Codec == nil {
|
|
|
|
m.RLock()
|
|
|
|
for _, msg := range msgs {
|
2021-07-23 12:03:18 +03:00
|
|
|
topic, _ := msg.Header.Get(metadata.HeaderTopic)
|
2021-07-22 22:53:44 +03:00
|
|
|
vs = append(vs, msgWrapper{topic: topic, body: m})
|
|
|
|
}
|
|
|
|
m.RUnlock()
|
|
|
|
} else {
|
|
|
|
m.RLock()
|
|
|
|
for _, msg := range msgs {
|
2021-07-23 12:03:18 +03:00
|
|
|
topic, _ := msg.Header.Get(metadata.HeaderTopic)
|
2021-07-22 22:53:44 +03:00
|
|
|
buf, err := m.opts.Codec.Marshal(msg)
|
|
|
|
if err != nil {
|
|
|
|
m.RUnlock()
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
vs = append(vs, msgWrapper{topic: topic, body: buf})
|
|
|
|
}
|
|
|
|
m.RUnlock()
|
|
|
|
}
|
|
|
|
|
|
|
|
if len(m.batchsubscribers) > 0 {
|
|
|
|
eh := m.opts.BatchErrorHandler
|
|
|
|
|
|
|
|
msgTopicMap := make(map[string]Events)
|
|
|
|
for _, v := range vs {
|
|
|
|
p := &memoryEvent{
|
|
|
|
topic: v.topic,
|
|
|
|
message: v.body,
|
|
|
|
opts: m.opts,
|
|
|
|
}
|
|
|
|
msgTopicMap[p.topic] = append(msgTopicMap[p.topic], p)
|
|
|
|
}
|
|
|
|
|
|
|
|
for t, ms := range msgTopicMap {
|
|
|
|
m.RLock()
|
|
|
|
subs, ok := m.batchsubscribers[t]
|
|
|
|
m.RUnlock()
|
|
|
|
if !ok {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
for _, sub := range subs {
|
|
|
|
if err := sub.handler(ms); err != nil {
|
|
|
|
ms.SetError(err)
|
|
|
|
if sub.opts.BatchErrorHandler != nil {
|
|
|
|
eh = sub.opts.BatchErrorHandler
|
|
|
|
}
|
|
|
|
if eh != nil {
|
|
|
|
eh(ms)
|
|
|
|
} else if m.opts.Logger.V(logger.ErrorLevel) {
|
|
|
|
m.opts.Logger.Error(m.opts.Context, err.Error())
|
|
|
|
}
|
|
|
|
} else if sub.opts.AutoAck {
|
|
|
|
if err := ms.Ack(); err != nil {
|
|
|
|
m.opts.Logger.Errorf(m.opts.Context, "ack failed: %v", err)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
eh := m.opts.ErrorHandler
|
|
|
|
|
|
|
|
for _, v := range vs {
|
|
|
|
p := &memoryEvent{
|
|
|
|
topic: v.topic,
|
|
|
|
message: v.body,
|
|
|
|
opts: m.opts,
|
|
|
|
}
|
|
|
|
|
|
|
|
m.RLock()
|
|
|
|
subs, ok := m.subscribers[p.topic]
|
|
|
|
m.RUnlock()
|
|
|
|
if !ok {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
for _, sub := range subs {
|
|
|
|
if err := sub.handler(p); err != nil {
|
|
|
|
p.SetError(err)
|
|
|
|
if sub.opts.ErrorHandler != nil {
|
|
|
|
eh = sub.opts.ErrorHandler
|
|
|
|
}
|
|
|
|
if eh != nil {
|
|
|
|
eh(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.Errorf(m.opts.Context, "ack failed: %v", err)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-02-12 16:33:16 +03:00
|
|
|
func (m *memoryBroker) Publish(ctx context.Context, topic string, msg *Message, opts ...PublishOption) error {
|
|
|
|
m.RLock()
|
|
|
|
if !m.connected {
|
|
|
|
m.RUnlock()
|
2021-07-22 22:53:44 +03:00
|
|
|
return ErrNotConnected
|
2021-02-12 16:33:16 +03:00
|
|
|
}
|
|
|
|
|
2021-07-22 22:53:44 +03:00
|
|
|
subs, ok := m.subscribers[topic]
|
2021-02-12 16:33:16 +03:00
|
|
|
m.RUnlock()
|
|
|
|
if !ok {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
var v interface{}
|
|
|
|
if m.opts.Codec != nil {
|
|
|
|
buf, err := m.opts.Codec.Marshal(msg)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
v = buf
|
|
|
|
} else {
|
|
|
|
v = msg
|
|
|
|
}
|
|
|
|
|
|
|
|
p := &memoryEvent{
|
|
|
|
topic: topic,
|
|
|
|
message: v,
|
|
|
|
opts: m.opts,
|
|
|
|
}
|
|
|
|
|
|
|
|
eh := m.opts.ErrorHandler
|
|
|
|
|
|
|
|
for _, sub := range subs {
|
|
|
|
if err := sub.handler(p); err != nil {
|
|
|
|
p.err = err
|
|
|
|
if sub.opts.ErrorHandler != nil {
|
|
|
|
eh = sub.opts.ErrorHandler
|
|
|
|
}
|
|
|
|
if eh != nil {
|
|
|
|
eh(p)
|
2021-02-13 01:46:16 +03:00
|
|
|
} else if m.opts.Logger.V(logger.ErrorLevel) {
|
|
|
|
m.opts.Logger.Error(m.opts.Context, err.Error())
|
2021-02-12 16:33:16 +03:00
|
|
|
}
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-07-22 22:53:44 +03:00
|
|
|
func (m *memoryBroker) BatchSubscribe(ctx context.Context, topic string, handler BatchHandler, opts ...SubscribeOption) (Subscriber, error) {
|
|
|
|
m.RLock()
|
|
|
|
if !m.connected {
|
|
|
|
m.RUnlock()
|
|
|
|
return nil, ErrNotConnected
|
|
|
|
}
|
|
|
|
m.RUnlock()
|
|
|
|
|
|
|
|
options := NewSubscribeOptions(opts...)
|
|
|
|
|
|
|
|
id, err := uuid.NewRandom()
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
|
|
|
sub := &memoryBatchSubscriber{
|
|
|
|
exit: make(chan bool, 1),
|
|
|
|
id: id.String(),
|
|
|
|
topic: topic,
|
|
|
|
handler: handler,
|
|
|
|
opts: options,
|
|
|
|
ctx: ctx,
|
|
|
|
}
|
|
|
|
|
|
|
|
m.Lock()
|
|
|
|
m.batchsubscribers[topic] = append(m.batchsubscribers[topic], sub)
|
|
|
|
m.Unlock()
|
|
|
|
|
|
|
|
go func() {
|
|
|
|
<-sub.exit
|
|
|
|
m.Lock()
|
|
|
|
var newSubscribers []*memoryBatchSubscriber
|
|
|
|
for _, sb := range m.batchsubscribers[topic] {
|
|
|
|
if sb.id == sub.id {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
newSubscribers = append(newSubscribers, sb)
|
|
|
|
}
|
|
|
|
m.batchsubscribers[topic] = newSubscribers
|
|
|
|
m.Unlock()
|
|
|
|
}()
|
|
|
|
|
|
|
|
return sub, nil
|
|
|
|
}
|
|
|
|
|
2021-02-12 16:33:16 +03:00
|
|
|
func (m *memoryBroker) Subscribe(ctx context.Context, topic string, handler Handler, opts ...SubscribeOption) (Subscriber, error) {
|
|
|
|
m.RLock()
|
|
|
|
if !m.connected {
|
|
|
|
m.RUnlock()
|
2021-07-22 22:53:44 +03:00
|
|
|
return nil, ErrNotConnected
|
2021-02-12 16:33:16 +03:00
|
|
|
}
|
|
|
|
m.RUnlock()
|
|
|
|
|
|
|
|
options := NewSubscribeOptions(opts...)
|
|
|
|
|
|
|
|
id, err := uuid.NewRandom()
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
|
|
|
sub := &memorySubscriber{
|
|
|
|
exit: make(chan bool, 1),
|
|
|
|
id: id.String(),
|
|
|
|
topic: topic,
|
|
|
|
handler: handler,
|
|
|
|
opts: options,
|
|
|
|
ctx: ctx,
|
|
|
|
}
|
|
|
|
|
|
|
|
m.Lock()
|
2021-07-22 22:53:44 +03:00
|
|
|
m.subscribers[topic] = append(m.subscribers[topic], sub)
|
2021-02-12 16:33:16 +03:00
|
|
|
m.Unlock()
|
|
|
|
|
|
|
|
go func() {
|
|
|
|
<-sub.exit
|
|
|
|
m.Lock()
|
|
|
|
var newSubscribers []*memorySubscriber
|
2021-07-22 22:53:44 +03:00
|
|
|
for _, sb := range m.subscribers[topic] {
|
2021-02-12 16:33:16 +03:00
|
|
|
if sb.id == sub.id {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
newSubscribers = append(newSubscribers, sb)
|
|
|
|
}
|
2021-07-22 22:53:44 +03:00
|
|
|
m.subscribers[topic] = newSubscribers
|
2021-02-12 16:33:16 +03:00
|
|
|
m.Unlock()
|
|
|
|
}()
|
|
|
|
|
|
|
|
return sub, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryBroker) String() string {
|
|
|
|
return "memory"
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryBroker) Name() string {
|
|
|
|
return m.opts.Name
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryEvent) Topic() string {
|
|
|
|
return m.topic
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryEvent) Message() *Message {
|
|
|
|
switch v := m.message.(type) {
|
|
|
|
case *Message:
|
|
|
|
return v
|
|
|
|
case []byte:
|
|
|
|
msg := &Message{}
|
|
|
|
if err := m.opts.Codec.Unmarshal(v, msg); err != nil {
|
|
|
|
if m.opts.Logger.V(logger.ErrorLevel) {
|
|
|
|
m.opts.Logger.Error(m.opts.Context, "[memory]: failed to unmarshal: %v", err)
|
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
return msg
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryEvent) Ack() error {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryEvent) Error() error {
|
|
|
|
return m.err
|
|
|
|
}
|
|
|
|
|
2021-07-22 22:53:44 +03:00
|
|
|
func (m *memoryEvent) SetError(err error) {
|
|
|
|
m.err = err
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryBatchSubscriber) Options() SubscribeOptions {
|
|
|
|
return m.opts
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryBatchSubscriber) Topic() string {
|
|
|
|
return m.topic
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *memoryBatchSubscriber) Unsubscribe(ctx context.Context) error {
|
|
|
|
m.exit <- true
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-02-12 16:33:16 +03:00
|
|
|
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
|
|
|
|
}
|
|
|
|
|
2021-02-13 15:35:56 +03:00
|
|
|
// NewBroker return new memory broker
|
2021-07-23 12:55:36 +03:00
|
|
|
func NewBroker(opts ...Option) BatchBroker {
|
2021-02-12 16:33:16 +03:00
|
|
|
return &memoryBroker{
|
2021-07-22 22:53:44 +03:00
|
|
|
opts: NewOptions(opts...),
|
|
|
|
subscribers: make(map[string][]*memorySubscriber),
|
|
|
|
batchsubscribers: make(map[string][]*memoryBatchSubscriber),
|
2021-02-12 16:33:16 +03:00
|
|
|
}
|
|
|
|
}
|