diff --git a/src/index.rs b/src/index.rs index 86ea82d..e0ec108 100644 --- a/src/index.rs +++ b/src/index.rs @@ -15,15 +15,15 @@ pub struct Index { is_changed: bool, /// Store the index value in memory only when it is in use by the init options has_name: bool, - has_length: bool, + has_size: bool, } impl Index { - pub fn init(capacity: usize, timeout: Option, has_name: bool, has_length: bool) -> Self { + pub fn init(capacity: usize, timeout: Option, has_name: bool, has_size: bool) -> Self { Self { index: HashMap::with_capacity(capacity), timeout: timeout.map(Duration::seconds), - has_length, + has_size, has_name, is_changed: false, } @@ -49,21 +49,15 @@ impl Index { self.index.values().map(|i| i.node).sum::() } - pub fn insert( - &mut self, - infohash: String, - node: u64, - name: Option, - length: Option, - ) { + pub fn insert(&mut self, infohash: String, node: u64, size: u64, name: Option) { if self .index .insert( infohash, Value::new( node, + if self.has_size { Some(size) } else { None }, if self.has_name { name } else { None }, - if self.has_length { length } else { None }, ), ) .is_none() @@ -88,9 +82,9 @@ fn test() { // test values auto-clean by timeout let mut i = Index::init(2, Some(3), false, false); - i.insert("h1".to_string(), 0, None, None); + i.insert("h1".to_string(), 0, 0, None); sleep(Duration::from_secs(1)); - i.insert("h2".to_string(), 0, None, None); + i.insert("h2".to_string(), 0, 0, None); i.refresh(); assert_eq!(i.len(), 2); diff --git a/src/index/value.rs b/src/index/value.rs index c7286d6..b721336 100644 --- a/src/index/value.rs +++ b/src/index/value.rs @@ -7,17 +7,17 @@ pub struct Value { pub time: DateTime, pub node: u64, // Isolate by applying internal filter on value set - length: Option, + size: Option, name: Option, } impl Value { /// Create new `Self` with current timestamp - pub fn new(node: u64, name: Option, length: Option) -> Self { + pub fn new(node: u64, size: Option, name: Option) -> Self { Self { time: Utc::now(), node, - length, + size, name: filter_name(name), } } @@ -26,8 +26,8 @@ impl Value { self.name.as_ref() } /// Get reference to the safely constructed `length` member - pub fn length(&self) -> Option { - self.length + pub fn size(&self) -> Option { + self.size } } diff --git a/src/main.rs b/src/main.rs index a32e8d3..73ff9f2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -150,7 +150,7 @@ async fn main() -> Result<()> { config.preload_max_filecount.unwrap_or_default(), ); mt.wait_until_initialized().await?; - let (name, length) = mt.with_metadata(|m| { + let (name, size) = mt.with_metadata(|m| { // init preload files list if let Some(ref p) = preload { for (id, info) in m.file_infos.iter().enumerate() { @@ -183,7 +183,8 @@ async fn main() -> Result<()> { if let Some(ref t) = torrent { save_torrent_file(t, &debug, &i, &m.torrent_bytes) } - (m.info.name.as_ref().map(|n| n.to_string()), m.info.length) + + (m.info.name.as_ref().map(|n| n.to_string()), size(&m.info)) })?; session.update_only_files(&mt, &only_files).await?; session.unpause(&mt).await?; @@ -198,7 +199,7 @@ async fn main() -> Result<()> { p.cleanup(&i, Some(only_files_keep))? } - index.insert(i, only_files_size, name, length) + index.insert(i, only_files_size, size, name) } Ok(AddTorrentResponse::ListOnly(r)) => { if let Some(ref t) = torrent { @@ -209,7 +210,7 @@ async fn main() -> Result<()> { // use `r.info` for Memory, SQLite, // Manticore and other alternative storage type - index.insert(i, 0, r.info.name.map(|n| n.to_string()), r.info.length) + index.insert(i, 0, size(&r.info), r.info.name.map(|n| n.to_string())) } // unexpected as should be deleted Ok(AddTorrentResponse::AlreadyManaged(..)) => panic!(), @@ -238,7 +239,7 @@ async fn main() -> Result<()> { rss.push( k, v.name().unwrap_or(k), - v.length().map(|l| l.bytes()), + v.size().map(|l| l.bytes()), Some(&v.time.to_rfc2822()), )? } @@ -285,3 +286,17 @@ fn magnet(infohash: &str, trackers: Option<&HashSet>) -> String { } m } + +/// Count total size, including torrent files +fn size(info: &librqbit::TorrentMetaV1Info) -> u64 { + let mut t = 0; + if let Some(l) = info.length { + t += l + } + if let Some(ref files) = info.files { + for f in files { + t += f.length + } + } + t +}