webtorrent/
store.rs

1use crate::error::Result;
2use bytes::Bytes;
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8/// Chunk store trait for storing torrent pieces
9#[async_trait]
10pub trait ChunkStore: Send + Sync {
11    /// Get a piece from the store
12    async fn get(&self, index: usize, offset: usize, length: usize) -> Result<Bytes>;
13
14    /// Put a piece into the store
15    async fn put(&self, index: usize, data: Bytes) -> Result<()>;
16
17    /// Close the store
18    async fn close(&self) -> Result<()>;
19
20    /// Destroy the store (delete all data)
21    async fn destroy(&self) -> Result<()>;
22}
23
24/// Memory-based chunk store
25pub struct MemoryChunkStore {
26    pieces: Arc<RwLock<HashMap<usize, Bytes>>>,
27}
28
29impl MemoryChunkStore {
30    pub fn new() -> Self {
31        Self {
32            pieces: Arc::new(RwLock::new(HashMap::new())),
33        }
34    }
35}
36
37#[async_trait]
38impl ChunkStore for MemoryChunkStore {
39    async fn get(&self, index: usize, offset: usize, length: usize) -> Result<Bytes> {
40        let pieces = self.pieces.read().await;
41        if let Some(piece) = pieces.get(&index) {
42            let end = (offset + length).min(piece.len());
43            Ok(piece.slice(offset..end))
44        } else {
45            Err(crate::error::WebTorrentError::Store(
46                format!("Piece {} not found", index),
47            ))
48        }
49    }
50
51    async fn put(&self, index: usize, data: Bytes) -> Result<()> {
52        let mut pieces = self.pieces.write().await;
53        pieces.insert(index, data);
54        Ok(())
55    }
56
57    async fn close(&self) -> Result<()> {
58        Ok(())
59    }
60
61    async fn destroy(&self) -> Result<()> {
62        let mut pieces = self.pieces.write().await;
63        pieces.clear();
64        Ok(())
65    }
66}
67
webtorrent/store.rs

webtorrent/
store.rs
