1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
use lru::LruCache;
use std::sync::Arc;
use tokio::sync::RwLock;
use std::num::NonZeroUsize;
use tokio::sync::mpsc::{channel, Receiver, Sender};
use twilight_model::channel::Message as TwiMessage;
use super::{MemberId, Message as GatewayMessage, MessageEvent, MessageId, SystemEvent};
pub struct AggregatorState {
rx: Receiver<MessageEvent>,
tx: Sender<MessageEvent>,
message_cache: lru::LruCache<MessageId, (TwiMessage, MemberId)>,
system_emitter: Option<Sender<SystemEvent>>,
}
pub struct MessageAggregator {
state: Arc<RwLock<AggregatorState>>,
}
impl MessageAggregator {
pub fn new(system_size: usize) -> Self {
let buf_size = std::cmp::max(system_size * 2, 1);
let (tx, rx) = channel::<MessageEvent>(buf_size);
Self {
state: Arc::new(RwLock::new( AggregatorState {
tx,
rx,
message_cache: LruCache::new(NonZeroUsize::new(buf_size).unwrap()),
system_emitter: None,
}))
}
}
pub async fn get_sender(&self) -> Sender<MessageEvent> {
self.state.read().await.tx.clone()
}
pub async fn set_system_handler(&mut self, emitter: Sender<SystemEvent>) -> () {
self.state.write().await.system_emitter = Some(emitter);
}
// We probably don't actully need this since we've got a separate sent-cache by channel
// pub async fn lookup_message(&self, message_id: MessageId) -> Option<TwiMessage> {
// self.state.write().await.message_cache.get(&message_id).map(|m| m.clone())
// }
pub fn start(&self) -> () {
let state = self.state.clone();
tokio::spawn(async move {
loop {
let system_emitter = { state.read().await.system_emitter.clone().expect("No system emitter") };
let self_emitter = { state.read().await.tx.clone() };
let next_event = { state.write().await.rx.recv().await };
match next_event {
None => (),
Some((timestamp, message)) => {
match message {
GatewayMessage::Partial(current_partial, member_id) => {
let cache_content = { state.write().await.message_cache.get(¤t_partial.id).map(|m| m.clone()) };
match cache_content {
Some((original_message, member_id)) => {
let mut updated_message = original_message.clone();
if let Some(edited_time) = current_partial.edited_timestamp {
updated_message.edited_timestamp = Some(edited_time);
}
if let Some(content) = ¤t_partial.content {
updated_message.content = content.clone()
}
self_emitter.send((timestamp, GatewayMessage::Complete(updated_message, member_id))).await;
},
None => {
system_emitter.send(
SystemEvent::RefetchMessage(member_id, current_partial.id, current_partial.channel_id)
).await;
},
};
},
GatewayMessage::Complete(message, member_id) => {
let previous_message = { state.write().await.message_cache.get(&message.id).map(|m| m.clone()) };
if let Some((previous_message, _last_seen_by)) = previous_message {
let previous_timestamp = previous_message.edited_timestamp.unwrap_or(previous_message.timestamp);
let current_timestamp = message.edited_timestamp.unwrap_or(message.timestamp);
// Should we skip sending
if previous_timestamp.as_micros() >= current_timestamp.as_micros() {
continue
}
// If not, fall through to update stored message
}
{ state.write().await.message_cache.put(message.id, (message.clone(), member_id)); };
system_emitter
.send(SystemEvent::NewMessage(timestamp, message, member_id))
.await;
},
};
}
}
}
});
}
}
|
