zeyn 0.0.5

Rust web framework for data heavy applications, written by AI agents and reviewed by people
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

//! Memory of an upload stays flat as the file grows: the body is streamed to the store in
//! bounded chunks. Its own test binary, so no other test allocates while it measures.

use axum::body::{Body, Bytes};
use axum::http::{HeaderMap, Request, StatusCode, header};
use http_body_util::BodyExt;
use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use tower::ServiceExt;
use zeyn::prelude::*;
use zeyn::{BoxFuture, FileChunks, FileMeta, FileStore, I18n, Identity, StoredFile};

struct Counting;

static LIVE: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(layout) };
        if !p.is_null() {
            let now = LIVE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
            PEAK.fetch_max(now, Ordering::Relaxed);
        }
        p
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) };
        LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
    }
}

#[global_allocator]
static A: Counting = Counting;

const TOKEN: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
const BIG: FileRules = FileRules::new("big").max_bytes(1 << 30).accept(&[FileType::Png]);
const CHUNK: usize = 64 * 1024;

/// Counts the bytes it is given and keeps none of them.
#[derive(Clone, Default)]
struct Counter(Arc<AtomicU64>);

impl FileStore for Counter {
    fn put(&self, _meta: FileMeta, mut chunks: FileChunks) -> BoxFuture<Result<()>> {
        let n = self.0.clone();
        Box::pin(async move {
            while let Some(c) = chunks.next().await {
                n.fetch_add(c.map_err(|e| Error::Internal(e.to_string()))?.len() as u64, Ordering::Relaxed);
            }
            Ok(())
        })
    }
    fn get(&self, _id: String) -> BoxFuture<Result<StoredFile>> {
        Box::pin(async { Err(Error::NotFound) })
    }
    fn meta(&self, _id: String) -> BoxFuture<Result<FileMeta>> {
        Box::pin(async { Err(Error::NotFound) })
    }
    fn delete(&self, _id: String) -> BoxFuture<Result<()>> {
        Box::pin(async { Ok(()) })
    }
}

/// A PNG header followed by zeros, made one chunk at a time.
fn body(total: usize) -> Body {
    let stream = futures_util::stream::unfold(0usize, move |sent| async move {
        if sent >= total {
            return None;
        }
        let n = CHUNK.min(total - sent);
        let mut chunk = vec![0u8; n];
        if sent == 0 {
            chunk[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n");
        }
        Some((Ok::<_, std::io::Error>(Bytes::from(chunk)), sent + n))
    });
    Body::from_stream(stream)
}

async fn upload(router: &axum::Router, total: usize) -> (StatusCode, usize) {
    let base = LIVE.load(Ordering::Relaxed);
    PEAK.store(base, Ordering::Relaxed);
    let req = Request::post("/_zeyn/files/big")
        .header("x-zeyn-csrf", TOKEN)
        .header(header::COOKIE, format!("zeyn_csrf={TOKEN}"))
        .header("x-zeyn-file-name", "huge.png")
        .body(body(total))
        .unwrap();
    let res = router.clone().oneshot(req).await.unwrap();
    let status = res.status();
    let _ = res.into_body().collect().await;
    (status, PEAK.load(Ordering::Relaxed).saturating_sub(base))
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_large_upload_keeps_memory_flat() {
    let counter = Counter::default();
    let dir = std::env::temp_dir().join(format!("zeyn-files-memory-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    let router = |store: Box<dyn Fn(Feature) -> Feature>| {
        App::new("Test")
            .i18n(I18n::new("en"))
            .auth(|_: &HeaderMap| Some(Identity::user("ada").with_permissions(["files.write"])))
            .feature(store(Feature::new("files")))
            .router()
            .unwrap()
    };
    let counted = {
        let c = counter.clone();
        router(Box::new(move |f| f.files(BIG, "files.write", c.clone())))
    };
    assert_eq!(upload(&counted, 1 << 20).await.0, StatusCode::CREATED, "warm up");
    let (status, small) = upload(&counted, 8 << 20).await;
    assert_eq!(status, StatusCode::CREATED);
    let (status, large) = upload(&counted, 256 << 20).await;
    assert_eq!(status, StatusCode::CREATED);
    assert_eq!(counter.0.load(Ordering::Relaxed), (1 << 20) + (8 << 20) + (256 << 20));
    println!("peak heap growth: 8 MiB upload {} KiB, 256 MiB upload {} KiB", small / 1024, large / 1024);
    assert!(large < 4 << 20, "a 256 MiB upload must not hold more than 4 MiB, held {large} bytes");

    let disk = {
        let d = dir.clone();
        router(Box::new(move |f| f.files(BIG, "files.write", zeyn::DiskStore::new(d.clone()))))
    };
    let (status, on_disk) = upload(&disk, 64 << 20).await;
    assert_eq!(status, StatusCode::CREATED);
    println!("peak heap growth: 64 MiB upload to disk {} KiB", on_disk / 1024);
    assert!(on_disk < 4 << 20, "a 64 MiB upload to disk must not hold more than 4 MiB, held {on_disk} bytes");
    let _ = std::fs::remove_dir_all(&dir);
}