micro/broker/memory/memory.go

228 lines
3.8 KiB
Go
Raw Normal View History

2019-01-14 18:27:25 +03:00
// Package memory provides a memory broker
package memory
2016-03-18 01:46:20 +03:00
import (
"context"
2016-03-18 01:46:20 +03:00
"errors"
"math/rand"
2016-03-18 01:46:20 +03:00
"sync"
"time"
2016-03-18 01:46:20 +03:00
"github.com/google/uuid"
"github.com/micro/go-micro/v2/broker"
log "github.com/micro/go-micro/v2/logger"
maddr "github.com/micro/go-micro/v2/util/addr"
mnet "github.com/micro/go-micro/v2/util/net"
2016-03-18 01:46:20 +03:00
)
2019-01-14 18:27:25 +03:00
type memoryBroker struct {
2016-03-18 01:46:20 +03:00
opts broker.Options
addr string
2016-03-18 01:46:20 +03:00
sync.RWMutex
connected bool
2019-01-14 18:27:25 +03:00
Subscribers map[string][]*memorySubscriber
2016-03-18 01:46:20 +03:00
}
2019-07-07 14:44:09 +03:00
type memoryEvent struct {
opts broker.Options
2016-03-18 01:46:20 +03:00
topic string
message interface{}
2016-03-18 01:46:20 +03:00
}
2019-01-14 18:27:25 +03:00
type memorySubscriber struct {
2016-03-18 01:46:20 +03:00
id string
topic string
exit chan bool
handler broker.Handler
opts broker.SubscribeOptions
}
2019-01-14 18:27:25 +03:00
func (m *memoryBroker) Options() broker.Options {
2016-03-18 01:46:20 +03:00
return m.opts
}
2019-01-14 18:27:25 +03:00
func (m *memoryBroker) Address() string {
return m.addr
2016-03-18 01:46:20 +03:00
}
2019-01-14 18:27:25 +03:00
func (m *memoryBroker) Connect() error {
2016-03-18 01:46:20 +03:00
m.Lock()
defer m.Unlock()
if m.connected {
return nil
}
addr, err := maddr.Extract("::")
if err != nil {
return err
}
i := rand.Intn(20000)
// set addr with port
addr = mnet.HostPort(addr, 10000+i)
m.addr = addr
2016-03-18 01:46:20 +03:00
m.connected = true
return nil
}
2019-01-14 18:27:25 +03:00
func (m *memoryBroker) Disconnect() error {
2016-03-18 01:46:20 +03:00
m.Lock()
defer m.Unlock()
if !m.connected {
return nil
}
m.connected = false
return nil
}
2019-01-14 18:27:25 +03:00
func (m *memoryBroker) Init(opts ...broker.Option) error {
2016-03-18 01:46:20 +03:00
for _, o := range opts {
o(&m.opts)
}
return nil
}
func (m *memoryBroker) Publish(topic string, msg *broker.Message, opts ...broker.PublishOption) error {
m.RLock()
2016-03-18 01:46:20 +03:00
if !m.connected {
m.RUnlock()
2016-03-18 01:46:20 +03:00
return errors.New("not connected")
}
subs, ok := m.Subscribers[topic]
m.RUnlock()
2016-03-18 01:46:20 +03:00
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
}
2019-07-07 14:44:09 +03:00
p := &memoryEvent{
2016-03-18 01:46:20 +03:00
topic: topic,
message: v,
opts: m.opts,
2016-03-18 01:46:20 +03:00
}
for _, sub := range subs {
if err := sub.handler(p); err != nil {
return err
}
}
return nil
}
2019-01-14 18:27:25 +03:00
func (m *memoryBroker) Subscribe(topic string, handler broker.Handler, opts ...broker.SubscribeOption) (broker.Subscriber, error) {
m.RLock()
2016-03-18 01:46:20 +03:00
if !m.connected {
m.RUnlock()
2016-03-18 01:46:20 +03:00
return nil, errors.New("not connected")
}
m.RUnlock()
2016-03-18 01:46:20 +03:00
var options broker.SubscribeOptions
for _, o := range opts {
o(&options)
}
2019-01-14 18:27:25 +03:00
sub := &memorySubscriber{
2016-03-18 01:46:20 +03:00
exit: make(chan bool, 1),
id: uuid.New().String(),
2016-03-18 01:46:20 +03:00
topic: topic,
handler: handler,
opts: options,
}
m.Lock()
2016-03-18 01:46:20 +03:00
m.Subscribers[topic] = append(m.Subscribers[topic], sub)
m.Unlock()
2016-03-18 01:46:20 +03:00
go func() {
<-sub.exit
m.Lock()
2019-01-14 18:27:25 +03:00
var newSubscribers []*memorySubscriber
2016-03-18 01:46:20 +03:00
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
}
2019-01-14 18:27:25 +03:00
func (m *memoryBroker) String() string {
return "memory"
2016-03-18 01:46:20 +03:00
}
2019-07-07 14:44:09 +03:00
func (m *memoryEvent) Topic() string {
2016-03-18 01:46:20 +03:00
return m.topic
}
2019-07-07 14:44:09 +03:00
func (m *memoryEvent) Message() *broker.Message {
switch v := m.message.(type) {
case *broker.Message:
return v
case []byte:
msg := &broker.Message{}
if err := m.opts.Codec.Unmarshal(v, msg); err != nil {
log.Errorf("[memory]: failed to unmarshal: %v\n", err)
return nil
}
return msg
}
return nil
2016-03-18 01:46:20 +03:00
}
2019-07-07 14:44:09 +03:00
func (m *memoryEvent) Ack() error {
2016-03-18 01:46:20 +03:00
return nil
}
2019-01-14 18:27:25 +03:00
func (m *memorySubscriber) Options() broker.SubscribeOptions {
2016-03-18 01:46:20 +03:00
return m.opts
}
2019-01-14 18:27:25 +03:00
func (m *memorySubscriber) Topic() string {
2016-03-18 01:46:20 +03:00
return m.topic
}
2019-01-14 18:27:25 +03:00
func (m *memorySubscriber) Unsubscribe() error {
2016-03-18 01:46:20 +03:00
m.exit <- true
return nil
}
func NewBroker(opts ...broker.Option) broker.Broker {
options := broker.Options{
Context: context.Background(),
}
rand.Seed(time.Now().UnixNano())
2016-03-18 01:46:20 +03:00
for _, o := range opts {
o(&options)
}
2019-01-14 18:27:25 +03:00
return &memoryBroker{
2016-03-18 01:46:20 +03:00
opts: options,
2019-01-14 18:27:25 +03:00
Subscribers: make(map[string][]*memorySubscriber),
2016-03-18 01:46:20 +03:00
}
}