atagen

macltok

public
Star0
Watch0
fbeaaa098107ea3a2693391aaa175f1c17cf2899
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/engine/context.rsRust546 lines16.6 KB
1
use std::{
2
    collections::{HashMap, VecDeque},
3
    sync::Arc,
4
};
5
6
use futures_util::future::{join_all, try_join};
7
use matrix_sdk::{
8
    Client, Room, RoomState, deserialized_responses::TimelineEvent as SdkTimelineEvent,
9
    room::MessagesOptions,
10
};
11
use ruma::{OwnedEventId, OwnedRoomId, OwnedUserId, UInt, push::Action};
12
use serde_json::Value;
13
use tokio::sync::{Mutex, broadcast::error::RecvError};
14
15
use crate::{
16
    Error, ErrorKind, Event, EventContext, EventContextCloseReason, EventContextCommand,
17
    EventContextPage, EventContextSession, Sender, TimelineEvent, TimelineItem,
18
    engine::MediaService,
19
    normalise::{
20
        MediaRegistration, MediaScope, apply_profiles, member_sender, normalise_context,
21
        normalise_observed, resolve_replies,
22
    },
23
    transport::EventSender,
24
};
25
26
const MAX_CONTEXT_LIMIT: u16 = 50;
27
/// a consumer that never closes its jump views still holds at most this many
28
const MAX_SESSIONS: usize = 32;
29
30
/// open event contexts, oldest first, each owning the media its pages registered
31
pub(crate) struct Sessions {
32
    open: Mutex<VecDeque<Session>>,
33
    events: EventSender,
34
}
35
36
struct Session {
37
    id: String,
38
    room_id: OwnedRoomId,
39
    scope: MediaScope,
40
}
41
42
impl Session {
43
    /// media ids carry the session, so two owners of one event never share an entry
44
    fn media_prefix(scope: &MediaScope) -> String {
45
        format!("{}#", scope.nonce())
46
    }
47
}
48
49
impl Sessions {
50
    pub(crate) fn new(events: EventSender) -> Self {
51
        Self {
52
            open: Mutex::default(),
53
            events,
54
        }
55
    }
56
57
    async fn open(
58
        &self,
59
        media: &MediaService,
60
        room_id: OwnedRoomId,
61
        scope: MediaScope,
62
        registrations: Vec<MediaRegistration>,
63
    ) -> String {
64
        let id = format!("context-{}", scope.nonce());
65
        let mut evicted = Vec::new();
66
        {
67
            let mut open = self.open.lock().await;
68
            media.register(&scope, registrations).await;
69
            open.push_back(Session {
70
                id: id.clone(),
71
                room_id,
72
                scope,
73
            });
74
            while open.len() > MAX_SESSIONS {
75
                if let Some(oldest) = open.pop_front() {
76
                    media.release(&oldest.scope).await;
77
                    evicted.push(oldest.id);
78
                }
79
            }
80
        }
81
        self.announce(evicted, EventContextCloseReason::Evicted)
82
            .await;
83
        id
84
    }
85
86
    /// the room and the media prefix the session's pages use
87
    async fn find(&self, id: &str) -> Result<(OwnedRoomId, String), Error> {
88
        let open = self.open.lock().await;
89
        let session = open
90
            .iter()
91
            .find(|session| session.id == id)
92
            .ok_or_else(not_open)?;
93
        Ok((
94
            session.room_id.clone(),
95
            Session::media_prefix(&session.scope),
96
        ))
97
    }
98
99
    /// under the lock, so a page that lands after its close registers nothing
100
    async fn register(
101
        &self,
102
        media: &MediaService,
103
        id: &str,
104
        registrations: Vec<MediaRegistration>,
105
    ) -> Result<(), Error> {
106
        let open = self.open.lock().await;
107
        let session = open
108
            .iter()
109
            .find(|session| session.id == id)
110
            .ok_or_else(not_open)?;
111
        media.register(&session.scope, registrations).await;
112
        Ok(())
113
    }
114
115
    async fn close(&self, media: &MediaService, id: &str) -> Result<(), Error> {
116
        let mut open = self.open.lock().await;
117
        let index = open
118
            .iter()
119
            .position(|session| session.id == id)
120
            .ok_or_else(not_open)?;
121
        if let Some(session) = open.remove(index) {
122
            media.release(&session.scope).await;
123
        }
124
        Ok(())
125
    }
126
127
    async fn close_left(&self, media: &MediaService, client: &Client) {
128
        let mut closed = Vec::new();
129
        {
130
            let mut open = self.open.lock().await;
131
            let mut kept = VecDeque::with_capacity(open.len());
132
            for session in open.drain(..) {
133
                let left = client
134
                    .get_room(&session.room_id)
135
                    .is_none_or(|room| matches!(room.state(), RoomState::Left | RoomState::Banned));
136
                if left {
137
                    media.release(&session.scope).await;
138
                    closed.push(session.id);
139
                } else {
140
                    kept.push_back(session);
141
                }
142
            }
143
            *open = kept;
144
        }
145
        self.announce(closed, EventContextCloseReason::RoomLeft)
146
            .await;
147
    }
148
149
    async fn announce(&self, ids: Vec<String>, reason: EventContextCloseReason) {
150
        for context_id in ids {
151
            let event = Event::EventContextClosed {
152
                context_id,
153
                reason: reason.clone(),
154
            };
155
            if self.events.send(event).await.is_err() {
156
                return;
157
            }
158
        }
159
    }
160
161
    /// the media service is cleared beside this, so nothing is released one by one
162
    pub(crate) async fn clear(&self) {
163
        self.open.lock().await.clear();
164
    }
165
}
166
167
fn not_open() -> Error {
168
    Error::new(
169
        ErrorKind::EventContextNotOpen,
170
        "event context is closed or was never opened",
171
    )
172
}
173
174
/// leaving a room closes its event contexts
175
pub(super) async fn close_on_leave(
176
    client: Client,
177
    media: Arc<MediaService>,
178
    sessions: Arc<Sessions>,
179
) {
180
    let mut info = client.room_info_notable_update_receiver();
181
    loop {
182
        match info.recv().await {
183
            Ok(_) | Err(RecvError::Lagged(_)) => sessions.close_left(&media, &client).await,
184
            Err(RecvError::Closed) => return,
185
        }
186
    }
187
}
188
189
pub(crate) async fn execute(
190
    client: Client,
191
    media: Arc<MediaService>,
192
    sessions: Arc<Sessions>,
193
    command: EventContextCommand,
194
) -> Result<crate::CommandOutput, Error> {
195
    let current_user_id = client.user_id().map(ToOwned::to_owned);
196
    let room = |room_id: &OwnedRoomId| client.get_room(room_id).ok_or_else(Error::room_unknown);
197
    let (context_id, token, limit, backwards) = match command {
198
        EventContextCommand::Fetch {
199
            room_id,
200
            event_id,
201
            before_limit,
202
            after_limit,
203
        } => {
204
            let scope = media.context_scope();
205
            let (context, registrations, _) = capture_in(
206
                room(&room_id)?,
207
                current_user_id,
208
                event_id,
209
                before_limit,
210
                after_limit,
211
                &Session::media_prefix(&scope),
212
            )
213
            .await?;
214
            let context_id = sessions.open(&media, room_id, scope, registrations).await;
215
            return Ok(crate::CommandOutput::EventContext(EventContextSession {
216
                context_id,
217
                context,
218
            }));
219
        }
220
        EventContextCommand::Close { context_id } => {
221
            sessions.close(&media, &context_id).await?;
222
            return Ok(crate::CommandOutput::Done);
223
        }
224
        EventContextCommand::Before {
225
            context_id,
226
            token,
227
            limit,
228
        } => (context_id, token, limit, true),
229
        EventContextCommand::After {
230
            context_id,
231
            token,
232
            limit,
233
        } => (context_id, token, limit, false),
234
    };
235
    if limit == 0 {
236
        return Err(Error::new(
237
            ErrorKind::Invalid,
238
            "context page limit must be at least 1",
239
        ));
240
    }
241
    validate_limit(limit)?;
242
    if token.is_empty() {
243
        return Err(Error::new(
244
            ErrorKind::Invalid,
245
            "context pagination token is empty",
246
        ));
247
    }
248
    let (room_id, media_prefix) = sessions.find(&context_id).await?;
249
    let room = room(&room_id)?;
250
    let side = fetch_side(&room, Some(&token), limit, backwards).await?;
251
    let mut page = Assembly::new(
252
        &room,
253
        current_user_id.as_deref(),
254
        &side.values,
255
        &side.highlights,
256
        None,
257
        &media_prefix,
258
    )
259
    .await;
260
    let events = page.take(0..side.values.len());
261
    sessions.register(&media, &context_id, page.media).await?;
262
    Ok(crate::CommandOutput::EventContextPage(EventContextPage {
263
        events,
264
        next_token: side.token,
265
    }))
266
}
267
268
pub(crate) async fn capture(
269
    room: Room,
270
    current_user_id: Option<OwnedUserId>,
271
    event_id: OwnedEventId,
272
    before_limit: u16,
273
    after_limit: u16,
274
) -> Result<(EventContext, Vec<MediaRegistration>, Vec<Value>), Error> {
275
    capture_in(
276
        room,
277
        current_user_id,
278
        event_id,
279
        before_limit,
280
        after_limit,
281
        "",
282
    )
283
    .await
284
}
285
286
async fn capture_in(
287
    room: Room,
288
    current_user_id: Option<OwnedUserId>,
289
    event_id: OwnedEventId,
290
    before_limit: u16,
291
    after_limit: u16,
292
    media_prefix: &str,
293
) -> Result<(EventContext, Vec<MediaRegistration>, Vec<Value>), Error> {
294
    validate_limit(before_limit)?;
295
    validate_limit(after_limit)?;
296
    let raw = room
297
        .event_with_context(&event_id, false, UInt::from(0_u16), None)
298
        .await
299
        .map_err(|source| {
300
            let error = Error::from_sdk("fetching event context failed", source);
301
            if error
302
                .server
303
                .as_ref()
304
                .is_some_and(|server| server.status == 404)
305
            {
306
                Error::new(ErrorKind::EventNotFound, "event was not found")
307
            } else {
308
                error
309
            }
310
        })?;
311
    let (target, target_highlight) = raw
312
        .event
313
        .as_ref()
314
        .and_then(raw_event)
315
        .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "event was not found"))?;
316
    let (before, after) = try_join(
317
        fetch_side(&room, raw.prev_batch_token.as_deref(), before_limit, true),
318
        fetch_side(&room, raw.next_batch_token.as_deref(), after_limit, false),
319
    )
320
    .await?;
321
    let target_index = before.values.len();
322
    let mut values = before.values;
323
    values.push(target);
324
    values.extend(after.values);
325
    let mut highlights = before.highlights;
326
    highlights.push(target_highlight);
327
    highlights.extend(after.highlights);
328
    let mut page = Assembly::new(
329
        &room,
330
        current_user_id.as_deref(),
331
        &values,
332
        &highlights,
333
        Some(target_index),
334
        media_prefix,
335
    )
336
    .await;
337
    let events_before = page.take(0..target_index);
338
    let event = page
339
        .take(target_index..target_index + 1)
340
        .pop()
341
        .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "event could not be read"))?;
342
    let events_after = page.take(target_index + 1..values.len());
343
    Ok((
344
        EventContext {
345
            events_before,
346
            event,
347
            events_after,
348
            start_token: before.token,
349
            end_token: after.token,
350
        },
351
        page.media,
352
        values,
353
    ))
354
}
355
356
struct Side {
357
    /// in timeline order
358
    values: Vec<Value>,
359
    /// beside `values`, from the push actions the sdk worked out as it fetched
360
    highlights: Vec<bool>,
361
    token: Option<String>,
362
}
363
364
fn raw_event(event: &SdkTimelineEvent) -> Option<(Value, bool)> {
365
    let value = serde_json::from_str(event.raw().json().get()).ok()?;
366
    let highlight = event
367
        .push_actions()
368
        .is_some_and(|actions| actions.iter().any(Action::is_highlight));
369
    Some((value, highlight))
370
}
371
372
async fn fetch_side(
373
    room: &Room,
374
    token: Option<&str>,
375
    limit: u16,
376
    backwards: bool,
377
) -> Result<Side, Error> {
378
    let Some(token) = token else {
379
        return Ok(Side {
380
            values: Vec::new(),
381
            highlights: Vec::new(),
382
            token: None,
383
        });
384
    };
385
    if limit == 0 {
386
        return Ok(Side {
387
            values: Vec::new(),
388
            highlights: Vec::new(),
389
            token: Some(token.to_owned()),
390
        });
391
    }
392
    let mut options = if backwards {
393
        MessagesOptions::backward()
394
    } else {
395
        MessagesOptions::forward()
396
    };
397
    options = options.from(token);
398
    options.limit = UInt::from(limit);
399
    let messages = room
400
        .messages(options)
401
        .await
402
        .map_err(|source| Error::from_sdk("fetching event context events failed", source))?;
403
    let mut events: Vec<_> = messages.chunk.iter().filter_map(raw_event).collect();
404
    if backwards {
405
        events.reverse();
406
    }
407
    let (values, highlights) = events.into_iter().unzip();
408
    Ok(Side {
409
        values,
410
        highlights,
411
        token: messages.end,
412
    })
413
}
414
415
/// a window of raw events normalised together, so relations inside it land on their targets
416
struct Assembly {
417
    items: Vec<(usize, TimelineItem)>,
418
    media: Vec<MediaRegistration>,
419
}
420
421
impl Assembly {
422
    async fn new(
423
        room: &Room,
424
        current_user_id: Option<&ruma::UserId>,
425
        values: &[Value],
426
        highlights: &[bool],
427
        keep: Option<usize>,
428
        media_prefix: &str,
429
    ) -> Self {
430
        let page = normalise_context(values, keep, current_user_id, room.room_id(), media_prefix);
431
        let mut items = page.items;
432
        for (index, item) in &mut items {
433
            if let TimelineItem::Event(event) = item {
434
                event.highlight = highlights[*index];
435
            }
436
        }
437
        let mut media = page.media;
438
        let mut targets: HashMap<OwnedEventId, TimelineEvent> = HashMap::new();
439
        let mut missing = Vec::new();
440
        for (_, item) in &items {
441
            if let TimelineItem::Event(event) = item
442
                && let Some(event_id) = &event.event_id
443
            {
444
                targets.insert(event_id.clone(), (**event).clone());
445
            }
446
        }
447
        for event in events(&mut items) {
448
            if let Some(reply) = &event.reply_to
449
                && !targets.contains_key(&reply.event_id)
450
                && !missing.contains(&reply.event_id)
451
            {
452
                missing.push(reply.event_id.clone());
453
            }
454
        }
455
        let fetched = join_all(
456
            missing
457
                .into_iter()
458
                .map(|event_id| fetch_reply(room, current_user_id, event_id, media_prefix)),
459
        )
460
        .await;
461
        for (event, registrations) in fetched.into_iter().flatten() {
462
            media.extend(registrations);
463
            if let Some(event_id) = event.event_id.clone() {
464
                targets.insert(event_id, event);
465
            }
466
        }
467
        let profiles = profiles(
468
            room,
469
            items
470
                .iter()
471
                .filter_map(|(_, item)| match item {
472
                    TimelineItem::Event(event) => Some(&event.sender.user_id),
473
                    _ => None,
474
                })
475
                .chain(targets.values().map(|event| &event.sender.user_id)),
476
        )
477
        .await;
478
        apply_profiles(targets.values_mut(), &profiles);
479
        apply_profiles(events(&mut items), &profiles);
480
        resolve_replies(events(&mut items), &targets);
481
        Self { items, media }
482
    }
483
484
    /// the items that came from raw events in `range`
485
    fn take(&mut self, range: std::ops::Range<usize>) -> Vec<TimelineItem> {
486
        let (taken, kept) = std::mem::take(&mut self.items)
487
            .into_iter()
488
            .partition(|(index, _)| range.contains(index));
489
        self.items = kept;
490
        taken.into_iter().map(|(_, item)| item).collect()
491
    }
492
}
493
494
fn events(items: &mut [(usize, TimelineItem)]) -> impl Iterator<Item = &mut TimelineEvent> {
495
    items.iter_mut().filter_map(|(_, item)| match item {
496
        TimelineItem::Event(event) => Some(&mut **event),
497
        _ => None,
498
    })
499
}
500
501
async fn fetch_reply(
502
    room: &Room,
503
    current_user_id: Option<&ruma::UserId>,
504
    event_id: OwnedEventId,
505
    media_prefix: &str,
506
) -> Option<(TimelineEvent, Vec<MediaRegistration>)> {
507
    let event = room
508
        .load_or_fetch_event(&event_id, None)
509
        .await
510
        .inspect_err(|error| tracing::debug!(%event_id, %error, "fetching a reply target failed"))
511
        .ok()?;
512
    let value: Value = serde_json::from_str(event.raw().json().get()).ok()?;
513
    let identity = format!("{media_prefix}{}\n{event_id}", room.room_id());
514
    match normalise_observed(&value, current_user_id, &identity)? {
515
        (TimelineItem::Event(event), media) => Some((*event, media)),
516
        _ => None,
517
    }
518
}
519
520
/// under lazy-loaded members a sender may be unknown locally, and `get_member` fetches the
521
/// member list once for the room when it is
522
async fn profiles<'a>(
523
    room: &Room,
524
    user_ids: impl Iterator<Item = &'a OwnedUserId>,
525
) -> HashMap<OwnedUserId, Sender> {
526
    let mut unique: Vec<&OwnedUserId> = user_ids.collect();
527
    unique.sort();
528
    unique.dedup();
529
    join_all(
530
        unique.into_iter().map(|user_id| async move {
531
            Some((user_id.clone(), member_sender(room, user_id).await?))
532
        }),
533
    )
534
    .await
535
    .into_iter()
536
    .flatten()
537
    .collect()
538
}
539
540
fn validate_limit(limit: u16) -> Result<(), Error> {
541
    if limit > MAX_CONTEXT_LIMIT {
542
        Err(Error::new(ErrorKind::Invalid, "context limit exceeds 50"))
543
    } else {
544
        Ok(())
545
    }
546
}