geosuggest_core/
storage.rs

1use crate::ArchivedEngineMetadata;
2use crate::EngineMetadata;
3use rkyv;
4use rkyv::{deserialize, option::ArchivedOption, rancor::Error};
5use std::fs::OpenOptions;
6use std::io::Read;
7use std::io::SeekFrom;
8use std::path::Path;
9
10#[cfg(feature = "tracing")]
11use std::time::Instant;
12
13/// rkyv storage in len-prefix format `<4-bytes metadata length><metadata><payload>`
14pub struct Storage {}
15
16impl Storage {
17    pub fn new() -> Self {
18        Self {}
19    }
20}
21
22impl Default for Storage {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl Storage {
29    /// Serialize
30    pub fn dump<W>(
31        &self,
32        buf: &mut W,
33        engine_data: &crate::EngineData,
34    ) -> Result<(), Box<dyn std::error::Error>>
35    where
36        W: std::io::Write,
37    {
38        let metadata = rkyv::to_bytes::<Error>(&engine_data.metadata)?;
39
40        buf.write_all(&(metadata.len() as u32).to_be_bytes())?;
41        #[cfg(feature = "tracing")]
42        buf.write_all(&metadata)?;
43
44        buf.write_all(&engine_data.data)?;
45        Ok(())
46    }
47
48    /// Deserialize
49    pub fn load<R>(&self, buf: &mut R) -> Result<crate::EngineData, Box<dyn std::error::Error>>
50    where
51        R: std::io::Read + std::io::Seek,
52    {
53        // skip metadata
54        let mut metadata_len = [0; 4];
55        buf.read_exact(&mut metadata_len)?;
56        let metadata_len = u32::from_be_bytes(metadata_len);
57        let _ = buf.seek(SeekFrom::Current(metadata_len as i64))?;
58
59        let mut bytes = rkyv::util::AlignedVec::new();
60        bytes.extend_from_reader(buf)?;
61
62        Ok(bytes.try_into()?)
63    }
64
65    /// Read engine metadata and don't load whole engine
66    pub fn read_metadata<P: AsRef<Path>>(
67        &self,
68        path: P,
69    ) -> Result<Option<EngineMetadata>, Box<dyn std::error::Error>> {
70        let mut file = OpenOptions::new()
71            .create(false)
72            .read(true)
73            .truncate(false)
74            .open(&path)?;
75
76        let mut metadata_len = [0; 4];
77        file.read_exact(&mut metadata_len)?;
78
79        let metadata_len = u32::from_be_bytes(metadata_len);
80        if metadata_len == 0 {
81            return Ok(None);
82        }
83        let mut raw_metadata = vec![0; metadata_len as usize];
84        file.read_exact(&mut raw_metadata)?;
85
86        let archived =
87            rkyv::access::<ArchivedOption<ArchivedEngineMetadata>, Error>(&raw_metadata[..])?;
88
89        Ok(deserialize::<Option<EngineMetadata>, Error>(archived)?)
90    }
91
92    /// Dump whole index to file
93    pub fn dump_to<P: AsRef<Path>>(
94        &self,
95        path: P,
96        engine_data: &crate::EngineData,
97    ) -> Result<(), Box<dyn std::error::Error>> {
98        #[cfg(feature = "tracing")]
99        tracing::info!("Start dump index to file...");
100        #[cfg(feature = "tracing")]
101        let now = Instant::now();
102
103        let mut file = OpenOptions::new()
104            .create(true)
105            .write(true)
106            .truncate(true)
107            .open(&path)?;
108
109        self.dump(&mut file, engine_data)?;
110
111        #[cfg(feature = "tracing")]
112        tracing::info!("Dump index to file. took {}ms", now.elapsed().as_millis(),);
113
114        Ok(())
115    }
116    /// Load whole index from file
117    pub fn load_from<P: AsRef<std::path::Path>>(
118        &self,
119        path: P,
120    ) -> Result<crate::EngineData, Box<dyn std::error::Error>> {
121        #[cfg(feature = "tracing")]
122        tracing::info!("Loading index...");
123        #[cfg(feature = "tracing")]
124        let now = Instant::now();
125
126        let mut file = OpenOptions::new()
127            .create(false)
128            .read(true)
129            .truncate(false)
130            .open(&path)?;
131
132        let index = self.load(&mut file)?;
133
134        #[cfg(feature = "tracing")]
135        tracing::info!(
136            "Loaded from file done. took {}ms",
137            now.elapsed().as_millis(),
138        );
139
140        Ok(index)
141    }
142}
geosuggest_core/storage.rs

geosuggest_core/
storage.rs
