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))
