util/eventbus: add a function-based subscriber type (#17432)

Originally proposed by @bradfitz in #17413.

In practice, a lot of subscribers have only one event type of interest, or a
small number of mostly independent ones. In that case, the overhead of running
and maintaining a goroutine to select on multiple channels winds up being more
noisy than we'd like for the user of the API.

For this common case, add a new SubscriberFunc[T] type that delivers events to
a callback owned by the subscriber, directly on the goroutine belonging to the
client itself. This frees the consumer from the need to maintain their own
goroutine to pull events from the channel, and to watch for closure of the
subscriber.

Before:

     s := eventbus.Subscribe[T](eventClient)
     go func() {
       for {
          select {
          case <-s.Done():
            return
          case e := <-s.Events():
            doSomethingWith(e)
          }
       }
     }()
     // ...
     s.Close()

After:

     func doSomethingWithT(e T) { ... }
     s := eventbus.SubscribeFunc(eventClient, doSomethingWithT)
     // ...
     s.Close()

Moreover, unless the caller wants to explicitly stop the subscriber separately
from its governing client, it need not capture the SubscriberFunc value at all.

One downside of this approach is that a slow or deadlocked callback could block
client's service routine and thus stall all other subscriptions on that client,
However, this can already happen more broadly if a subscriber fails to service
its delivery channel in a timely manner, it just feeds back more immediately.

Updates #17487

Change-Id: I64592d786005177aa9fd445c263178ed415784d5
Signed-off-by: M. J. Fromberger <fromberger@tailscale.com>

util/eventbus/subscribe.go

@@ -61,45 +61,45 @@ func newSubscribeState(c *Client) *subscribeState {
 	return ret
 }
 
-func (q *subscribeState) pump(ctx context.Context) {
+func (s *subscribeState) pump(ctx context.Context) {
 	var vals queue[DeliveredEvent]
 	acceptCh := func() chan DeliveredEvent {
 		if vals.Full() {
 			return nil
 		}
-		return q.write
+		return s.write
 	}
 	for {
 		if !vals.Empty() {
 			val := vals.Peek()
-			sub := q.subscriberFor(val.Event)
+			sub := s.subscriberFor(val.Event)
 			if sub == nil {
 				// Raced with unsubscribe.
 				vals.Drop()
 				continue
 			}
-			if !sub.dispatch(ctx, &vals, acceptCh, q.snapshot) {
+			if !sub.dispatch(ctx, &vals, acceptCh, s.snapshot) {
 				return
 			}
 
-			if q.debug.active() {
-				q.debug.run(DeliveredEvent{
+			if s.debug.active() {
+				s.debug.run(DeliveredEvent{
 					Event: val.Event,
 					From:  val.From,
-					To:    q.client,
+					To:    s.client,
 				})
 			}
 		} else {
 			// Keep the cases in this select in sync with
-			// Subscriber.dispatch below. The only difference should be
-			// that this select doesn't deliver queued values to
-			// anyone, and unconditionally accepts new values.
+			// Subscriber.dispatch and SubscriberFunc.dispatch below.
+			// The only difference should be that this select doesn't deliver
+			// queued values to anyone, and unconditionally accepts new values.
 			select {
-			case val := <-q.write:
+			case val := <-s.write:
 				vals.Add(val)
 			case <-ctx.Done():
 				return
-			case ch := <-q.snapshot:
+			case ch := <-s.snapshot:
 				ch <- vals.Snapshot()
 			}
 		}
@@ -152,10 +152,10 @@ func (s *subscribeState) deleteSubscriber(t reflect.Type) {
 	s.client.deleteSubscriber(t, s)
 }
 
-func (q *subscribeState) subscriberFor(val any) subscriber {
-	q.outputsMu.Lock()
-	defer q.outputsMu.Unlock()
-	return q.outputs[reflect.TypeOf(val)]
+func (s *subscribeState) subscriberFor(val any) subscriber {
+	s.outputsMu.Lock()
+	defer s.outputsMu.Unlock()
+	return s.outputs[reflect.TypeOf(val)]
 }
 
 // Close closes the subscribeState. It implicitly closes all Subscribers
@@ -177,6 +177,7 @@ func (s *subscribeState) closed() <-chan struct{} {
 }
 
 // A Subscriber delivers one type of event from a [Client].
+// Events are sent to the [Subscriber.Events] channel.
 type Subscriber[T any] struct {
 	stop       stopFlag
 	read       chan T
@@ -252,3 +253,49 @@ func (s *Subscriber[T]) Close() {
 	s.stop.Stop() // unblock receivers
 	s.unregister()
 }
+
+// A SubscriberFunc delivers one type of event from a [Client].
+// Events are forwarded synchronously to a function provided at construction.
+type SubscriberFunc[T any] struct {
+	stop       stopFlag
+	read       func(T)
+	unregister func()
+}
+
+func newSubscriberFunc[T any](r *subscribeState, f func(T)) *SubscriberFunc[T] {
+	return &SubscriberFunc[T]{
+		read:       f,
+		unregister: func() { r.deleteSubscriber(reflect.TypeFor[T]()) },
+	}
+}
+
+// Close closes the SubscriberFunc, indicating the caller no longer wishes to
+// receive this event type.  After Close, no further events will be passed to
+// the callback.
+//
+// If the [Bus] from which s was created is closed, s is implicitly closed and
+// does not need to be closed separately.
+func (s *SubscriberFunc[T]) Close() { s.stop.Stop(); s.unregister() }
+
+// subscribeType implements part of the subscriber interface.
+func (s *SubscriberFunc[T]) subscribeType() reflect.Type { return reflect.TypeFor[T]() }
+
+// dispatch implements part of the subscriber interface.
+func (s *SubscriberFunc[T]) dispatch(ctx context.Context, vals *queue[DeliveredEvent], acceptCh func() chan DeliveredEvent, snapshot chan chan []DeliveredEvent) bool {
+	// Keep the cases in this select in sync with subscribeState.pump
+	// above. The only different should be that this select
+	// delivers a value by calling s.read.
+	select {
+	case val := <-acceptCh():
+		vals.Add(val)
+	case <-ctx.Done():
+		return false
+	case ch := <-snapshot:
+		ch <- vals.Snapshot()
+	default:
+	}
+	t := vals.Peek().Event.(T)
+	s.read(t)
+	vals.Drop()
+	return true
+}