Optimize the room-worker add-member hot path

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

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

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}
 }