Skip to main content

classicube_relay/events/
store.rs

1use std::{
2    collections::HashMap,
3    time::{Duration, Instant},
4};
5
6use tracing::{debug, error, warn};
7
8use super::{CallbackFn, PartialStream, PartialStreamError};
9use crate::packet::{ContinuePacket, Packet, PlayerScope, Scope, StartPacket};
10
11#[derive(Debug, thiserror::Error)]
12pub enum StoreError {
13    #[error("got non-Player scope")]
14    Thing,
15
16    #[error("found continue packet before start")]
17    Thing2,
18
19    #[error(transparent)]
20    PartialStream(#[from] PartialStreamError),
21
22    #[error(transparent)]
23    Io(#[from] std::io::Error),
24}
25type Result<T> = std::result::Result<T, StoreError>;
26
27#[derive(Default)]
28pub(crate) struct Store {
29    pub(crate) event_handlers: Vec<CallbackFn>,
30    streams: HashMap<u8, PartialStream>,
31    cleanup_times: HashMap<u8, Instant>,
32}
33
34impl Store {
35    const STREAM_TIMEOUT: Duration = Duration::from_secs(10);
36
37    pub(crate) fn process_packet(&mut self, packet: Packet) -> Result<()> {
38        debug!("process_packet {:?}", packet);
39
40        let finished_stream = match packet {
41            Packet::Start(StartPacket {
42                stream_id,
43                scope,
44                data_length,
45                data_part,
46            }) => {
47                if let Scope::Player(PlayerScope { player_id }) = scope {
48                    debug!(stream_id, player_id, data_length, "new stream");
49
50                    let mut stream = PartialStream {
51                        player_id,
52                        data_length,
53                        data_buffer: Vec::with_capacity(data_length as usize),
54                    };
55                    stream.write_part(data_part)?;
56
57                    if let Some(old_stream) = self.streams.remove(&stream_id) {
58                        warn!("restarting stream {:?}", old_stream);
59                    }
60                    if stream.is_finished() {
61                        Some(stream)
62                    } else {
63                        self.streams.insert(stream_id, stream);
64                        self.cleanup_times
65                            .insert(stream_id, Instant::now() + Self::STREAM_TIMEOUT);
66                        None
67                    }
68                } else {
69                    return Err(StoreError::Thing);
70                }
71            }
72
73            Packet::Continue(ContinuePacket {
74                stream_id,
75                data_part,
76            }) => {
77                let is_finished = if let Some(stream) = self.streams.get_mut(&stream_id) {
78                    stream.write_part(data_part)?;
79                    debug!(
80                        stream_id,
81                        player_id = stream.player_id,
82                        current_length = stream.data_buffer.len(),
83                        data_length = stream.data_length,
84                        "continue stream"
85                    );
86
87                    stream.is_finished()
88                } else {
89                    return Err(StoreError::Thing2);
90                };
91
92                if is_finished {
93                    self.cleanup_times.remove(&stream_id);
94                    Some(self.streams.remove(&stream_id).unwrap())
95                } else {
96                    None
97                }
98            }
99        };
100
101        if let Some(stream) = finished_stream {
102            debug!("finished_stream");
103            for f in &self.event_handlers {
104                f(stream.player_id, &stream.data_buffer);
105            }
106        }
107
108        Ok(())
109    }
110
111    pub(crate) fn tick(&mut self) {
112        let now = Instant::now();
113        let mut stream_ids_to_removes = self
114            .cleanup_times
115            .iter()
116            .filter_map(|(stream_id, cleanup_time)| {
117                if &now > cleanup_time {
118                    Some(*stream_id)
119                } else {
120                    None
121                }
122            })
123            .collect::<Vec<_>>();
124
125        for stream_id in stream_ids_to_removes.drain(..) {
126            debug!(stream_id, "timed out, removing");
127            self.cleanup_times.remove(&stream_id);
128            self.streams.remove(&stream_id);
129        }
130    }
131}