libsql_wal/io/
mod.rs

1use std::path::Path;
2use std::sync::Arc;
3use std::{future::Future, io};
4
5use chrono::{DateTime, Utc};
6use rand::{rngs::ThreadRng, thread_rng, Rng};
7use uuid::Uuid;
8
9pub use self::file::FileExt;
10
11pub mod buf;
12pub mod compat;
13pub mod file;
14
15pub trait Io: Send + Sync + 'static {
16    type File: FileExt;
17    type TempFile: FileExt;
18    type Rng: Rng;
19
20    fn create_dir_all(&self, path: &Path) -> io::Result<()>;
21    /// TODO: when adding an async variant make sure all places where async is needed are replaced
22    fn open(
23        &self,
24        create_new: bool,
25        read: bool,
26        write: bool,
27        path: &Path,
28    ) -> io::Result<Self::File>;
29
30    // todo: create an async counterpart
31    fn tempfile(&self) -> io::Result<Self::TempFile>;
32    fn now(&self) -> DateTime<Utc>;
33    fn hard_link(&self, src: &Path, dst: &Path) -> io::Result<()>;
34    fn with_rng<F, R>(&self, f: F) -> R
35    where
36        F: FnOnce(&mut Self::Rng) -> R;
37    fn uuid(&self) -> uuid::Uuid {
38        self.with_rng(|rng| {
39            let n: u128 = rng.gen();
40            Uuid::from_u128(n)
41        })
42    }
43
44    fn remove_file_async(&self, path: &Path) -> impl Future<Output = io::Result<()>> + Send;
45}
46
47#[derive(Default, Debug, Clone, Copy)]
48pub struct StdIO(pub(crate) ());
49
50impl Io for StdIO {
51    type File = std::fs::File;
52    type TempFile = std::fs::File;
53    type Rng = ThreadRng;
54
55    fn create_dir_all(&self, path: &Path) -> io::Result<()> {
56        std::fs::create_dir_all(path)
57    }
58
59    fn open(
60        &self,
61        create_new: bool,
62        read: bool,
63        write: bool,
64        path: &Path,
65    ) -> io::Result<Self::File> {
66        std::fs::OpenOptions::new()
67            .create_new(create_new)
68            .create(write)
69            .read(read)
70            .write(write)
71            .open(path)
72    }
73
74    fn tempfile(&self) -> io::Result<Self::TempFile> {
75        tempfile::tempfile()
76    }
77
78    fn now(&self) -> DateTime<Utc> {
79        Utc::now()
80    }
81
82    fn uuid(&self) -> Uuid {
83        Uuid::new_v4()
84    }
85
86    fn hard_link(&self, src: &Path, dst: &Path) -> io::Result<()> {
87        std::fs::hard_link(src, dst)
88    }
89
90    fn with_rng<F, R>(&self, f: F) -> R
91    where
92        F: FnOnce(&mut Self::Rng) -> R,
93    {
94        f(&mut thread_rng())
95    }
96
97    async fn remove_file_async(&self, path: &Path) -> io::Result<()> {
98        tokio::fs::remove_file(path).await
99    }
100}
101
102impl<T: Io> Io for Arc<T> {
103    type File = T::File;
104    type TempFile = T::TempFile;
105    type Rng = T::Rng;
106
107    fn create_dir_all(&self, path: &Path) -> io::Result<()> {
108        self.as_ref().create_dir_all(path)
109    }
110
111    fn open(
112        &self,
113        create_new: bool,
114        read: bool,
115        write: bool,
116        path: &Path,
117    ) -> io::Result<Self::File> {
118        self.as_ref().open(create_new, read, write, path)
119    }
120
121    fn tempfile(&self) -> io::Result<Self::TempFile> {
122        self.as_ref().tempfile()
123    }
124
125    fn now(&self) -> DateTime<Utc> {
126        self.as_ref().now()
127    }
128
129    fn uuid(&self) -> Uuid {
130        self.as_ref().uuid()
131    }
132
133    fn hard_link(&self, src: &Path, dst: &Path) -> io::Result<()> {
134        self.as_ref().hard_link(src, dst)
135    }
136
137    fn with_rng<F, R>(&self, f: F) -> R
138    where
139        F: FnOnce(&mut Self::Rng) -> R,
140    {
141        self.as_ref().with_rng(f)
142    }
143
144    async fn remove_file_async(&self, path: &Path) -> io::Result<()> {
145        self.as_ref().remove_file_async(path).await
146    }
147}
148
149pub struct Inspect<W, F> {
150    inner: W,
151    f: F,
152}
153
154impl<W, F> Inspect<W, F> {
155    pub fn new(inner: W, f: F) -> Self {
156        Self { inner, f }
157    }
158
159    pub(crate) fn into_inner(self) -> W {
160        self.inner
161    }
162}
163
164impl<W, F> io::Write for Inspect<W, F>
165where
166    W: io::Write,
167    for<'a> F: FnMut(&'a [u8]),
168{
169    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
170        (self.f)(buf);
171        self.inner.write(buf)
172    }
173
174    fn flush(&mut self) -> io::Result<()> {
175        self.inner.flush()
176    }
177}