Tweak speed test buffer size

james58899 · Feb 29, 2024 · 3a6e507 · 3a6e507
1 parent 73e330b
commit 3a6e507
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/src/route/speed_test.rs b/src/route/speed_test.rs
@@ -7,11 +7,13 @@ use axum::{
     http::header::CONTENT_LENGTH,
     response::Response,
 };
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use rand::{prelude::SmallRng, RngCore, SeedableRng};
 
 use crate::{route::forbidden, util::string_to_hash, AppState, MAX_KEY_TIME_DRIFT};
 
+const BUFFER_SIZE: usize = 16384;
+
 // example: /t/5242880/1645930666/bce541b2a97788319e53a754b47e1801204ae7bf/43432228
 pub(super) async fn speedtest(Path((size, time, hash)): Path<(u64, i64, String)>, data: State<Arc<AppState>>) -> Response<Body> {
     // Check time & hash
@@ -27,14 +29,17 @@ pub(super) fn random_response(size: u64) -> Response<Body> {
     Response::builder()
         .header(CONTENT_LENGTH, size)
         .body(Body::from_stream(stream! {
-            let mut buffer = [0; 8192];
+            let mut buffer = BytesMut::zeroed(BUFFER_SIZE);
             SmallRng::from_entropy().fill_bytes(&mut buffer);
-            let buffer = Bytes::copy_from_slice(&buffer);
+            let buffer = buffer.freeze();
 
             let mut filled = 0;
             while filled < size {
-                let size = cmp::min(size - filled, 8192) as usize;
-                yield Result::Ok::<Bytes, Infallible>(buffer.slice(0..size));
+                let size = cmp::min(size - filled, BUFFER_SIZE as u64) as usize;
+                if size == BUFFER_SIZE {
+                    yield Ok(buffer.clone());
+                }
+                yield Ok::<Bytes, Infallible>(buffer.slice(0..size));
                 filled += size as u64;
             }
         }))