diff --git a/roomserver/api/api.go b/roomserver/api/api.go index ab56529c5..c29406a1a 100644 --- a/roomserver/api/api.go +++ b/roomserver/api/api.go @@ -227,6 +227,7 @@ type UserRoomserverAPI interface { QueryMembershipsForRoom(ctx context.Context, req *QueryMembershipsForRoomRequest, res *QueryMembershipsForRoomResponse) error PerformAdminEvacuateUser(ctx context.Context, userID string) (affected []string, err error) PerformJoin(ctx context.Context, req *PerformJoinRequest) (roomID string, joinedVia spec.ServerName, err error) + JoinedUserCount(ctx context.Context, roomID string) (int, error) } type FederationRoomserverAPI interface { diff --git a/roomserver/internal/query/query.go b/roomserver/internal/query/query.go index 626d3c13e..39e3bd0ec 100644 --- a/roomserver/internal/query/query.go +++ b/roomserver/internal/query/query.go @@ -974,6 +974,20 @@ func (r *Queryer) LocallyJoinedUsers(ctx context.Context, roomVersion gomatrixse return joinedUsers, nil } +func (r *Queryer) JoinedUserCount(ctx context.Context, roomID string) (int, error) { + info, err := r.DB.RoomInfo(ctx, roomID) + if err != nil { + return 0, err + } + if info == nil { + return 0, nil + } + + // TODO: this can be further optimised by just using a SELECT COUNT query + nids, err := r.DB.GetMembershipEventNIDsForRoom(ctx, info.RoomNID, true, false) + return len(nids), err +} + // nolint:gocyclo func (r *Queryer) QueryRestrictedJoinAllowed(ctx context.Context, roomID spec.RoomID, senderID spec.SenderID) (string, error) { // Look up if we know anything about the room. If it doesn't exist diff --git a/userapi/consumers/roomserver.go b/userapi/consumers/roomserver.go index 9a9a407ce..1f866ef4d 100644 --- a/userapi/consumers/roomserver.go +++ b/userapi/consumers/roomserver.go @@ -405,18 +405,25 @@ func newLocalMembership(event *synctypes.ClientEvent) (*localMembership, error) // localRoomMembers fetches the current local members of a room, and // the total number of members. func (s *OutputRoomEventConsumer) localRoomMembers(ctx context.Context, roomID string) ([]*localMembership, int, error) { + // Get only locally joined users to avoid unmarshalling and caching + // membership events we only use to calculate the room size. req := &rsapi.QueryMembershipsForRoomRequest{ RoomID: roomID, JoinedOnly: true, + LocalOnly: true, } var res rsapi.QueryMembershipsForRoomResponse - - // XXX: This could potentially race if the state for the event is not known yet - // e.g. the event came over federation but we do not have the full state persisted. if err := s.rsAPI.QueryMembershipsForRoom(ctx, req, &res); err != nil { return nil, 0, err } + // Since we only queried locally joined users above, + // we also need to ask the roomserver about the joined user count. + totalCount, err := s.rsAPI.JoinedUserCount(ctx, roomID) + if err != nil { + return nil, 0, err + } + var members []*localMembership for _, event := range res.JoinEvents { // Filter out invalid join events @@ -426,31 +433,18 @@ func (s *OutputRoomEventConsumer) localRoomMembers(ctx context.Context, roomID s if *event.StateKey == "" { continue } - _, serverName, err := gomatrixserverlib.SplitID('@', *event.StateKey) - if err != nil { - log.WithError(err).Error("failed to get servername from statekey") - continue - } - // Only get memberships for our server - if serverName != s.serverName { - continue - } + // We're going to trust the Query from above to really just return + // local users member, err := newLocalMembership(&event) if err != nil { log.WithError(err).Errorf("Parsing MemberContent") continue } - if member.Membership != spec.Join { - continue - } - if member.Domain != s.cfg.Matrix.ServerName { - continue - } members = append(members, member) } - return members, len(res.JoinEvents), nil + return members, totalCount, nil } // roomName returns the name in the event (if type==m.room.name), or diff --git a/userapi/consumers/roomserver_test.go b/userapi/consumers/roomserver_test.go index 4dc81e74a..49dd5b238 100644 --- a/userapi/consumers/roomserver_test.go +++ b/userapi/consumers/roomserver_test.go @@ -2,16 +2,22 @@ package consumers import ( "context" + "crypto/ed25519" "reflect" "sync" "testing" "time" + "github.com/matrix-org/dendrite/internal/caching" "github.com/matrix-org/dendrite/internal/sqlutil" + "github.com/matrix-org/dendrite/roomserver" "github.com/matrix-org/dendrite/roomserver/types" + "github.com/matrix-org/dendrite/setup/jetstream" + "github.com/matrix-org/dendrite/test/testrig" "github.com/matrix-org/gomatrixserverlib" "github.com/matrix-org/gomatrixserverlib/spec" "github.com/stretchr/testify/assert" + "golang.org/x/crypto/bcrypt" "github.com/matrix-org/dendrite/internal/pushrules" rsapi "github.com/matrix-org/dendrite/roomserver/api" @@ -139,6 +145,42 @@ func Test_evaluatePushRules(t *testing.T) { }) } +func TestLocalRoomMembers(t *testing.T) { + alice := test.NewUser(t) + _, sk, err := ed25519.GenerateKey(nil) + assert.NoError(t, err) + bob := test.NewUser(t, test.WithSigningServer("notlocalhost", "ed25519:abc", sk)) + charlie := test.NewUser(t, test.WithSigningServer("notlocalhost", "ed25519:abc", sk)) + + room := test.NewRoom(t, alice) + room.CreateAndInsert(t, bob, spec.MRoomMember, map[string]string{"membership": spec.Join}, test.WithStateKey(bob.ID)) + room.CreateAndInsert(t, charlie, spec.MRoomMember, map[string]string{"membership": spec.Join}, test.WithStateKey(charlie.ID)) + + test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) { + cfg, processCtx, close := testrig.CreateConfig(t, dbType) + defer close() + + cm := sqlutil.NewConnectionManager(processCtx, cfg.Global.DatabaseOptions) + natsInstance := &jetstream.NATSInstance{} + caches := caching.NewRistrettoCache(8*1024*1024, time.Hour, caching.DisableMetrics) + rsAPI := roomserver.NewInternalAPI(processCtx, cfg, cm, natsInstance, caches, caching.DisableMetrics) + rsAPI.SetFederationAPI(nil, nil) + db, err := storage.NewUserDatabase(processCtx.Context(), cm, &cfg.UserAPI.AccountDatabase, cfg.Global.ServerName, bcrypt.MinCost, 1000, 1000, "") + assert.NoError(t, err) + + err = rsapi.SendEvents(processCtx.Context(), rsAPI, rsapi.KindNew, room.Events(), "", "test", "test", nil, false) + assert.NoError(t, err) + + consumer := OutputRoomEventConsumer{db: db, rsAPI: rsAPI, serverName: "test", cfg: &cfg.UserAPI} + members, count, err := consumer.localRoomMembers(processCtx.Context(), room.ID) + assert.NoError(t, err) + assert.Equal(t, 3, count) + expectedLocalMember := &localMembership{UserID: alice.ID, Localpart: alice.Localpart, Domain: "test", MemberContent: gomatrixserverlib.MemberContent{Membership: spec.Join}} + assert.Equal(t, expectedLocalMember, members[0]) + }) + +} + func TestMessageStats(t *testing.T) { type args struct { eventType string @@ -257,3 +299,42 @@ func TestMessageStats(t *testing.T) { } }) } + +func BenchmarkLocalRoomMembers(b *testing.B) { + t := &testing.T{} + + cfg, processCtx, close := testrig.CreateConfig(t, test.DBTypePostgres) + defer close() + cm := sqlutil.NewConnectionManager(processCtx, cfg.Global.DatabaseOptions) + natsInstance := &jetstream.NATSInstance{} + caches := caching.NewRistrettoCache(8*1024*1024, time.Hour, caching.DisableMetrics) + rsAPI := roomserver.NewInternalAPI(processCtx, cfg, cm, natsInstance, caches, caching.DisableMetrics) + rsAPI.SetFederationAPI(nil, nil) + db, err := storage.NewUserDatabase(processCtx.Context(), cm, &cfg.UserAPI.AccountDatabase, cfg.Global.ServerName, bcrypt.MinCost, 1000, 1000, "") + assert.NoError(b, err) + + consumer := OutputRoomEventConsumer{db: db, rsAPI: rsAPI, serverName: "test", cfg: &cfg.UserAPI} + _, sk, err := ed25519.GenerateKey(nil) + assert.NoError(b, err) + + alice := test.NewUser(t) + room := test.NewRoom(t, alice) + + for i := 0; i < 100; i++ { + user := test.NewUser(t, test.WithSigningServer("notlocalhost", "ed25519:abc", sk)) + room.CreateAndInsert(t, user, spec.MRoomMember, map[string]string{"membership": spec.Join}, test.WithStateKey(user.ID)) + } + + err = rsapi.SendEvents(processCtx.Context(), rsAPI, rsapi.KindNew, room.Events(), "", "test", "test", nil, false) + assert.NoError(b, err) + + expectedLocalMember := &localMembership{UserID: alice.ID, Localpart: alice.Localpart, Domain: "test", MemberContent: gomatrixserverlib.MemberContent{Membership: spec.Join}} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + members, count, err := consumer.localRoomMembers(processCtx.Context(), room.ID) + assert.NoError(b, err) + assert.Equal(b, 101, count) + assert.Equal(b, expectedLocalMember, members[0]) + } +}