From 4675685ed1d2ee0695afae58bd0487a2deea45f3 Mon Sep 17 00:00:00 2001 From: Logan Pratt Date: Sun, 18 Oct 2026 08:51:04 -0400 Subject: [PATCH] fix(watch,wait): register update handlers before connecting watch and wait added their handlers to the update dispatcher inside the run callback, i.e. after the client had connected and could already be dispatching updates. tg.UpdateDispatcher keeps handlers in a plain map, so registering concurrently with dispatch is a data race (and can panic with a concurrent map write). runParams gains a register hook that connectWith calls on the fresh dispatcher before the client is created. The message stream is registered there and stays inert until armed with the resolved peer filter, which keeps the existing behavior: events before the filter is known were never delivered. Co-Authored-By: Claude Opus 5.5 --- cmd/tg/app.go | 7 +++++ cmd/tg/wait.go | 26 +++++++++--------- cmd/tg/watch.go | 48 +++++++++++++++++++-------------- cmd/tg/watch_test.go | 63 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 110 insertions(+), 34 deletions(-) create mode 100644 cmd/tg/watch_test.go diff --git a/cmd/tg/app.go b/cmd/tg/app.go index 9a8529c..d253ff4 100644 --- a/cmd/tg/app.go +++ b/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) diff --git a/cmd/tg/wait.go b/cmd/tg/wait.go index 2d962f4..366a480 100644 --- a/cmd/tg/wait.go +++ b/cmd/tg/wait.go @@ -24,8 +24,17 @@ 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 } @@ -33,18 +42,7 @@ func (a *app) newWaitCmd() *cobra.Command { 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 { diff --git a/cmd/tg/watch.go b/cmd/tg/watch.go index 012de71..a7779b4 100644 --- a/cmd/tg/watch.go +++ b/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,8 +170,17 @@ 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 } @@ -169,18 +188,7 @@ func (a *app) watchWith(ctx context.Context, st *accountState, header string, ar 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() diff --git a/cmd/tg/watch_test.go b/cmd/tg/watch_test.go new file mode 100644 index 0000000..00f97e2 --- /dev/null +++ b/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() +}