fix(watch,wait): register update handlers before connecting

cmd/tg/app.go

@@ -173,6 +173,10 @@ func skipConfig(cmd *cobra.Command) bool {
 type runParams struct {
 	auth    authKind
 	updates bool
+	// register attaches update handlers to the dispatcher (updates only). It
+	// runs before the client connects: the dispatcher's handler map is not
+	// safe to modify while updates are being dispatched.
+	register func(d tg.UpdateDispatcher)
 	// authorize performs interactive authentication when the session is not yet
 	// authorized. If nil, user sessions error with errNotAuthorized and bot
 	// sessions authenticate with the configured token.
@@ -217,6 +221,9 @@ func (a *app) connectWith(
 	var d tg.UpdateDispatcher
 	if rp.updates {
 		d = tg.NewUpdateDispatcher()
+		if rp.register != nil {
+			rp.register(d)
+		}
 	}
 
 	appID, appHash, err := effectiveCreds(st.acc)

cmd/tg/wait.go

@@ -24,27 +24,25 @@ func (a *app) newWaitCmd() *cobra.Command {
 		Args:              cobra.MaximumNArgs(1),
 		ValidArgsFunction: peerArgCompletion,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return a.connect(cmd.Context(), runParams{auth: authUser, updates: true},
-				func(ctx context.Context, client *telegram.Client, d tg.UpdateDispatcher) error {
+			events := make(chan watchEvent, 1)
+			stream := &messageStream{
+				onEvent: func(ev watchEvent) {
+					select {
+					case events <- ev:
+					default:
+					}
+				},
+			}
+			return a.connect(cmd.Context(), runParams{auth: authUser, updates: true, register: stream.register},
+				func(ctx context.Context, client *telegram.Client, _ tg.UpdateDispatcher) error {
 					if err := requireAuth(ctx, client); err != nil {
 						return err
 					}
 					filterID, err := a.resolveFilterFor(ctx, client.API(), a.active, args)
 					if err != nil {
 						return err
 					}
-
-					events := make(chan watchEvent, 1)
-					stream := &messageStream{
-						filterID: filterID,
-						onEvent: func(ev watchEvent) {
-							select {
-							case events <- ev:
-							default:
-							}
-						},
-					}
-					stream.register(d)
+					stream.arm(filterID)
 
 					var timer <-chan time.Time
 					if timeout > 0 {

cmd/tg/watch.go

@@ -25,11 +25,21 @@ type watchEvent struct {
 }
 
 // messageStream turns incoming new-message updates into watchEvents, filtered by
-// an optional peer id, delivered to onEvent. Safe for concurrent dispatch.
+// an optional peer id, delivered to onEvent. It is registered before the client
+// connects and drops events until armed. Safe for concurrent dispatch.
 type messageStream struct {
-	filterID int64
-	onEvent  func(watchEvent)
+	onEvent func(watchEvent)
+
 	mu       sync.Mutex
+	armed    bool
+	filterID int64 // 0 = all peers
+}
+
+// arm starts delivering events, limited to filterID when non-zero.
+func (s *messageStream) arm(filterID int64) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.armed, s.filterID = true, filterID
 }
 
 func (s *messageStream) handle(msg tg.MessageClass, e tg.Entities) {
@@ -39,11 +49,11 @@ func (s *messageStream) handle(msg tg.MessageClass, e tg.Entities) {
 	}
 	ent := peer.EntitiesFromUpdate(e)
 	ev := watchEvent{Peer: describePeer(m.PeerID, ent), Message: buildMessageItem(m, ent)}
-	if s.filterID != 0 && ev.Peer.ID != s.filterID {
-		return
-	}
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	if !s.armed || (s.filterID != 0 && ev.Peer.ID != s.filterID) {
+		return
+	}
 	s.onEvent(ev)
 }
 
@@ -160,27 +170,25 @@ func (a *app) watchAll(ctx context.Context, labels, args []string) error {
 // serializes stdout across concurrent accounts.
 func (a *app) watchWith(ctx context.Context, st *accountState, header string, args []string, mu *sync.Mutex) error {
 	format := a.printer.Format()
-	return a.connectWith(ctx, st, runParams{auth: authUser, updates: true},
-		func(ctx context.Context, client *telegram.Client, d tg.UpdateDispatcher) error {
+	stream := &messageStream{
+		onEvent: func(ev watchEvent) {
+			if mu != nil {
+				mu.Lock()
+				defer mu.Unlock()
+			}
+			emitLine(format, header, ev)
+		},
+	}
+	return a.connectWith(ctx, st, runParams{auth: authUser, updates: true, register: stream.register},
+		func(ctx context.Context, client *telegram.Client, _ tg.UpdateDispatcher) error {
 			if err := requireAuth(ctx, client); err != nil {
 				return err
 			}
 			filterID, err := a.resolveFilterFor(ctx, client.API(), st, args)
 			if err != nil {
 				return err
 			}
-
-			stream := &messageStream{
-				filterID: filterID,
-				onEvent: func(ev watchEvent) {
-					if mu != nil {
-						mu.Lock()
-						defer mu.Unlock()
-					}
-					emitLine(format, header, ev)
-				},
-			}
-			stream.register(d)
+			stream.arm(filterID)
 
 			_, _ = fmt.Fprintf(os.Stderr, "Watching %s for new messages (Ctrl-C to stop)…\n", st.label)
 			<-ctx.Done()

cmd/tg/watch_test.go

@@ -0,0 +1,63 @@
+package main
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/gotd/td/tg"
+)
+
+func pushNewMessage(t *testing.T, d tg.UpdateDispatcher, peerID int64, id int) {
+	t.Helper()
+	upd := &tg.Updates{Updates: []tg.UpdateClass{&tg.UpdateNewMessage{
+		Message: &tg.Message{ID: id, PeerID: &tg.PeerUser{UserID: peerID}, Message: "hi"},
+	}}}
+	if err := d.Handle(context.Background(), upd); err != nil {
+		t.Error(err)
+	}
+}
+
+func TestMessageStreamArm(t *testing.T) {
+	var (
+		mu  sync.Mutex
+		got []int
+	)
+	s := &messageStream{onEvent: func(ev watchEvent) {
+		mu.Lock()
+		defer mu.Unlock()
+		got = append(got, ev.Message.ID)
+	}}
+	d := tg.NewUpdateDispatcher()
+	s.register(d) // registered before any update is dispatched
+
+	pushNewMessage(t, d, 7, 1) // dropped: not armed yet
+	s.arm(7)
+	pushNewMessage(t, d, 7, 2)
+	pushNewMessage(t, d, 8, 3) // dropped: other peer
+	s.arm(0)
+	pushNewMessage(t, d, 8, 4)
+
+	if len(got) != 2 || got[0] != 2 || got[1] != 4 {
+		t.Errorf("delivered %v, want [2 4]", got)
+	}
+}
+
+func TestMessageStreamConcurrentArm(t *testing.T) {
+	// Updates are dispatched from the client's goroutines while the command
+	// arms the stream; with -race this fails if the two are not synchronized.
+	s := &messageStream{onEvent: func(watchEvent) {}}
+	d := tg.NewUpdateDispatcher()
+	s.register(d)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < 100; i++ {
+			pushNewMessage(t, d, 7, i)
+		}
+	}()
+	s.arm(7)
+	wg.Wait()
+}