From c341cd177e5fcfa5d2a3487831664cdff83e0fb1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yago=20P=C3=A9rez?= <5002453+yagop@users.noreply.github.com> Date: Sun, 28 Jun 2026 14:21:12 +0200 Subject: [PATCH] fix(chats,drafts): batch getDialogs requests instead of one RPC per dialog The dialogs iterator's default batch size is 1, so it issued one messages.getDialogs RPC per dialog. `tg chats list` was slow (one round trip per chat up to --limit) and `tg drafts` was worse: it scans every dialog in the account with no limit, so it walked the whole list one RPC at a time. Set BatchSize: min(limit, 100) for chats (respecting --limit) and a flat 100 for drafts (full scan, no user limit). Telegram does not reject a per-request limit above 100 but silently returns fewer/incorrect results, so 100 is the safe maximum. chats clamps to at least 1, since a non-positive batch size panics the iterator (it preallocates a slice with that capacity) and a negative --limit would otherwise crash the process. Co-Authored-By: Claude Opus 4.8 --- cmd/tg/chats.go | 9 ++++++++- cmd/tg/chats_test.go | 37 +++++++++++++++++++++++++++++++++++ cmd/tg/drafts.go | 6 +++++- cmd/tg/drafts_test.go | 45 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 95 insertions(+), 2 deletions(-) create mode 100644 cmd/tg/drafts_test.go diff --git a/cmd/tg/chats.go b/cmd/tg/chats.go index 687260b..187142b 100644 --- a/cmd/tg/chats.go +++ b/cmd/tg/chats.go @@ -74,7 +74,14 @@ func listChats(ctx context.Context, api *tg.Client, m *peerManager, limit int, a folder = 1 } - iter := query.GetDialogs(api).FolderID(folder).Iter() + // Fetch in large batches to minimize round trips. The dialogs iterator's + // default batch size is 1, so it issues one messages.getDialogs RPC per + // dialog, which makes listing slow. Telegram does not reject a per-request + // limit above 100 but silently returns fewer/incorrect results, so cap at + // 100. Clamp to at least 1: a non-positive batch size would panic the + // iterator (it preallocates a slice with that capacity). + batch := max(1, min(limit, 100)) + iter := query.GetDialogs(api).FolderID(folder).BatchSize(batch).Iter() now := time.Now().Unix() var ( diff --git a/cmd/tg/chats_test.go b/cmd/tg/chats_test.go index 3eeb5c3..ff7cc53 100644 --- a/cmd/tg/chats_test.go +++ b/cmd/tg/chats_test.go @@ -75,3 +75,40 @@ func TestListChatsLimit(t *testing.T) { t.Fatalf("limit not respected: got %d", len(list.Chats)) } } + +// TestListChatsBatchSize asserts the per-request limit is batched (capped at +// 100) instead of the iterator default of 1, and never goes below 1 (a +// non-positive batch size panics the iterator). +func TestListChatsBatchSize(t *testing.T) { + for _, tt := range []struct { + limit int + wantBatch int + }{ + {limit: 5, wantBatch: 5}, + {limit: 100, wantBatch: 100}, + {limit: 350, wantBatch: 100}, + {limit: 0, wantBatch: 1}, + {limit: -1, wantBatch: 1}, + } { + var gotBatch int + api := newFuncAPI(t, func(req bin.Encoder) (bin.Encoder, error) { + r, ok := req.(*tg.MessagesGetDialogsRequest) + if !ok { + return nil, errors.Errorf("unexpected request %T", req) + } + gotBatch = r.Limit + return &tg.MessagesDialogs{ + Dialogs: []tg.DialogClass{&tg.Dialog{Peer: &tg.PeerUser{UserID: 1}, TopMessage: 1}}, + Messages: []tg.MessageClass{&tg.Message{ID: 1, PeerID: &tg.PeerUser{UserID: 1}, Date: 1}}, + Users: []tg.UserClass{&tg.User{ID: 1}}, + }, nil + }) + + if _, err := listChats(context.Background(), api, nil, tt.limit, false); err != nil { + t.Fatalf("limit %d: %v", tt.limit, err) + } + if gotBatch != tt.wantBatch { + t.Errorf("limit %d: request limit = %d, want %d", tt.limit, gotBatch, tt.wantBatch) + } + } +} diff --git a/cmd/tg/drafts.go b/cmd/tg/drafts.go index 038bce9..7bb16b9 100644 --- a/cmd/tg/drafts.go +++ b/cmd/tg/drafts.go @@ -47,7 +47,11 @@ func saveDraft(ctx context.Context, api *tg.Client, peer tg.InputPeerClass, text // listDrafts collects drafts from the dialog list. func listDrafts(ctx context.Context, api *tg.Client) (draftsResult, error) { - iter := query.GetDialogs(api).Iter() + // Scan every dialog to collect drafts. The iterator's default batch size is + // 1 (one messages.getDialogs RPC per dialog), so without batching this walks + // the whole account one round trip at a time. 100 is the per-request maximum + // Telegram serves. + iter := query.GetDialogs(api).BatchSize(100).Iter() var out draftsResult for iter.Next(ctx) { elem := iter.Value() diff --git a/cmd/tg/drafts_test.go b/cmd/tg/drafts_test.go new file mode 100644 index 0000000..cde7340 --- /dev/null +++ b/cmd/tg/drafts_test.go @@ -0,0 +1,45 @@ +package main + +import ( + "context" + "testing" + + "github.com/go-faster/errors" + + "github.com/gotd/td/bin" + "github.com/gotd/td/tg" +) + +// TestListDraftsBatchSize asserts the dialog scan is batched at 100 per request +// instead of the iterator default of 1: listDrafts walks every dialog, so a +// batch size of 1 would issue one messages.getDialogs RPC per dialog. +func TestListDraftsBatchSize(t *testing.T) { + var gotBatch int + api := newFuncAPI(t, func(req bin.Encoder) (bin.Encoder, error) { + r, ok := req.(*tg.MessagesGetDialogsRequest) + if !ok { + return nil, errors.Errorf("unexpected request %T", req) + } + gotBatch = r.Limit + return &tg.MessagesDialogs{ + Dialogs: []tg.DialogClass{&tg.Dialog{ + Peer: &tg.PeerUser{UserID: 1}, + TopMessage: 1, + Draft: &tg.DraftMessage{Message: "wip"}, + }}, + Messages: []tg.MessageClass{&tg.Message{ID: 1, PeerID: &tg.PeerUser{UserID: 1}, Date: 1}}, + Users: []tg.UserClass{&tg.User{ID: 1}}, + }, nil + }) + + res, err := listDrafts(context.Background(), api) + if err != nil { + t.Fatal(err) + } + if gotBatch != 100 { + t.Errorf("request limit = %d, want 100", gotBatch) + } + if len(res.Drafts) != 1 || res.Drafts[0].Message != "wip" { + t.Errorf("drafts = %+v", res.Drafts) + } +}