sensenet/pubsub/subscriber.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
}