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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
|
//! Reference service for Pure's HTTP benchmarks: the same contract as Pure's
//! document server (GET/PUT /objects/{key}, If-Match / If-None-Match: *,
//! a PUT body of u32le start, u32le remove, then the bytes to insert; the
//! reply is the inserted bytes; ETag is the revision).
//!
//! axum-baseline --port P [--durable DIR | --atomic DIR]
//! Without either, documents live in memory. --durable writes the document to
//! DIR/key and fdatasyncs before answering, as a straightforward durable
//! handler would (not crash-safe: a crash after truncation loses it).
//! --atomic writes DIR/key.tmp, fsyncs it, renames it over DIR/key and fsyncs
//! the directory: the usual crash-safe replacement.
use axum::{
body::Bytes,
extract::{Path, State},
http::{HeaderMap, HeaderValue, StatusCode},
response::{IntoResponse, Response},
routing::get,
Router,
};
use std::{collections::HashMap, path::PathBuf, sync::Arc};
use tokio::sync::Mutex;
#[derive(Default)]
struct Document {
bytes: Vec<u8>,
revision: u64,
}
struct Service {
documents: Mutex<HashMap<String, Arc<Mutex<Document>>>>,
durable: Option<PathBuf>,
atomic: bool,
}
fn valid_key(key: &str) -> bool {
!key.is_empty()
&& key.len() <= 64
&& key.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
}
fn tagged(status: StatusCode, revision: u64, error: &'static str, body: Bytes) -> Response {
let mut response = (status, body).into_response();
let headers = response.headers_mut();
headers.insert("X-Pure-Error", HeaderValue::from_static(error));
if revision > 0 {
headers.insert("ETag", HeaderValue::from_str(&format!("\"{revision}\"")).unwrap());
}
response
}
async fn document(service: &Service, key: &str) -> Arc<Mutex<Document>> {
let mut documents = service.documents.lock().await;
documents.entry(key.to_string()).or_default().clone()
}
async fn show(State(service): State<Arc<Service>>, Path(key): Path<String>) -> Response {
if !valid_key(&key) {
return tagged(StatusCode::BAD_REQUEST, 0, "request", Bytes::new());
}
let document = document(&service, &key).await;
let document = document.lock().await;
if document.revision == 0 {
return tagged(StatusCode::NOT_FOUND, 0, "not_found", Bytes::new());
}
tagged(StatusCode::OK, document.revision, "none", Bytes::copy_from_slice(&document.bytes))
}
async fn edit(
State(service): State<Arc<Service>>,
Path(key): Path<String>,
headers: HeaderMap,
body: Bytes,
) -> Response {
let expected = match (headers.get("if-match"), headers.get("if-none-match")) {
(Some(tag), None) => tag
.to_str()
.ok()
.and_then(|t| t.strip_prefix('"')?.strip_suffix('"')?.parse::<u64>().ok()),
(None, Some(star)) if star == "*" => Some(0),
_ => None,
};
let Some(expected) = expected.filter(|_| valid_key(&key) && body.len() >= 8) else {
return tagged(StatusCode::BAD_REQUEST, 0, "request", Bytes::new());
};
let start = u32::from_le_bytes(body[0..4].try_into().unwrap()) as usize;
let remove = u32::from_le_bytes(body[4..8].try_into().unwrap()) as usize;
let document = document(&service, &key).await;
let mut document = document.lock().await;
if document.revision != expected {
return tagged(StatusCode::PRECONDITION_FAILED, document.revision, "conflict", Bytes::new());
}
if start > document.bytes.len() || remove > document.bytes.len() - start {
return tagged(StatusCode::PAYLOAD_TOO_LARGE, 0, "limit", Bytes::new());
}
let insert = body.slice(8..);
document.bytes.splice(start..start + remove, insert.iter().copied());
if let Some(directory) = &service.durable {
let (directory, path, atomic) = (directory.clone(), directory.join(&key), service.atomic);
let bytes = document.bytes.clone();
let written = tokio::task::spawn_blocking(move || -> std::io::Result<()> {
use std::io::Write;
if !atomic {
let mut file = std::fs::File::create(path)?;
file.write_all(&bytes)?;
return file.sync_data();
}
let temporary = path.with_extension("tmp");
let mut file = std::fs::File::create(&temporary)?;
file.write_all(&bytes)?;
file.sync_all()?;
std::fs::rename(&temporary, &path)?;
std::fs::File::open(&directory)?.sync_all()
})
.await;
if !matches!(written, Ok(Ok(()))) {
return tagged(StatusCode::INTERNAL_SERVER_ERROR, 0, "storage", Bytes::new());
}
}
document.revision += 1;
tagged(StatusCode::OK, document.revision, "none", insert)
}
#[tokio::main]
async fn main() {
let arguments: Vec<String> = std::env::args().collect();
let mut port = 8768u16;
let mut durable = None;
let mut atomic = false;
let mut at = 1;
while at + 1 < arguments.len() {
match arguments[at].as_str() {
"--port" => port = arguments[at + 1].parse().expect("port"),
"--durable" => durable = Some(PathBuf::from(&arguments[at + 1])),
"--atomic" => {
durable = Some(PathBuf::from(&arguments[at + 1]));
atomic = true;
}
other => panic!("unknown argument {other}"),
}
at += 2;
}
if let Some(directory) = &durable {
std::fs::create_dir_all(directory).expect("durable directory");
}
let service = Arc::new(Service { documents: Mutex::new(HashMap::new()), durable, atomic });
let app = Router::new()
.route("/objects/{key}", get(show).put(edit))
.with_state(service);
let listener = tokio::net::TcpListener::bind(("127.0.0.1", port)).await.expect("bind");
println!("{{\"port\":{}}}", listener.local_addr().unwrap().port());
use std::io::Write;
std::io::stdout().flush().unwrap();
axum::serve(listener, app)
.with_graceful_shutdown(async {
tokio::signal::ctrl_c().await.ok();
})
.await
.unwrap();
}
|
