relay_server/services/buffer/envelope_stack/
sqlite.rs

1use std::fmt::Debug;
2use std::num::NonZeroUsize;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use chrono::{DateTime, Utc};
7use relay_base_schema::project::ProjectKey;
8
9use crate::envelope::Envelope;
10use crate::services::buffer::envelope_stack::EnvelopeStack;
11use crate::services::buffer::envelope_store::sqlite::{
12    DatabaseBatch, DatabaseEnvelope, InsertEnvelopeError, SqliteEnvelopeStore,
13    SqliteEnvelopeStoreError,
14};
15use crate::services::buffer::throttle::Throttle;
16use crate::statsd::RelayTimers;
17
18/// An error returned when doing an operation on [`SqliteEnvelopeStack`].
19#[derive(Debug, thiserror::Error)]
20pub enum SqliteEnvelopeStackError {
21    #[error("envelope store error: {0}")]
22    EnvelopeStoreError(#[from] SqliteEnvelopeStoreError),
23    #[error("envelope encode error: {0}")]
24    Envelope(#[from] InsertEnvelopeError),
25}
26
27#[derive(Debug)]
28/// An [`EnvelopeStack`] that is implemented on an SQLite database.
29///
30/// For efficiency reasons, the implementation has an in-memory buffer that is periodically spooled
31/// to disk in a batched way.
32pub struct SqliteEnvelopeStack {
33    /// Shared SQLite database pool which will be used to read and write from disk.
34    envelope_store: SqliteEnvelopeStore,
35    /// Maximum number of bytes in the in-memory cache before we write to disk.
36    batch_size_bytes: NonZeroUsize,
37    /// The project key of the project to which all the envelopes belong.
38    own_key: ProjectKey,
39    /// The project key of the root project of the trace to which all the envelopes belong.
40    sampling_key: ProjectKey,
41    /// In-memory stack containing a batch of envelopes that either have not been written to disk yet, or have been read from disk recently.
42    batch: Vec<DatabaseEnvelope>,
43    /// Boolean representing whether calls to `push()` and `peek()` check disk in case not enough
44    /// elements are available in the `batches_buffer`.
45    check_disk: bool,
46    /// The tag value of this partition which is used for reporting purposes.
47    partition_tag: String,
48    /// Time after which to flush the buffer to disk, regardless of batch size.
49    flush_timeout: Option<Duration>,
50    /// Time of last flush to disk (or creation time of the envelope stack).
51    last_flush: Instant,
52    /// Optional throttle which paces how quickly envelopes are unspooled from disk.
53    unspool_throttle: Arc<Throttle>,
54}
55
56/// Configuration for a [`SqliteEnvelopeStack`].
57#[derive(Debug, Clone)]
58pub struct SqliteEnvelopeStackConfig {
59    /// The partition of the envelope buffer to which the stack belongs.
60    pub partition_id: u8,
61    /// Maximum number of bytes in the in-memory cache before writing to disk.
62    pub batch_size_bytes: usize,
63    /// Time after which to flush the buffer to disk.
64    pub flush_timeout: Option<Duration>,
65    /// Optional throttle which paces how quickly envelopes are unspooled from disk.
66    pub unspool_throttle: Arc<Throttle>,
67}
68
69impl SqliteEnvelopeStack {
70    /// Creates a new empty [`SqliteEnvelopeStack`].
71    pub fn new(
72        config: SqliteEnvelopeStackConfig,
73        envelope_store: SqliteEnvelopeStore,
74        own_key: ProjectKey,
75        sampling_key: ProjectKey,
76        check_disk: bool,
77    ) -> Self {
78        Self {
79            envelope_store,
80            batch_size_bytes: NonZeroUsize::new(config.batch_size_bytes)
81                .expect("batch bytes should be > 0"),
82            own_key,
83            sampling_key,
84            batch: vec![],
85            check_disk,
86            partition_tag: config.partition_id.to_string(),
87            flush_timeout: config.flush_timeout,
88            last_flush: Instant::now(),
89            unspool_throttle: config.unspool_throttle,
90        }
91    }
92
93    /// Threshold above which the [`SqliteEnvelopeStack`] will spool data from the `buffer` to disk.
94    ///
95    /// Returns the reason why the batch should be spooled.
96    fn should_spool_to_disk(&self) -> Option<&'static str> {
97        let batch_size = self.batch.iter().map(|e| e.len()).sum::<usize>();
98        if batch_size > self.batch_size_bytes.get() {
99            return Some("size");
100        }
101
102        if let Some(timeout) = self.flush_timeout
103            && self.last_flush.elapsed() > timeout
104        {
105            return Some("timeout");
106        }
107
108        None
109    }
110
111    /// Spools to disk a batch of envelopes from the `batch`.
112    ///
113    /// In case there is a failure while writing envelopes, all the envelopes that were enqueued
114    /// to be written to disk are lost. The explanation for this behavior can be found in the body
115    /// of the method.
116    async fn spool_to_disk(&mut self, why: &'static str) -> Result<(), SqliteEnvelopeStackError> {
117        let batch = std::mem::take(&mut self.batch);
118        let Ok(batch) = DatabaseBatch::try_from(batch) else {
119            return Ok(());
120        };
121
122        // When early return here, we are acknowledging that the elements that we popped from
123        // the buffer are lost in case of failure. We are doing this on purpose, since if we were
124        // to have a database corruption during runtime, and we were to put the values back into
125        // the buffer we will end up with an infinite cycle.
126        relay_statsd::metric!(
127            timer(RelayTimers::BufferSpool),
128            partition_id = &self.partition_tag,
129            reason = why,
130            {
131                self.envelope_store
132                    .insert_batch(batch)
133                    .await
134                    .map_err(SqliteEnvelopeStackError::EnvelopeStoreError)?;
135            }
136        );
137
138        // If we successfully spooled to disk, we know that data should be there.
139        self.check_disk = true;
140        self.last_flush = Instant::now();
141
142        Ok(())
143    }
144
145    /// Unspools from disk a batch of envelopes and appends them to the `batch`.
146    ///
147    /// In case there is a failure while deleting envelopes, the envelopes will be lost.
148    async fn unspool_from_disk(&mut self) -> Result<(), SqliteEnvelopeStackError> {
149        debug_assert!(self.batch.is_empty());
150        let batch = relay_statsd::metric!(
151            timer(RelayTimers::BufferUnspool),
152            partition_id = &self.partition_tag,
153            {
154                self.envelope_store
155                    .delete_batch(self.own_key, self.sampling_key)
156                    .await
157                    .map_err(SqliteEnvelopeStackError::EnvelopeStoreError)?
158            }
159        );
160
161        match batch {
162            Some(batch) => {
163                self.batch = batch.into();
164                self.unspool_throttle.acquire(self.batch.len()).await;
165            }
166            None => self.check_disk = false,
167        }
168
169        Ok(())
170    }
171
172    /// Validates that the incoming [`Envelope`] has the same project keys at the
173    /// [`SqliteEnvelopeStack`].
174    fn validate_envelope(&self, envelope: &Envelope) -> bool {
175        let own_key = envelope.meta().public_key();
176        let sampling_key = envelope.sampling_key().unwrap_or(own_key);
177
178        self.own_key == own_key && self.sampling_key == sampling_key
179    }
180}
181
182impl EnvelopeStack for SqliteEnvelopeStack {
183    type Error = SqliteEnvelopeStackError;
184
185    async fn push(&mut self, envelope: Box<Envelope>) -> Result<(), Self::Error> {
186        debug_assert!(self.validate_envelope(&envelope));
187
188        if let Some(reason) = self.should_spool_to_disk() {
189            self.spool_to_disk(reason).await?;
190        }
191
192        let encoded_envelope = relay_statsd::metric!(
193            timer(RelayTimers::BufferEnvelopesSerialization),
194            partition_id = &self.partition_tag,
195            { DatabaseEnvelope::try_from(envelope.as_ref())? }
196        );
197        self.batch.push(encoded_envelope);
198
199        Ok(())
200    }
201
202    async fn peek(&mut self) -> Result<Option<DateTime<Utc>>, Self::Error> {
203        if self.batch.is_empty() && self.check_disk {
204            self.unspool_from_disk().await?
205        }
206
207        let Some(envelope) = self.batch.last() else {
208            return Ok(None);
209        };
210
211        Ok(Some(envelope.received_at()))
212    }
213
214    async fn pop(&mut self) -> Result<Option<Box<Envelope>>, Self::Error> {
215        if self.batch.is_empty() && self.check_disk {
216            self.unspool_from_disk().await?
217        }
218
219        let Some(envelope) = self.batch.pop() else {
220            return Ok(None);
221        };
222        let envelope = envelope.try_into()?;
223
224        Ok(Some(envelope))
225    }
226
227    async fn flush(mut self) {
228        if let Err(e) = self.spool_to_disk("flush").await {
229            relay_log::error!(error = &e as &dyn std::error::Error, "flush error");
230        }
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use chrono::Utc;
237    use relay_base_schema::project::ProjectKey;
238    use std::time::Duration;
239
240    use super::*;
241    use crate::services::buffer::testutils::utils::{mock_envelope, mock_envelopes, setup_db};
242
243    /// Helper function to calculate the total size of a slice of envelopes after compression
244    fn calculate_compressed_size(envelopes: &[Box<Envelope>]) -> usize {
245        envelopes
246            .iter()
247            .map(|e| DatabaseEnvelope::try_from(e.as_ref()).unwrap().len())
248            .sum()
249    }
250
251    #[tokio::test]
252    #[should_panic]
253    async fn test_push_with_mismatching_project_keys() {
254        let db = setup_db(false).await;
255        let envelope_store = SqliteEnvelopeStore::new(0, db, Duration::from_millis(100));
256        let mut stack = SqliteEnvelopeStack::new(
257            SqliteEnvelopeStackConfig {
258                partition_id: 0,
259                batch_size_bytes: 10,
260                flush_timeout: None,
261                unspool_throttle: Arc::new(Throttle::new(None)),
262            },
263            envelope_store,
264            ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
265            ProjectKey::parse("c25ae32be2584e0bbd7a4cbb95971fe1").unwrap(),
266            true,
267        );
268
269        let envelope = mock_envelope(Utc::now());
270        let _ = stack.push(envelope).await;
271    }
272
273    const COMPRESSED_ENVELOPE_SIZE: usize = 313;
274
275    #[tokio::test]
276    async fn test_push_when_db_is_not_valid() {
277        let db = setup_db(false).await;
278        let envelope_store = SqliteEnvelopeStore::new(0, db, Duration::from_millis(100));
279
280        // Create envelopes first so we can calculate actual size
281        let envelopes = mock_envelopes(4);
282        let threshold_size = calculate_compressed_size(&envelopes) - 1;
283
284        let mut stack = SqliteEnvelopeStack::new(
285            SqliteEnvelopeStackConfig {
286                partition_id: 0,
287                batch_size_bytes: threshold_size,
288                flush_timeout: None,
289                unspool_throttle: Arc::new(Throttle::new(None)),
290            },
291            envelope_store,
292            ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
293            ProjectKey::parse("b81ae32be2584e0bbd7a4cbb95971fe1").unwrap(),
294            true,
295        );
296
297        // We push the 4 envelopes without errors because they are below the threshold.
298        for envelope in envelopes.clone() {
299            assert!(stack.push(envelope).await.is_ok());
300        }
301
302        // We push 1 more envelope which results in spooling, which fails because of a database
303        // problem.
304        let envelope = mock_envelope(Utc::now());
305        assert!(matches!(
306            stack.push(envelope).await,
307            Err(SqliteEnvelopeStackError::EnvelopeStoreError(_))
308        ));
309
310        // The stack now contains the last of the 1 elements that were added. If we add a new one
311        // we will end up with 2.
312        let envelope = mock_envelope(Utc::now());
313        assert!(stack.push(envelope.clone()).await.is_ok());
314        assert_eq!(stack.batch.len(), 1);
315
316        // We pop the remaining elements, expecting the last added envelope to be on top.
317        let popped_envelope_1 = stack.pop().await.unwrap().unwrap();
318        assert_eq!(
319            popped_envelope_1.event_id().unwrap(),
320            envelope.event_id().unwrap()
321        );
322        assert_eq!(stack.batch.len(), 0);
323    }
324
325    #[tokio::test]
326    async fn test_pop_when_db_is_not_valid() {
327        let db = setup_db(false).await;
328        let envelope_store = SqliteEnvelopeStore::new(0, db, Duration::from_millis(100));
329        let mut stack = SqliteEnvelopeStack::new(
330            SqliteEnvelopeStackConfig {
331                partition_id: 0,
332                batch_size_bytes: 2,
333                flush_timeout: None,
334                unspool_throttle: Arc::new(Throttle::new(None)),
335            },
336            envelope_store,
337            ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
338            ProjectKey::parse("b81ae32be2584e0bbd7a4cbb95971fe1").unwrap(),
339            true,
340        );
341
342        // We pop with an invalid db.
343        assert!(matches!(
344            stack.pop().await,
345            Err(SqliteEnvelopeStackError::EnvelopeStoreError(_))
346        ));
347    }
348
349    #[tokio::test]
350    async fn test_pop_when_stack_is_empty() {
351        let db = setup_db(true).await;
352        let envelope_store = SqliteEnvelopeStore::new(0, db, Duration::from_millis(100));
353        let mut stack = SqliteEnvelopeStack::new(
354            SqliteEnvelopeStackConfig {
355                partition_id: 0,
356                batch_size_bytes: 2,
357                flush_timeout: None,
358                unspool_throttle: Arc::new(Throttle::new(None)),
359            },
360            envelope_store,
361            ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
362            ProjectKey::parse("b81ae32be2584e0bbd7a4cbb95971fe1").unwrap(),
363            true,
364        );
365
366        // We pop with no elements.
367        // We pop with no elements.
368        assert!(stack.pop().await.unwrap().is_none());
369    }
370
371    #[tokio::test]
372    async fn test_push_below_threshold_and_pop() {
373        let db = setup_db(true).await;
374        let envelope_store = SqliteEnvelopeStore::new(0, db, Duration::from_millis(100));
375        let mut stack = SqliteEnvelopeStack::new(
376            SqliteEnvelopeStackConfig {
377                partition_id: 0,
378                batch_size_bytes: 9999,
379                flush_timeout: None,
380                unspool_throttle: Arc::new(Throttle::new(None)),
381            },
382            envelope_store,
383            ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
384            ProjectKey::parse("b81ae32be2584e0bbd7a4cbb95971fe1").unwrap(),
385            true,
386        );
387
388        let envelopes = mock_envelopes(5);
389
390        // We push 5 envelopes.
391        for envelope in envelopes.clone() {
392            assert!(stack.push(envelope).await.is_ok());
393        }
394        assert_eq!(stack.batch.len(), 5);
395
396        // We peek the top element.
397        let peeked = stack.peek().await.unwrap().unwrap();
398        assert_eq!(
399            peeked.timestamp_millis(),
400            envelopes.clone()[4].received_at().timestamp_millis()
401        );
402
403        // We pop 5 envelopes.
404        for envelope in envelopes.iter().rev() {
405            let popped_envelope = stack.pop().await.unwrap().unwrap();
406            assert_eq!(
407                popped_envelope.event_id().unwrap(),
408                envelope.event_id().unwrap()
409            );
410        }
411
412        assert_eq!(stack.batch.len(), 0);
413    }
414
415    #[tokio::test]
416    async fn test_push_with_flush_timeout() {
417        let db = setup_db(true).await;
418        let envelope_store = SqliteEnvelopeStore::new(0, db, Duration::from_millis(100));
419        let envelopes = mock_envelopes(4);
420        let timeout = Duration::from_secs(3600);
421        let mut stack = SqliteEnvelopeStack::new(
422            SqliteEnvelopeStackConfig {
423                partition_id: 0,
424                batch_size_bytes: calculate_compressed_size(&envelopes),
425                flush_timeout: Some(timeout),
426                unspool_throttle: Arc::new(Throttle::new(None)),
427            },
428            envelope_store.clone(),
429            ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
430            ProjectKey::parse("b81ae32be2584e0bbd7a4cbb95971fe1").unwrap(),
431            false,
432        );
433
434        // First push: no spool
435        assert_eq!(stack.batch.len(), 0);
436        stack.push(envelopes[0].clone()).await.unwrap();
437        assert_eq!(stack.batch.len(), 1);
438        assert_eq!(envelope_store.total_count().await.unwrap(), 0);
439
440        // Second push: no spool
441        stack.push(envelopes[1].clone()).await.unwrap();
442        assert_eq!(stack.batch.len(), 2);
443        assert_eq!(envelope_store.total_count().await.unwrap(), 0);
444
445        // Third push (after timeout): spool
446        stack.last_flush = Instant::now() - timeout - Duration::from_secs(1);
447        stack.push(envelopes[2].clone()).await.unwrap();
448        assert_eq!(stack.batch.len(), 1);
449        assert_eq!(envelope_store.total_count().await.unwrap(), 2);
450
451        // Fourth push: no spool
452        stack.push(envelopes[3].clone()).await.unwrap();
453        assert_eq!(stack.batch.len(), 2);
454        assert_eq!(envelope_store.total_count().await.unwrap(), 2);
455
456        // Envelopes from memory and disk are still returned in stack order.
457        for envelope in envelopes.iter().rev() {
458            let popped = stack.pop().await.unwrap().unwrap();
459            assert_eq!(popped.event_id().unwrap(), envelope.event_id().unwrap());
460        }
461        assert!(stack.pop().await.unwrap().is_none());
462    }
463
464    #[tokio::test]
465    async fn test_push_above_threshold_and_pop() {
466        let db = setup_db(true).await;
467        let envelope_store = SqliteEnvelopeStore::new(0, db, Duration::from_millis(100));
468
469        // Create envelopes first so we can calculate actual size
470        let envelopes = mock_envelopes(7);
471        let threshold_size = calculate_compressed_size(&envelopes[..5]) - 1;
472
473        // Create stack with threshold just below the size of first 5 envelopes
474        let mut stack = SqliteEnvelopeStack::new(
475            SqliteEnvelopeStackConfig {
476                partition_id: 0,
477                batch_size_bytes: threshold_size,
478                flush_timeout: None,
479                unspool_throttle: Arc::new(Throttle::new(None)),
480            },
481            envelope_store,
482            ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
483            ProjectKey::parse("b81ae32be2584e0bbd7a4cbb95971fe1").unwrap(),
484            true,
485        );
486
487        // We push 7 envelopes.
488        for envelope in envelopes.clone() {
489            assert!(stack.push(envelope).await.is_ok());
490        }
491        assert_eq!(stack.batch.len(), 2);
492
493        // We peek the top element.
494        let peeked = stack.peek().await.unwrap().unwrap();
495        assert_eq!(
496            peeked.timestamp_millis(),
497            envelopes[6].received_at().timestamp_millis()
498        );
499
500        // We pop envelopes, and we expect that the last 2 are in memory, since the first 5
501        // should have been spooled to disk.
502        for envelope in envelopes[5..7].iter().rev() {
503            let popped_envelope = stack.pop().await.unwrap().unwrap();
504            assert_eq!(
505                popped_envelope.event_id().unwrap(),
506                envelope.event_id().unwrap()
507            );
508        }
509        assert_eq!(stack.batch.len(), 0);
510
511        // We peek the top element, which since the buffer is empty should result in a disk load.
512        let peeked = stack.peek().await.unwrap().unwrap();
513        assert_eq!(
514            peeked.timestamp_millis(),
515            envelopes[4].received_at().timestamp_millis()
516        );
517
518        // We insert a new envelope, to test the load from disk happening during `peek()` gives
519        // priority to this envelope in the stack.
520        let envelope = mock_envelope(Utc::now());
521        assert!(stack.push(envelope.clone()).await.is_ok());
522
523        // We pop and expect the newly inserted element.
524        let popped_envelope = stack.pop().await.unwrap().unwrap();
525        assert_eq!(
526            popped_envelope.event_id().unwrap(),
527            envelope.event_id().unwrap()
528        );
529
530        // We pop 5 envelopes, which should not result in a disk load since `peek()` already should
531        // have caused it.
532        for envelope in envelopes[0..5].iter().rev() {
533            let popped_envelope = stack.pop().await.unwrap().unwrap();
534            assert_eq!(
535                popped_envelope.event_id().unwrap(),
536                envelope.event_id().unwrap()
537            );
538        }
539        assert_eq!(stack.batch.len(), 0);
540    }
541
542    #[tokio::test]
543    async fn test_drain() {
544        let db = setup_db(true).await;
545        let envelope_store = SqliteEnvelopeStore::new(0, db, Duration::from_millis(100));
546        let mut stack = SqliteEnvelopeStack::new(
547            SqliteEnvelopeStackConfig {
548                partition_id: 0,
549                batch_size_bytes: 10 * COMPRESSED_ENVELOPE_SIZE,
550                flush_timeout: None,
551                unspool_throttle: Arc::new(Throttle::new(None)),
552            },
553            envelope_store.clone(),
554            ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
555            ProjectKey::parse("b81ae32be2584e0bbd7a4cbb95971fe1").unwrap(),
556            true,
557        );
558
559        let envelopes = mock_envelopes(5);
560
561        // We push 5 envelopes and check that there is nothing on disk.
562        for envelope in envelopes.clone() {
563            assert!(stack.push(envelope).await.is_ok());
564        }
565        assert_eq!(stack.batch.len(), 5);
566        assert_eq!(envelope_store.total_count().await.unwrap(), 0);
567
568        // We drain the stack and make sure everything was spooled to disk.
569        stack.flush().await;
570        assert_eq!(envelope_store.total_count().await.unwrap(), 5);
571    }
572}
relay_server/services/buffer/envelope_stack/sqlite.rs

relay_server/services/buffer/envelope_stack/
sqlite.rs
