diff --git a/sync3/tracker.go b/sync3/tracker.go index e33fabfc..13bb6eb3 100644 --- a/sync3/tracker.go +++ b/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) + } + } +} diff --git a/sync3/tracker_test.go b/sync3/tracker_test.go index 0bc58009..fb294873 100644 --- a/sync3/tracker_test.go +++ b/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"