From 1782f22d307f007c25ab62bb83dd264f4f2cda9c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 07:49:08 +0000 Subject: [PATCH 1/4] pipelines: de-aggregate per-add candidate resolution into indexed finds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The add-member candidate/capacity resolution ran a correlated-$lookup aggregation on the users collection on every add (room-service CountNewMembers, room-worker ListAddMemberCandidates/ListNewMembersForNewRoom). Even with an optimal plan (IXSCAN, ~10 docs), the aggregation framework plus per-candidate $lookup sub-pipelines are markedly more CPU-heavy in mongod than plain finds, and this query runs once per add — so under the members load test it tops the slow log purely by volume, queueing behind itself on a saturated server. Replace the shared GetNewMembersPipeline/GetAddMemberCandidatesPipeline builders with a plain MatchCandidatesFilter predicate. Each caller now resolves candidates with one indexed find on users, then derives membership state with indexed reads scoped to the candidate set ($in): subscriptions on (roomId, u.account) and room_members on (rid, member.type, member.id), with the set difference computed in Go. Same results and examined-doc counts, none of the aggregation/$lookup overhead. Store interfaces are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_018ynZvBuRnWgG9moeqhSwkM --- pkg/pipelines/integration_test.go | 51 ++++++++ pkg/pipelines/main_test.go | 11 ++ pkg/pipelines/member.go | 122 ++++-------------- pkg/pipelines/member_test.go | 197 ++++++------------------------ pkg/pipelines/subscription.go | 38 ++++++ room-service/store.go | 3 +- room-service/store_mongo.go | 43 +++++-- room-worker/store_mongo.go | 104 ++++++++++++---- 8 files changed, 272 insertions(+), 297 deletions(-) create mode 100644 pkg/pipelines/integration_test.go create mode 100644 pkg/pipelines/main_test.go create mode 100644 pkg/pipelines/subscription.go diff --git a/pkg/pipelines/integration_test.go b/pkg/pipelines/integration_test.go new file mode 100644 index 000000000..cf63dce27 --- /dev/null +++ b/pkg/pipelines/integration_test.go @@ -0,0 +1,51 @@ +//go:build integration + +package pipelines + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.mongodb.org/mongo-driver/v2/bson" + + "github.com/hmchangw/chat/pkg/testutil" +) + +func TestSubscribedAccounts(t *testing.T) { + subs := testutil.MongoDB(t, "pipelines_test").Collection("subscriptions") + ctx := context.Background() + + // alice/bob subscribed to r1; carol subscribed to a different room. + _, err := subs.InsertMany(ctx, []any{ + bson.M{"_id": "s1", "roomId": "r1", "u": bson.M{"account": "alice"}}, + bson.M{"_id": "s2", "roomId": "r1", "u": bson.M{"account": "bob"}}, + bson.M{"_id": "s3", "roomId": "r2", "u": bson.M{"account": "carol"}}, + }) + require.NoError(t, err) + + t.Run("returns only the subset subscribed to the room", func(t *testing.T) { + got, err := SubscribedAccounts(ctx, subs, "r1", []string{"alice", "bob", "carol", "dave"}) + require.NoError(t, err) + assert.Equal(t, map[string]struct{}{"alice": {}, "bob": {}}, got) + }) + + t.Run("empty when none of the candidates are subscribed", func(t *testing.T) { + got, err := SubscribedAccounts(ctx, subs, "r1", []string{"carol", "dave"}) + require.NoError(t, err) + assert.Empty(t, got) + }) + + t.Run("does not match a different room", func(t *testing.T) { + got, err := SubscribedAccounts(ctx, subs, "r2", []string{"alice", "bob"}) + require.NoError(t, err) + assert.Empty(t, got) + }) + + t.Run("empty when the candidate set is empty", func(t *testing.T) { + got, err := SubscribedAccounts(ctx, subs, "r1", []string{}) + require.NoError(t, err) + assert.Empty(t, got) + }) +} diff --git a/pkg/pipelines/main_test.go b/pkg/pipelines/main_test.go new file mode 100644 index 000000000..1a3ce4d59 --- /dev/null +++ b/pkg/pipelines/main_test.go @@ -0,0 +1,11 @@ +//go:build integration + +package pipelines + +import ( + "testing" + + "github.com/hmchangw/chat/pkg/testutil" +) + +func TestMain(m *testing.M) { testutil.RunTests(m) } diff --git a/pkg/pipelines/member.go b/pkg/pipelines/member.go index 8e46ceabd..58b8c10ba 100644 --- a/pkg/pipelines/member.go +++ b/pkg/pipelines/member.go @@ -1,66 +1,34 @@ -// Package pipelines holds shared MongoDB aggregation pipelines used by more -// than one service. Putting them here lets each service append its own -// terminal stage (e.g. $count vs. $group) without duplicating the leading -// stages. +// Package pipelines holds shared MongoDB candidate-resolution helpers used by +// more than one service — both the query predicate (MatchCandidatesFilter) and +// the companion read (SubscribedAccounts). Centralizing them keeps room-service +// (capacity check) and room-worker (candidate resolution) in lock-step on org +// expansion, bot exclusion, and the already-subscribed subtraction. package pipelines import ( - "errors" - "go.mongodb.org/mongo-driver/v2/bson" ) +// botOrPseudoAccountRegex matches bot (".bot" suffix) and pseudo ("p_" prefix) +// accounts. It is the wire-side equivalent of model.IsBot / model.IsPlatformAdminAccount, +// applied as a residual filter so org-expanded candidates — whose accounts the +// caller does not know up front — are excluded server-side. const botOrPseudoAccountRegex = `(\.bot$|^p_)` -// ErrRoomIDRequired is returned by pipeline builders that require a non-empty roomID. -var ErrRoomIDRequired = errors.New("roomID required") - -// GetNewMembersPipeline returns the common stages for finding the unique, -// non-bot, not-already-subscribed users that an add-members request would -// add to roomID, given org IDs and direct account names. -// -// Pipeline target: the users collection. +// MatchCandidatesFilter returns the query predicate selecting the users an +// add/create-member request resolves to: members of any of orgIDs (by sectId or +// deptId) OR any of directAccounts, excluding bot/pseudo accounts and — when +// non-empty — excludeAccount. // -// Stages: $match (org/account filter, exclude bots, optionally exclude one -// account), then (when roomID != "") $lookup + $match to filter out -// already-subscribed accounts. Empty roomID returns the $match stage only -// (used by capacity-check at create time, where the room doesn't exist yet -// so there are no subscriptions to filter against). +// It is a plain find/count predicate, not an aggregation stage. Callers resolve +// each candidate's membership state with separate indexed reads (subscriptions +// keyed on (roomId, u.account), room_members on (rid, member.type, member.id)) +// rather than a correlated $lookup, which keeps every per-add query a bounded +// indexed lookup instead of an aggregation pipeline. // -// excludeAccount is empty string to disable, or an account that must be -// dropped from the candidate set. Create-room callers pass the requester's -// account so the requester (who joins as owner separately) is not double- -// counted via org expansion. -// -// Callers MUST append a terminal stage that fits their need: -// - room-service: bson.M{"$count": "n"} (capacity check) -// - room-worker: bson.M{"$group": {"_id": nil, "accounts": {"$addToSet": "$account"}}} -func GetNewMembersPipeline(orgIDs, directAccounts []string, roomID, excludeAccount string) bson.A { - stages := bson.A{matchCandidates(orgIDs, directAccounts, excludeAccount)} - if roomID != "" { - stages = append(stages, - bson.M{"$lookup": bson.M{ - "from": "subscriptions", - "let": bson.M{"userAccount": "$account"}, - "pipeline": bson.A{ - bson.M{"$match": bson.M{"$expr": bson.M{"$and": bson.A{ - bson.M{"$eq": bson.A{"$roomId", roomID}}, - bson.M{"$eq": bson.A{"$u.account", "$$userAccount"}}, - }}}}, - bson.M{"$limit": 1}, - // Outer stage only checks emptiness — drop the full sub doc. - bson.M{"$project": bson.M{"_id": 1}}, - }, - "as": "existingSub", - }}, - bson.M{"$match": bson.M{"existingSub": bson.M{"$eq": bson.A{}}}}, - ) - } - return stages -} - -// matchCandidates: $match users by (sectId|deptId IN orgIDs) OR (account IN directAccounts), bot/excludeAccount filtered. -func matchCandidates(orgIDs, directAccounts []string, excludeAccount string) bson.M { +// At least one of orgIDs/directAccounts must be non-empty: an all-empty call +// produces a {$or: []} predicate that MongoDB rejects, so callers guard first. +func MatchCandidatesFilter(orgIDs, directAccounts []string, excludeAccount string) bson.M { orFilter := bson.A{} if len(orgIDs) > 0 { orFilter = append(orFilter, bson.M{"sectId": bson.M{"$in": orgIDs}}, bson.M{"deptId": bson.M{"$in": orgIDs}}) @@ -72,51 +40,5 @@ func matchCandidates(orgIDs, directAccounts []string, excludeAccount string) bso if excludeAccount != "" { accountFilter["$ne"] = excludeAccount } - return bson.M{"$match": bson.M{"$or": orFilter, "account": accountFilter}} -} - -// GetAddMemberCandidatesPipeline returns per-candidate {account, hasSubscription, hasIndividualRoomMember} for the worker. -// Returns ErrRoomIDRequired if roomID is empty. -func GetAddMemberCandidatesPipeline(orgIDs, directAccounts []string, roomID, excludeAccount string) (bson.A, error) { - if roomID == "" { - return nil, ErrRoomIDRequired - } - return bson.A{ - matchCandidates(orgIDs, directAccounts, excludeAccount), - bson.M{"$lookup": bson.M{ - "from": "subscriptions", - "let": bson.M{"acct": "$account"}, - "pipeline": bson.A{ - bson.M{"$match": bson.M{"$expr": bson.M{"$and": bson.A{ - bson.M{"$eq": bson.A{"$roomId", roomID}}, - bson.M{"$eq": bson.A{"$u.account", "$$acct"}}, - }}}}, - bson.M{"$limit": 1}, - // Outer $project only reads $size of _sub — drop everything else. - bson.M{"$project": bson.M{"_id": 1}}, - }, - "as": "_sub", - }}, - bson.M{"$lookup": bson.M{ - "from": "room_members", - "let": bson.M{"uid": "$_id"}, - "pipeline": bson.A{ - bson.M{"$match": bson.M{"$expr": bson.M{"$and": bson.A{ - bson.M{"$eq": bson.A{"$rid", roomID}}, - bson.M{"$eq": bson.A{"$member.type", "individual"}}, - bson.M{"$eq": bson.A{"$member.id", "$$uid"}}, - }}}}, - bson.M{"$limit": 1}, - // Same rationale as the _sub lookup above. - bson.M{"$project": bson.M{"_id": 1}}, - }, - "as": "_irm", - }}, - bson.M{"$project": bson.M{ - "_id": 0, - "account": "$account", - "hasSubscription": bson.M{"$gt": bson.A{bson.M{"$size": "$_sub"}, 0}}, - "hasIndividualRoomMember": bson.M{"$gt": bson.A{bson.M{"$size": "$_irm"}, 0}}, - }}, - }, nil + return bson.M{"$or": orFilter, "account": accountFilter} } diff --git a/pkg/pipelines/member_test.go b/pkg/pipelines/member_test.go index d3b449e14..01195dae0 100644 --- a/pkg/pipelines/member_test.go +++ b/pkg/pipelines/member_test.go @@ -8,163 +8,44 @@ import ( "go.mongodb.org/mongo-driver/v2/bson" ) -func TestGetNewMembersPipeline(t *testing.T) { - t.Run("three stages returned", func(t *testing.T) { - got := GetNewMembersPipeline([]string{"org1"}, []string{"alice"}, "room1", "") - assert.Len(t, got, 3) - }) - - t.Run("$or filter with both orgIDs and directAccounts", func(t *testing.T) { - got := GetNewMembersPipeline([]string{"org1", "org2"}, []string{"alice"}, "room1", "") - stage0 := got[0].(bson.M) - match := stage0["$match"].(bson.M) - orFilter := match["$or"].(bson.A) - - // sectId + deptId for the one orgID group, plus account = 3 clauses. - assert.Len(t, orFilter, 3) - assert.NotNil(t, orFilter[0]) - assert.NotNil(t, orFilter[1]) - assert.NotNil(t, orFilter[2]) - }) - - t.Run("bot exclusion via $not regex", func(t *testing.T) { - got := GetNewMembersPipeline([]string{"org1"}, []string{"alice"}, "room1", "") - stage0 := got[0].(bson.M) - match := stage0["$match"].(bson.M) - - notFilter := match["account"] - assert.NotNil(t, notFilter) - }) - - t.Run("$or filter contains orgIDs when provided", func(t *testing.T) { - got := GetNewMembersPipeline([]string{"org1"}, nil, "room1", "") - stage0 := got[0].(bson.M) - match := stage0["$match"].(bson.M) - orFilter := match["$or"].(bson.A) - - // sectId + deptId clauses for the org group. - assert.Len(t, orFilter, 2) - sectIdFilter := orFilter[0].(bson.M) - assert.Contains(t, sectIdFilter, "sectId") - deptIdFilter := orFilter[1].(bson.M) - assert.Contains(t, deptIdFilter, "deptId") - }) - - t.Run("$or filter contains directAccounts when provided", func(t *testing.T) { - got := GetNewMembersPipeline(nil, []string{"alice"}, "room1", "") - stage0 := got[0].(bson.M) - match := stage0["$match"].(bson.M) - orFilter := match["$or"].(bson.A) - - assert.Len(t, orFilter, 1) - accountFilter := orFilter[0].(bson.M) - assert.Contains(t, accountFilter, "account") - }) - - t.Run("$lookup stage correct", func(t *testing.T) { - got := GetNewMembersPipeline([]string{"org1"}, []string{"alice"}, "room1", "") - stage1 := got[1].(bson.M) - lookup := stage1["$lookup"].(bson.M) - - assert.Equal(t, "subscriptions", lookup["from"]) - assert.Equal(t, "existingSub", lookup["as"]) - assert.NotNil(t, lookup["let"]) - assert.NotNil(t, lookup["pipeline"]) - }) - - t.Run("$match stage filters empty existingSub array", func(t *testing.T) { - got := GetNewMembersPipeline([]string{"org1"}, []string{"alice"}, "room1", "") - stage2 := got[2].(bson.M) - match := stage2["$match"].(bson.M) - existingSub := match["existingSub"].(bson.M) - eqA := existingSub["$eq"].(bson.A) - assert.Len(t, eqA, 0, "compare against the empty array literal") - }) -} - -func TestGetNewMembersPipelineEmptyRoomID(t *testing.T) { - pipe := GetNewMembersPipeline([]string{"org-fx"}, []string{"bob"}, "", "") - require.NotEmpty(t, pipe) - - hasLookup := false - for _, stage := range pipe { - m, ok := stage.(bson.M) - if !ok { - continue - } - if _, found := m["$lookup"]; found { - hasLookup = true - } - } - assert.False(t, hasLookup, "empty-roomID branch must drop the subscriptions $lookup") -} - -func TestGetNewMembersPipelineWithRoomIDStillHasLookup(t *testing.T) { - pipe := GetNewMembersPipeline([]string{"org-fx"}, []string{"bob"}, "r1", "") - require.NotEmpty(t, pipe) - - hasLookup := false - for _, stage := range pipe { - m, ok := stage.(bson.M) - if !ok { - continue - } - if _, found := m["$lookup"]; found { - hasLookup = true - } - } - assert.True(t, hasLookup, "non-empty roomID must keep the subscriptions $lookup") -} - -func TestGetAddMemberCandidatesPipeline(t *testing.T) { - t.Run("stage count is 4", func(t *testing.T) { - got, err := GetAddMemberCandidatesPipeline([]string{"org1"}, []string{"alice"}, "room1", "") - require.NoError(t, err) - assert.Len(t, got, 4) - }) - - t.Run("$match shape with orgIDs only — 2 OR clauses: sectId, deptId", func(t *testing.T) { - got, err := GetAddMemberCandidatesPipeline([]string{"org1"}, nil, "room1", "") - require.NoError(t, err) - stage0 := got[0].(bson.M) - match := stage0["$match"].(bson.M) - orFilter := match["$or"].(bson.A) - assert.Len(t, orFilter, 2, "sectId + deptId clauses for orgIDs only") - assert.Contains(t, orFilter[0].(bson.M), "sectId") - assert.Contains(t, orFilter[1].(bson.M), "deptId") - }) - - t.Run("$match shape with directAccounts only — 1 OR clause for account", func(t *testing.T) { - got, err := GetAddMemberCandidatesPipeline(nil, []string{"alice"}, "room1", "") - require.NoError(t, err) - stage0 := got[0].(bson.M) - match := stage0["$match"].(bson.M) - orFilter := match["$or"].(bson.A) - assert.Len(t, orFilter, 1) - assert.Contains(t, orFilter[0].(bson.M), "account") - }) - - t.Run("$match shape with both — 3 elements in $or: sectId, deptId, account", func(t *testing.T) { - got, err := GetAddMemberCandidatesPipeline([]string{"org1"}, []string{"alice"}, "room1", "") - require.NoError(t, err) - stage0 := got[0].(bson.M) - match := stage0["$match"].(bson.M) - orFilter := match["$or"].(bson.A) - assert.Len(t, orFilter, 3, "sectId + deptId + account") - }) - - t.Run("excludeAccount adds $ne to the account filter", func(t *testing.T) { - got, err := GetAddMemberCandidatesPipeline([]string{"org1"}, []string{"alice"}, "room1", "exclude-me") - require.NoError(t, err) - stage0 := got[0].(bson.M) - match := stage0["$match"].(bson.M) - accountFilter := match["account"].(bson.M) - assert.Equal(t, "exclude-me", accountFilter["$ne"]) - }) - - t.Run("error on empty roomID", func(t *testing.T) { - got, err := GetAddMemberCandidatesPipeline([]string{"org1"}, nil, "", "") - assert.Nil(t, got) - assert.ErrorIs(t, err, ErrRoomIDRequired) +func TestMatchCandidatesFilter(t *testing.T) { + t.Run("orgs and accounts produce three $or branches", func(t *testing.T) { + f := MatchCandidatesFilter([]string{"org1", "org2"}, []string{"alice", "bob"}, "") + or, ok := f["$or"].(bson.A) + require.True(t, ok, "$or must be a bson.A") + assert.Equal(t, bson.A{ + bson.M{"sectId": bson.M{"$in": []string{"org1", "org2"}}}, + bson.M{"deptId": bson.M{"$in": []string{"org1", "org2"}}}, + bson.M{"account": bson.M{"$in": []string{"alice", "bob"}}}, + }, or) + }) + + t.Run("orgs only omits the account branch", func(t *testing.T) { + f := MatchCandidatesFilter([]string{"org1"}, nil, "") + assert.Len(t, f["$or"], 2) + }) + + t.Run("accounts only omits the org branches", func(t *testing.T) { + f := MatchCandidatesFilter(nil, []string{"alice"}, "") + assert.Equal(t, bson.A{bson.M{"account": bson.M{"$in": []string{"alice"}}}}, f["$or"]) + }) + + t.Run("always excludes bots via a $not regex on account", func(t *testing.T) { + f := MatchCandidatesFilter(nil, []string{"alice"}, "") + acct, ok := f["account"].(bson.M) + require.True(t, ok) + rx, ok := acct["$not"].(bson.Regex) + require.True(t, ok, "account.$not must be a regex") + assert.Equal(t, botOrPseudoAccountRegex, rx.Pattern) + _, hasNe := acct["$ne"] + assert.False(t, hasNe, "no $ne when excludeAccount is empty") + }) + + t.Run("excludeAccount adds a $ne on account", func(t *testing.T) { + f := MatchCandidatesFilter([]string{"org1"}, nil, "exclude-me") + acct := f["account"].(bson.M) + assert.Equal(t, "exclude-me", acct["$ne"]) + _, hasNot := acct["$not"] + assert.True(t, hasNot, "bot filter is retained alongside $ne") }) } diff --git a/pkg/pipelines/subscription.go b/pkg/pipelines/subscription.go new file mode 100644 index 000000000..4c1e3bda9 --- /dev/null +++ b/pkg/pipelines/subscription.go @@ -0,0 +1,38 @@ +package pipelines + +import ( + "context" + "fmt" + + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" + "go.mongodb.org/mongo-driver/v2/mongo/options" +) + +// SubscribedAccounts returns the subset of accounts that already have a +// subscription to roomID, via an indexed point read on (roomId, u.account). +// It is the "subtract already-subscribed" half of candidate resolution, shared +// by room-service (CountNewMembers) and room-worker (ListAddMemberCandidates) +// so it stays in lock-step with MatchCandidatesFilter — exactly the drift this +// package exists to prevent. +func SubscribedAccounts(ctx context.Context, subscriptions *mongo.Collection, roomID string, accounts []string) (map[string]struct{}, error) { + cursor, err := subscriptions.Find(ctx, + bson.M{"roomId": roomID, "u.account": bson.M{"$in": accounts}}, + options.Find().SetProjection(bson.M{"u.account": 1, "_id": 0})) + if err != nil { + return nil, fmt.Errorf("find existing subscriptions for room %q: %w", roomID, err) + } + var rows []struct { + User struct { + Account string `bson:"account"` + } `bson:"u"` + } + if err := cursor.All(ctx, &rows); err != nil { + return nil, fmt.Errorf("decode existing subscriptions: %w", err) + } + set := make(map[string]struct{}, len(rows)) + for _, r := range rows { + set[r.User.Account] = struct{}{} + } + return set, nil +} diff --git a/room-service/store.go b/room-service/store.go index ade81e84b..80d6fb489 100644 --- a/room-service/store.go +++ b/room-service/store.go @@ -93,7 +93,8 @@ type RoomStore interface { // account so an org-expanded requester is not double-counted against the // cap (the requester is added separately as the owner). // Used by addMembers and handleCreateRoomChannel for capacity validation. - // Delegates to pkg/pipelines.GetNewMembersPipeline + a $count terminal stage. + // Resolves candidates via pkg/pipelines.MatchCandidatesFilter, then (for a + // non-empty roomID) subtracts already-subscribed accounts via an indexed read. CountNewMembers(ctx context.Context, orgIDs, directAccounts []string, roomID, excludeAccount string) (int, error) // ListRoomMembers returns the members of roomID. When enrich=true, the // returned RoomMember.Member entries carry display fields populated via diff --git a/room-service/store_mongo.go b/room-service/store_mongo.go index a2a7b53a2..583f05acd 100644 --- a/room-service/store_mongo.go +++ b/room-service/store_mongo.go @@ -461,23 +461,44 @@ func (s *MongoStore) CountNewMembers(ctx context.Context, orgIDs, directAccounts if len(orgIDs) == 0 && len(directAccounts) == 0 { return 0, nil } - pipeline := pipelines.GetNewMembersPipeline(orgIDs, directAccounts, roomID, excludeAccount) - pipeline = append(pipeline, bson.M{"$count": "n"}) - - cursor, err := s.users.Aggregate(ctx, pipeline) + filter := pipelines.MatchCandidatesFilter(orgIDs, directAccounts, excludeAccount) + // Create path (no room yet): every resolved candidate is a new member, so a + // plain indexed count suffices — there are no subscriptions to subtract. + if roomID == "" { + n, err := s.users.CountDocuments(ctx, filter) + if err != nil { + return 0, fmt.Errorf("count new members: %w", err) + } + return int(n), nil + } + // Add path: resolve candidate accounts, then subtract those already + // subscribed via one indexed read scoped to the candidate set — instead of + // a correlated-$lookup aggregation. + cursor, err := s.users.Find(ctx, filter, options.Find().SetProjection(bson.M{"account": 1, "_id": 0})) if err != nil { - return 0, fmt.Errorf("count new members: %w", err) + return 0, fmt.Errorf("find candidate accounts: %w", err) } - var results []struct { - Count int `bson:"n"` + var rows []struct { + Account string `bson:"account"` } - if err := cursor.All(ctx, &results); err != nil { - return 0, fmt.Errorf("decode count new members: %w", err) + if err := cursor.All(ctx, &rows); err != nil { + return 0, fmt.Errorf("decode candidate accounts: %w", err) } - if len(results) == 0 { + if len(rows) == 0 { return 0, nil } - return results[0].Count, nil + // account is unique in the users collection, so rows carry no duplicates. + accounts := make([]string, len(rows)) + for i, r := range rows { + accounts[i] = r.Account + } + subbed, err := pipelines.SubscribedAccounts(ctx, s.subscriptions, roomID, accounts) + if err != nil { + return 0, fmt.Errorf("resolve subscribed accounts: %w", err) + } + // subbed is a subset of the candidate accounts (the $in query bounds it), so + // the new-member count is just the candidates minus those already subscribed. + return len(accounts) - len(subbed), nil } // ListRoomMembers returns the members of a room. It prefers the room_members diff --git a/room-worker/store_mongo.go b/room-worker/store_mongo.go index 1dfb6a46b..aa8356f9a 100644 --- a/room-worker/store_mongo.go +++ b/room-worker/store_mongo.go @@ -157,32 +157,30 @@ func (s *MongoStore) CreateRoom(ctx context.Context, room *model.Room, key *room } func (s *MongoStore) ListNewMembersForNewRoom(ctx context.Context, orgIDs, accounts []string, excludeAccount string) ([]string, error) { - pipe := pipelines.GetNewMembersPipeline(orgIDs, accounts, "", excludeAccount) - pipe = append(pipe, bson.M{"$group": bson.M{ - "_id": nil, - "accounts": bson.M{"$addToSet": "$account"}, - }}) - cur, err := s.users.Aggregate(ctx, pipe) + if len(orgIDs) == 0 && len(accounts) == 0 { + return nil, nil + } + filter := pipelines.MatchCandidatesFilter(orgIDs, accounts, excludeAccount) + cur, err := s.users.Find(ctx, filter, options.Find().SetProjection(bson.M{"account": 1, "_id": 0})) if err != nil { return nil, fmt.Errorf("list new members for new room: %w", err) } defer cur.Close(ctx) - if !cur.Next(ctx) { - // Distinguish a true empty result from a cursor/read failure — the - // caller must not proceed to room creation when membership resolution - // silently failed. - if err := cur.Err(); err != nil { - return nil, fmt.Errorf("iterate aggregation result: %w", err) - } - return nil, nil + var rows []struct { + Account string `bson:"account"` } - var doc struct { - Accounts []string `bson:"accounts"` + if err := cur.All(ctx, &rows); err != nil { + return nil, fmt.Errorf("decode new members for new room: %w", err) } - if err := cur.Decode(&doc); err != nil { - return nil, fmt.Errorf("decode aggregation result: %w", err) + if len(rows) == 0 { + return nil, nil + } + // account is unique in the users collection, so rows carry no duplicates. + out := make([]string, len(rows)) + for i, r := range rows { + out[i] = r.Account } - return doc.Accounts, nil + return out, nil } func (s *MongoStore) GetSubscription(ctx context.Context, account, roomID string) (*model.Subscription, error) { @@ -438,22 +436,74 @@ func (s *MongoStore) ListAddMemberCandidates(ctx context.Context, orgIDs, direct if len(orgIDs) == 0 && len(directAccounts) == 0 { return nil, nil } - pipeline, err := pipelines.GetAddMemberCandidatesPipeline(orgIDs, directAccounts, roomID, "") + // 1. Resolve the candidate users (account + id) with one indexed find on + // users, instead of a correlated-$lookup aggregation. + filter := pipelines.MatchCandidatesFilter(orgIDs, directAccounts, "") + cursor, err := s.users.Find(ctx, filter, options.Find().SetProjection(bson.M{"account": 1, "_id": 1})) if err != nil { - return nil, fmt.Errorf("build add-member candidates pipeline: %w", err) + return nil, fmt.Errorf("find add-member candidates: %w", err) } - cursor, err := s.users.Aggregate(ctx, pipeline) - if err != nil { - return nil, fmt.Errorf("aggregate add-member candidates: %w", err) + var candidates []struct { + ID string `bson:"_id"` + Account string `bson:"account"` } - defer cursor.Close(ctx) - var out []AddMemberCandidate - if err := cursor.All(ctx, &out); err != nil { + if err := cursor.All(ctx, &candidates); err != nil { return nil, fmt.Errorf("decode add-member candidates: %w", err) } + if len(candidates) == 0 { + return nil, nil + } + accounts := make([]string, len(candidates)) + ids := make([]string, len(candidates)) + for i, c := range candidates { + accounts[i] = c.Account + ids[i] = c.ID + } + // 2 & 3. Per-candidate membership state via two indexed reads scoped to the + // candidate set ($in), so examined keys stay bounded by the request size + // rather than the room size. + subbed, err := pipelines.SubscribedAccounts(ctx, s.subscriptions, roomID, accounts) + if err != nil { + return nil, err + } + irm, err := s.individualMemberIDs(ctx, roomID, ids) + if err != nil { + return nil, err + } + out := make([]AddMemberCandidate, len(candidates)) + for i, c := range candidates { + _, hasSub := subbed[c.Account] + _, hasIRM := irm[c.ID] + out[i] = AddMemberCandidate{Account: c.Account, HasSubscription: hasSub, HasIndividualRoomMember: hasIRM} + } return out, nil } +// individualMemberIDs returns the subset of user ids that already have an +// individual room_members row for roomID. Indexed on +// (rid, member.type, member.id). +func (s *MongoStore) individualMemberIDs(ctx context.Context, roomID string, ids []string) (map[string]struct{}, error) { + cursor, err := s.roomMembers.Find(ctx, + bson.M{"rid": roomID, "member.type": string(model.RoomMemberIndividual), "member.id": bson.M{"$in": ids}}, + options.Find().SetProjection(bson.M{"member.id": 1, "_id": 0})) + if err != nil { + return nil, fmt.Errorf("find existing room members for room %q: %w", roomID, err) + } + var rows []struct { + Member struct { + ID string `bson:"id"` + } `bson:"member"` + } + if err := cursor.All(ctx, &rows); err != nil { + return nil, fmt.Errorf("decode existing room members: %w", err) + } + set := make(map[string]struct{}, len(rows)) + for _, r := range rows { + set[r.Member.ID] = struct{}{} + } + return set, nil +} + func (s *MongoStore) GetSubscriptionAccounts(ctx context.Context, roomID string) ([]string, error) { cursor, err := s.subscriptions.Find(ctx, bson.M{"roomId": roomID}, options.Find().SetProjection(bson.M{"u.account": 1, "_id": 0})) From ebd53435ddbb123dada22b0d893e0ca72f055f8c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 00:00:00 +0000 Subject: [PATCH 2/4] room-worker: make add-member count update O(1) with TTL-bounded reconcile MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit processAddMembers ran a full ReconcileMemberCounts on every add, recomputing userCount/appCount from two whole-room CountDocuments. That scales with room size (O(members), ~1000 index entries/add as rooms fill toward MAX_ROOM_SIZE) and dominates MongoDB CPU on the members load test, capping throughput. The hot path now applies the count delta incrementally via an atomic $inc (ApplyMemberCountDelta), and runs the authoritative O(room) recompute only once per room per MEMBER_COUNT_RECONCILE_TTL (default 60s) as a drift safety net. The delta is the net-new subscription set (needSub, recomputed from live state each delivery), so it stays correct under JetStream redelivery; $inc is atomic under concurrent adds. The only residual drift — a crash between the subscription write and the $inc, always an undercount — converges within the TTL. ReconcileMemberCounts stamps countsReconciledAt to reset the TTL clock; the rarer remove/create paths keep the authoritative recompute. TTL=0 forces a recompute on every add (legacy behaviour). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_018ynZvBuRnWgG9moeqhSwkM --- room-worker/handler.go | 31 ++++++++++++++++---- room-worker/handler_test.go | 52 ++++++++++++++++++--------------- room-worker/integration_test.go | 37 +++++++++++++++++++++++ room-worker/main.go | 8 +++++ room-worker/mock_store_test.go | 15 ++++++++++ room-worker/store.go | 9 +++++- room-worker/store_mongo.go | 41 ++++++++++++++++++++++++-- 7 files changed, 160 insertions(+), 33 deletions(-) diff --git a/room-worker/handler.go b/room-worker/handler.go index 0ffb9e81c..4dc4225e7 100644 --- a/room-worker/handler.go +++ b/room-worker/handler.go @@ -67,6 +67,10 @@ type Handler struct { // after authoritative writes. nil disables invalidation (best-effort). valkey valkeyutil.Client keyFanoutWorkers int + // reconcileTTL bounds how often the add-member hot path runs a full + // O(room) member-count recompute; see config.MemberCountReconcileTTL. + // Zero means recompute on every add (the pre-optimisation behaviour). + reconcileTTL time.Duration } func NewHandler(store SubscriptionStore, siteID string, publish PublishFunc, keyStore RoomKeyStore, keySender *roomkeysender.Sender) *Handler { @@ -995,11 +999,28 @@ func (h *Handler) processAddMembers(ctx context.Context, data []byte) (err error } } - // 6. Reconcile userCount. Idempotent $set converges to the correct value - // even under JetStream redelivery; an upstream log-and-continue would - // leave the counter drifted forever, so we propagate the error. - if err := h.store.ReconcileMemberCounts(ctx, req.RoomID); err != nil { - return fmt.Errorf("reconcile member counts: %w", err) + // 6. Update member counts. The hot path applies the delta incrementally + // ($inc, O(1)); a full O(room) recompute runs only once the per-room TTL + // elapses (drift safety net). subs is exactly the net-new subscriptions + // (needSub, recomputed from live state each delivery), so the delta is + // redelivery-safe. Errors propagate — log-and-continue would drift the + // counter. + var userDelta, appDelta int + for _, sub := range subs { + if sub.User.IsBot { + appDelta++ + } else { + userDelta++ + } + } + reconcileDue, err := h.store.ApplyMemberCountDelta(ctx, req.RoomID, userDelta, appDelta, h.reconcileTTL) + if err != nil { + return fmt.Errorf("apply member count delta: %w", err) + } + if reconcileDue { + if err := h.store.ReconcileMemberCounts(ctx, req.RoomID); err != nil { + return fmt.Errorf("reconcile member counts: %w", err) + } } h.bustRoomMeta(ctx, req.RoomID) diff --git a/room-worker/handler_test.go b/room-worker/handler_test.go index ab819a966..57b529447 100644 --- a/room-worker/handler_test.go +++ b/room-worker/handler_test.go @@ -378,6 +378,10 @@ func TestHandler_ProcessAddMembers(t *testing.T) { } return nil }) + // Two non-bot members are created, so the count delta is (userDelta=2, + // appDelta=0). Returning reconcileDue=true exercises the TTL safety-net + // path: the full ReconcileMemberCounts must follow the incremental $inc. + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", 2, 0, gomock.Any()).Return(true, nil) store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) @@ -441,7 +445,7 @@ func TestHandler_ProcessAddMembers_PublishesSubscriptionUpdateBeforeRoomKey(t *t ID: "u1", Account: "alice", SiteID: "site-a", EngName: "Alice", ChineseName: "愛", }, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) req := model.AddMembersRequest{ @@ -496,7 +500,7 @@ func TestHandler_ProcessAddMembers_HistoryAll(t *testing.T) { assert.Nil(t, subs[0].HistorySharedSince, "HistorySharedSince should be nil for mode all") return nil }) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) req := model.AddMembersRequest{ @@ -558,7 +562,7 @@ func TestHandler_ProcessAddMembers_RestrictedPropagatesPointer(t *testing.T) { ID: "u1", Account: "alice", SiteID: "site-a", EngName: "Alice", ChineseName: "愛", }, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) const reqTS int64 = 1744300000000 @@ -621,7 +625,7 @@ func TestHandler_ProcessAddMembers_UnrestrictedOmitsFieldFromWire(t *testing.T) ID: "u1", Account: "alice", SiteID: "site-a", EngName: "Alice", ChineseName: "愛", }, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) req := model.AddMembersRequest{ @@ -660,7 +664,7 @@ func TestHandler_ProcessAddMembers_WithOrgs(t *testing.T) { ID: "u1", Account: "alice", SiteID: "site-a", EngName: "Alice", ChineseName: "愛", }, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) // HasOrgRoomMembers is now called unconditionally (Task 2). Return false to // preserve this test's first-org-transition semantics so backfill fires. store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) @@ -803,7 +807,7 @@ func TestHandler_ProcessAddMembers_MultipleSiteInbox(t *testing.T) { ID: "u1", Account: "alice", SiteID: "site-b", EngName: "Alice", ChineseName: "愛", }, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) req := model.AddMembersRequest{ @@ -1114,7 +1118,7 @@ func TestHandler_ProcessAddMembers_ExistingOrgsWritesIndividuals(t *testing.T) { ID: "u1", Account: "alice", SiteID: "site-a", EngName: "Alice", ChineseName: "愛", }, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(true, nil) store.EXPECT().BulkCreateRoomMembers(gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, members []*model.RoomMember) error { @@ -1177,7 +1181,7 @@ func TestHandler_ProcessAddMembers_OrgToIndividualUpgrade(t *testing.T) { assert.Equal(t, "u_alice", members[0].Member.ID) return nil }) - store.EXPECT().ReconcileMemberCounts(ctx, roomID).Return(nil) + store.EXPECT().ApplyMemberCountDelta(ctx, roomID, gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) var published []publishedMsg h := &Handler{ @@ -1322,7 +1326,7 @@ func TestHandler_processAddMembers_PublishesSuccessEventToRequesterSubject(t *te ID: "u1", Account: "alice", SiteID: "site1", EngName: "Alice", ChineseName: "愛", }, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) ctx := natsutil.WithRequestID(context.Background(), testRequestID) @@ -1494,7 +1498,7 @@ func setupAddMembersHappyPath(t *testing.T, mockStore *MockSubscriptionStore, ac }, nil) mockStore.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) mockStore.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) - mockStore.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + mockStore.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) } // The legacy TestProcessAddMembers_RequiresRequestID test pinned the old @@ -1550,7 +1554,7 @@ func TestProcessAddMembers_PopulatesSubName(t *testing.T) { return nil }) mockStore.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) - mockStore.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + mockStore.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) body, err := json.Marshal(model.AddMembersRequest{ RoomID: "r1", Users: []string{"bob"}, @@ -1590,7 +1594,7 @@ func TestProcessAddMembers_HistoryNone_NoTimestamp(t *testing.T) { return nil }) mockStore.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) - mockStore.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + mockStore.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) // No SharedSince in HistoryConfig — falls back to req.Timestamp (acceptedAt). body, err := json.Marshal(model.AddMembersRequest{ @@ -1630,7 +1634,7 @@ func TestProcessAddMembers_NoHistoryConfig_LeavesNil(t *testing.T) { return nil }) mockStore.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) - mockStore.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + mockStore.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) // No History.Mode — HistorySharedSince must remain nil. body, err := json.Marshal(model.AddMembersRequest{ @@ -1665,7 +1669,7 @@ func TestProcessAddMembers_InboxCarriesRoomName(t *testing.T) { }, nil) mockStore.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) mockStore.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) - mockStore.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + mockStore.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) body, err := json.Marshal(model.AddMembersRequest{ RoomID: "r1", Users: []string{"bob"}, @@ -3518,7 +3522,7 @@ func TestProcessAddMembers_FansOutKeyToNewAccountsOnly(t *testing.T) { }, nil) mockStore.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) mockStore.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) - mockStore.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + mockStore.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) pair := &roomkeystore.VersionedKeyPair{ Version: 1, @@ -3560,7 +3564,7 @@ func TestProcessAddMembers_PermanentErrorWhenKeyMissing(t *testing.T) { }, nil) mockStore.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) mockStore.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) - mockStore.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + mockStore.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) keyStore.EXPECT().Get(gomock.Any(), "r1").Return(nil, nil) // key missing h := NewHandler(mockStore, "site-a", func(_ context.Context, _ string, _ []byte, _ string) error { return nil }, keyStore, roomkeysender.NewSender(&mockPublisher{})) @@ -3596,7 +3600,7 @@ func TestProcessAddMembers_TransientErrorWhenValkeyFails(t *testing.T) { }, nil) mockStore.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) mockStore.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) - mockStore.EXPECT().ReconcileMemberCounts(gomock.Any(), "r1").Return(nil) + mockStore.EXPECT().ApplyMemberCountDelta(gomock.Any(), "r1", gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) keyStore.EXPECT().Get(gomock.Any(), "r1").Return(nil, fmt.Errorf("valkey timeout")) h := NewHandler(mockStore, "site-a", func(_ context.Context, _ string, _ []byte, _ string) error { return nil }, keyStore, roomkeysender.NewSender(&mockPublisher{})) @@ -3697,7 +3701,7 @@ func TestHandler_ProcessAddMembers_BackfillRunsOnFirstOrgTransition(t *testing.T store.EXPECT().FindUsersByAccounts(gomock.Any(), []string{"existing_user"}). Return([]model.User{{ID: "u_e", Account: "existing_user", SiteID: "site-a", EngName: "Ex", ChineseName: "存"}}, nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), roomID).Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), roomID, gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) h := &Handler{store: store, siteID: "site-a", publish: func(_ context.Context, _ string, _ []byte, _ string) error { return nil }, keyStore: testKeyStore, keySender: testKeySender} @@ -3796,7 +3800,7 @@ func TestHandler_ProcessAddMembers_BackfillSkippedWhenRoomAlreadyHasOrgs(t *test store.EXPECT().BulkCreateRoomMembers(gomock.Any(), gomock.Any()).Return(nil) // NO GetSubscriptionAccounts expectation — backfill must be skipped. - store.EXPECT().ReconcileMemberCounts(gomock.Any(), roomID).Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), roomID, gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) h := &Handler{store: store, siteID: "site-a", publish: func(_ context.Context, _ string, _ []byte, _ string) error { return nil }, keyStore: testKeyStore, keySender: testKeySender} @@ -3841,7 +3845,7 @@ func TestHandler_ProcessAddMembers_IndividualFilter_DirectAndOrgOverlap(t *testi captured = m return nil }) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), roomID).Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), roomID, gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) h := &Handler{store: store, siteID: "site-a", publish: func(_ context.Context, _ string, _ []byte, _ string) error { return nil }, keyStore: testKeyStore, keySender: testKeySender} @@ -3892,7 +3896,7 @@ func TestHandler_ProcessAddMembers_IndividualFilter_OrgOnly(t *testing.T) { captured = m return nil }) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), roomID).Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), roomID, gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) h := &Handler{store: store, siteID: "site-a", publish: func(_ context.Context, _ string, _ []byte, _ string) error { return nil }, keyStore: testKeyStore, keySender: testKeySender} @@ -4039,7 +4043,7 @@ func TestHandler_ProcessAddMembers_Content_Single(t *testing.T) { Return(&model.User{ID: "u_a", Account: "alice", SiteID: "site-a", EngName: "Alice", ChineseName: "愛"}, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), roomID).Return(false, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), roomID).Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), roomID, gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) // No BulkCreateRoomMembers expected (no orgs, no pre-existing orgs → lite-mode add). var published []publishedMsg @@ -4079,7 +4083,7 @@ func TestHandler_ProcessAddMembers_Content_Multi(t *testing.T) { Return(&model.User{ID: "u_a", Account: "alice", SiteID: "site-a", EngName: "Alice", ChineseName: "愛"}, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), roomID).Return(false, nil) store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), roomID).Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), roomID, gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) var published []publishedMsg h := &Handler{store: store, siteID: "site-a", publish: func(_ context.Context, subj string, data []byte, _ string) error { @@ -4469,7 +4473,7 @@ func TestHandler_ProcessAddMembers_Content_OrgAddWithOneMember_UsesMulti(t *test store.EXPECT().BulkCreateSubscriptions(gomock.Any(), gomock.Any()).Return(nil) store.EXPECT().BulkCreateRoomMembers(gomock.Any(), gomock.Any()).Return(nil) store.EXPECT().GetSubscriptionAccounts(gomock.Any(), roomID).Return([]string{}, nil) - store.EXPECT().ReconcileMemberCounts(gomock.Any(), roomID).Return(nil) + store.EXPECT().ApplyMemberCountDelta(gomock.Any(), roomID, gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil) var published []publishedMsg h := &Handler{store: store, siteID: "site-a", publish: func(_ context.Context, subj string, data []byte, _ string) error { diff --git a/room-worker/integration_test.go b/room-worker/integration_test.go index acbd724fd..1b64bc3ca 100644 --- a/room-worker/integration_test.go +++ b/room-worker/integration_test.go @@ -128,6 +128,43 @@ func TestMongoStore_Integration(t *testing.T) { } } +// TestMongoStore_ApplyMemberCountDelta covers the add-member hot path: the $inc +// applies the delta atomically, and reconcileDue follows the per-room TTL — +// true when never reconciled or once the TTL has elapsed, false within it. +func TestMongoStore_ApplyMemberCountDelta(t *testing.T) { + db := setupMongo(t) + store := NewMongoStore(db) + ctx := context.Background() + + _, err := db.Collection("rooms").InsertOne(ctx, model.Room{ID: "rd", Name: "delta", UserCount: 5, AppCount: 1}) + require.NoError(t, err) + + // Never reconciled → due; $inc applies userDelta/appDelta. + due, err := store.ApplyMemberCountDelta(ctx, "rd", 3, 2, time.Minute) + require.NoError(t, err) + assert.True(t, due, "a room never reconciled should be due") + room, err := store.GetRoom(ctx, "rd") + require.NoError(t, err) + assert.Equal(t, 8, room.UserCount, "userCount incremented by 3") + assert.Equal(t, 3, room.AppCount, "appCount incremented by 2") + + // ReconcileMemberCounts stamps countsReconciledAt → not due within the TTL. + require.NoError(t, store.ReconcileMemberCounts(ctx, "rd")) + due, err = store.ApplyMemberCountDelta(ctx, "rd", 0, 0, time.Minute) + require.NoError(t, err) + assert.False(t, due, "within TTL the recompute should not be due") + + // ttl=0 forces a recompute every time (legacy behaviour). + due, err = store.ApplyMemberCountDelta(ctx, "rd", 0, 0, 0) + require.NoError(t, err) + assert.True(t, due, "ttl=0 is always due") + + // Missing room: no-op, not due, no error (matches ReconcileMemberCounts). + due, err = store.ApplyMemberCountDelta(ctx, "missing", 1, 0, time.Minute) + require.NoError(t, err) + assert.False(t, due) +} + func TestMongoStore_GetSubscription_Integration(t *testing.T) { db := setupMongo(t) store := NewMongoStore(db) diff --git a/room-worker/main.go b/room-worker/main.go index 660a479c7..a47a7018a 100644 --- a/room-worker/main.go +++ b/room-worker/main.go @@ -52,6 +52,13 @@ type config struct { // Grace window during which a rotated-out previous key remains valid for decrypt. RoomKeyGracePeriod time.Duration `env:"ROOM_KEY_GRACE_PERIOD" envDefault:"24h"` + // MemberCountReconcileTTL bounds how often the add-member hot path runs a + // full O(room) recompute of userCount/appCount. Between recomputes the + // counts are maintained incrementally ($inc by the actual delta); a full + // reconcile runs at most once per room per TTL as a drift safety net. 0 + // forces a recompute on every add (legacy behaviour). + MemberCountReconcileTTL time.Duration `env:"MEMBER_COUNT_RECONCILE_TTL" envDefault:"60s"` + // Valkey backs best-effort room-meta L2 cache invalidation. Optional: when // VALKEY_ADDRS is empty the bust is a no-op (the L2 TTL reconciles). ValkeyAddrs []string `env:"VALKEY_ADDRS" envSeparator:","` @@ -166,6 +173,7 @@ func main() { handler.SetKeyFanoutWorkers(cfg.KeyFanoutWorkers) handler.dekProvisioner = dekProvisioner handler.valkey = metaValkey + handler.reconcileTTL = cfg.MemberCountReconcileTTL router := natsrouter.New(nc, "room-worker") router.Use(natsrouter.Recovery(), natsrouter.RequestID(), natsrouter.Logging()) diff --git a/room-worker/mock_store_test.go b/room-worker/mock_store_test.go index 275b8509d..577446f13 100644 --- a/room-worker/mock_store_test.go +++ b/room-worker/mock_store_test.go @@ -43,6 +43,21 @@ func (m *MockSubscriptionStore) EXPECT() *MockSubscriptionStoreMockRecorder { return m.recorder } +// ApplyMemberCountDelta mocks base method. +func (m *MockSubscriptionStore) ApplyMemberCountDelta(ctx context.Context, roomID string, userDelta, appDelta int, ttl time.Duration) (bool, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ApplyMemberCountDelta", ctx, roomID, userDelta, appDelta, ttl) + ret0, _ := ret[0].(bool) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ApplyMemberCountDelta indicates an expected call of ApplyMemberCountDelta. +func (mr *MockSubscriptionStoreMockRecorder) ApplyMemberCountDelta(ctx, roomID, userDelta, appDelta, ttl any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ApplyMemberCountDelta", reflect.TypeOf((*MockSubscriptionStore)(nil).ApplyMemberCountDelta), ctx, roomID, userDelta, appDelta, ttl) +} + // BulkCreateRoomMembers mocks base method. func (m *MockSubscriptionStore) BulkCreateRoomMembers(ctx context.Context, members []*model.RoomMember) error { m.ctrl.T.Helper() diff --git a/room-worker/store.go b/room-worker/store.go index 8cebb718e..fdb92e678 100644 --- a/room-worker/store.go +++ b/room-worker/store.go @@ -67,8 +67,15 @@ type SubscriptionStore interface { // ReconcileMemberCounts recomputes Room.UserCount (non-bot subs) and // 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. + // update. It also stamps countsReconciledAt, resetting the per-room TTL + // clock that ApplyMemberCountDelta consults. ReconcileMemberCounts(ctx context.Context, roomID string) error + // ApplyMemberCountDelta atomically applies userDelta/appDelta to the room's + // counters ($inc, O(1)) and reports whether a full ReconcileMemberCounts is + // now due — i.e. the room has never been reconciled or its countsReconciledAt + // is older than ttl. This keeps the add-member hot path O(1) while a bounded + // periodic recompute (the drift safety net) restores convergence. + ApplyMemberCountDelta(ctx context.Context, roomID string, userDelta, appDelta int, ttl time.Duration) (reconcileDue bool, err error) GetRoom(ctx context.Context, roomID string) (*model.Room, error) GetSubscription(ctx context.Context, account, roomID string) (*model.Subscription, error) GetUser(ctx context.Context, account string) (*model.User, error) diff --git a/room-worker/store_mongo.go b/room-worker/store_mongo.go index aa8356f9a..f8332a075 100644 --- a/room-worker/store_mongo.go +++ b/room-worker/store_mongo.go @@ -70,11 +70,13 @@ func (s *MongoStore) ReconcileMemberCounts(ctx context.Context, roomID string) e return fmt.Errorf("count app subscriptions: %w", err) } + now := time.Now().UTC() if _, err := s.rooms.UpdateOne(ctx, bson.M{"_id": roomID}, bson.M{ "$set": bson.M{ - "userCount": total - appCount, - "appCount": appCount, - "updatedAt": time.Now().UTC(), + "userCount": total - appCount, + "appCount": appCount, + "countsReconciledAt": now, + "updatedAt": now, }, }); err != nil { return fmt.Errorf("update room counts: %w", err) @@ -82,6 +84,39 @@ func (s *MongoStore) ReconcileMemberCounts(ctx context.Context, roomID string) e return nil } +// ApplyMemberCountDelta atomically $inc's userCount/appCount and reports whether +// a full recompute is now due. The $inc is O(1); the delta is the actual number +// of subscriptions created/removed by the caller, so it stays correct under +// JetStream redelivery (net-new is recomputed from live state each delivery) and +// concurrent adds ($inc is atomic). The returned countsReconciledAt drives the +// per-room TTL: a crash between the subscription write and this $inc, or a rare +// racing duplicate, can leave the counter slightly drifted until the next +// reconcile, which this method schedules once the TTL elapses. +func (s *MongoStore) ApplyMemberCountDelta(ctx context.Context, roomID string, userDelta, appDelta int, ttl time.Duration) (bool, error) { + var doc struct { + CountsReconciledAt time.Time `bson:"countsReconciledAt"` + } + err := s.rooms.FindOneAndUpdate(ctx, + bson.M{"_id": roomID}, + bson.M{ + "$inc": bson.M{"userCount": userDelta, "appCount": appDelta}, + "$set": bson.M{"updatedAt": time.Now().UTC()}, + }, + options.FindOneAndUpdate(). + SetReturnDocument(options.After). + SetProjection(bson.M{"countsReconciledAt": 1}), + ).Decode(&doc) + if err != nil { + if errors.Is(err, mongo.ErrNoDocuments) { + // No such room — nothing to count, nothing to reconcile (matches + // ReconcileMemberCounts' no-op-on-missing UpdateOne). + return false, nil + } + return false, fmt.Errorf("apply member count delta for room %q: %w", roomID, err) + } + return doc.CountsReconciledAt.IsZero() || time.Since(doc.CountsReconciledAt) > ttl, nil +} + func (s *MongoStore) GetRoom(ctx context.Context, roomID string) (*model.Room, error) { var room model.Room if err := s.rooms.FindOne(ctx, bson.M{"_id": roomID}).Decode(&room); err != nil { From 6f7ea9cd15a57ab488b4e8b1e03a80d95dd23c59 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 12:21:40 +0000 Subject: [PATCH 3/4] room-worker: cache user reads and add direct-account candidate fast path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every add resolves users several times (GetUser for the requester, FindUsersByAccounts for the new members, plus candidate resolution), all uncached Mongo round-trips — and under the members load test those queue behind a saturated mongod, throttling room-worker's ROOMS-consumer throughput (the E2 backlog). Front the store's user reads with the existing pkg/userstore.Cache (in-process LRU+TTL, already used by broadcast-worker/message-worker). The store now holds a userReader, defaulting to an uncached userstore read (behaviour-identical, so the integration tests are unchanged); EnableUserCache wraps it in the cache at startup, configured via USER_CACHE_SIZE/USER_CACHE_TTL (10000 / 5m). Safe because the fields room-worker consumes (id, account, siteId, names) are within userstore's projection and are effectively immutable; ErrUserNotFound is not negatively cached, and is mapped back to the store's sentinel. Additionally, the add-member candidate resolution now takes a fast path when no orgs are involved: resolve directly from the (cached) user reader and filter bots in Go via model.IsBotAccount, eliminating both the users query and the $not regex for the common direct-add case. Org adds keep the sectId/deptId query, which a by-account reader cannot serve. room-service is intentionally left for a follow-up: its authorizeRoomAppRead reads user.Roles via IsPlatformAdmin, which userstore's projection omits, so caching there first needs the projection widened. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_018ynZvBuRnWgG9moeqhSwkM --- room-worker/integration_test.go | 31 ++++++++++++ room-worker/main.go | 13 +++++ room-worker/store_mongo.go | 88 ++++++++++++++++++++++----------- 3 files changed, 103 insertions(+), 29 deletions(-) diff --git a/room-worker/integration_test.go b/room-worker/integration_test.go index 1b64bc3ca..eb4976f94 100644 --- a/room-worker/integration_test.go +++ b/room-worker/integration_test.go @@ -511,6 +511,37 @@ func mustInsertUser(t *testing.T, db *mongo.Database, u *model.User) { require.NoError(t, err) } +// TestMongoStore_UserCache covers the cached user-read path: EnableUserCache +// wraps the reader without breaking reads, and a missing account still maps to +// the store's ErrUserNotFound sentinel that handlers branch on. +func TestMongoStore_UserCache(t *testing.T) { + db := setupMongo(t) + store := NewMongoStore(db) + require.NoError(t, store.EnableUserCache(100, time.Minute)) + ctx := context.Background() + + _, err := db.Collection("users").InsertOne(ctx, model.User{ID: "u1", Account: "alice", SiteID: "site-a", EngName: "Alice"}) + require.NoError(t, err) + + // Read through the cache twice; both return the seeded user. + for range 2 { + u, err := store.GetUser(ctx, "alice") + require.NoError(t, err) + assert.Equal(t, "u1", u.ID) + assert.Equal(t, "alice", u.Account) + } + + // Missing account maps to the store's ErrUserNotFound (not userstore's). + _, err = store.GetUser(ctx, "ghost") + assert.ErrorIs(t, err, ErrUserNotFound) + + // Batch lookup returns only existing users (negative results not cached). + users, err := store.FindUsersByAccounts(ctx, []string{"alice", "ghost"}) + require.NoError(t, err) + require.Len(t, users, 1) + assert.Equal(t, "alice", users[0].Account) +} + func TestMongoStore_ListAddMemberCandidates_Integration(t *testing.T) { ctx := context.Background() db := setupMongo(t) diff --git a/room-worker/main.go b/room-worker/main.go index a47a7018a..eff1ae270 100644 --- a/room-worker/main.go +++ b/room-worker/main.go @@ -42,6 +42,8 @@ type config struct { MongoPassword string `env:"MONGO_PASSWORD" envDefault:""` MaxWorkers int `env:"MAX_WORKERS" envDefault:"100"` KeyFanoutWorkers int `env:"KEY_FANOUT_WORKERS" envDefault:"32"` // see defaultKeyFanoutWorkers in handler.go + UserCacheSize int `env:"USER_CACHE_SIZE" envDefault:"10000"` + UserCacheTTL time.Duration `env:"USER_CACHE_TTL" envDefault:"5m"` Consumer stream.ConsumerSettings `envPrefix:"CONSUMER_"` Bootstrap bootstrapConfig `envPrefix:"BOOTSTRAP_"` HealthAddr string `env:"HEALTH_ADDR" envDefault:":8081"` @@ -155,6 +157,17 @@ func main() { streamCfg := stream.Rooms(cfg.SiteID) store := NewMongoStore(mongoClient.Database(cfg.MongoDB)) + // User cache is on by default; a non-positive size disables it cleanly rather + // than failing startup (the LRU constructor rejects size<=0). + if cfg.UserCacheSize > 0 { + if err := store.EnableUserCache(cfg.UserCacheSize, cfg.UserCacheTTL); err != nil { + slog.Error("failed to enable user cache", "error", err) + os.Exit(1) + } + slog.Info("user-cache enabled", "size", cfg.UserCacheSize, "ttl", cfg.UserCacheTTL) + } else { + slog.Info("user-cache disabled", "size", cfg.UserCacheSize) + } handler := NewHandler(store, cfg.SiteID, func(ctx context.Context, subj string, data []byte, msgID string) error { msg := natsutil.NewMsg(ctx, subj, data) if msgID == "" { diff --git a/room-worker/store_mongo.go b/room-worker/store_mongo.go index f8332a075..d3de7bb9d 100644 --- a/room-worker/store_mongo.go +++ b/room-worker/store_mongo.go @@ -14,6 +14,7 @@ import ( "github.com/hmchangw/chat/pkg/mongoutil" "github.com/hmchangw/chat/pkg/pipelines" "github.com/hmchangw/chat/pkg/roomkeystore" + "github.com/hmchangw/chat/pkg/userstore" ) type MongoStore struct { @@ -22,6 +23,12 @@ type MongoStore struct { roomMembers *mongo.Collection users *mongo.Collection apps *mongo.Collection + // userReader serves point lookups (GetUser, FindUsersByAccounts) and the + // direct-account candidate fast path. Defaults to an uncached read of the + // users collection; EnableUserCache wraps it in an LRU+TTL cache. The + // `users` collection above is still used directly for the org-expansion + // candidate query, which the by-account reader cannot serve. + userReader userstore.UserStore } func NewMongoStore(db *mongo.Database) *MongoStore { @@ -31,9 +38,23 @@ func NewMongoStore(db *mongo.Database) *MongoStore { roomMembers: db.Collection("room_members"), users: db.Collection("users"), apps: db.Collection("apps"), + userReader: userstore.NewMongoStore(db.Collection("users")), } } +// EnableUserCache wraps the store's user reader in an in-process LRU+TTL cache. +// Call once at startup; user docs are near-static (account→id and isBot are +// immutable, org fields change rarely) so a short TTL bounds staleness. Reads +// fall through to Mongo on miss; ErrUserNotFound is not negatively cached. +func (s *MongoStore) EnableUserCache(size int, ttl time.Duration) error { + cache, err := userstore.NewCache(s.userReader, size, ttl) + if err != nil { + return fmt.Errorf("enable user cache: %w", err) + } + s.userReader = cache + return nil +} + // ListByRoom returns all subscriptions for roomID across every site. Not part // of SubscriptionStore — the handler's hot paths only need accounts (see // GetSubscriptionAccounts); this full-document read is retained for integration @@ -126,15 +147,14 @@ func (s *MongoStore) GetRoom(ctx context.Context, roomID string) (*model.Room, e } func (s *MongoStore) GetUser(ctx context.Context, account string) (*model.User, error) { - var u model.User - err := s.users.FindOne(ctx, bson.M{"account": account}).Decode(&u) - if errors.Is(err, mongo.ErrNoDocuments) { + u, err := s.userReader.FindUserByAccount(ctx, account) + if errors.Is(err, userstore.ErrUserNotFound) { return nil, ErrUserNotFound } if err != nil { return nil, fmt.Errorf("get user %q: %w", account, err) } - return &u, nil + return u, nil } // GetApp reads the apps collection, which room-service owns and indexes @@ -445,18 +465,7 @@ func (s *MongoStore) BulkCreateRoomMembers(ctx context.Context, members []*model } func (s *MongoStore) FindUsersByAccounts(ctx context.Context, accounts []string) ([]model.User, error) { - if len(accounts) == 0 { - return nil, nil - } - cursor, err := s.users.Find(ctx, bson.M{"account": bson.M{"$in": accounts}}) - if err != nil { - return nil, fmt.Errorf("find users by accounts: %w", err) - } - var users []model.User - if err := cursor.All(ctx, &users); err != nil { - return nil, fmt.Errorf("decode users: %w", err) - } - return users, nil + return s.userReader.FindUsersByAccounts(ctx, accounts) } func (s *MongoStore) HasOrgRoomMembers(ctx context.Context, roomID string) (bool, error) { @@ -471,19 +480,40 @@ func (s *MongoStore) ListAddMemberCandidates(ctx context.Context, orgIDs, direct if len(orgIDs) == 0 && len(directAccounts) == 0 { return nil, nil } - // 1. Resolve the candidate users (account + id) with one indexed find on - // users, instead of a correlated-$lookup aggregation. - filter := pipelines.MatchCandidatesFilter(orgIDs, directAccounts, "") - cursor, err := s.users.Find(ctx, filter, options.Find().SetProjection(bson.M{"account": 1, "_id": 1})) - if err != nil { - return nil, fmt.Errorf("find add-member candidates: %w", err) - } - var candidates []struct { - ID string `bson:"_id"` - Account string `bson:"account"` - } - if err := cursor.All(ctx, &candidates); err != nil { - return nil, fmt.Errorf("decode add-member candidates: %w", err) + // 1. Resolve the candidate users (account + id). + type candidate struct{ ID, Account string } + var candidates []candidate + if len(orgIDs) == 0 { + // Direct accounts only: resolve via the user reader (cache-friendly) and + // filter bots in Go — avoids the users query and the $not regex entirely. + users, err := s.userReader.FindUsersByAccounts(ctx, directAccounts) + if err != nil { + return nil, fmt.Errorf("find add-member candidate users: %w", err) + } + for i := range users { + if model.IsBot(users[i].Account) || model.IsPlatformAdminAccount(users[i].Account) { + continue + } + candidates = append(candidates, candidate{ID: users[i].ID, Account: users[i].Account}) + } + } else { + // Org expansion needs the sectId/deptId query on the users collection, + // which the by-account reader cannot serve. + filter := pipelines.MatchCandidatesFilter(orgIDs, directAccounts, "") + cursor, err := s.users.Find(ctx, filter, options.Find().SetProjection(bson.M{"account": 1, "_id": 1})) + if err != nil { + return nil, fmt.Errorf("find add-member candidates: %w", err) + } + var rows []struct { + ID string `bson:"_id"` + Account string `bson:"account"` + } + if err := cursor.All(ctx, &rows); err != nil { + return nil, fmt.Errorf("decode add-member candidates: %w", err) + } + for _, r := range rows { + candidates = append(candidates, candidate{ID: r.ID, Account: r.Account}) + } } if len(candidates) == 0 { return nil, nil From 3b4cc3c0c5c54c0911872d9b065240e591541125 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 20 Jun 2026 07:13:16 +0000 Subject: [PATCH 4/4] room-worker: add a room-metadata cache for the add path and skip the no-org room_members read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Profiling showed room-worker pinned ~1 core, with time spread evenly across the per-add Mongo calls — i.e. it's CPU-bound on driver overhead (BSON + syscalls) per round-trip, not on any single query. Cutting round-trips per add directly cuts that CPU. Two reductions on the add hot path: - Room metadata: add GetRoomMeta, an add-path read of just the stable fields (ID/Type/Name/SiteID/UserCount), fronted by an in-process LRU+TTL roommetacache (EnableRoomMetaCache, ROOM_META_CACHE_SIZE/TTL). loadAddMemberInputs now calls GetRoomMeta; a rename reflects within the TTL. GetRoom is deliberately left returning the full document and is never cache-served: serverCreateDM's idempotency path reads CreatedAt for sub.JoinedAt and the outbox dedup seed, so a partial cache-served room would zero it. An integration test guards that GetRoom still returns the real CreatedAt with the cache enabled. Default-off in the constructor (tests/ integration read Mongo directly), enabled in main. - individualMemberIDs: skip this room_members read when the request has no orgs. The handler only consults HasIndividualRoomMember when writeIndividuals is true (orgs in the request or pre-existing org rows); the no-org path never reads it, and BulkCreateRoomMembers ignores dup-key, so the rare no-org-add- to-a-formerly-org-room case is covered by idempotent re-inserts. Together these remove ~2 of the ~7 per-add Mongo round-trips for the common direct-add workload, lowering the per-core driver CPU that caps throughput. HasOrgRoomMembers stays — it is load-bearing for writeIndividuals/backfill. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_018ynZvBuRnWgG9moeqhSwkM --- room-worker/handler.go | 4 +- room-worker/handler_test.go | 68 ++++++++++++++-------------- room-worker/integration_test.go | 78 +++++++++++++++++++++++++++++++++ room-worker/loadinputs_test.go | 10 ++--- room-worker/main.go | 47 +++++++++++++------- room-worker/mock_store_test.go | 15 +++++++ room-worker/store.go | 5 +++ room-worker/store_mongo.go | 67 +++++++++++++++++++++++++--- 8 files changed, 230 insertions(+), 64 deletions(-) diff --git a/room-worker/handler.go b/room-worker/handler.go index 4dc4225e7..62d08c82d 100644 --- a/room-worker/handler.go +++ b/room-worker/handler.go @@ -716,7 +716,7 @@ type addMemberInputs struct { hadOrgsBefore bool } -// loadAddMemberInputs runs GetRoom, ListAddMemberCandidates, and +// loadAddMemberInputs runs GetRoomMeta, ListAddMemberCandidates, and // HasOrgRoomMembers concurrently, collapsing three serial Mongo round trips into // one. A plain errgroup.Group (not WithContext) is used deliberately: these are // independent reads, so a failure in one need not cancel the others — matching @@ -730,7 +730,7 @@ func (h *Handler) loadAddMemberInputs(ctx context.Context, req *model.AddMembers g errgroup.Group ) g.Go(func() error { - room, err := h.store.GetRoom(ctx, req.RoomID) + room, err := h.store.GetRoomMeta(ctx, req.RoomID) if err != nil { return fmt.Errorf("get room: %w", err) } diff --git a/room-worker/handler_test.go b/room-worker/handler_test.go index 57b529447..6f47ee67f 100644 --- a/room-worker/handler_test.go +++ b/room-worker/handler_test.go @@ -321,7 +321,7 @@ func TestHandler_ProcessAddMembers_FallsBackToNowOnInvalidTimestamp(t *testing.T // rejection. ctrl := gomock.NewController(t) store := NewMockSubscriptionStore(ctrl) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(nil, fmt.Errorf("db error")) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(nil, fmt.Errorf("db error")) // The three up-front reads now run concurrently, so candidates/has-orgs are // issued alongside GetRoom; their results are discarded once GetRoom errors. store.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1").Return(nil, nil).AnyTimes() @@ -353,7 +353,7 @@ func TestHandler_ProcessAddMembers(t *testing.T) { } h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), nil, []string{"bob", "charlie"}, "r1"). Return([]AddMemberCandidate{ {Account: "bob", HasSubscription: false, HasIndividualRoomMember: false}, @@ -431,7 +431,7 @@ func TestHandler_ProcessAddMembers_PublishesSubscriptionUpdateBeforeRoomKey(t *t } h := NewHandler(store, "site-a", publish, testKeyStore, roomkeysender.NewSender(pub)) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), nil, []string{"bob", "charlie"}, "r1"). Return([]AddMemberCandidate{ {Account: "bob", HasSubscription: false, HasIndividualRoomMember: false}, @@ -485,7 +485,7 @@ func TestHandler_ProcessAddMembers_HistoryAll(t *testing.T) { publish := func(_ context.Context, _ string, _ []byte, _ string) error { return nil } h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), nil, []string{"bob"}, "r1"). Return([]AddMemberCandidate{{Account: "bob"}}, nil) store.EXPECT().FindUsersByAccounts(gomock.Any(), []string{"bob"}).Return([]model.User{ @@ -549,7 +549,7 @@ func TestHandler_ProcessAddMembers_RestrictedPropagatesPointer(t *testing.T) { } h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), nil, []string{"bob", "charlie"}, "r1"). Return([]AddMemberCandidate{ {Account: "bob"}, {Account: "charlie"}, @@ -615,7 +615,7 @@ func TestHandler_ProcessAddMembers_UnrestrictedOmitsFieldFromWire(t *testing.T) } h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), nil, []string{"bob"}, "r1"). Return([]AddMemberCandidate{{Account: "bob"}}, nil) store.EXPECT().FindUsersByAccounts(gomock.Any(), []string{"bob"}).Return([]model.User{ @@ -654,7 +654,7 @@ func TestHandler_ProcessAddMembers_WithOrgs(t *testing.T) { publish := func(_ context.Context, _ string, _ []byte, _ string) error { return nil } h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"eng"}, []string{"bob"}, "r1"). Return([]AddMemberCandidate{{Account: "bob"}}, nil) store.EXPECT().FindUsersByAccounts(gomock.Any(), []string{"bob"}).Return([]model.User{ @@ -713,7 +713,7 @@ func TestHandler_ProcessAddMembers_BackfillUserMissing(t *testing.T) { h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) // Org-only add on a room with no prior orgs → backfill fires. - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"eng"}, nil, "r1"). Return([]AddMemberCandidate{}, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) @@ -755,7 +755,7 @@ func TestHandler_ProcessAddMembers_UserNotFound(t *testing.T) { publish := func(_ context.Context, _ string, _ []byte, _ string) error { return nil } h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), nil, []string{"bob", "ghost"}, "r1"). Return([]AddMemberCandidate{{Account: "bob"}, {Account: "ghost"}}, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) @@ -793,7 +793,7 @@ func TestHandler_ProcessAddMembers_MultipleSiteInbox(t *testing.T) { } h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), nil, []string{"alice", "bob", "charlie"}, "r1"). Return([]AddMemberCandidate{ {Account: "alice"}, {Account: "bob"}, {Account: "charlie"}, @@ -1108,7 +1108,7 @@ func TestHandler_ProcessAddMembers_ExistingOrgsWritesIndividuals(t *testing.T) { publish := func(_ context.Context, _ string, _ []byte, _ string) error { return nil } h := NewHandler(store, "site-a", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site-a"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), nil, []string{"bob"}, "r1"). Return([]AddMemberCandidate{{Account: "bob"}}, nil) store.EXPECT().FindUsersByAccounts(gomock.Any(), []string{"bob"}).Return([]model.User{ @@ -1157,7 +1157,7 @@ func TestHandler_ProcessAddMembers_OrgToIndividualUpgrade(t *testing.T) { requestID := idgen.GenerateRequestID() ctx := natsutil.WithRequestID(context.Background(), requestID) - store.EXPECT().GetRoom(ctx, roomID).Return(&model.Room{ + store.EXPECT().GetRoomMeta(ctx, roomID).Return(&model.Room{ ID: roomID, Type: model.RoomTypeChannel, SiteID: "site-a", Name: "Room 1", }, nil) // Alice has a subscription (added earlier via org) but no individual row. @@ -1316,7 +1316,7 @@ func TestHandler_processAddMembers_PublishesSuccessEventToRequesterSubject(t *te } h := NewHandler(store, "site1", publish, testKeyStore, testKeySender) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site1"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site1"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), []string{"bob"}, "r1"). Return([]AddMemberCandidate{{Account: "bob"}}, nil) store.EXPECT().FindUsersByAccounts(gomock.Any(), []string{"bob"}).Return([]model.User{ @@ -1366,7 +1366,7 @@ func TestHandler_processAddMembers_PublishesFailureEventOnError(t *testing.T) { h := NewHandler(store, "site1", publish, testKeyStore, testKeySender) // Mock store to fail on FindUsersByAccounts (first store operation after candidate resolution) - store.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site1"}, nil) + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel, SiteID: "site1"}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), []string{"bob"}, "r1"). Return([]AddMemberCandidate{{Account: "bob"}}, nil) store.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) @@ -1479,7 +1479,7 @@ func newAddMembersTestHandler(t *testing.T) (*Handler, *MockSubscriptionStore, f // All users are on site-A (no cross-site inbox). HasOrgRoomMembers returns false. func setupAddMembersHappyPath(t *testing.T, mockStore *MockSubscriptionStore, accounts []string) { t.Helper() - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Name: "deal team", Type: model.RoomTypeChannel, SiteID: "site-A", }, nil) candidates := make([]AddMemberCandidate, len(accounts)) @@ -1536,7 +1536,7 @@ func TestProcessAddMembers_PopulatesSubName(t *testing.T) { ctx := natsutil.WithRequestID(context.Background(), reqID) // Use a custom BulkCreateSubscriptions expectation to capture subs. - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Name: "deal team", Type: model.RoomTypeChannel, SiteID: "site-A", }, nil) mockStore.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1"). @@ -1576,7 +1576,7 @@ func TestProcessAddMembers_HistoryNone_NoTimestamp(t *testing.T) { const reqTimestampMs = int64(1740000000000) acceptedAt := time.UnixMilli(reqTimestampMs).UTC() - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Name: "deal team", Type: model.RoomTypeChannel, SiteID: "site-A", }, nil) mockStore.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1"). @@ -1616,7 +1616,7 @@ func TestProcessAddMembers_NoHistoryConfig_LeavesNil(t *testing.T) { h, mockStore, _ := newAddMembersTestHandler(t) ctx := natsutil.WithRequestID(context.Background(), "0193abcd-0193-7abc-89ab-0193abcd0003") - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Name: "deal team", Type: model.RoomTypeChannel, SiteID: "site-A", }, nil) mockStore.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1"). @@ -1656,7 +1656,7 @@ func TestProcessAddMembers_InboxCarriesRoomName(t *testing.T) { ctx := natsutil.WithRequestID(context.Background(), reqID) // Cross-site member: bob lives on site-B. - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Name: "deal team", Type: model.RoomTypeChannel, SiteID: "site-A", }, nil) mockStore.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1"). @@ -3509,7 +3509,7 @@ func TestProcessAddMembers_FansOutKeyToNewAccountsOnly(t *testing.T) { pub := &mockPublisher{} keySender := roomkeysender.NewSender(pub) - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Name: "deal team", Type: model.RoomTypeChannel, SiteID: "site-a", }, nil) mockStore.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1"). @@ -3551,7 +3551,7 @@ func TestProcessAddMembers_PermanentErrorWhenKeyMissing(t *testing.T) { mockStore := NewMockSubscriptionStore(ctrl) keyStore := NewMockRoomKeyStore(ctrl) - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Name: "deal team", Type: model.RoomTypeChannel, SiteID: "site-a", }, nil) mockStore.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1"). @@ -3587,7 +3587,7 @@ func TestProcessAddMembers_TransientErrorWhenValkeyFails(t *testing.T) { mockStore := NewMockSubscriptionStore(ctrl) keyStore := NewMockRoomKeyStore(ctrl) - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Name: "deal team", Type: model.RoomTypeChannel, SiteID: "site-a", }, nil) mockStore.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1"). @@ -3619,7 +3619,7 @@ func TestProcessAddMembers_TransientErrorWhenValkeyFails(t *testing.T) { func TestProcessAddMembers_RejectsNonChannel(t *testing.T) { ctrl := gomock.NewController(t) mockStore := NewMockSubscriptionStore(ctrl) - mockStore.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ + mockStore.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ ID: "r1", Type: model.RoomTypeDM, SiteID: "site-a", }, nil) // The up-front reads run concurrently, so candidates/has-orgs fire before the @@ -3683,7 +3683,7 @@ func TestHandler_ProcessAddMembers_BackfillRunsOnFirstOrgTransition(t *testing.T store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"o1"}, []string(nil), roomID). Return([]AddMemberCandidate{{Account: "u_new"}}, nil) @@ -3722,7 +3722,7 @@ func TestHandler_ProcessAddMembers_BackfillSubscriptionAccountsErrorFailsHard(t store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"o1"}, []string(nil), roomID). Return([]AddMemberCandidate{{Account: "u_new"}}, nil) @@ -3753,7 +3753,7 @@ func TestHandler_ProcessAddMembers_BackfillFindUsersErrorFailsHard(t *testing.T) store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"o1"}, []string(nil), roomID). Return([]AddMemberCandidate{{Account: "u_new"}}, nil) @@ -3785,7 +3785,7 @@ func TestHandler_ProcessAddMembers_BackfillSkippedWhenRoomAlreadyHasOrgs(t *test store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"o_new"}, []string(nil), roomID). Return([]AddMemberCandidate{{Account: "u_new"}}, nil) @@ -3823,7 +3823,7 @@ func TestHandler_ProcessAddMembers_IndividualFilter_DirectAndOrgOverlap(t *testi store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"o1"}, []string{"u1"}, roomID). Return([]AddMemberCandidate{{Account: "u1"}, {Account: "u2"}}, nil) @@ -3877,7 +3877,7 @@ func TestHandler_ProcessAddMembers_IndividualFilter_OrgOnly(t *testing.T) { store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"o1"}, []string(nil), roomID). Return([]AddMemberCandidate{{Account: "u1"}}, nil) @@ -3973,7 +3973,7 @@ func TestHandler_ProcessAddMembers_RequesterNotFound(t *testing.T) { store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string(nil), []string{"u1"}, roomID). Return([]AddMemberCandidate{{Account: "u1"}}, nil) @@ -4033,7 +4033,7 @@ func TestHandler_ProcessAddMembers_Content_Single(t *testing.T) { store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string(nil), []string{"u1"}, roomID). Return([]AddMemberCandidate{{Account: "u1"}}, nil) @@ -4070,7 +4070,7 @@ func TestHandler_ProcessAddMembers_Content_Multi(t *testing.T) { store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string(nil), []string{"u1", "u2"}, roomID). Return([]AddMemberCandidate{{Account: "u1"}, {Account: "u2"}}, nil) @@ -4460,7 +4460,7 @@ func TestHandler_ProcessAddMembers_Content_OrgAddWithOneMember_UsesMulti(t *test store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) // req.Users is empty; org "eng" expands to one user "u1". store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string{"eng"}, []string(nil), roomID). @@ -4500,7 +4500,7 @@ func TestHandler_ProcessAddMembers_HasOrgRoomMembersError_FailsClosed(t *testing store := NewMockSubscriptionStore(ctrl) roomID := "r1" - store.EXPECT().GetRoom(gomock.Any(), roomID). + store.EXPECT().GetRoomMeta(gomock.Any(), roomID). Return(&model.Room{ID: roomID, Name: "Chan", SiteID: "site-a", Type: model.RoomTypeChannel}, nil) store.EXPECT().ListAddMemberCandidates(gomock.Any(), []string(nil), []string{"u1"}, roomID). Return([]AddMemberCandidate{{Account: "u1"}}, nil) diff --git a/room-worker/integration_test.go b/room-worker/integration_test.go index eb4976f94..5710b88b4 100644 --- a/room-worker/integration_test.go +++ b/room-worker/integration_test.go @@ -542,6 +542,84 @@ func TestMongoStore_UserCache(t *testing.T) { assert.Equal(t, "alice", users[0].Account) } +// TestMongoStore_RoomMetaCache covers the cached GetRoomMeta path: it serves the +// room's stable fields from the in-process cache and still errors on a miss. +func TestMongoStore_RoomMetaCache(t *testing.T) { + db := setupMongo(t) + store := NewMongoStore(db) + require.NoError(t, store.EnableRoomMetaCache(100, time.Minute)) + ctx := context.Background() + + _, err := db.Collection("rooms").InsertOne(ctx, model.Room{ + ID: "rm", Name: "general", Type: model.RoomTypeChannel, SiteID: "site-a", UserCount: 7, + }) + require.NoError(t, err) + + // Two reads (miss then hit) both return the stable fields. + for range 2 { + room, err := store.GetRoomMeta(ctx, "rm") + require.NoError(t, err) + assert.Equal(t, "rm", room.ID) + assert.Equal(t, "general", room.Name) + assert.Equal(t, model.RoomTypeChannel, room.Type) + assert.Equal(t, "site-a", room.SiteID) + } + + _, err = store.GetRoomMeta(ctx, "nope") + require.Error(t, err) +} + +// TestMongoStore_GetRoom_FullDocumentWithCacheEnabled guards the regression +// where enabling the meta cache made GetRoom serve a partial (Meta-only) room +// with a zero CreatedAt. GetRoom must always return the full document — the DM +// idempotency path (serverCreateDM) depends on existing.CreatedAt — regardless +// of whether the add-path meta cache is enabled. +func TestMongoStore_GetRoom_FullDocumentWithCacheEnabled(t *testing.T) { + db := setupMongo(t) + store := NewMongoStore(db) + require.NoError(t, store.EnableRoomMetaCache(100, time.Minute)) + ctx := context.Background() + + createdAt := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + _, err := db.Collection("rooms").InsertOne(ctx, model.Room{ + ID: "rfull", Name: "general", Type: model.RoomTypeDM, SiteID: "site-a", + UserCount: 2, CreatedAt: createdAt, UpdatedAt: createdAt, + }) + require.NoError(t, err) + + room, err := store.GetRoom(ctx, "rfull") + require.NoError(t, err) + assert.Equal(t, "rfull", room.ID) + assert.Equal(t, model.RoomTypeDM, room.Type) + assert.Equal(t, createdAt, room.CreatedAt, "GetRoom must return the persisted CreatedAt, not a zero value") +} + +// TestMongoStore_ListAddMemberCandidates_SkipsIRMForDirectAdds asserts the +// no-org optimization: the room_members read is skipped, so +// HasIndividualRoomMember is false even when an individual row exists. +func TestMongoStore_ListAddMemberCandidates_SkipsIRMForDirectAdds(t *testing.T) { + db := setupMongo(t) + store := NewMongoStore(db) + ctx := context.Background() + + _, err := db.Collection("users").InsertOne(ctx, model.User{ID: "u9", Account: "dave", SiteID: "site-a"}) + require.NoError(t, err) + // An individual room_members row exists for dave in room rx. + _, err = db.Collection("room_members").InsertOne(ctx, model.RoomMember{ + ID: idgen.GenerateUUIDv7(), RoomID: "rx", + Member: model.RoomMemberEntry{ID: "u9", Type: model.RoomMemberIndividual, Account: "dave"}, + }) + require.NoError(t, err) + + // Direct (no-org) add: the IRM read is skipped, so the flag is false despite + // the existing row (the handler never reads it on the no-org path). + cands, err := store.ListAddMemberCandidates(ctx, nil, []string{"dave"}, "rx") + require.NoError(t, err) + require.Len(t, cands, 1) + assert.Equal(t, "dave", cands[0].Account) + assert.False(t, cands[0].HasIndividualRoomMember, "no-org add skips the room_members read") +} + func TestMongoStore_ListAddMemberCandidates_Integration(t *testing.T) { ctx := context.Background() db := setupMongo(t) diff --git a/room-worker/loadinputs_test.go b/room-worker/loadinputs_test.go index 92746daff..a2ad3b6b6 100644 --- a/room-worker/loadinputs_test.go +++ b/room-worker/loadinputs_test.go @@ -14,7 +14,7 @@ import ( ) // TestLoadAddMemberInputs_RunsReadsConcurrently proves the three independent -// up-front reads (GetRoom, ListAddMemberCandidates, HasOrgRoomMembers) are +// up-front reads (GetRoomMeta, ListAddMemberCandidates, HasOrgRoomMembers) are // issued concurrently: each mock blocks on a shared release channel after // signalling arrival, so serial execution would only ever reach one before the // others are unblocked and the test would time out. @@ -27,7 +27,7 @@ func TestLoadAddMemberInputs_RunsReadsConcurrently(t *testing.T) { release := make(chan struct{}) block := func() { arrived <- struct{}{}; <-release } - store.EXPECT().GetRoom(gomock.Any(), "r1").DoAndReturn( + store.EXPECT().GetRoomMeta(gomock.Any(), "r1").DoAndReturn( func(_ context.Context, _ string) (*model.Room, error) { block() return &model.Room{ID: "r1", Type: model.RoomTypeChannel}, nil @@ -86,7 +86,7 @@ func TestLoadAddMemberInputs_PropagatesErrors(t *testing.T) { { name: "get room fails", setup: func(s *MockSubscriptionStore) { - s.EXPECT().GetRoom(gomock.Any(), "r1").Return(nil, sentinel) + s.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(nil, sentinel) s.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1").Return(nil, nil) s.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) }, @@ -95,7 +95,7 @@ func TestLoadAddMemberInputs_PropagatesErrors(t *testing.T) { { name: "list candidates fails", setup: func(s *MockSubscriptionStore) { - s.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel}, nil) + s.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel}, nil) s.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1").Return(nil, sentinel) s.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, nil) }, @@ -104,7 +104,7 @@ func TestLoadAddMemberInputs_PropagatesErrors(t *testing.T) { { name: "has-org-members fails", setup: func(s *MockSubscriptionStore) { - s.EXPECT().GetRoom(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel}, nil) + s.EXPECT().GetRoomMeta(gomock.Any(), "r1").Return(&model.Room{ID: "r1", Type: model.RoomTypeChannel}, nil) s.EXPECT().ListAddMemberCandidates(gomock.Any(), gomock.Any(), gomock.Any(), "r1").Return(nil, nil) s.EXPECT().HasOrgRoomMembers(gomock.Any(), "r1").Return(false, sentinel) }, diff --git a/room-worker/main.go b/room-worker/main.go index eff1ae270..6dab5808b 100644 --- a/room-worker/main.go +++ b/room-worker/main.go @@ -33,23 +33,25 @@ import ( ) type config struct { - NatsURL string `env:"NATS_URL" envDefault:"nats://localhost:4222"` - NatsCredsFile string `env:"NATS_CREDS_FILE" envDefault:""` - SiteID string `env:"SITE_ID" envDefault:"site-local"` - MongoURI string `env:"MONGO_URI" envDefault:"mongodb://localhost:27017"` - MongoDB string `env:"MONGO_DB" envDefault:"chat"` - MongoUsername string `env:"MONGO_USERNAME" envDefault:""` - MongoPassword string `env:"MONGO_PASSWORD" envDefault:""` - MaxWorkers int `env:"MAX_WORKERS" envDefault:"100"` - KeyFanoutWorkers int `env:"KEY_FANOUT_WORKERS" envDefault:"32"` // see defaultKeyFanoutWorkers in handler.go - UserCacheSize int `env:"USER_CACHE_SIZE" envDefault:"10000"` - UserCacheTTL time.Duration `env:"USER_CACHE_TTL" envDefault:"5m"` - Consumer stream.ConsumerSettings `envPrefix:"CONSUMER_"` - Bootstrap bootstrapConfig `envPrefix:"BOOTSTRAP_"` - HealthAddr string `env:"HEALTH_ADDR" envDefault:":8081"` - PProfEnabled bool `env:"PPROF_ENABLED" envDefault:"false"` - MetricsAddr string `env:"METRICS_ADDR" envDefault:":9090"` - DebugLog logctx.Config `envPrefix:"DEBUG_LOG_"` + NatsURL string `env:"NATS_URL" envDefault:"nats://localhost:4222"` + NatsCredsFile string `env:"NATS_CREDS_FILE" envDefault:""` + SiteID string `env:"SITE_ID" envDefault:"site-local"` + MongoURI string `env:"MONGO_URI" envDefault:"mongodb://localhost:27017"` + MongoDB string `env:"MONGO_DB" envDefault:"chat"` + MongoUsername string `env:"MONGO_USERNAME" envDefault:""` + MongoPassword string `env:"MONGO_PASSWORD" envDefault:""` + MaxWorkers int `env:"MAX_WORKERS" envDefault:"100"` + KeyFanoutWorkers int `env:"KEY_FANOUT_WORKERS" envDefault:"32"` // see defaultKeyFanoutWorkers in handler.go + UserCacheSize int `env:"USER_CACHE_SIZE" envDefault:"10000"` + UserCacheTTL time.Duration `env:"USER_CACHE_TTL" envDefault:"5m"` + RoomMetaCacheSize int `env:"ROOM_META_CACHE_SIZE" envDefault:"10000"` + RoomMetaCacheTTL time.Duration `env:"ROOM_META_CACHE_TTL" envDefault:"60s"` + Consumer stream.ConsumerSettings `envPrefix:"CONSUMER_"` + Bootstrap bootstrapConfig `envPrefix:"BOOTSTRAP_"` + HealthAddr string `env:"HEALTH_ADDR" envDefault:":8081"` + PProfEnabled bool `env:"PPROF_ENABLED" envDefault:"false"` + MetricsAddr string `env:"METRICS_ADDR" envDefault:":9090"` + DebugLog logctx.Config `envPrefix:"DEBUG_LOG_"` // Grace window during which a rotated-out previous key remains valid for decrypt. RoomKeyGracePeriod time.Duration `env:"ROOM_KEY_GRACE_PERIOD" envDefault:"24h"` @@ -168,6 +170,17 @@ func main() { } else { slog.Info("user-cache disabled", "size", cfg.UserCacheSize) } + // Room-meta cache is on by default; a non-positive size disables it cleanly + // rather than failing startup (the LRU constructor rejects size<=0). + if cfg.RoomMetaCacheSize > 0 { + if err := store.EnableRoomMetaCache(cfg.RoomMetaCacheSize, cfg.RoomMetaCacheTTL); err != nil { + slog.Error("failed to enable room-meta cache", "error", err) + os.Exit(1) + } + slog.Info("room-meta-cache enabled", "size", cfg.RoomMetaCacheSize, "ttl", cfg.RoomMetaCacheTTL) + } else { + slog.Info("room-meta-cache disabled", "size", cfg.RoomMetaCacheSize) + } handler := NewHandler(store, cfg.SiteID, func(ctx context.Context, subj string, data []byte, msgID string) error { msg := natsutil.NewMsg(ctx, subj, data) if msgID == "" { diff --git a/room-worker/mock_store_test.go b/room-worker/mock_store_test.go index 577446f13..c5ee64c0a 100644 --- a/room-worker/mock_store_test.go +++ b/room-worker/mock_store_test.go @@ -221,6 +221,21 @@ func (mr *MockSubscriptionStoreMockRecorder) GetRoom(ctx, roomID any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRoom", reflect.TypeOf((*MockSubscriptionStore)(nil).GetRoom), ctx, roomID) } +// GetRoomMeta mocks base method. +func (m *MockSubscriptionStore) GetRoomMeta(ctx context.Context, roomID string) (*model.Room, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetRoomMeta", ctx, roomID) + ret0, _ := ret[0].(*model.Room) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetRoomMeta indicates an expected call of GetRoomMeta. +func (mr *MockSubscriptionStoreMockRecorder) GetRoomMeta(ctx, roomID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRoomMeta", reflect.TypeOf((*MockSubscriptionStore)(nil).GetRoomMeta), ctx, roomID) +} + // GetSubscription mocks base method. func (m *MockSubscriptionStore) GetSubscription(ctx context.Context, account, roomID string) (*model.Subscription, error) { m.ctrl.T.Helper() diff --git a/room-worker/store.go b/room-worker/store.go index fdb92e678..8a39f4dc5 100644 --- a/room-worker/store.go +++ b/room-worker/store.go @@ -77,6 +77,11 @@ type SubscriptionStore interface { // periodic recompute (the drift safety net) restores convergence. ApplyMemberCountDelta(ctx context.Context, roomID string, userDelta, appDelta int, ttl time.Duration) (reconcileDue bool, err error) GetRoom(ctx context.Context, roomID string) (*model.Room, error) + // GetRoomMeta returns a room populated with only its stable fields + // (ID/Type/Name/SiteID/UserCount); CreatedAt/UpdatedAt are zero. It is the + // add-member hot path's read and is served from the meta cache when enabled. + // Callers needing time-sensitive fields (e.g. CreatedAt) must use GetRoom. + GetRoomMeta(ctx context.Context, roomID string) (*model.Room, error) GetSubscription(ctx context.Context, account, roomID string) (*model.Subscription, error) GetUser(ctx context.Context, account string) (*model.User, error) // GetApp returns the app whose Assistant.Name == botAccount, or ErrAppNotFound. diff --git a/room-worker/store_mongo.go b/room-worker/store_mongo.go index d3de7bb9d..77283bd80 100644 --- a/room-worker/store_mongo.go +++ b/room-worker/store_mongo.go @@ -14,6 +14,7 @@ import ( "github.com/hmchangw/chat/pkg/mongoutil" "github.com/hmchangw/chat/pkg/pipelines" "github.com/hmchangw/chat/pkg/roomkeystore" + "github.com/hmchangw/chat/pkg/roommetacache" "github.com/hmchangw/chat/pkg/userstore" ) @@ -29,6 +30,10 @@ type MongoStore struct { // `users` collection above is still used directly for the org-expansion // candidate query, which the by-account reader cannot serve. userReader userstore.UserStore + // roomMeta, when non-nil (EnableRoomMetaCache), fronts GetRoomMeta with an + // in-process LRU+TTL cache of the room's stable fields. Nil means GetRoomMeta + // reads Mongo directly. + roomMeta *roommetacache.Cache } func NewMongoStore(db *mongo.Database) *MongoStore { @@ -55,6 +60,23 @@ func (s *MongoStore) EnableUserCache(size int, ttl time.Duration) error { return nil } +// EnableRoomMetaCache fronts GetRoomMeta with an in-process LRU+TTL cache. The +// add-member hot path (loadAddMemberInputs) reads only the stable fields +// (Type/Name/SiteID/ID/UserCount), so caching them is safe; a rename is +// reflected after at most the TTL. GetRoom is deliberately left uncached so the +// DM idempotency path still sees the room's real CreatedAt. Call once at startup. +func (s *MongoStore) EnableRoomMetaCache(size int, ttl time.Duration) error { + rooms := s.rooms // capture only the collection, not the whole store + cache, err := roommetacache.New(size, ttl, func(ctx context.Context, roomID string) (roommetacache.Meta, error) { + return roommetacache.FetchFromMongo(ctx, rooms, roomID) + }) + if err != nil { + return fmt.Errorf("enable room meta cache: %w", err) + } + s.roomMeta = cache + return nil +} + // ListByRoom returns all subscriptions for roomID across every site. Not part // of SubscriptionStore — the handler's hot paths only need accounts (see // GetSubscriptionAccounts); this full-document read is retained for integration @@ -138,6 +160,10 @@ func (s *MongoStore) ApplyMemberCountDelta(ctx context.Context, roomID string, u return doc.CountsReconciledAt.IsZero() || time.Since(doc.CountsReconciledAt) > ttl, nil } +// GetRoom returns the full room document from Mongo. It is never cache-served: +// callers such as serverCreateDM's idempotency path read time-sensitive fields +// (CreatedAt) that the meta cache does not carry. The add-member hot path uses +// GetRoomMeta instead, which only needs the stable fields. func (s *MongoStore) GetRoom(ctx context.Context, roomID string) (*model.Room, error) { var room model.Room if err := s.rooms.FindOne(ctx, bson.M{"_id": roomID}).Decode(&room); err != nil { @@ -146,6 +172,27 @@ func (s *MongoStore) GetRoom(ctx context.Context, roomID string) (*model.Room, e return &room, nil } +// GetRoomMeta returns a room populated with only its stable fields +// (ID/Type/Name/SiteID/UserCount) — the subset the add-member hot path reads. +// When the meta cache is enabled it is served from the in-process LRU+TTL cache; +// otherwise it falls through to a direct Mongo read. The returned room's +// CreatedAt/UpdatedAt are zero — callers needing those must use GetRoom. +func (s *MongoStore) GetRoomMeta(ctx context.Context, roomID string) (*model.Room, error) { + var ( + meta roommetacache.Meta + err error + ) + if s.roomMeta != nil { + meta, err = s.roomMeta.Get(ctx, roomID) + } else { + meta, err = roommetacache.FetchFromMongo(ctx, s.rooms, roomID) + } + if err != nil { + return nil, err + } + return &model.Room{ID: meta.ID, Type: meta.Type, Name: meta.Name, SiteID: meta.SiteID, UserCount: meta.UserCount}, nil +} + func (s *MongoStore) GetUser(ctx context.Context, account string) (*model.User, error) { u, err := s.userReader.FindUserByAccount(ctx, account) if errors.Is(err, userstore.ErrUserNotFound) { @@ -524,16 +571,24 @@ func (s *MongoStore) ListAddMemberCandidates(ctx context.Context, orgIDs, direct accounts[i] = c.Account ids[i] = c.ID } - // 2 & 3. Per-candidate membership state via two indexed reads scoped to the - // candidate set ($in), so examined keys stay bounded by the request size - // rather than the room size. + // 2. Already-subscribed candidates (always needed). subbed, err := pipelines.SubscribedAccounts(ctx, s.subscriptions, roomID, accounts) if err != nil { return nil, err } - irm, err := s.individualMemberIDs(ctx, roomID, ids) - if err != nil { - return nil, err + // 3. Existing individual room_members rows — only consulted when the room + // tracks individuals, which the worker gates on writeIndividuals (orgs in + // the request OR pre-existing org rows). With no orgs the handler never + // reads HasIndividualRoomMember, so skip this room_members read entirely; + // for an org-bearing request it always runs. (A no-org add to a room that + // formerly had orgs may re-attempt a few individual inserts, which + // BulkCreateRoomMembers absorbs as idempotent dup-key no-ops.) + var irm map[string]struct{} + if len(orgIDs) > 0 { + irm, err = s.individualMemberIDs(ctx, roomID, ids) + if err != nil { + return nil, err + } } out := make([]AddMemberCandidate, len(candidates)) for i, c := range candidates {