move pubsub inmem to subpackage (#210)

Split types from implementation to support multiple pubsub types in the near future.

blueprint/listener.go

@@ -62,7 +62,7 @@ func (db *DB) ApplyToDevice(ctx context.Context, svc command.Service, bp *Bluepr
 }
 
 func (db *DB) StartListener(sub pubsub.Subscriber, cmdSvc command.Service) error {
-	tokenUpdateEvents, err := sub.Subscribe("applyAtEnroll", device.DeviceEnrolledTopic)
+	tokenUpdateEvents, err := sub.Subscribe(context.TODO(), "applyAtEnroll", device.DeviceEnrolledTopic)
 	if err != nil {
 		return errors.Wrapf(err,
 			"subscribing devices to %s topic", device.DeviceEnrolledTopic)

checkin/checkin.go

@@ -92,7 +92,7 @@ func (svc *Checkin) archiveAndPublish(topic string, cmd mdm.CheckinCommand) erro
 	if err := svc.archiveFn(event.Time.UnixNano(), msg); err != nil {
 		return errors.Wrap(err, "archive checkin")
 	}
-	if err := svc.publisher.Publish(topic, msg); err != nil {
+	if err := svc.publisher.Publish(context.TODO(), topic, msg); err != nil {
 		return errors.Wrapf(err, "publish checkin on topic: %s", topic)
 	}
 	return nil

checkin/checkin_test.go

@@ -47,7 +47,7 @@ type mockPublisher struct {
 	PublishFn func(string, []byte) error
 }
 
-func (m *mockPublisher) Publish(s string, b []byte) error {
+func (m *mockPublisher) Publish(ctx context.Context, s string, b []byte) error {
 	m.Invoked = true
 	return m.PublishFn(s, b)
 }

command/command.go

@@ -57,7 +57,7 @@ func (svc *Command) NewCommand(ctx context.Context, request *mdm.CommandRequest)
 	if err := svc.archive(event.Time.UnixNano(), msg); err != nil {
 		return nil, errors.Wrap(err, "archive mdm command")
 	}
-	if err := svc.publisher.Publish(CommandTopic, msg); err != nil {
+	if err := svc.publisher.Publish(context.TODO(), CommandTopic, msg); err != nil {
 		return nil, errors.Wrapf(err, "publish mdm command on topic: %s", CommandTopic)
 	}
 	return payload, nil

command/command_test.go

@@ -1,6 +1,7 @@
 package command
 
 import (
+	"context"
 	"errors"
 	"io/ioutil"
 	"os"
@@ -8,7 +9,6 @@ import (
 
 	"github.com/boltdb/bolt"
 	"github.com/micromdm/mdm"
-	"golang.org/x/net/context"
 )
 
 func TestService_NewCommand(t *testing.T) {
@@ -82,7 +82,7 @@ type mockPublisher struct {
 	PublishFn func(string, []byte) error
 }
 
-func (m *mockPublisher) Publish(s string, b []byte) error {
+func (m *mockPublisher) Publish(ctx context.Context, s string, b []byte) error {
 	return m.PublishFn(s, b)
 }
 

core/apply/service.go

@@ -57,7 +57,7 @@ func (svc *ApplyService) UploadApp(ctx context.Context, manifestName string, man
 }
 
 func (svc *ApplyService) WatchTokenUpdates(pubsub pubsub.Subscriber) error {
-	tokenAdded, err := pubsub.Subscribe("apply-token-events", deptoken.DEPTokenTopic)
+	tokenAdded, err := pubsub.Subscribe(context.TODO(), "apply-token-events", deptoken.DEPTokenTopic)
 	if err != nil {
 		return err
 	}

core/list/service.go

@@ -81,7 +81,7 @@ func (svc *ListService) ListApplications(ctx context.Context, opts ListAppsOptio
 }
 
 func (svc *ListService) WatchTokenUpdates(pubsub pubsub.Subscriber) error {
-	tokenAdded, err := pubsub.Subscribe("list-token-events", deptoken.DEPTokenTopic)
+	tokenAdded, err := pubsub.Subscribe(context.TODO(), "list-token-events", deptoken.DEPTokenTopic)
 	if err != nil {
 		return err
 	}

depsync/depsync.go

@@ -1,6 +1,7 @@
 package depsync
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"log"
@@ -103,7 +104,7 @@ func New(pub pubsub.PublishSubscriber, db *bolt.DB, opts ...Option) (Syncer, err
 }
 
 func (w *watcher) updateClient(pubsub pubsub.Subscriber) error {
-	tokenAdded, err := pubsub.Subscribe("token-events", deptoken.DEPTokenTopic)
+	tokenAdded, err := pubsub.Subscribe(context.TODO(), "token-events", deptoken.DEPTokenTopic)
 	if err != nil {
 		return err
 	}
@@ -168,7 +169,7 @@ FETCH:
 		if err != nil {
 			return err
 		}
-		if err := w.publisher.Publish(SyncTopic, data); err != nil {
+		if err := w.publisher.Publish(context.TODO(), SyncTopic, data); err != nil {
 			return err
 		}
 		if !resp.MoreToFollow {
@@ -198,7 +199,7 @@ SYNC:
 			if err != nil {
 				return err
 			}
-			if err := w.publisher.Publish(SyncTopic, data); err != nil {
+			if err := w.publisher.Publish(context.TODO(), SyncTopic, data); err != nil {
 				return err
 			}
 		}

deptoken/deptopken.go

@@ -2,6 +2,7 @@ package deptoken
 
 import (
 	"bytes"
+	"context"
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/json"
@@ -56,7 +57,7 @@ func (db *DB) AddToken(consumerKey string, json []byte) error {
 	if err != nil {
 		return err
 	}
-	if err := db.Publisher.Publish(DEPTokenTopic, json); err != nil {
+	if err := db.Publisher.Publish(context.TODO(), DEPTokenTopic, json); err != nil {
 		return err
 	}
 	return nil

device/db.go

@@ -1,6 +1,7 @@
 package device
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -165,22 +166,22 @@ func isNotFound(err error) bool {
 }
 
 func (db *DB) pollCheckin(pubsubSvc pubsub.PublishSubscriber) error {
-	authenticateEvents, err := pubsubSvc.Subscribe("devices", checkin.AuthenticateTopic)
+	authenticateEvents, err := pubsubSvc.Subscribe(context.TODO(), "devices", checkin.AuthenticateTopic)
 	if err != nil {
 		return errors.Wrapf(err,
 			"subscribing devices to %s topic", checkin.AuthenticateTopic)
 	}
-	tokenUpdateEvents, err := pubsubSvc.Subscribe("devices", checkin.TokenUpdateTopic)
+	tokenUpdateEvents, err := pubsubSvc.Subscribe(context.TODO(), "devices", checkin.TokenUpdateTopic)
 	if err != nil {
 		return errors.Wrapf(err,
 			"subscribing devices to %s topic", checkin.TokenUpdateTopic)
 	}
-	checkoutEvents, err := pubsubSvc.Subscribe("devices", checkin.CheckoutTopic)
+	checkoutEvents, err := pubsubSvc.Subscribe(context.TODO(), "devices", checkin.CheckoutTopic)
 	if err != nil {
 		return errors.Wrapf(err,
 			"subscribing devices to %s topic", checkin.CheckoutTopic)
 	}
-	depSyncEvents, err := pubsubSvc.Subscribe("devices", depsync.SyncTopic)
+	depSyncEvents, err := pubsubSvc.Subscribe(context.TODO(), "devices", depsync.SyncTopic)
 	if err != nil {
 		return errors.Wrapf(err,
 			"subscribing devices to %s topic", depsync.SyncTopic)
@@ -265,7 +266,7 @@ func (db *DB) pollCheckin(pubsubSvc pubsub.PublishSubscriber) error {
 				}
 				if newlyEnrolled {
 					fmt.Printf("device %s enrolled\n", ev.Command.UDID)
-					err := pubsubSvc.Publish(DeviceEnrolledTopic, event.Message)
+					err := pubsubSvc.Publish(context.TODO(), DeviceEnrolledTopic, event.Message)
 					if err != nil {
 						fmt.Println(err)
 					}

pubsub/inmem/consumer.go

@@ -1,16 +1,13 @@
-package pubsub
+package inmem
 
-type Subscriber interface {
-	Subscribe(name, topic string) (<-chan Event, error)
-}
+import (
+	"context"
 
-type PublishSubscriber interface {
-	Publisher
-	Subscriber
-}
+	"github.com/micromdm/micromdm/pubsub"
+)
 
-func (p *Inmem) Subscribe(name, topic string) (<-chan Event, error) {
-	events := make(chan Event)
+func (p *Inmem) Subscribe(_ context.Context, name, topic string) (<-chan pubsub.Event, error) {
+	events := make(chan pubsub.Event)
 	sub := subscription{
 		name:      name,
 		topic:     topic,

pubsub/inmem/inmem_test.go

@@ -1,33 +1,35 @@
-package pubsub
+package inmem
 
 import (
+	"context"
 	"fmt"
 	"testing"
 	"time"
 )
 
 func TestPubSub(t *testing.T) {
-	inmem := NewInmemPubsub()
+	ctx := context.Background()
+	inmem := NewPubSub()
 	tests := []string{"a", "b", "c"}
 	for _, tt := range tests {
-		if err := inmem.Publish(tt, []byte(tt+tt+tt)); err != nil {
+		if err := inmem.Publish(ctx, tt, []byte(tt+tt+tt)); err != nil {
 			t.Fatal(err)
 		}
-		if err := inmem.Publish(tt, []byte(tt+tt)); err != nil {
+		if err := inmem.Publish(ctx, tt, []byte(tt+tt)); err != nil {
 			t.Fatal(err)
 		}
 	}
 
-	subA, err := inmem.Subscribe("asub", "a")
+	subA, err := inmem.Subscribe(ctx, "asub", "a")
 	if err != nil {
 		t.Fatal(err)
 	}
-	subA1, err := inmem.Subscribe("asub1", "a")
+	subA1, err := inmem.Subscribe(ctx, "asub1", "a")
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	subB, err := inmem.Subscribe("bsub", "b")
+	subB, err := inmem.Subscribe(ctx, "bsub", "b")
 	if err != nil {
 		t.Fatal(err)
 	}

pubsub/inmem/publisher.go

@@ -1,18 +1,14 @@
-package pubsub
+package inmem
 
-import "sync"
+import (
+	"context"
+	"sync"
 
-type Publisher interface {
-	Publish(topic string, msg []byte) error
-}
+	"github.com/micromdm/micromdm/pubsub"
+)
 
-type Event struct {
-	Topic   string
-	Message []byte
-}
-
-func NewInmemPubsub() *Inmem {
-	publish := make(chan Event)
+func NewPubSub() *Inmem {
+	publish := make(chan pubsub.Event)
 	subscriptions := make(map[string][]subscription)
 	inmem := &Inmem{
 		publish:       publish,
@@ -26,17 +22,17 @@ type Inmem struct {
 	mtx           sync.RWMutex
 	subscriptions map[string][]subscription
 
-	publish chan Event
+	publish chan pubsub.Event
 }
 
 type subscription struct {
 	name      string
 	topic     string
-	eventChan chan<- Event
+	eventChan chan<- pubsub.Event
 }
 
-func (p *Inmem) Publish(topic string, msg []byte) error {
-	event := Event{Topic: topic, Message: msg}
+func (p *Inmem) Publish(_ context.Context, topic string, msg []byte) error {
+	event := pubsub.Event{Topic: topic, Message: msg}
 	go func() { p.publish <- event }()
 	return nil
 }

pubsub/pubsub.go

@@ -0,0 +1,21 @@
+package pubsub
+
+import "context"
+
+type Event struct {
+	Topic   string
+	Message []byte
+}
+
+type Publisher interface {
+	Publish(ctx context.Context, topic string, msg []byte) error
+}
+
+type Subscriber interface {
+	Subscribe(ctx context.Context, name, topic string) (<-chan Event, error)
+}
+
+type PublishSubscriber interface {
+	Publisher
+	Subscriber
+}

push/db.go

@@ -1,6 +1,7 @@
 package push
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/boltdb/bolt"
@@ -78,7 +79,7 @@ func (db *DB) Save(info *PushInfo) error {
 }
 
 func (db *DB) pollCheckin(sub pubsub.Subscriber) error {
-	tokenUpdateEvents, err := sub.Subscribe("push-info", checkin.TokenUpdateTopic)
+	tokenUpdateEvents, err := sub.Subscribe(context.TODO(), "push-info", checkin.TokenUpdateTopic)
 	if err != nil {
 		return errors.Wrapf(err,
 			"subscribing push to %s topic", checkin.TokenUpdateTopic)

push/service.go

@@ -28,7 +28,7 @@ func New(db *DB, push *push.Service, sub pubsub.Subscriber) (*Push, error) {
 }
 
 func (svc *Push) startQueuedSubscriber(push *push.Service, sub pubsub.Subscriber) error {
-	commandQueuedEvents, err := sub.Subscribe("push-info", queue.CommandQueuedTopic)
+	commandQueuedEvents, err := sub.Subscribe(context.TODO(), "push-info", queue.CommandQueuedTopic)
 	if err != nil {
 		return errors.Wrapf(err,
 			"subscribing push to %s topic", queue.CommandQueuedTopic)

queue/queue.go

@@ -175,7 +175,7 @@ func (e *notFound) Error() string {
 }
 
 func (db *Store) pollCommands(pubsub pubsub.PublishSubscriber) error {
-	commandEvents, err := pubsub.Subscribe("command-queue", command.CommandTopic)
+	commandEvents, err := pubsub.Subscribe(context.TODO(), "command-queue", command.CommandTopic)
 	if err != nil {
 		return errors.Wrapf(err,
 			"subscribing push to %s topic", command.CommandTopic)
@@ -222,7 +222,7 @@ func (db *Store) pollCommands(pubsub pubsub.PublishSubscriber) error {
 					continue
 				}
 
-				pubsub.Publish(CommandQueuedTopic, msgBytes)
+				pubsub.Publish(context.TODO(), CommandQueuedTopic, msgBytes)
 			}
 		}
 	}()

serve.go

@@ -55,6 +55,7 @@ import (
 	"github.com/micromdm/micromdm/enroll"
 	"github.com/micromdm/micromdm/profile"
 	"github.com/micromdm/micromdm/pubsub"
+	"github.com/micromdm/micromdm/pubsub/inmem"
 	nanopush "github.com/micromdm/micromdm/push"
 	"github.com/micromdm/micromdm/queue"
 )
@@ -473,7 +474,7 @@ func tlsConfig() *tls.Config {
 type config struct {
 	configPath          string
 	depsim              bool
-	pubclient           *pubsub.Inmem
+	pubclient           pubsub.PublishSubscriber
 	db                  *bolt.DB
 	pushCert            pushServiceCert
 	ServerPublicURL     string
@@ -504,7 +505,7 @@ func (c *config) setupPubSub() {
 	if c.err != nil {
 		return
 	}
-	c.pubclient = pubsub.NewInmemPubsub()
+	c.pubclient = inmem.NewPubSub()
 }
 
 func (c *config) setupCommandService() {