/
subscription.go
96 lines (81 loc) · 1.69 KB
/
subscription.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
// Copyright (c) 2013 ActiveState Software Inc. All rights reserved.
package zmqpubsub
import (
zmq "github.com/alecthomas/gozmq"
"gopkg.in/tomb.v1"
"time"
)
// Subscription provides channel abstraction over zmq SUB sockets
type Subscription struct {
addr string
filters []string
Ch chan Message // Channel to read messages from
tomb.Tomb
}
func newSubscription(addr string, filters []string) *Subscription {
sub := new(Subscription)
sub.addr = addr
sub.filters = filters
sub.Ch = make(chan Message)
go sub.loop()
return sub
}
func (sub *Subscription) loop() {
defer sub.Done()
defer close(sub.Ch)
ctx, err := GetGlobalContext()
if err != nil {
sub.Kill(err)
return
}
// Establish a connection and subscription filter
socket, err := ctx.NewSocket(zmq.SUB)
if err != nil {
sub.Kill(err)
return
}
for _, filter := range sub.filters {
err = socket.SetSockOptString(zmq.SUBSCRIBE, filter)
if err != nil {
sub.Kill(err)
return
}
}
err = socket.Connect(sub.addr)
if err != nil {
sub.Killf("Couldn't connect to %s: %s", sub.addr, err)
return
}
defer socket.Close()
// Read and stream the results in a channel
pollItems := []zmq.PollItem{zmq.PollItem{socket, 0, zmq.POLLIN, 0}}
for {
n, err := zmq.Poll(pollItems, time.Duration(1)*time.Second)
if err != nil {
sub.Kill(err)
return
}
select {
case <-sub.Dying():
return
default:
}
if n > 0 {
data, err := socket.Recv(zmq.DONTWAIT)
if err != nil {
sub.Kill(err)
return
}
select {
case sub.Ch <- NewMessage(data):
case <-sub.Dying():
return
}
}
}
}
// Stop stops this Subscription
func (sub *Subscription) Stop() error {
sub.Kill(nil)
return sub.Wait()
}