From b4e4e0350e484beb9e0ef2f245d456598eb5da50 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 29 May 2026 03:22:57 +0000 Subject: [PATCH] perf(room-worker): count members off denormalized u.isBot instead of per-doc regex ReconcileMemberCounts ran a $regexMatch on $u.account for every subscription in the room, on every create/add/remove, to split user vs app counts. The regex is not index-usable and runs once per document, so the cost scaled with room size on every membership change. model.Subscription already carries a u.isBot field that was never populated. This wires it up: - Add model.IsBotAccount as the single source of truth for the bot rule (".bot" suffix or "p_" prefix; equivalent to the (\.bot$|^p_) regex). - Stamp u.isBot at sub-creation in newSub, and route the add-member inline build through newSub so there is one construction site. Also reuse the predicate in determineRoomTypeFromPayload. - Rewrite ReconcileMemberCounts to two index-backed counts: total subs and bot subs ({roomId, u.isBot}); UserCount = total - bots. Deriving by subtraction keeps legacy docs missing the field counted as users. Recompute-and-$set preserves JetStream-redelivery idempotency. - Add the {roomId, u.isBot} index in room-service EnsureIndexes. Deploy ordering: ship this (writers stamp u.isBot) and backfill existing docs before relying on the new counts, otherwise pre-existing bots count as users until backfilled. Backfill (run once): db.subscriptions.updateMany({}, [ {$set: {"u.isBot": {$regexMatch: {input: "$u.account", regex: "(\\.bot$|^p_)"}}}} ]) https://claude.ai/code/session_01KyEPakZnVkZKrPL5j9cjpw --- pkg/model/account.go | 13 ++++++ pkg/model/account_test.go | 34 ++++++++++++++ room-service/store_mongo.go | 8 ++++ room-worker/handler.go | 24 ++++------ room-worker/integration_test.go | 5 +- room-worker/store.go | 5 +- room-worker/store_mongo.go | 63 +++++++------------------- room-worker/subscription_isbot_test.go | 27 +++++++++++ 8 files changed, 114 insertions(+), 65 deletions(-) create mode 100644 pkg/model/account.go create mode 100644 pkg/model/account_test.go create mode 100644 room-worker/subscription_isbot_test.go diff --git a/pkg/model/account.go b/pkg/model/account.go new file mode 100644 index 000000000..0c3654de3 --- /dev/null +++ b/pkg/model/account.go @@ -0,0 +1,13 @@ +package model + +import "strings" + +// IsBotAccount reports whether an account name denotes a bot/pseudo user. +// The rule — a ".bot" suffix or a "p_" prefix — is the single source of truth +// for bot classification and is equivalent to the regex `(\.bot$|^p_)` used by +// pkg/pipelines and room-service. Subscriptions store the result in u.isBot so +// member-count reconciliation can split user vs app counts off an indexed field +// instead of evaluating a regex per document on every read. +func IsBotAccount(account string) bool { + return strings.HasSuffix(account, ".bot") || strings.HasPrefix(account, "p_") +} diff --git a/pkg/model/account_test.go b/pkg/model/account_test.go new file mode 100644 index 000000000..d1119ce7c --- /dev/null +++ b/pkg/model/account_test.go @@ -0,0 +1,34 @@ +package model_test + +import ( + "testing" + + "github.com/hmchangw/chat/pkg/model" +) + +func TestIsBotAccount(t *testing.T) { + tests := []struct { + name string + account string + want bool + }{ + {name: "dot-bot suffix", account: "weather.bot", want: true}, + {name: "bare dot-bot", account: ".bot", want: true}, + {name: "p_ prefix webhook", account: "p_webhook", want: true}, + {name: "bare p_ prefix", account: "p_", want: true}, + {name: "plain human account", account: "alice", want: false}, + {name: "ends with bot but no dot", account: "robot", want: false}, + {name: "dot-bot not at end", account: "alice.bot.com", want: false}, + {name: "p underscore not at start", account: "alice_p_x", want: false}, + {name: "case sensitive prefix", account: "P_webhook", want: false}, + {name: "case sensitive suffix", account: "weather.BOT", want: false}, + {name: "empty", account: "", want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := model.IsBotAccount(tt.account); got != tt.want { + t.Errorf("IsBotAccount(%q) = %v, want %v", tt.account, got, tt.want) + } + }) + } +} diff --git a/room-service/store_mongo.go b/room-service/store_mongo.go index b647580db..4410e3ccc 100644 --- a/room-service/store_mongo.go +++ b/room-service/store_mongo.go @@ -91,6 +91,14 @@ func (s *MongoStore) EnsureIndexes(ctx context.Context) error { }); err != nil { return fmt.Errorf("ensure subscriptions (roomId,lastSeenAt) index: %w", err) } + // Backs room-worker's ReconcileMemberCounts, which counts bot vs non-bot + // subs per room off u.isBot — keeps both CountDocuments index-only instead + // of scanning every subscription in the room. + if _, err := s.subscriptions.Indexes().CreateOne(ctx, mongo.IndexModel{ + Keys: bson.D{{Key: "roomId", Value: 1}, {Key: "u.isBot", Value: 1}}, + }); err != nil { + return fmt.Errorf("ensure subscriptions (roomId,u.isBot) index: %w", err) + } // Lookup index for FindDMSubscription (filters on u.account+name). // Without this index, FindDMSubscription falls back to a collection scan. if _, err := s.subscriptions.Indexes().CreateOne(ctx, mongo.IndexModel{ diff --git a/room-worker/handler.go b/room-worker/handler.go index 466bdc996..87b8366c0 100644 --- a/room-worker/handler.go +++ b/room-worker/handler.go @@ -866,16 +866,9 @@ func (h *Handler) processAddMembers(ctx context.Context, data []byte) (err error actualAccounts := make([]string, 0, len(needSub)) for _, c := range needSub { user := userMap[c.Account] - sub := &model.Subscription{ - ID: idgen.GenerateUUIDv7(), - User: model.SubscriptionUser{ID: user.ID, Account: user.Account}, - RoomID: req.RoomID, - Name: room.Name, - RoomType: model.RoomTypeChannel, - SiteID: room.SiteID, - Roles: []model.Role{model.RoleMember}, - JoinedAt: acceptedAt, - } + // newSub stamps u.isBot from the account; room is the channel fetched by + // req.RoomID so RoomType/SiteID/Name/ID all match the prior inline build. + sub := newSub(idgen.GenerateUUIDv7(), &user, room, []model.Role{model.RoleMember}, room.Name, false, acceptedAt) // Resolve once via the shared helper so the local sub, the per-user // SubscriptionUpdateEvent fan-out, and the cross-site MemberAddEvent // all carry the same HistorySharedSince value. @@ -1188,7 +1181,7 @@ func newSub(id string, user *model.User, room *model.Room, roles []model.Role, name string, isSubscribed bool, joinedAt time.Time) *model.Subscription { return &model.Subscription{ ID: id, - User: model.SubscriptionUser{ID: user.ID, Account: user.Account}, + User: model.SubscriptionUser{ID: user.ID, Account: user.Account, IsBot: model.IsBotAccount(user.Account)}, RoomID: room.ID, SiteID: room.SiteID, Roles: roles, @@ -1311,13 +1304,12 @@ func (h *Handler) processCreateRoom(ctx context.Context, data []byte) (err error } } -// determineRoomTypeFromPayload mirrors room-service's determineRoomType on the canonical payload. -// botPattern matches both ".bot" suffix and "p_" prefix to classify webhook-style bots -// consistently with room-service/helper.go and pkg/pipelines. +// determineRoomTypeFromPayload mirrors room-service's determineRoomType on the +// canonical payload. model.IsBotAccount classifies webhook-style bots (".bot" +// suffix or "p_" prefix) consistently with room-service/helper.go and pkg/pipelines. func determineRoomTypeFromPayload(req *model.CreateRoomRequest) model.RoomType { if req.Name == "" && len(req.Orgs) == 0 && len(req.Channels) == 0 && len(req.Users) == 1 { - acct := req.Users[0] - if strings.HasSuffix(acct, ".bot") || strings.HasPrefix(acct, "p_") { + if model.IsBotAccount(req.Users[0]) { return model.RoomTypeBotDM } return model.RoomTypeDM diff --git a/room-worker/integration_test.go b/room-worker/integration_test.go index d7f7e5379..929279fc3 100644 --- a/room-worker/integration_test.go +++ b/room-worker/integration_test.go @@ -470,8 +470,11 @@ func TestReconcileMemberCountsSplitsBots(t *testing.T) { mustInsertSub(t, db, &model.Subscription{ ID: "s3", User: model.SubscriptionUser{Account: "carol"}, RoomID: "r1", }) + // Bot subs carry u.isBot=true in production (stamped by newSub via + // model.IsBotAccount); set it explicitly here since the test inserts the + // document directly. ReconcileMemberCounts now counts off this flag. mustInsertSub(t, db, &model.Subscription{ - ID: "s4", User: model.SubscriptionUser{Account: "weather.bot"}, RoomID: "r1", + ID: "s4", User: model.SubscriptionUser{Account: "weather.bot", IsBot: true}, RoomID: "r1", }) mustInsertRoom(t, db, &model.Room{ID: "r1", Type: model.RoomTypeChannel}) diff --git a/room-worker/store.go b/room-worker/store.go index 33671ae17..3f75466e6 100644 --- a/room-worker/store.go +++ b/room-worker/store.go @@ -59,8 +59,9 @@ type SubscriptionStore interface { // ListByRoom returns all subscriptions for roomID across every site. ListByRoom(ctx context.Context, roomID string) ([]model.Subscription, error) // ReconcileMemberCounts recomputes Room.UserCount (non-bot subs) and - // Room.AppCount (bot subs) by scanning the subscriptions collection, - // then writes both back to the rooms collection in a single update. + // Room.AppCount (bot subs) via index-backed counts on the denormalized + // u.isBot flag, then writes both back to the rooms collection in a single + // update. ReconcileMemberCounts(ctx context.Context, roomID string) error GetRoom(ctx context.Context, roomID string) (*model.Room, error) GetSubscription(ctx context.Context, account, roomID string) (*model.Subscription, error) diff --git a/room-worker/store_mongo.go b/room-worker/store_mongo.go index d3ccc0a76..ec10bc707 100644 --- a/room-worker/store_mongo.go +++ b/room-worker/store_mongo.go @@ -43,59 +43,30 @@ func (s *MongoStore) ListByRoom(ctx context.Context, roomID string) ([]model.Sub return subs, nil } -// ReconcileMemberCounts counts the room's subscriptions, splitting on -// the bot account naming pattern to produce both UserCount (non-bot) and -// AppCount (bot). A single $group aggregation does both buckets in one -// collection scan (was: two CountDocuments queries). Writes both fields -// to the rooms collection in a single updateOne. The regex must stay in -// lockstep with pkg/pipelines.GetNewMembersPipeline — both classify -// accounts matching `.bot$|^p_` as bots. +// ReconcileMemberCounts recomputes the room's AppCount (bot subs) and UserCount +// (everyone else) and writes both back in a single updateOne. AppCount is an +// index-backed CountDocuments on {roomId, u.isBot} (the flag is stamped at +// sub-creation via model.IsBotAccount) and UserCount is total minus bots — both +// counts use the index and no per-document regex runs. Deriving UserCount by +// subtraction also means legacy docs written before u.isBot existed (and any +// missing the field) correctly fall into UserCount rather than being dropped. +// Recompute-and-$set keeps the counts idempotent under JetStream redelivery. func (s *MongoStore) ReconcileMemberCounts(ctx context.Context, roomID string) error { - const botRegex = `(\.bot$|^p_)` - pipe := []bson.M{ - {"$match": bson.M{"roomId": roomID}}, - {"$group": bson.M{ - "_id": nil, - "appCount": bson.M{"$sum": bson.M{ - "$cond": []any{ - bson.M{"$regexMatch": bson.M{"input": "$u.account", "regex": botRegex}}, - 1, 0, - }, - }}, - "userCount": bson.M{"$sum": bson.M{ - "$cond": []any{ - bson.M{"$regexMatch": bson.M{"input": "$u.account", "regex": botRegex}}, - 0, 1, - }, - }}, - }}, - } - cur, err := s.subscriptions.Aggregate(ctx, pipe) + // A transient count error must not fall through to an UpdateOne with zero + // counts, which would clobber the rooms doc. + total, err := s.subscriptions.CountDocuments(ctx, bson.M{"roomId": roomID}) if err != nil { - return fmt.Errorf("aggregate member counts: %w", err) - } - defer cur.Close(ctx) - - var counts struct { - UserCount int64 `bson:"userCount"` - AppCount int64 `bson:"appCount"` + return fmt.Errorf("count subscriptions: %w", err) } - if cur.Next(ctx) { - if err := cur.Decode(&counts); err != nil { - return fmt.Errorf("decode member counts: %w", err) - } - } else if err := cur.Err(); err != nil { - // A cursor failure must not silently fall through to an UpdateOne with - // zero counts, which would clobber the rooms doc on a transient error. - return fmt.Errorf("iterate member counts: %w", err) + appCount, err := s.subscriptions.CountDocuments(ctx, bson.M{"roomId": roomID, "u.isBot": true}) + if err != nil { + return fmt.Errorf("count app subscriptions: %w", err) } - // No rows match → both counts stay 0, which is the correct reset behavior - // for a room whose last subscription was just removed. if _, err := s.rooms.UpdateOne(ctx, bson.M{"_id": roomID}, bson.M{ "$set": bson.M{ - "userCount": counts.UserCount, - "appCount": counts.AppCount, + "userCount": total - appCount, + "appCount": appCount, "updatedAt": time.Now().UTC(), }, }); err != nil { diff --git a/room-worker/subscription_isbot_test.go b/room-worker/subscription_isbot_test.go new file mode 100644 index 000000000..bc61a8c47 --- /dev/null +++ b/room-worker/subscription_isbot_test.go @@ -0,0 +1,27 @@ +package main + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/hmchangw/chat/pkg/model" +) + +// TestNewSub_SetsIsBotFromAccount asserts subscriptions are stamped with the +// bot classification at creation time, so member-count reconciliation can read +// u.isBot directly instead of re-deriving it with a per-document regex. +func TestNewSub_SetsIsBotFromAccount(t *testing.T) { + room := &model.Room{ID: "r1", SiteID: "site-a", Type: model.RoomTypeChannel} + ts := time.Now().UTC() + + bot := newSub("s1", &model.User{ID: "u_bot", Account: "helper.bot"}, room, nil, "n", false, ts) + assert.True(t, bot.User.IsBot, "helper.bot must be flagged as a bot") + + pseudo := newSub("s2", &model.User{ID: "u_p", Account: "p_webhook"}, room, nil, "n", false, ts) + assert.True(t, pseudo.User.IsBot, "p_ prefix must be flagged as a bot") + + human := newSub("s3", &model.User{ID: "u_h", Account: "alice"}, room, nil, "n", false, ts) + assert.False(t, human.User.IsBot, "human account must not be flagged") +}