relay_server/services/buffer/stack_provider/
sqlite.rs

1use std::error::Error;
2use std::time::Duration;
3
4use relay_config::ConfigSnapshot;
5
6use crate::services::buffer::common::ProjectKeyPair;
7use crate::services::buffer::envelope_stack::caching::CachingEnvelopeStack;
8use crate::services::buffer::envelope_store::sqlite::{
9    SqliteEnvelopeStore, SqliteEnvelopeStoreError,
10};
11use crate::services::buffer::stack_provider::{
12    InitializationState, StackCreationType, StackProvider,
13};
14use crate::statsd::RelayTimers;
15use crate::{EnvelopeStack, SqliteEnvelopeStack};
16
17#[derive(Debug)]
18pub struct SqliteStackProvider {
19    envelope_store: SqliteEnvelopeStore,
20    batch_size_bytes: usize,
21    flush_timeout: Option<Duration>,
22    max_disk_size: usize,
23    partition_id: u8,
24    ephemeral: bool,
25}
26
27#[warn(dead_code)]
28impl SqliteStackProvider {
29    /// Creates a new [`SqliteStackProvider`] from the provided [`ConfigSnapshot`].
30    pub async fn new(
31        partition_id: u8,
32        config: &ConfigSnapshot,
33    ) -> Result<Self, SqliteEnvelopeStoreError> {
34        let envelope_store = SqliteEnvelopeStore::prepare(partition_id, config).await?;
35        Ok(Self {
36            envelope_store,
37            batch_size_bytes: config.spool_envelopes_batch_size_bytes(),
38            flush_timeout: config.spool_envelopes_flush_timeout(),
39            max_disk_size: config.spool_envelopes_max_disk_size(),
40            partition_id,
41            ephemeral: config.spool_ephemeral(),
42        })
43    }
44
45    /// Returns `true` if data is stored on non-persistent disks.
46    pub fn ephemeral(&self) -> bool {
47        self.ephemeral
48    }
49
50    /// Returns `true` when there might be data residing on disk, `false` otherwise.
51    fn assume_data_on_disk(stack_creation_type: StackCreationType) -> bool {
52        matches!(stack_creation_type, StackCreationType::Initialization)
53    }
54}
55
56impl StackProvider for SqliteStackProvider {
57    type Stack = CachingEnvelopeStack<SqliteEnvelopeStack>;
58
59    async fn initialize(&self) -> InitializationState {
60        match self.envelope_store.project_key_pairs().await {
61            Ok(project_key_pairs) => InitializationState::new(project_key_pairs),
62            Err(error) => {
63                relay_log::error!(
64                    error = &error as &dyn Error,
65                    "failed to initialize the sqlite stack provider"
66                );
67                InitializationState::empty()
68            }
69        }
70    }
71
72    fn create_stack(
73        &self,
74        stack_creation_type: StackCreationType,
75        project_key_pair: ProjectKeyPair,
76    ) -> Self::Stack {
77        let inner = SqliteEnvelopeStack::new(
78            self.partition_id,
79            self.envelope_store.clone(),
80            self.batch_size_bytes,
81            project_key_pair.own_key,
82            project_key_pair.sampling_key,
83            // We want to check the disk by default if we are creating the stack for the first time,
84            // since we might have some data on disk.
85            // On the other hand, if we are recreating a stack, it means that we popped it because
86            // it was empty, or we never had data on disk for that stack, so we assume by default
87            // that there is no need to check disk until some data is spooled.
88            Self::assume_data_on_disk(stack_creation_type),
89            self.flush_timeout,
90        );
91
92        CachingEnvelopeStack::new(inner)
93    }
94
95    fn has_store_capacity(&self) -> bool {
96        (self.envelope_store.usage() as usize) < self.max_disk_size
97    }
98
99    async fn store_total_count(&self) -> u64 {
100        self.envelope_store
101            .total_count()
102            .await
103            .unwrap_or_else(|error| {
104                relay_log::error!(
105                    error = &error as &dyn Error,
106                    "failed to get the total count of envelopes for the sqlite envelope store",
107                );
108                // In case we have an error, we default to communicating a total count of 0.
109                0
110            })
111    }
112
113    fn total_size(&self) -> Option<u64> {
114        Some(self.envelope_store.usage())
115    }
116
117    fn stack_type<'a>(&self) -> &'a str {
118        "sqlite"
119    }
120
121    async fn flush(&mut self, envelope_stacks: impl IntoIterator<Item = Self::Stack>) {
122        relay_log::trace!("Flushing sqlite envelope buffer");
123
124        let partition_tag = self.partition_id.to_string();
125        relay_statsd::metric!(
126            timer(RelayTimers::BufferDrain),
127            partition_id = &partition_tag,
128            {
129                for envelope_stack in envelope_stacks {
130                    envelope_stack.flush().await;
131                }
132            }
133        );
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use std::sync::Arc;
140
141    use relay_base_schema::project::ProjectKey;
142    use relay_config::Config;
143    use uuid::Uuid;
144
145    use crate::EnvelopeStack;
146    use crate::services::buffer::common::ProjectKeyPair;
147    use crate::services::buffer::stack_provider::sqlite::SqliteStackProvider;
148    use crate::services::buffer::stack_provider::{StackCreationType, StackProvider};
149    use crate::services::buffer::testutils::utils::mock_envelopes;
150
151    fn mock_config() -> Arc<Config> {
152        let path = std::env::temp_dir()
153            .join(Uuid::new_v4().to_string())
154            .into_os_string()
155            .into_string()
156            .unwrap();
157
158        Config::from_json_value(serde_json::json!({
159            "spool": {
160                "envelopes": {
161                    "path": path,
162                    "disk_batch_size": 100,
163                    "max_batches": 1,
164                }
165            }
166        }))
167        .unwrap()
168        .into()
169    }
170
171    #[tokio::test]
172    async fn test_flush() {
173        let config = mock_config();
174        let mut stack_provider = SqliteStackProvider::new(0, &config.current())
175            .await
176            .unwrap();
177
178        let own_key = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
179        let sampling_key = ProjectKey::parse("b81ae32be2584e0bbd7a4cbb95971fe1").unwrap();
180
181        let mut envelope_stack = stack_provider.create_stack(
182            StackCreationType::New,
183            ProjectKeyPair::new(own_key, sampling_key),
184        );
185
186        let envelopes = mock_envelopes(10);
187        for envelope in envelopes {
188            envelope_stack.push(envelope).await.unwrap();
189        }
190
191        let envelope_store = stack_provider.envelope_store.clone();
192
193        // We make sure that no data is on disk since we will spool when more than 100 elements are
194        // in the in-memory stack.
195        assert_eq!(envelope_store.total_count().await.unwrap(), 0);
196
197        // We drain the stack provider, and we expect all in-memory envelopes to be spooled to disk.
198        stack_provider.flush(vec![envelope_stack]).await;
199        assert_eq!(envelope_store.total_count().await.unwrap(), 10);
200    }
201}
relay_server/services/buffer/stack_provider/sqlite.rs

relay_server/services/buffer/stack_provider/
sqlite.rs
