Merge pull request #361 from matrix-org/dmr/reload-jrt

sync3/tracker.go

@@ -193,3 +193,40 @@ func (t *JoinedRoomsTracker) NumInvitedUsersForRoom(roomID string) int {
 	defer t.mu.RUnlock()
 	return len(t.roomIDToInvitedUsers[roomID])
 }
+
+// ReloadMembershipsForRoom overwrites the JoinedRoomsTracker state for one room to the
+// given list of joined and invited users.
+func (t *JoinedRoomsTracker) ReloadMembershipsForRoom(roomID string, joined, invited []string) {
+	newJoined := make(set, len(joined))
+	newInvited := make(set, len(invited))
+	for _, member := range joined {
+		newJoined[member] = struct{}{}
+	}
+	for _, member := range invited {
+		newInvited[member] = struct{}{}
+	}
+
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	// 1. Overwrite the room's memberships with the given arguments.
+	oldJoined := t.roomIDToJoinedUsers[roomID]
+	t.roomIDToJoinedUsers[roomID] = newJoined
+	t.roomIDToInvitedUsers[roomID] = newInvited
+
+	// 2. Mark the joined users as being joined to this room.
+	for userID := range newJoined {
+		if t.userIDToJoinedRooms[userID] == nil {
+			t.userIDToJoinedRooms[userID] = make(set)
+		}
+		t.userIDToJoinedRooms[userID][roomID] = struct{}{}
+	}
+
+	// 3. Scan the old joined list for users who are no longer joined, and mark them as such.
+	for userID := range oldJoined {
+		_, stillJoined := newJoined[userID]
+		if !stillJoined {
+			delete(t.userIDToJoinedRooms[userID], roomID)
+		}
+	}
+}

sync3/tracker_test.go

@@ -83,6 +83,45 @@ func TestTrackerStartup(t *testing.T) {
 	assertInt(t, jrt.NumInvitedUsersForRoom(roomC), 0)
 }
 
+func TestTrackerReload(t *testing.T) {
+	roomA := "!a"
+	roomB := "!b"
+	roomC := "!c"
+	alice := "@alice"
+	bob := "@bob"
+	chris := "@chris"
+	jrt := NewJoinedRoomsTracker()
+	jrt.Startup(map[string][]string{
+		roomA: {alice, bob},
+		roomB: {bob},
+		roomC: {alice},
+	})
+
+	t.Log("Chris joins room C.")
+	jrt.ReloadMembershipsForRoom(roomC, []string{alice, chris}, nil)
+	members, _ := jrt.JoinedUsersForRoom(roomC, nil)
+	assertEqualSlices(t, "roomC joined members", members, []string{alice, chris})
+	assertEqualSlices(t, "alice's rooms", jrt.JoinedRoomsForUser(alice), []string{roomA, roomC})
+	assertEqualSlices(t, "chris's rooms", jrt.JoinedRoomsForUser(chris), []string{roomC})
+	assertInt(t, jrt.NumInvitedUsersForRoom(roomC), 0)
+
+	t.Log("Bob leaves room B.")
+	jrt.ReloadMembershipsForRoom(roomB, nil, nil)
+	members, _ = jrt.JoinedUsersForRoom(roomB, nil)
+	assertEqualSlices(t, "roomB joined members", members, nil)
+	assertEqualSlices(t, "bob's rooms", jrt.JoinedRoomsForUser(bob), []string{roomA})
+	assertInt(t, jrt.NumInvitedUsersForRoom(roomB), 0)
+
+	t.Log("Chris joins room A. Alice and Bob leave it, but Chris reinvites Bob.")
+	jrt.ReloadMembershipsForRoom(roomA, []string{chris}, []string{bob})
+	members, _ = jrt.JoinedUsersForRoom(roomA, nil)
+	assertEqualSlices(t, "roomA joined members", members, []string{chris})
+	assertEqualSlices(t, "alice's rooms", jrt.JoinedRoomsForUser(alice), []string{roomC})
+	assertEqualSlices(t, "bob's rooms", jrt.JoinedRoomsForUser(bob), nil)
+	assertEqualSlices(t, "chris's rooms", jrt.JoinedRoomsForUser(chris), []string{roomA, roomC})
+	assertInt(t, jrt.NumInvitedUsersForRoom(roomA), 1)
+}
+
 func TestJoinedRoomsTracker_UserLeftRoom_ReturnValue(t *testing.T) {
 	alice := "@alice"
 	bob := "@bob"