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() +}