util/eventbus: initial implementation of an in-process event bus

Updates #15160

Signed-off-by: David Anderson <dave@tailscale.com>
Co-authored-by: M. J. Fromberger <fromberger@tailscale.com>

util/eventbus/bus.go

@@ -0,0 +1,223 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package eventbus
+
+import (
+	"context"
+	"reflect"
+	"slices"
+	"sync"
+
+	"tailscale.com/util/set"
+)
+
+// Bus is an event bus that distributes published events to interested
+// subscribers.
+type Bus struct {
+	write    chan any
+	stop     goroutineShutdownControl
+	snapshot chan chan []any
+
+	topicsMu sync.Mutex // guards everything below.
+	topics   map[reflect.Type][]*Queue
+
+	// Used for introspection/debugging only, not in the normal event
+	// publishing path.
+	publishers set.Set[publisher]
+	queues     set.Set[*Queue]
+}
+
+// New returns a new bus. Use [PublisherOf] to make event publishers,
+// and [Bus.Queue] and [Subscribe] to make event subscribers.
+func New() *Bus {
+	stopCtl, stopWorker := newGoroutineShutdown()
+	ret := &Bus{
+		write:      make(chan any),
+		stop:       stopCtl,
+		snapshot:   make(chan chan []any),
+		topics:     map[reflect.Type][]*Queue{},
+		publishers: set.Set[publisher]{},
+		queues:     set.Set[*Queue]{},
+	}
+	go ret.pump(stopWorker)
+	return ret
+}
+
+func (b *Bus) pump(stop goroutineShutdownWorker) {
+	defer stop.Done()
+	var vals queue
+	acceptCh := func() chan any {
+		if vals.Full() {
+			return nil
+		}
+		return b.write
+	}
+	for {
+		// Drain all pending events. Note that while we're draining
+		// events into subscriber queues, we continue to
+		// opportunistically accept more incoming events, if we have
+		// queue space for it.
+		for !vals.Empty() {
+			val := vals.Peek()
+			dests := b.dest(reflect.ValueOf(val).Type())
+			for _, d := range dests {
+			deliverOne:
+				for {
+					select {
+					case d.write <- val:
+						break deliverOne
+					case <-d.stop.WaitChan():
+						// Queue closed, don't block but continue
+						// delivering to others.
+						break deliverOne
+					case in := <-acceptCh():
+						vals.Add(in)
+					case <-stop.Stop():
+						return
+					case ch := <-b.snapshot:
+						ch <- vals.Snapshot()
+					}
+				}
+			}
+			vals.Drop()
+		}
+
+		// Inbound queue empty, wait for at least 1 work item before
+		// resuming.
+		for vals.Empty() {
+			select {
+			case <-stop.Stop():
+				return
+			case val := <-b.write:
+				vals.Add(val)
+			case ch := <-b.snapshot:
+				ch <- nil
+			}
+		}
+	}
+}
+
+func (b *Bus) dest(t reflect.Type) []*Queue {
+	b.topicsMu.Lock()
+	defer b.topicsMu.Unlock()
+	return b.topics[t]
+}
+
+func (b *Bus) subscribe(t reflect.Type, q *Queue) (cancel func()) {
+	b.topicsMu.Lock()
+	defer b.topicsMu.Unlock()
+	b.topics[t] = append(b.topics[t], q)
+	return func() {
+		b.unsubscribe(t, q)
+	}
+}
+
+func (b *Bus) unsubscribe(t reflect.Type, q *Queue) {
+	b.topicsMu.Lock()
+	defer b.topicsMu.Unlock()
+	// Topic slices are accessed by pump without holding a lock, so we
+	// have to replace the entire slice when unsubscribing.
+	// Unsubscribing should be infrequent enough that this won't
+	// matter.
+	i := slices.Index(b.topics[t], q)
+	if i < 0 {
+		return
+	}
+	b.topics[t] = slices.Delete(slices.Clone(b.topics[t]), i, i+1)
+}
+
+func (b *Bus) Close() {
+	b.stop.StopAndWait()
+}
+
+// Queue returns a new queue with no subscriptions. Use [Subscribe] to
+// atach subscriptions to it.
+//
+// The queue's name should be a short, human-readable string that
+// identifies this queue. The name is only visible through debugging
+// APIs.
+func (b *Bus) Queue(name string) *Queue {
+	return newQueue(b, name)
+}
+
+func (b *Bus) addQueue(q *Queue) {
+	b.topicsMu.Lock()
+	defer b.topicsMu.Unlock()
+	b.queues.Add(q)
+}
+
+func (b *Bus) deleteQueue(q *Queue) {
+	b.topicsMu.Lock()
+	defer b.topicsMu.Unlock()
+	b.queues.Delete(q)
+}
+
+func (b *Bus) addPublisher(p publisher) {
+	b.topicsMu.Lock()
+	defer b.topicsMu.Unlock()
+	b.publishers.Add(p)
+}
+
+func (b *Bus) deletePublisher(p publisher) {
+	b.topicsMu.Lock()
+	defer b.topicsMu.Unlock()
+	b.publishers.Delete(p)
+}
+
+func newGoroutineShutdown() (goroutineShutdownControl, goroutineShutdownWorker) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	ctl := goroutineShutdownControl{
+		startShutdown:    cancel,
+		shutdownFinished: make(chan struct{}),
+	}
+	work := goroutineShutdownWorker{
+		startShutdown:    ctx.Done(),
+		shutdownFinished: ctl.shutdownFinished,
+	}
+
+	return ctl, work
+}
+
+// goroutineShutdownControl is a helper type to manage the shutdown of
+// a worker goroutine. The worker goroutine should use the
+// goroutineShutdownWorker related to this controller.
+type goroutineShutdownControl struct {
+	startShutdown    context.CancelFunc
+	shutdownFinished chan struct{}
+}
+
+func (ctl *goroutineShutdownControl) Stop() {
+	ctl.startShutdown()
+}
+
+func (ctl *goroutineShutdownControl) Wait() {
+	<-ctl.shutdownFinished
+}
+
+func (ctl *goroutineShutdownControl) WaitChan() <-chan struct{} {
+	return ctl.shutdownFinished
+}
+
+func (ctl *goroutineShutdownControl) StopAndWait() {
+	ctl.Stop()
+	ctl.Wait()
+}
+
+// goroutineShutdownWorker is a helper type for a worker goroutine to
+// be notified that it should shut down, and to report that shutdown
+// has completed. The notification is triggered by the related
+// goroutineShutdownControl.
+type goroutineShutdownWorker struct {
+	startShutdown    <-chan struct{}
+	shutdownFinished chan struct{}
+}
+
+func (work *goroutineShutdownWorker) Stop() <-chan struct{} {
+	return work.startShutdown
+}
+
+func (work *goroutineShutdownWorker) Done() {
+	close(work.shutdownFinished)
+}