diff options
author | Rutger Broekhoff | 2024-07-11 19:11:51 +0200 |
---|---|---|
committer | Rutger Broekhoff | 2024-07-11 19:11:51 +0200 |
commit | 3e67a3486eed22522f4352503ef7067ca81a8050 (patch) | |
tree | 971f0a84c073731c944069bdefdb848099428bdf /gitolfs3-server/src/dlimit.rs | |
parent | 0822bdf658e03bdb5b31a90b8c64f3d0b6a6cff7 (diff) | |
download | gitolfs3-3e67a3486eed22522f4352503ef7067ca81a8050.tar.gz gitolfs3-3e67a3486eed22522f4352503ef7067ca81a8050.zip |
Split server code into multiple smaller modules
Diffstat (limited to 'gitolfs3-server/src/dlimit.rs')
-rw-r--r-- | gitolfs3-server/src/dlimit.rs | 73 |
1 files changed, 73 insertions, 0 deletions
diff --git a/gitolfs3-server/src/dlimit.rs b/gitolfs3-server/src/dlimit.rs new file mode 100644 index 0000000..f68bec1 --- /dev/null +++ b/gitolfs3-server/src/dlimit.rs | |||
@@ -0,0 +1,73 @@ | |||
1 | use std::sync::Arc; | ||
2 | use std::time::Duration; | ||
3 | use tokio::io::AsyncWriteExt; | ||
4 | use tokio::sync::Mutex; | ||
5 | |||
6 | // I know that this is pretty bad, but it's good enough (??) for now. | ||
7 | pub struct DownloadLimiter { | ||
8 | current: u64, | ||
9 | limit: u64, | ||
10 | } | ||
11 | |||
12 | impl DownloadLimiter { | ||
13 | pub async fn new(limit: u64) -> Arc<Mutex<DownloadLimiter>> { | ||
14 | let dlimit_str = match tokio::fs::read_to_string(".gitolfs3-dlimit").await { | ||
15 | Ok(dlimit_str) => dlimit_str, | ||
16 | Err(e) => { | ||
17 | println!("Failed to read download counter, assuming 0: {e}"); | ||
18 | return DownloadLimiter { current: 0, limit }.auto_resetting(); | ||
19 | } | ||
20 | }; | ||
21 | let current: u64 = match dlimit_str | ||
22 | .parse() | ||
23 | .map_err(|e| tokio::io::Error::new(tokio::io::ErrorKind::InvalidData, e)) | ||
24 | { | ||
25 | Ok(current) => current, | ||
26 | Err(e) => { | ||
27 | println!("Failed to read download counter, assuming 0: {e}"); | ||
28 | return DownloadLimiter { current: 0, limit }.auto_resetting(); | ||
29 | } | ||
30 | }; | ||
31 | |||
32 | DownloadLimiter { current, limit }.auto_resetting() | ||
33 | } | ||
34 | |||
35 | fn auto_resetting(self) -> Arc<Mutex<Self>> { | ||
36 | let limiter_ref = Arc::new(Mutex::new(self)); | ||
37 | let limiter_ref_cloned = limiter_ref.clone(); | ||
38 | tokio::spawn(async move { | ||
39 | loop { | ||
40 | println!("Resetting download counter in one hour"); | ||
41 | tokio::time::sleep(Duration::from_secs(3600)).await; | ||
42 | println!("Resetting download counter"); | ||
43 | limiter_ref_cloned.lock().await.reset().await; | ||
44 | } | ||
45 | }); | ||
46 | limiter_ref | ||
47 | } | ||
48 | |||
49 | pub async fn request(&mut self, n: u64) -> tokio::io::Result<bool> { | ||
50 | if self.current + n > self.limit { | ||
51 | return Ok(false); | ||
52 | } | ||
53 | self.current += n; | ||
54 | self.write_new_count().await?; | ||
55 | Ok(true) | ||
56 | } | ||
57 | |||
58 | pub async fn reset(&mut self) { | ||
59 | self.current = 0; | ||
60 | if let Err(e) = self.write_new_count().await { | ||
61 | println!("Failed to reset download counter: {e}"); | ||
62 | } | ||
63 | } | ||
64 | |||
65 | async fn write_new_count(&self) -> tokio::io::Result<()> { | ||
66 | let cwd = tokio::fs::File::open(std::env::current_dir()?).await?; | ||
67 | let mut file = tokio::fs::File::create(".gitolfs3-dlimit.tmp").await?; | ||
68 | file.write_all(self.current.to_string().as_bytes()).await?; | ||
69 | file.sync_all().await?; | ||
70 | tokio::fs::rename(".gitolfs3-dlimit.tmp", ".gitolfs3-dlimit").await?; | ||
71 | cwd.sync_all().await | ||
72 | } | ||
73 | } | ||