ri/queue/backends/
memory_backend.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18#![allow(non_snake_case)]
19
20//! # In-Memory Queue Implementation
21//!
22//! This file implements an in-memory queue backend for the Ri queue system. The in-memory queue
23//! provides a lightweight, fast queue implementation suitable for testing, development, and
24//! scenarios where durability is not a strict requirement. It also supports optional persistence
25//! to disk for basic durability.
26//!
27//! ## Key Components
28//!
29//! - **RiMemoryQueue**: Main in-memory queue implementation
30//! - **MemoryQueueState**: Internal state management for the queue
31//! - **MemoryQueueProducer**: Producer implementation for sending messages
32//! - **MemoryQueueConsumer**: Consumer implementation for receiving messages
33//!
34//! ## Design Principles
35//!
36//! 1. **Lightweight**: Minimal dependencies and overhead
37//! 2. **Fast Performance**: In-memory operations for low latency
38//! 3. **Optional Persistence**: Can be configured to persist messages to disk
39//! 4. **Consumer Groups**: Supports multiple consumer groups with message distribution
40//! 5. **Async-First**: All operations are asynchronous
41//! 6. **Thread-safe**: Uses Arc and RwLock for safe concurrent access
42//! 7. **Durable Option**: Optional disk persistence for message durability
43//! 8. **Simple API**: Implements the standard RiQueue interfaces
44//! 9. **Non-blocking**: Uses tokio's spawn_blocking for file I/O operations
45//! 10. **Message Retry**: Supports message requeueing with retry count increment
46//!
47//! ## Usage
48//!
49//! ```rust
50//! use ri::queue::{RiQueue, RiQueueMessage, RiQueueProducer, RiQueueConsumer};
51//! use ri::queue::backends::RiMemoryQueue;
52//! use ri::core::RiResult;
53//! use serde_json::json;
54//!
55//! async fn example() -> RiResult<()> {
56//!     // Create a basic in-memory queue
57//!     let queue = RiMemoryQueue::new("example_queue");
58//!     
59//!     // Or create a queue with disk persistence
60//!     // let queue = RiMemoryQueue::with_persistence("example_queue", "/tmp/queue_persistence");
61//!     
62//!     // Create a producer
63//!     let producer = queue.create_producer().await?;
64//!     
65//!     // Create a message
66//!     let payload = json!({ "key": "value" }).to_string().into_bytes();
67//!     let message = RiQueueMessage::new(payload);
68//!     
69//!     // Send the message
70//!     producer.send(message).await?;
71//!     
72//!     // Create a consumer
73//!     let consumer = queue.create_consumer("consumer_group_1").await?;
74//!     
75//!     // Receive a message
76//!     if let Some(message) = consumer.receive().await? {
77//!         // Process the message
78//!         let payload = String::from_utf8_lossy(&message.payload);
79//!         println!("Received message: {}", payload);
80//!         
81//!         // Acknowledge the message
82//!         consumer.ack(&message.id).await?;
83//!     }
84//!     
85//!     Ok(())
86//! }
87//! ```
88
89use crate::core::RiResult;
90use crate::queue::{RiQueue, RiQueueConsumer, RiQueueMessage, RiQueueProducer, RiQueueStats};
91use async_trait::async_trait;
92use std::collections::{HashMap, VecDeque};
93use std::fs::{File, OpenOptions};
94use std::io::{Read, Write};
95use std::path::Path;
96use std::sync::Arc;
97use tokio::sync::{Mutex, RwLock};
98use tokio::task::spawn_blocking;
99
100/// Internal state management for the in-memory queue.
101///
102/// This struct holds the queue's messages and consumer-specific queues. It is protected by a
103/// RwLock to ensure thread-safe access.
104struct MemoryQueueState {
105    /// Main queue of messages waiting to be consumed
106    messages: VecDeque<RiQueueMessage>,
107    /// Map of consumer group names to their respective message queues
108    consumers: HashMap<String, VecDeque<RiQueueMessage>>,
109}
110
111impl MemoryQueueState {
112    /// Creates a new MemoryQueueState with empty queues.
113    ///
114    /// # Returns
115    ///
116    /// A new MemoryQueueState instance
117    fn new() -> Self {
118        Self {
119            messages: VecDeque::new(),
120            consumers: HashMap::new(),
121        }
122    }
123}
124
125/// In-memory queue implementation.
126///
127/// This struct implements the RiQueue trait for an in-memory queue backend. It supports optional
128/// disk persistence for message durability.
129pub struct RiMemoryQueue {
130    /// Name of the queue
131    name: String,
132    /// Internal queue state protected by a RwLock
133    state: Arc<RwLock<MemoryQueueState>>,
134    /// Optional path for disk persistence
135    persistence_path: Option<String>,
136}
137
138#[allow(dead_code)]
139impl RiMemoryQueue {
140    /// Creates a new in-memory queue without persistence.
141    ///
142    /// # Parameters
143    ///
144    /// - `name`: The name of the queue
145    ///
146    /// # Returns
147    ///
148    /// A new RiMemoryQueue instance
149    pub fn new(name: &str) -> Self {
150        Self {
151            name: name.to_string(),
152            state: Arc::new(RwLock::new(MemoryQueueState::new())),
153            persistence_path: None,
154        }
155    }
156
157    /// Creates a new in-memory queue with disk persistence.
158    ///
159    /// # Parameters
160    ///
161    /// - `name`: The name of the queue
162    /// - `persistence_path`: Path to the file where messages will be persisted
163    ///
164    /// # Returns
165    ///
166    /// A new RiMemoryQueue instance with persistence enabled
167    pub fn with_persistence(name: &str, persistence_path: &str) -> Self {
168        let queue = Self {
169            name: name.to_string(),
170            state: Arc::new(RwLock::new(MemoryQueueState::new())),
171            persistence_path: Some(persistence_path.to_string()),
172        };
173
174        // Load messages from disk if persistence is enabled
175        if let Err(e) = queue.load_messages() {
176            log::warn!("Failed to load persisted messages for queue '{name}': {e}");
177        }
178
179        queue
180    }
181
182    /// Loads messages from disk if persistence is enabled.
183    ///
184    /// # Returns
185    ///
186    /// A `RiResult<()>` indicating success or failure
187    fn load_messages(&self) -> RiResult<()> {
188        if let Some(path) = &self.persistence_path {
189            if Path::new(path).exists() {
190                let mut file = File::open(path)?;
191                
192                // Security: Check file size before reading to prevent memory exhaustion
193                const MAX_PERSISTENCE_FILE_SIZE: u64 = 100 * 1024 * 1024; // 100 MB
194                let metadata = file.metadata()?;
195                if metadata.len() > MAX_PERSISTENCE_FILE_SIZE {
196                    log::warn!(
197                        "[Ri.MemoryQueue] Persistence file too large: {} bytes (max {} bytes)",
198                        metadata.len(), MAX_PERSISTENCE_FILE_SIZE
199                    );
200                    return Err(crate::core::RiError::Other(format!(
201                        "Persistence file too large: {} bytes (max {} bytes)",
202                        metadata.len(), MAX_PERSISTENCE_FILE_SIZE
203                    )));
204                }
205                
206                let mut content = String::new();
207                file.read_to_string(&mut content)?;
208
209                if !content.is_empty() {
210                    // Security: Use bounded deserialization
211                    let messages: VecDeque<RiQueueMessage> = serde_json::from_str(&content)
212                        .map_err(|e| {
213                            log::warn!("[Ri.MemoryQueue] Failed to deserialize messages: {}", e);
214                            crate::core::RiError::Other(format!("Failed to deserialize messages: {}", e))
215                        })?;
216                    
217                    // Security: Limit number of messages loaded
218                    const MAX_MESSAGES: usize = 100000;
219                    let messages: VecDeque<RiQueueMessage> = messages.into_iter().take(MAX_MESSAGES).collect();
220                    
221                    let mut state = self.state.blocking_write();
222                    state.messages = messages;
223                }
224            }
225        }
226        Ok(())
227    }
228
229    /// Saves messages to disk if persistence is enabled.
230    ///
231    /// # Returns
232    ///
233    /// A `RiResult<()>` indicating success or failure
234    fn save_messages(&self) -> RiResult<()> {
235        if let Some(path) = &self.persistence_path {
236            let state = self.state.blocking_read();
237            let content = serde_json::to_string(&state.messages)?;
238
239            let mut file = OpenOptions::new()
240                .write(true)
241                .create(true)
242                .truncate(true)
243                .open(path)?;
244
245            file.write_all(content.as_bytes())?;
246        }
247        Ok(())
248    }
249}
250
251#[async_trait]
252impl RiQueue for RiMemoryQueue {
253    /// Creates a new producer for this queue.
254    ///
255    /// # Returns
256    ///
257    /// A `RiResult<Box<dyn RiQueueProducer>>` containing the producer
258    async fn create_producer(&self) -> RiResult<Box<dyn RiQueueProducer>> {
259        Ok(Box::new(MemoryQueueProducer {
260            state: self.state.clone(),
261            persistence_path: self.persistence_path.clone(),
262        }))
263    }
264
265    /// Creates a new consumer for this queue with the given consumer group.
266    ///
267    /// # Parameters
268    ///
269    /// - `consumer_group`: The name of the consumer group
270    ///
271    /// # Returns
272    ///
273    /// A `RiResult<Box<dyn RiQueueConsumer>>` containing the consumer
274    async fn create_consumer(
275        &self,
276        consumer_group: &str,
277    ) -> RiResult<Box<dyn RiQueueConsumer>> {
278        Ok(Box::new(MemoryQueueConsumer {
279            state: self.state.clone(),
280            consumer_group: consumer_group.to_string(),
281            paused: Arc::new(Mutex::new(false)),
282            persistence_path: self.persistence_path.clone(),
283        }))
284    }
285
286    /// Gets statistics for this queue.
287    ///
288    /// # Returns
289    ///
290    /// A `RiResult<RiQueueStats>` containing the queue statistics
291    async fn get_stats(&self) -> RiResult<RiQueueStats> {
292        let state = self.state.read().await;
293        Ok(RiQueueStats {
294            queue_name: self.name.clone(),
295            message_count: state.messages.len() as u64,
296            consumer_count: state.consumers.len() as u32,
297            producer_count: 1,
298            processed_messages: 0,
299            failed_messages: 0,
300            avg_processing_time_ms: 0.0,
301            total_bytes_sent: 0,
302            total_bytes_received: 0,
303            last_message_time: 0,
304        })
305    }
306
307    /// Purges all messages from this queue.
308    ///
309    /// # Returns
310    ///
311    /// A `RiResult<()>` indicating success or failure
312    async fn purge(&self) -> RiResult<()> {
313        let mut state = self.state.write().await;
314        state.messages.clear();
315        state.consumers.clear();
316
317        // Clear persistence file if enabled
318        if let Some(path) = &self.persistence_path {
319            let path_clone = path.clone();
320            spawn_blocking(move || {
321                if Path::new(&path_clone).exists() {
322                    if let Err(e) = std::fs::remove_file(&path_clone) {
323                        log::warn!("Failed to remove persistence file '{path_clone}': {e}");
324                    }
325                }
326            })
327            .await
328            .map_err(|e| {
329                log::error!("Failed to execute persistence file removal: {e}");
330                crate::core::RiError::Other(format!("Failed to clear persistence: {e}"))
331            })?;
332        }
333
334        Ok(())
335    }
336
337    /// Deletes this queue.
338    ///
339    /// # Returns
340    ///
341    /// A `RiResult<()>` indicating success or failure
342    async fn delete(&self) -> RiResult<()> {
343        self.purge().await
344    }
345}
346
347/// Producer implementation for the in-memory queue.
348///
349/// This struct implements the RiQueueProducer trait for sending messages to the in-memory queue.
350struct MemoryQueueProducer {
351    /// Shared queue state
352    state: Arc<RwLock<MemoryQueueState>>,
353    /// Optional path for disk persistence
354    persistence_path: Option<String>,
355}
356
357#[async_trait]
358impl RiQueueProducer for MemoryQueueProducer {
359    /// Sends a single message to the queue.
360    ///
361    /// # Parameters
362    ///
363    /// - `message`: The message to send
364    ///
365    /// # Returns
366    ///
367    /// A `RiResult<()>` indicating success or failure
368    async fn send(&self, message: RiQueueMessage) -> RiResult<()> {
369        let mut state = self.state.write().await;
370        state.messages.push_back(message);
371
372        // Save to disk if persistence is enabled
373        if let Some(path) = &self.persistence_path {
374            let messages_clone = state.messages.clone();
375            let path_clone = path.clone();
376
377            let _ = spawn_blocking(move || {
378            let content = serde_json::to_string(&messages_clone)
379                .map_err(|e| {
380                    log::error!("Failed to serialize messages for persistence: {e}");
381                    crate::core::RiError::Serde(format!("Serialization failed: {e}"))
382                })?;
383            let mut file = OpenOptions::new()
384                .write(true)
385                .create(true)
386                .truncate(true)
387                .open(path_clone)
388                .map_err(|e| {
389                    log::error!("Failed to open persistence file: {e}");
390                    crate::core::RiError::Io(format!("File open failed: {e}"))
391                })?;
392            file.write_all(content.as_bytes())
393                .map_err(|e| {
394                    log::error!("Failed to write persistence file: {e}");
395                    crate::core::RiError::Io(format!("File write failed: {e}"))
396                })?;
397            Ok::<(), crate::core::RiError>(())
398        })
399        .await
400        .map_err(|e| {
401            log::error!("Failed to execute persistence task: {e}");
402            crate::core::RiError::Other(format!("Persistence task failed: {e}"))
403        });
404        }
405
406        Ok(())
407    }
408
409    /// Sends multiple messages to the queue in a batch.
410    ///
411    /// # Parameters
412    ///
413    /// - `messages`: A vector of messages to send
414    ///
415    /// # Returns
416    ///
417    /// A `RiResult<()>` indicating success or failure
418    async fn send_batch(&self, messages: Vec<RiQueueMessage>) -> RiResult<()> {
419        let mut state = self.state.write().await;
420        for message in messages {
421            state.messages.push_back(message);
422        }
423
424        // Save to disk if persistence is enabled
425        if let Some(path) = &self.persistence_path {
426            let messages_clone = state.messages.clone();
427            let path_clone = path.clone();
428
429            let _ = spawn_blocking(move || {
430            let content = serde_json::to_string(&messages_clone)
431                .map_err(|e| {
432                    log::error!("Failed to serialize messages for persistence: {e}");
433                    crate::core::RiError::Serde(format!("Serialization failed: {e}"))
434                })?;
435            let mut file = OpenOptions::new()
436                .write(true)
437                .create(true)
438                .truncate(true)
439                .open(path_clone)
440                .map_err(|e| {
441                    log::error!("Failed to open persistence file: {e}");
442                    crate::core::RiError::Io(format!("File open failed: {e}"))
443                })?;
444            file.write_all(content.as_bytes())
445                .map_err(|e| {
446                    log::error!("Failed to write persistence file: {e}");
447                    crate::core::RiError::Io(format!("File write failed: {e}"))
448                })?;
449            Ok::<(), crate::core::RiError>(())
450        })
451        .await
452        .map_err(|e| {
453            log::error!("Failed to execute persistence task: {e}");
454            crate::core::RiError::Other(format!("Persistence task failed: {e}"))
455        });
456        }
457
458        Ok(())
459    }
460}
461
462/// Consumer implementation for the in-memory queue.
463///
464/// This struct implements the RiQueueConsumer trait for receiving messages from the in-memory queue.
465struct MemoryQueueConsumer {
466    /// Shared queue state
467    state: Arc<RwLock<MemoryQueueState>>,
468    /// Name of the consumer group
469    consumer_group: String,
470    /// Flag indicating if the consumer is paused
471    paused: Arc<Mutex<bool>>,
472    /// Optional path for disk persistence
473    persistence_path: Option<String>,
474}
475
476#[async_trait]
477impl RiQueueConsumer for MemoryQueueConsumer {
478    /// Receives a message from the queue.
479    ///
480    /// # Returns
481    ///
482    /// A `RiResult<Option<RiQueueMessage>>` containing the message if available, or None if no message is available
483    async fn receive(&self) -> RiResult<Option<RiQueueMessage>> {
484        let paused = *self.paused.lock().await;
485        if paused {
486            return Ok(None);
487        }
488
489        let mut state = self.state.write().await;
490
491        // If consumer queue exists and has messages, return one
492        if let Some(consumer_queue) = state.consumers.get_mut(&self.consumer_group) {
493            if let Some(message) = consumer_queue.pop_front() {
494                return Ok(Some(message));
495            }
496        }
497
498        // If main queue has messages, move one to consumer queue
499        if let Some(message) = state.messages.pop_front() {
500            let mut consumer_queue = VecDeque::new();
501            consumer_queue.push_back(message.clone());
502            state
503                .consumers
504                .insert(self.consumer_group.clone(), consumer_queue);
505
506            // Save to disk if persistence is enabled (since main queue changed)
507            if let Some(path) = &self.persistence_path {
508                let messages_clone = state.messages.clone();
509                let path_clone = path.clone();
510
511                let _ = spawn_blocking(move || {
512                    let content = serde_json::to_string(&messages_clone)
513                        .map_err(|e| {
514                            log::error!("Failed to serialize messages for persistence: {e}");
515                            crate::core::RiError::Serde(format!("Serialization failed: {e}"))
516                        })?;
517                    let mut file = OpenOptions::new()
518                        .write(true)
519                        .create(true)
520                        .truncate(true)
521                        .open(path_clone)
522                        .map_err(|e| {
523                            log::error!("Failed to open persistence file: {e}");
524                            crate::core::RiError::Io(format!("File open failed: {e}"))
525                        })?;
526                    file.write_all(content.as_bytes())
527                        .map_err(|e| {
528                            log::error!("Failed to write persistence file: {e}");
529                            crate::core::RiError::Io(format!("File write failed: {e}"))
530                        })?;
531                    Ok::<(), crate::core::RiError>(())
532                })
533                .await
534                .map_err(|e| {
535                    log::error!("Failed to execute persistence task: {e}");
536                    crate::core::RiError::Other(format!("Persistence task failed: {e}"))
537                });
538            }
539
540            Ok(Some(message))
541        } else {
542            Ok(None)
543        }
544    }
545
546    /// Acknowledges a message, indicating it has been successfully processed.
547    ///
548    /// For in-memory queues, acknowledgment is implicit when the message is received.
549    ///
550    /// # Parameters
551    ///
552    /// - `_message_id`: The ID of the message to acknowledge (not used for in-memory queues)
553    ///
554    /// # Returns
555    ///
556    /// A `RiResult<()>` indicating success or failure
557    async fn ack(&self, _message_id: &str) -> RiResult<()> {
558        // In memory queue, acknowledgment is implicit when message is received
559        Ok(())
560    }
561
562    /// Negatively acknowledges a message, indicating it failed to process and should be retried.
563    ///
564    /// # Parameters
565    ///
566    /// - `message_id`: The ID of the message to negatively acknowledge
567    ///
568    /// # Returns
569    ///
570    /// A `RiResult<()>` indicating success or failure
571    async fn nack(&self, message_id: &str) -> RiResult<()> {
572        // Find the message in consumer queue and put it back in main queue
573        let mut state = self.state.write().await;
574
575        if let Some(consumer_queue) = state.consumers.get_mut(&self.consumer_group) {
576            // Find the message by ID
577            let mut message_to_requeue: Option<RiQueueMessage> = None;
578
579            // Iterate through consumer queue to find the message
580            let mut index = 0;
581            for (i, message) in consumer_queue.iter().enumerate() {
582                if message.id == message_id {
583                    message_to_requeue = Some(message.clone());
584                    index = i;
585                    break;
586                }
587            }
588
589            // If found, remove from consumer queue and put back in main queue
590            if let Some(mut message) = message_to_requeue {
591                consumer_queue.remove(index);
592                message.increment_retry();
593                state.messages.push_back(message);
594
595                // Save to disk if persistence is enabled
596                if let Some(path) = &self.persistence_path {
597                    let messages_clone = state.messages.clone();
598                    let path_clone = path.clone();
599
600                    let _ = spawn_blocking(move || {
601                        let content = serde_json::to_string(&messages_clone)
602                            .map_err(|e| {
603                                log::error!("Failed to serialize messages for persistence: {e}");
604                                crate::core::RiError::Serde(format!("Serialization failed: {e}"))
605                            })?;
606                        let mut file = OpenOptions::new()
607                            .write(true)
608                            .create(true)
609                            .truncate(true)
610                            .open(path_clone)
611                            .map_err(|e| {
612                                log::error!("Failed to open persistence file: {e}");
613                                crate::core::RiError::Io(format!("File open failed: {e}"))
614                            })?;
615                        file.write_all(content.as_bytes())
616                            .map_err(|e| {
617                                log::error!("Failed to write persistence file: {e}");
618                                crate::core::RiError::Io(format!("File write failed: {e}"))
619                            })?;
620                        Ok::<(), crate::core::RiError>(())
621                    })
622                    .await
623                    .map_err(|e| {
624                        log::error!("Failed to execute persistence task: {e}");
625                        crate::core::RiError::Other(format!("Persistence task failed: {e}"))
626                    });
627                }
628            }
629        }
630
631        Ok(())
632    }
633
634    /// Pauses message consumption.
635    ///
636    /// # Returns
637    ///
638    /// A `RiResult<()>` indicating success or failure
639    async fn pause(&self) -> RiResult<()> {
640        let mut paused = self.paused.lock().await;
641        *paused = true;
642        Ok(())
643    }
644
645    /// Resumes message consumption after pausing.
646    ///
647    /// # Returns
648    ///
649    /// A `RiResult<()>` indicating success or failure
650    async fn resume(&self) -> RiResult<()> {
651        let mut paused = self.paused.lock().await;
652        *paused = false;
653        Ok(())
654    }
655}
ri/queue/backends/memory_backend.rs

ri/queue/backends/
memory_backend.rs
