121 lines
2.1 KiB
Go
121 lines
2.1 KiB
Go
package pubsub
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"sort"
|
|
|
|
telemetrypb "git.wntrmute.dev/kyle/sensenet/proto"
|
|
"git.wntrmute.dev/kyle/sensenet/topic"
|
|
"google.golang.org/protobuf/proto"
|
|
"gopkg.in/zeromq/goczmq.v4"
|
|
)
|
|
|
|
func readTopic(b []byte) string {
|
|
size := int(b[1])
|
|
return string(b[2 : size+2])
|
|
}
|
|
|
|
func protoTopic(topic string) string {
|
|
return fmt.Sprintf("\x0a%c%s", len(topic), topic)
|
|
}
|
|
|
|
type Subscriber struct {
|
|
addr string
|
|
publisher string
|
|
sock *goczmq.Sock
|
|
topics map[string]bool
|
|
}
|
|
|
|
func NewSubscriber(addr, publisher string, topics ...string) (*Subscriber, error) {
|
|
sub := &Subscriber{
|
|
addr: addr,
|
|
publisher: publisher,
|
|
sock: goczmq.NewSock(goczmq.Sub),
|
|
topics: map[string]bool{},
|
|
}
|
|
|
|
err := sub.connect()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
sub.Conflate(1)
|
|
for _, topicName := range topics {
|
|
sub.Subscribe(topicName)
|
|
}
|
|
|
|
return sub, nil
|
|
}
|
|
|
|
func (sub *Subscriber) Conflate(n int) {
|
|
sub.sock.SetConflate(n)
|
|
}
|
|
|
|
func (sub *Subscriber) connect() error {
|
|
log.Printf("subscriber dialing %s", sub.addr)
|
|
return sub.sock.Connect(sub.addr)
|
|
}
|
|
|
|
func (sub *Subscriber) Subscribe(topic string) {
|
|
if sub.topics[topic] {
|
|
return
|
|
}
|
|
|
|
sub.topics[topic] = true
|
|
sub.sock.SetSubscribe(protoTopic(topic))
|
|
}
|
|
|
|
func (sub *Subscriber) Topics() []string {
|
|
var topics = make([]string, 0, len(sub.topics))
|
|
for k := range sub.topics {
|
|
topics = append(topics, k)
|
|
}
|
|
|
|
sort.Strings(topics)
|
|
return topics
|
|
}
|
|
|
|
func (sub *Subscriber) receive() ([]byte, error) {
|
|
var (
|
|
data []byte
|
|
err error
|
|
packet []byte
|
|
more = 1
|
|
)
|
|
for more > 0 {
|
|
packet, more, err = sub.sock.RecvFrame()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
data = append(data, packet...)
|
|
}
|
|
|
|
return data, nil
|
|
}
|
|
|
|
func (sub *Subscriber) Receive() (*topic.Packet, error) {
|
|
data, err := sub.receive()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
pbPacket := &telemetrypb.Packet{}
|
|
err = proto.Unmarshal(data, pbPacket)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
t := readTopic(data)
|
|
|
|
packet := &topic.Packet{
|
|
Topic: t,
|
|
Publisher: sub.publisher,
|
|
Received: int64(pbPacket.Timestamp),
|
|
Payload: pbPacket.Payload,
|
|
}
|
|
|
|
return packet, nil
|
|
}
|