Cleanup stale device lists for users we don't share a room with anymore (#2857)

The stale device lists table might contain entries for users we don't
share a room with anymore. This now asks the roomserver about left users
and removes those entries from the table.

Co-authored-by: Neil Alexander <neilalexander@users.noreply.github.com>

roomserver/storage/shared/storage.go

@@ -1365,6 +1365,43 @@ func (d *Database) JoinedUsersSetInRooms(ctx context.Context, roomIDs, userIDs [
 	return result, nil
 }
 
+// GetLeftUsers calculates users we (the server) don't share a room with anymore.
+func (d *Database) GetLeftUsers(ctx context.Context, userIDs []string) ([]string, error) {
+	// Get the userNID for all users with a stale device list
+	stateKeyNIDMap, err := d.EventStateKeyNIDs(ctx, userIDs)
+	if err != nil {
+		return nil, err
+	}
+
+	userNIDs := make([]types.EventStateKeyNID, 0, len(stateKeyNIDMap))
+	userNIDtoUserID := make(map[types.EventStateKeyNID]string, len(stateKeyNIDMap))
+	// Create a map from userNID -> userID
+	for userID, nid := range stateKeyNIDMap {
+		userNIDs = append(userNIDs, nid)
+		userNIDtoUserID[nid] = userID
+	}
+
+	// Get all users whose membership is still join, knock or invite.
+	stillJoinedUsersNIDs, err := d.MembershipTable.SelectJoinedUsers(ctx, nil, userNIDs)
+	if err != nil {
+		return nil, err
+	}
+
+	// Remove joined users from the "user with stale devices" list, which contains left AND joined users
+	for _, joinedUser := range stillJoinedUsersNIDs {
+		delete(userNIDtoUserID, joinedUser)
+	}
+
+	// The users still in our userNIDtoUserID map are the users we don't share a room with anymore,
+	// and the return value we are looking for.
+	leftUsers := make([]string, 0, len(userNIDtoUserID))
+	for _, userID := range userNIDtoUserID {
+		leftUsers = append(leftUsers, userID)
+	}
+
+	return leftUsers, nil
+}
+
 // GetLocalServerInRoom returns true if we think we're in a given room or false otherwise.
 func (d *Database) GetLocalServerInRoom(ctx context.Context, roomNID types.RoomNID) (bool, error) {
 	return d.MembershipTable.SelectLocalServerInRoom(ctx, nil, roomNID)

roomserver/storage/shared/storage_test.go

@@ -0,0 +1,96 @@
+package shared_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/roomserver/storage/postgres"
+	"github.com/matrix-org/dendrite/roomserver/storage/shared"
+	"github.com/matrix-org/dendrite/roomserver/storage/sqlite3"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
+	"github.com/matrix-org/dendrite/setup/config"
+	"github.com/matrix-org/dendrite/test"
+	"github.com/matrix-org/dendrite/test/testrig"
+)
+
+func mustCreateRoomserverDatabase(t *testing.T, dbType test.DBType) (*shared.Database, func()) {
+	t.Helper()
+
+	connStr, clearDB := test.PrepareDBConnectionString(t, dbType)
+	base, _, _ := testrig.Base(nil)
+	dbOpts := &config.DatabaseOptions{ConnectionString: config.DataSource(connStr)}
+
+	db, err := sqlutil.Open(dbOpts, sqlutil.NewExclusiveWriter())
+	assert.NoError(t, err)
+
+	var membershipTable tables.Membership
+	var stateKeyTable tables.EventStateKeys
+	switch dbType {
+	case test.DBTypePostgres:
+		err = postgres.CreateEventStateKeysTable(db)
+		assert.NoError(t, err)
+		err = postgres.CreateMembershipTable(db)
+		assert.NoError(t, err)
+		membershipTable, err = postgres.PrepareMembershipTable(db)
+		assert.NoError(t, err)
+		stateKeyTable, err = postgres.PrepareEventStateKeysTable(db)
+	case test.DBTypeSQLite:
+		err = sqlite3.CreateEventStateKeysTable(db)
+		assert.NoError(t, err)
+		err = sqlite3.CreateMembershipTable(db)
+		assert.NoError(t, err)
+		membershipTable, err = sqlite3.PrepareMembershipTable(db)
+		assert.NoError(t, err)
+		stateKeyTable, err = sqlite3.PrepareEventStateKeysTable(db)
+	}
+	assert.NoError(t, err)
+
+	return &shared.Database{
+			DB:                  db,
+			EventStateKeysTable: stateKeyTable,
+			MembershipTable:     membershipTable,
+			Writer:              sqlutil.NewExclusiveWriter(),
+		}, func() {
+			err := base.Close()
+			assert.NoError(t, err)
+			clearDB()
+			err = db.Close()
+			assert.NoError(t, err)
+		}
+}
+
+func Test_GetLeftUsers(t *testing.T) {
+	alice := test.NewUser(t)
+	bob := test.NewUser(t)
+	charlie := test.NewUser(t)
+
+	ctx := context.Background()
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		db, close := mustCreateRoomserverDatabase(t, dbType)
+		defer close()
+
+		// Create dummy entries
+		for _, user := range []*test.User{alice, bob, charlie} {
+			nid, err := db.EventStateKeysTable.InsertEventStateKeyNID(ctx, nil, user.ID)
+			assert.NoError(t, err)
+			err = db.MembershipTable.InsertMembership(ctx, nil, 1, nid, true)
+			assert.NoError(t, err)
+			// We must update the membership with a non-zero event NID or it will get filtered out in later queries
+			membershipNID := tables.MembershipStateLeaveOrBan
+			if user == alice {
+				membershipNID = tables.MembershipStateJoin
+			}
+			_, err = db.MembershipTable.UpdateMembership(ctx, nil, 1, nid, nid, membershipNID, 1, false)
+			assert.NoError(t, err)
+		}
+
+		// Now try to get the left users, this should be Bob and Charlie, since they have a "leave" membership
+		expectedUserIDs := []string{bob.ID, charlie.ID}
+		leftUsers, err := db.GetLeftUsers(context.Background(), []string{alice.ID, bob.ID, charlie.ID})
+		assert.NoError(t, err)
+		assert.ElementsMatch(t, expectedUserIDs, leftUsers)
+	})
+}