def test_partially_clear_queue(self): state1 = UserPresenceState.default("@user1:test") state2 = UserPresenceState.default("@user2:test") state3 = UserPresenceState.default("@user3:test") prev_token = self.queue.get_current_token(self.instance_name) self.queue.send_presence_to_destinations((state1, state2), ("dest1", "dest2")) self.reactor.advance(2 * 60 * 1000) self.queue.send_presence_to_destinations((state3, ), ("dest3", )) self.reactor.advance(4 * 60 * 1000) now_token = self.queue.get_current_token(self.instance_name) rows, upto_token, limited = self.get_success( self.queue.get_replication_rows("master", prev_token, now_token, 10)) self.assertEqual(upto_token, now_token) self.assertFalse(limited) expected_rows = [ (2, ("dest3", "@user3:test")), ] self.assertCountEqual(rows, []) prev_token = self.queue.get_current_token(self.instance_name) self.queue.send_presence_to_destinations((state1, state2), ("dest1", "dest2")) self.queue.send_presence_to_destinations((state3, ), ("dest3", )) now_token = self.queue.get_current_token(self.instance_name) rows, upto_token, limited = self.get_success( self.queue.get_replication_rows("master", prev_token, now_token, 10)) self.assertEqual(upto_token, now_token) self.assertFalse(limited) expected_rows = [ (3, ("dest1", "@user1:test")), (3, ("dest2", "@user1:test")), (3, ("dest1", "@user2:test")), (3, ("dest2", "@user2:test")), (4, ("dest3", "@user3:test")), ] self.assertCountEqual(rows, expected_rows)
def test_busy_no_idle(self): """ Tests that a user setting their presence to busy but idling doesn't turn their presence state into unavailable. """ user_id = "@foo:bar" status_msg = "I'm here!" now = 5000000 state = UserPresenceState.default(user_id) state = state.copy_and_replace( state=PresenceState.BUSY, last_active_ts=now - IDLE_TIMER - 1, last_user_sync_ts=now, status_msg=status_msg, ) new_state = handle_timeout(state, is_mine=True, syncing_user_ids=set(), now=now) self.assertIsNotNone(new_state) self.assertEquals(new_state.state, PresenceState.BUSY) self.assertEquals(new_state.status_msg, status_msg)
async def current_state_for_users( self, user_ids: Iterable[str] ) -> Dict[str, UserPresenceState]: """Get the current presence state for multiple users. Returns: dict: `user_id` -> `UserPresenceState` """ states = { user_id: self.user_to_current_state.get(user_id, None) for user_id in user_ids } missing = [user_id for user_id, state in states.items() if not state] if missing: # There are things not in our in memory cache. Lets pull them out of # the database. res = await self.store.get_presence_for_users(missing) states.update(res) missing = [user_id for user_id, state in states.items() if not state] if missing: new = { user_id: UserPresenceState.default(user_id) for user_id in missing } states.update(new) self.user_to_current_state.update(new) return states
def test_online_to_idle(self): wheel_timer = Mock() user_id = "@foo:bar" now = 5000000 prev_state = UserPresenceState.default(user_id) prev_state = prev_state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now, currently_active=True ) new_state = prev_state.copy_and_replace(state=PresenceState.UNAVAILABLE) state, persist_and_notify, federation_ping = handle_update( prev_state, new_state, is_mine=True, wheel_timer=wheel_timer, now=now ) self.assertTrue(persist_and_notify) self.assertEquals(new_state.state, state.state) self.assertEquals(state.last_federation_update_ts, now) self.assertEquals(new_state.state, state.state) self.assertEquals(new_state.status_msg, state.status_msg) self.assertEquals(wheel_timer.insert.call_count, 1) wheel_timer.insert.assert_has_calls( [ call( now=now, obj=user_id, then=new_state.last_user_sync_ts + SYNC_ONLINE_TIMEOUT, ) ], any_order=True, )
def test_remote_ping_timer(self): wheel_timer = Mock() user_id = "@foo:bar" now = 5000000 prev_state = UserPresenceState.default(user_id) prev_state = prev_state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now ) new_state = prev_state.copy_and_replace(state=PresenceState.ONLINE) state, persist_and_notify, federation_ping = handle_update( prev_state, new_state, is_mine=False, wheel_timer=wheel_timer, now=now ) self.assertFalse(persist_and_notify) self.assertFalse(federation_ping) self.assertFalse(state.currently_active) self.assertEquals(new_state.state, state.state) self.assertEquals(new_state.status_msg, state.status_msg) self.assertEquals(wheel_timer.insert.call_count, 1) wheel_timer.insert.assert_has_calls( [ call( now=now, obj=user_id, then=new_state.last_federation_update_ts + FEDERATION_TIMEOUT, ) ], any_order=True, )
async def get_presence_for_all_users( self, include_offline: bool = True, ) -> Dict[str, UserPresenceState]: """Retrieve the current presence state for all users. Note that the presence_stream table is culled frequently, so it should only contain the latest presence state for each user. Args: include_offline: Whether to include offline presence states Returns: A dict of user IDs to their current UserPresenceState. """ users_to_state = {} exclude_keyvalues = None if not include_offline: # Exclude offline presence state exclude_keyvalues = {"state": "offline"} # This may be a very heavy database query. # We paginate in order to not block a database connection. limit = 100 offset = 0 while True: rows = await self.db_pool.runInteraction( "get_presence_for_all_users", self.db_pool.simple_select_list_paginate_txn, "presence_stream", orderby="stream_id", start=offset, limit=limit, exclude_keyvalues=exclude_keyvalues, retcols=( "user_id", "state", "last_active_ts", "last_federation_update_ts", "last_user_sync_ts", "status_msg", "currently_active", ), order_direction="ASC", ) for row in rows: users_to_state[row["user_id"]] = UserPresenceState(**row) # We've run out of updates to query if len(rows) < limit: break offset += limit return users_to_state
def test_send_and_get_split(self): state1 = UserPresenceState.default("@user1:test") state2 = UserPresenceState.default("@user2:test") state3 = UserPresenceState.default("@user3:test") prev_token = self.queue.get_current_token(self.instance_name) self.queue.send_presence_to_destinations((state1, state2), ("dest1", "dest2")) now_token = self.queue.get_current_token(self.instance_name) self.queue.send_presence_to_destinations((state3, ), ("dest3", )) rows, upto_token, limited = self.get_success( self.queue.get_replication_rows("master", prev_token, now_token, 10)) self.assertEqual(upto_token, now_token) self.assertFalse(limited) expected_rows = [ (1, ("dest1", "@user1:test")), (1, ("dest2", "@user1:test")), (1, ("dest1", "@user2:test")), (1, ("dest2", "@user2:test")), ] self.assertCountEqual(rows, expected_rows) now_token = self.queue.get_current_token(self.instance_name) rows, upto_token, limited = self.get_success( self.queue.get_replication_rows("master", upto_token, now_token, 10)) self.assertEqual(upto_token, now_token) self.assertFalse(limited) expected_rows = [ (2, ("dest3", "@user3:test")), ] self.assertCountEqual(rows, expected_rows)
async def get_states( self, target_user_ids: Iterable[str] ) -> List[UserPresenceState]: """Get the presence state for users.""" updates_d = await self.current_state_for_users(target_user_ids) updates = list(updates_d.values()) for user_id in set(target_user_ids) - {u.user_id for u in updates}: updates.append(UserPresenceState.default(user_id)) return updates
async def _handle_timeouts(self): """Checks the presence of users that have timed out and updates as appropriate. """ logger.debug("Handling presence timeouts") now = self.clock.time_msec() # Fetch the list of users that *may* have timed out. Things may have # changed since the timeout was set, so we won't necessarily have to # take any action. users_to_check = set(self.wheel_timer.fetch(now)) # Check whether the lists of syncing processes from an external # process have expired. expired_process_ids = [ process_id for process_id, last_update in self.external_process_last_updated_ms.items() if now - last_update > EXTERNAL_PROCESS_EXPIRY ] for process_id in expired_process_ids: # For each expired process drop tracking info and check the users # that were syncing on that process to see if they need to be timed # out. users_to_check.update( self.external_process_to_current_syncs.pop(process_id, ()) ) self.external_process_last_updated_ms.pop(process_id) states = [ self.user_to_current_state.get(user_id, UserPresenceState.default(user_id)) for user_id in users_to_check ] timers_fired_counter.inc(len(states)) syncing_user_ids = { user_id for user_id, count in self.user_to_num_current_syncs.items() if count } for user_ids in self.external_process_to_current_syncs.values(): syncing_user_ids.update(user_ids) changes = handle_timeouts( states, is_mine_fn=self.is_mine_id, syncing_user_ids=syncing_user_ids, now=now, ) return await self._update_states(changes)
def test_no_timeout(self): user_id = "@foo:bar" now = 5000000 state = UserPresenceState.default(user_id) state = state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now, last_user_sync_ts=now, last_federation_update_ts=now, ) new_state = handle_timeout(state, is_mine=True, syncing_user_ids=set(), now=now) self.assertIsNone(new_state)
def test_sync_timeout(self): user_id = "@foo:bar" now = 5000000 state = UserPresenceState.default(user_id) state = state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=0, last_user_sync_ts=now - SYNC_ONLINE_TIMEOUT - 1, ) new_state = handle_timeout(state, is_mine=True, syncing_user_ids=set(), now=now) self.assertIsNotNone(new_state) self.assertEquals(new_state.state, PresenceState.OFFLINE)
def test_online_to_online_last_active_noop(self): wheel_timer = Mock() user_id = "@foo:bar" now = 5000000 prev_state = UserPresenceState.default(user_id) prev_state = prev_state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now - LAST_ACTIVE_GRANULARITY - 10, currently_active=True, ) new_state = prev_state.copy_and_replace(state=PresenceState.ONLINE, last_active_ts=now) state, persist_and_notify, federation_ping = handle_update( prev_state, new_state, is_mine=True, wheel_timer=wheel_timer, now=now) self.assertFalse(persist_and_notify) self.assertTrue(federation_ping) self.assertTrue(state.currently_active) self.assertEquals(new_state.state, state.state) self.assertEquals(new_state.status_msg, state.status_msg) self.assertEquals(state.last_federation_update_ts, now) self.assertEquals(wheel_timer.insert.call_count, 3) wheel_timer.insert.assert_has_calls( [ call(now=now, obj=user_id, then=new_state.last_active_ts + IDLE_TIMER), call( now=now, obj=user_id, then=new_state.last_user_sync_ts + SYNC_ONLINE_TIMEOUT, ), call( now=now, obj=user_id, then=new_state.last_active_ts + LAST_ACTIVE_GRANULARITY, ), ], any_order=True, )
def _get_active_presence(self, db_conn: Connection): """Fetch non-offline presence from the database so that we can register the appropriate time outs. """ sql = ( "SELECT user_id, state, last_active_ts, last_federation_update_ts," " last_user_sync_ts, status_msg, currently_active FROM presence_stream" " WHERE state != ?") txn = db_conn.cursor() txn.execute(sql, (PresenceState.OFFLINE, )) rows = self.db_pool.cursor_to_dict(txn) txn.close() for row in rows: row["currently_active"] = bool(row["currently_active"]) return [UserPresenceState(**row) for row in rows]
def test_online_to_offline(self): wheel_timer = Mock() user_id = "@foo:bar" now = 5000000 prev_state = UserPresenceState.default(user_id) prev_state = prev_state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now, currently_active=True ) new_state = prev_state.copy_and_replace(state=PresenceState.OFFLINE) state, persist_and_notify, federation_ping = handle_update( prev_state, new_state, is_mine=True, wheel_timer=wheel_timer, now=now ) self.assertTrue(persist_and_notify) self.assertEquals(new_state.state, state.state) self.assertEquals(state.last_federation_update_ts, now) self.assertEquals(wheel_timer.insert.call_count, 0)
def test_federation_ping(self): user_id = "@foo:bar" status_msg = "I'm here!" now = 5000000 state = UserPresenceState.default(user_id) state = state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now, last_user_sync_ts=now, last_federation_update_ts=now - FEDERATION_PING_INTERVAL - 1, status_msg=status_msg, ) new_state = handle_timeout(state, is_mine=True, syncing_user_ids=set(), now=now) self.assertIsNotNone(new_state) self.assertEqual(state, new_state)
def test_last_active(self): user_id = "@foo:bar" status_msg = "I'm here!" now = 5000000 state = UserPresenceState.default(user_id) state = state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now - LAST_ACTIVE_GRANULARITY - 1, last_user_sync_ts=now, last_federation_update_ts=now, status_msg=status_msg, ) new_state = handle_timeout(state, is_mine=True, syncing_user_ids=set(), now=now) self.assertIsNotNone(new_state) self.assertEquals(state, new_state)
def test_idle_timer(self): user_id = "@foo:bar" status_msg = "I'm here!" now = 5000000 state = UserPresenceState.default(user_id) state = state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now - IDLE_TIMER - 1, last_user_sync_ts=now, status_msg=status_msg, ) new_state = handle_timeout(state, is_mine=True, syncing_user_ids=set(), now=now) self.assertIsNotNone(new_state) self.assertEquals(new_state.state, PresenceState.UNAVAILABLE) self.assertEquals(new_state.status_msg, status_msg)
def test_persisting_presence_updates(self): """Tests that the latest presence state for each user is persisted correctly""" # Create some test users and presence states for them presence_states = [] for i in range(5): user_id = self.register_user(f"user_{i}", "password") presence_state = UserPresenceState( user_id=user_id, state="online", last_active_ts=1, last_federation_update_ts=1, last_user_sync_ts=1, status_msg="I'm online!", currently_active=True, ) presence_states.append(presence_state) # Persist these presence updates to the database self.get_success(self.store.update_presence(presence_states)) # Check that each update is present in the database db_presence_states = self.get_success( self.store.get_all_presence_updates( instance_name="master", last_id=0, current_id=len(presence_states) + 1, limit=len(presence_states), )) # Extract presence update user ID and state information into lists of tuples db_presence_states = [(ps[0], ps[1]) for _, ps in db_presence_states[0]] presence_states_compare = [(ps.user_id, ps.state) for ps in presence_states] # Compare what we put into the storage with what we got out. # They should be identical. self.assertEqual(presence_states_compare, db_presence_states)
async def get_presence_for_users(self, user_ids): rows = await self.db_pool.simple_select_many_batch( table="presence_stream", column="user_id", iterable=user_ids, keyvalues={}, retcols=( "user_id", "state", "last_active_ts", "last_federation_update_ts", "last_user_sync_ts", "status_msg", "currently_active", ), desc="get_presence_for_users", ) for row in rows: row["currently_active"] = bool(row["currently_active"]) return {row["user_id"]: UserPresenceState(**row) for row in rows}
def test_sync_online(self): user_id = "@foo:bar" status_msg = "I'm here!" now = 5000000 state = UserPresenceState.default(user_id) state = state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now - SYNC_ONLINE_TIMEOUT - 1, last_user_sync_ts=now - SYNC_ONLINE_TIMEOUT - 1, status_msg=status_msg, ) new_state = handle_timeout(state, is_mine=True, syncing_user_ids={user_id}, now=now) self.assertIsNotNone(new_state) assert new_state is not None self.assertEqual(new_state.state, PresenceState.ONLINE) self.assertEqual(new_state.status_msg, status_msg)
def test_federation_timeout(self): user_id = "@foo:bar" status_msg = "I'm here!" now = 5000000 state = UserPresenceState.default(user_id) state = state.copy_and_replace( state=PresenceState.ONLINE, last_active_ts=now, last_user_sync_ts=now, last_federation_update_ts=now - FEDERATION_TIMEOUT - 1, status_msg=status_msg, ) new_state = handle_timeout(state, is_mine=False, syncing_user_ids=set(), now=now) self.assertIsNotNone(new_state) assert new_state is not None self.assertEqual(new_state.state, PresenceState.OFFLINE) self.assertEqual(new_state.status_msg, status_msg)
async def _update_states(self, new_states): """Updates presence of users. Sets the appropriate timeouts. Pokes the notifier and federation if and only if the changed presence state should be sent to clients/servers. """ now = self.clock.time_msec() with Measure(self.clock, "presence_update_states"): # NOTE: We purposefully don't await between now and when we've # calculated what we want to do with the new states, to avoid races. to_notify = {} # Changes we want to notify everyone about to_federation_ping = {} # These need sending keep-alives # Only bother handling the last presence change for each user new_states_dict = {} for new_state in new_states: new_states_dict[new_state.user_id] = new_state new_state = new_states_dict.values() for new_state in new_states: user_id = new_state.user_id # Its fine to not hit the database here, as the only thing not in # the current state cache are OFFLINE states, where the only field # of interest is last_active which is safe enough to assume is 0 # here. prev_state = self.user_to_current_state.get( user_id, UserPresenceState.default(user_id) ) new_state, should_notify, should_ping = handle_update( prev_state, new_state, is_mine=self.is_mine_id(user_id), wheel_timer=self.wheel_timer, now=now, ) self.user_to_current_state[user_id] = new_state if should_notify: to_notify[user_id] = new_state elif should_ping: to_federation_ping[user_id] = new_state # TODO: We should probably ensure there are no races hereafter presence_updates_counter.inc(len(new_states)) if to_notify: notified_presence_counter.inc(len(to_notify)) await self._persist_and_notify(list(to_notify.values())) self.unpersisted_users_changes |= {s.user_id for s in new_states} self.unpersisted_users_changes -= set(to_notify.keys()) to_federation_ping = { user_id: state for user_id, state in to_federation_ping.items() if user_id not in to_notify } if to_federation_ping: federation_presence_out_counter.inc(len(to_federation_ping)) self._push_to_remotes(to_federation_ping.values())
def from_data(data): return PresenceDestinationsRow(state=UserPresenceState.from_dict( data["state"]), destinations=data["dests"])
def from_data(data): return PresenceRow(state=UserPresenceState.from_dict(data))