max / makenotwork

8.6 KB · 276 lines History Blame Raw

1	//! `ChatIdentity`: display name, avatar and flair for the people in a room.
2	//!
3	//! # Why there is a cache here at all
4	//!
5	//! `SessionUser` carries no `avatar_url`. It lives only on the `users` row, so
6	//! every rendered message needs a join the session cannot supply. Backlog
7	//! replay after a deploy asks for a whole window of messages at once, and every
8	//! client reconnects at the same moment, so the uncached shape is one query per
9	//! author per reconnecting client.
10	//!
11	//! The crate already batches the lookup (`ChatIdentity::attach` collects
12	//! distinct authors and calls `identify` once). This cache is the second half:
13	//! it makes the batch itself mostly free, because a room's authors are a small
14	//! set that barely changes between reconnects.
15	//!
16	//! # Bounding it
17	//!
18	//! Both units run under a hard 512M cgroup cap where an OOM restarts the whole
19	//! site, so an unbounded map keyed by user id is not acceptable: the key space
20	//! is every account that has ever spoken. Entries expire, and the map is
21	//! capped; past the cap the coldest entries go. Neither bound is tuning, both
22	//! are the memory budget.
23	//!
24	//! Invalidation is on profile update rather than only by expiry, so a rename
25	//! takes effect without waiting out the TTL.
26
27	use std::collections::HashMap;
28	use std::sync::Arc;
29	use std::time::{Duration, Instant};
30
31	use async_trait::async_trait;
32	use dashmap::DashMap;
33	use livechat::{ChatError, ChatIdentity, Identity, UserId};
34	use sqlx::PgPool;
35	use uuid::Uuid;
36
37	use super::host_error;
38
39	/// How long a cached identity is trusted.
40	///
41	/// Short, because the only cost of a miss is one row and the cost of a stale
42	/// hit is a user seeing their own rename not take. Profile updates invalidate
43	/// directly, so this is the backstop for changes that arrive another way.
44	const IDENTITY_TTL: Duration = Duration::from_mins(5);
45
46	/// Most identities held at once.
47	///
48	/// Sized for the realistic shape: a few busy rooms of a few hundred people
49	/// each. Past it the cache sheds rather than grows, because growing is what
50	/// trips the cgroup cap.
51	const IDENTITY_CAPACITY: usize = 5_000;
52
53	struct Cached {
54	identity: Identity,
55	fetched_at: Instant,
56	}
57
58	/// Process-wide identity cache. Cheap to clone; clones share one map.
59	#[derive(Clone)]
60	pub struct IdentityCache {
61	entries: Arc<DashMap<Uuid, Cached>>,
62	}
63
64	impl IdentityCache {
65	pub fn new() -> Self {
66	Self {
67	entries: Arc::new(DashMap::new()),
68	}
69	}
70
71	/// Drop a user's cached identity, so the next render re-reads the row.
72	///
73	/// Called from the profile-update path. Cheap and safe to call for a user
74	/// who was never cached.
75	pub fn invalidate(&self, user: Uuid) {
76	self.entries.remove(&user);
77	}
78
79	fn get_fresh(&self, user: Uuid, now: Instant) -> Option<Identity> {
80	let entry = self.entries.get(&user)?;
81	(now.duration_since(entry.fetched_at) < IDENTITY_TTL).then(\|\| entry.identity.clone())
82	}
83
84	fn insert(&self, user: Uuid, identity: Identity, now: Instant) {
85	// Shed before inserting, so the map cannot exceed the cap even briefly.
86	// Expired entries first: they are free to drop and usually enough.
87	if self.entries.len() >= IDENTITY_CAPACITY {
88	self.entries
89	.retain(\|_, e\| now.duration_since(e.fetched_at) < IDENTITY_TTL);
90	}
91	// Still full means every entry is live. Take the coldest rather than
92	// refusing to cache, or a room busy enough to fill the map would be the
93	// one room that never gets a cache.
94	if self.entries.len() >= IDENTITY_CAPACITY
95	&& let Some(coldest) = self
96	.entries
97	.iter()
98	.min_by_key(\|e\| e.fetched_at)
99	.map(\|e\| *e.key())
100	{
101	self.entries.remove(&coldest);
102	}
103
104	self.entries.insert(
105	user,
106	Cached {
107	identity,
108	fetched_at: now,
109	},
110	);
111	}
112
113	pub fn len(&self) -> usize {
114	self.entries.len()
115	}
116
117	pub fn is_empty(&self) -> bool {
118	self.entries.is_empty()
119	}
120	}
121
122	impl Default for IdentityCache {
123	fn default() -> Self {
124	Self::new()
125	}
126	}
127
128	pub struct MtChatIdentity {
129	db: PgPool,
130	cache: IdentityCache,
131	}
132
133	impl MtChatIdentity {
134	pub fn new(db: PgPool, cache: IdentityCache) -> Self {
135	Self { db, cache }
136	}
137	}
138
139	#[async_trait]
140	impl ChatIdentity for MtChatIdentity {
141	async fn identify(&self, users: &[UserId]) -> Result<HashMap<UserId, Identity>, ChatError> {
142	let now = Instant::now();
143	let mut resolved = HashMap::with_capacity(users.len());
144	let mut misses = Vec::new();
145
146	for id in users {
147	match self.cache.get_fresh(id.0, now) {
148	Some(identity) => {
149	resolved.insert(*id, identity);
150	}
151	None => misses.push(id.0),
152	}
153	}
154
155	if misses.is_empty() {
156	return Ok(resolved);
157	}
158
159	let rows = mt_db::queries::chat_identities(&self.db, &misses)
160	.await
161	.map_err(host_error)?;
162
163	for row in rows {
164	let identity = Identity {
165	// Display name falls back to the username, which is NOT NULL.
166	// A blank name in a chat room is worse than an ugly one.
167	display_name: row.display_name.unwrap_or(row.username),
168	avatar_url: row.avatar_url,
169	// Role flair is deliberately absent: it is per-community, and
170	// this cache is process-wide. Caching a role here would show a
171	// moderator's badge in a community they moderate nothing in.
172	flair: None,
173	};
174	self.cache.insert(row.id, identity.clone(), now);
175	resolved.insert(UserId(row.id), identity);
176	}
177
178	// Anyone still missing could not be resolved: a deleted account. Left
179	// out of the map rather than errored, so their messages stay readable
180	// instead of the room failing to render.
181	Ok(resolved)
182	}
183	}
184
185	#[cfg(test)]
186	mod tests {
187	use super::*;
188
189	fn identity(name: &str) -> Identity {
190	Identity {
191	display_name: name.to_owned(),
192	avatar_url: None,
193	flair: None,
194	}
195	}
196
197	#[test]
198	fn a_fresh_entry_is_returned_and_a_stale_one_is_not() {
199	let cache = IdentityCache::new();
200	let user = Uuid::new_v4();
201	let t0 = Instant::now();
202
203	cache.insert(user, identity("alice"), t0);
204
205	assert_eq!(
206	cache.get_fresh(user, t0).map(\|i\| i.display_name),
207	Some("alice".to_owned())
208	);
209	assert!(
210	cache.get_fresh(user, t0 + IDENTITY_TTL).is_none(),
211	"the TTL boundary is exclusive"
212	);
213	}
214
215	#[test]
216	fn invalidate_makes_a_rename_visible_without_waiting_out_the_ttl() {
217	let cache = IdentityCache::new();
218	let user = Uuid::new_v4();
219	let t0 = Instant::now();
220
221	cache.insert(user, identity("old name"), t0);
222	cache.invalidate(user);
223
224	assert!(cache.get_fresh(user, t0).is_none());
225	}
226
227	#[test]
228	fn invalidating_an_uncached_user_is_harmless() {
229	IdentityCache::new().invalidate(Uuid::new_v4());
230	}
231
232	#[test]
233	fn expired_entries_are_shed_before_the_cap_bites() {
234	let cache = IdentityCache::new();
235	let t0 = Instant::now();
236
237	for _ in 0..IDENTITY_CAPACITY {
238	cache.insert(Uuid::new_v4(), identity("old"), t0);
239	}
240	assert_eq!(cache.len(), IDENTITY_CAPACITY);
241
242	// Long enough later that everything already held is expired.
243	cache.insert(Uuid::new_v4(), identity("new"), t0 + IDENTITY_TTL * 2);
244
245	assert_eq!(cache.len(), 1, "the expired generation went in one sweep");
246	}
247
248	#[test]
249	fn a_full_cache_of_live_entries_still_accepts_a_new_one() {
250	// The room busy enough to fill the map must not be the one room that
251	// never gets a cache.
252	let cache = IdentityCache::new();
253	let t0 = Instant::now();
254
255	for i in 0..IDENTITY_CAPACITY {
256	// Staggered so there is a distinct coldest entry to evict.
257	cache.insert(
258	Uuid::new_v4(),
259	identity("live"),
260	t0 + Duration::from_millis(i as u64),
261	);
262	}
263
264	let newcomer = Uuid::new_v4();
265	cache.insert(newcomer, identity("newcomer"), t0 + Duration::from_secs(1));
266
267	assert!(cache.len() <= IDENTITY_CAPACITY, "the cap holds");
268	assert!(
269	cache
270	.get_fresh(newcomer, t0 + Duration::from_secs(1))
271	.is_some(),
272	"the new entry is the one kept"
273	);
274	}
275	}
276