-
Notifications
You must be signed in to change notification settings - Fork 211
/
pubsub.go
148 lines (133 loc) · 3.9 KB
/
pubsub.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
package events
import (
"fmt"
"github.com/spacemeshos/go-spacemesh/log"
"nanomsg.org/go-mangos"
"nanomsg.org/go-mangos/protocol/pub"
"nanomsg.org/go-mangos/protocol/sub"
"nanomsg.org/go-mangos/transport/ipc"
"nanomsg.org/go-mangos/transport/tcp"
"sync"
"time"
)
// channelBuffer defines the listening channel buffer size.
const channelBuffer = 100
// ChannelID is the ID on which subscribers must register in order to listen to messages on the channel.
type ChannelID byte
// Subscriber defines the struct of the receiving end of the pubsub messages.
type Subscriber struct {
sock mangos.Socket
output map[ChannelID]chan []byte
allOutput chan []byte
chanLock sync.RWMutex
}
// NewSubscriber received url string as input on which it will register to receive messages passed by server.
func NewSubscriber(url string) (*Subscriber, error) {
socket, err := sub.NewSocket()
if err != nil {
return nil, err
}
socket.AddTransport(ipc.NewTransport())
socket.AddTransport(tcp.NewTransport())
socket.SetOption(mangos.OptionBestEffort, false)
socket.SetOption(mangos.OptionReadQLen, 10000)
err = socket.Dial(url)
if err != nil {
return nil, err
}
return &Subscriber{
socket,
make(map[ChannelID]chan []byte, 100),
nil,
sync.RWMutex{},
}, nil
}
// StartListening runs in a go routine and listens to all channels this subscriber is registered to. it then passes the
// messages received to the appropriate reader channel.
func (sub *Subscriber) StartListening() {
go func() {
for {
data, err := sub.sock.Recv()
log.Debug("got msg %v", string(data))
if err != nil {
if err.Error() == "connection closed" {
log.Warning("connection closed on pubsub reader")
} else {
log.Error("error on recv: %v", err)
}
return
}
// cannot route empty message
if len(data) < 1 {
log.Error("got empty message")
}
sub.chanLock.RLock()
if sub.allOutput != nil {
sub.allOutput <- data
}
if c, ok := sub.output[ChannelID(data[0])]; ok {
c <- data
}
sub.chanLock.RUnlock()
}
}()
}
// Close closes the socket which in turn will Close the listener func
func (sub *Subscriber) Close() error {
return sub.sock.Close()
}
// Subscribe subscribes to the given topic, returns a channel on which data from the topic is received.
func (sub *Subscriber) Subscribe(topic ChannelID) (chan []byte, error) {
sub.chanLock.Lock()
defer sub.chanLock.Unlock()
if _, ok := sub.output[topic]; !ok {
sub.output[topic] = make(chan []byte, channelBuffer)
}
err := sub.sock.SetOption(mangos.OptionSubscribe, []byte{byte(topic)})
if err != nil {
return nil, err
}
return sub.output[topic], err
}
// SubscribeToAll subscribes to all available topics.
func (sub *Subscriber) SubscribeToAll() (chan []byte, error) {
err := sub.sock.SetOption(mangos.OptionSubscribe, []byte(""))
if err != nil {
return nil, err
}
sub.chanLock.Lock()
sub.allOutput = make(chan []byte)
sub.chanLock.Unlock()
return sub.allOutput, err
}
// Publisher is a wrapper for mangos pubsub publisher socket
type Publisher struct {
sock mangos.Socket
}
func newPublisher(url string) (*Publisher, error) {
var sock mangos.Socket
var err error
if sock, err = pub.NewSocket(); err != nil {
return nil, fmt.Errorf("can't get NewEventPublisher pub socket: %s", err)
}
sock.AddTransport(ipc.NewTransport())
sock.AddTransport(tcp.NewTransport())
sock.SetOption(mangos.OptionBestEffort, false)
sock.SetOption(mangos.OptionWriteQLen, 10000)
if err = sock.Listen(url); err != nil {
return nil, fmt.Errorf("can't listen on pub socket: %s", err.Error())
}
time.Sleep(time.Second)
p := &Publisher{sock: sock}
return p, nil
}
func (p *Publisher) publish(topic ChannelID, payload []byte) error {
msg := append([]byte{byte(topic)}, payload...)
log.Debug("sending message %v", string(msg))
err := p.sock.Send(msg)
return err
}
// Close closes the publishers output socket
func (p *Publisher) Close() {
p.sock.Close()
}