Skip to main content

grammers_session/storages/
memory.rs

1// Copyright 2020 - developers of the `grammers` project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9use std::fmt;
10use std::sync::{Mutex, MutexGuard};
11
12use crate::types::{ChannelState, DcOption, PeerId, PeerInfo, UpdateState, UpdatesState};
13use crate::{BoxFuture, Session, SessionData};
14
15/// In-memory session interface.
16///
17/// Does not actually offer direct ways to persist the state anywhere,
18/// so it should only be used in very few select cases.
19///
20/// Logging in has a very high cost in terms of flood wait errors,
21/// so the state really should be persisted by other means.
22#[derive(Default)]
23pub struct MemorySession(Mutex<SessionData>);
24
25impl From<SessionData> for MemorySession {
26    /// Constructs a memory session from the entirety of the session data,
27    /// unlike the blanket `From` implementation which cannot import all values
28    fn from(session_data: SessionData) -> Self {
29        Self(Mutex::new(session_data))
30    }
31}
32
33#[derive(Debug)]
34pub enum MemorySessionError {
35    Poisoned,
36}
37
38impl std::error::Error for MemorySessionError {}
39
40impl fmt::Display for MemorySessionError {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        match self {
43            MemorySessionError::Poisoned => write!(f, "session lock is poisoned"),
44        }
45    }
46}
47
48impl MemorySession {
49    /// Lock the session data, mapping a poisoned mutex into [`MemorySessionError`].
50    fn data(&self) -> Result<MutexGuard<'_, SessionData>, MemorySessionError> {
51        self.0.lock().map_err(|_| MemorySessionError::Poisoned)
52    }
53}
54
55impl Session for MemorySession {
56    type Error = MemorySessionError;
57
58    fn home_dc_id(&self) -> Result<i32, MemorySessionError> {
59        Ok(self.data()?.home_dc)
60    }
61
62    fn set_home_dc_id(&self, dc_id: i32) -> BoxFuture<'_, Result<(), MemorySessionError>> {
63        Box::pin(async move {
64            self.data()?.home_dc = dc_id;
65            Ok(())
66        })
67    }
68
69    fn dc_option(&self, dc_id: i32) -> Result<Option<DcOption>, MemorySessionError> {
70        Ok(self.data()?.dc_options.get(&dc_id).cloned())
71    }
72
73    fn set_dc_option(&self, dc_option: &DcOption) -> BoxFuture<'_, Result<(), MemorySessionError>> {
74        let dc_option = dc_option.clone();
75        Box::pin(async move {
76            self.data()?
77                .dc_options
78                .insert(dc_option.id, dc_option.clone());
79            Ok(())
80        })
81    }
82
83    fn peer(&self, peer: PeerId) -> BoxFuture<'_, Result<Option<PeerInfo>, MemorySessionError>> {
84        Box::pin(async move { Ok(self.data()?.peer_infos.get(&peer).cloned()) })
85    }
86
87    fn cache_peer(&self, peer: &PeerInfo) -> BoxFuture<'_, Result<(), MemorySessionError>> {
88        let peer = peer.clone();
89        Box::pin(async move {
90            self.data()?
91                .peer_infos
92                .entry(peer.id())
93                .or_insert_with(|| peer.clone())
94                .extend_info(&peer);
95            Ok(())
96        })
97    }
98
99    fn updates_state(&self) -> BoxFuture<'_, Result<UpdatesState, MemorySessionError>> {
100        Box::pin(async move { Ok(self.data()?.updates_state.clone()) })
101    }
102
103    fn set_update_state(
104        &self,
105        update: UpdateState,
106    ) -> BoxFuture<'_, Result<(), MemorySessionError>> {
107        Box::pin(async move {
108            let mut data = self.data()?;
109
110            match update {
111                UpdateState::All(updates_state) => {
112                    data.updates_state = updates_state;
113                }
114                UpdateState::Primary { pts, date, seq } => {
115                    data.updates_state.pts = pts;
116                    data.updates_state.date = date;
117                    data.updates_state.seq = seq;
118                }
119                UpdateState::Secondary { qts } => {
120                    data.updates_state.qts = qts;
121                }
122                UpdateState::Channel { id, pts } => {
123                    data.updates_state.channels.retain(|c| c.id != id);
124                    data.updates_state.channels.push(ChannelState { id, pts });
125                }
126            }
127
128            Ok(())
129        })
130    }
131}