axum/examples/key_value_store.rs

73 lines
1.8 KiB
Rust
Raw Normal View History

2021-05-30 14:33:20 +02:00
use bytes::Bytes;
2021-06-01 00:34:09 +02:00
use http::{Request, StatusCode};
2021-05-30 14:33:20 +02:00
use hyper::Server;
use std::{
collections::HashMap,
net::SocketAddr,
sync::{Arc, Mutex},
time::Duration,
};
use tower::{make::Shared, ServiceBuilder};
use tower_http::{
add_extension::AddExtensionLayer, compression::CompressionLayer, trace::TraceLayer,
};
2021-06-01 15:07:16 +02:00
use tower_web::{
body::Body,
extract::{BytesMaxLength, Extension, UrlParams},
2021-06-04 01:00:48 +02:00
get, route, Handler,
2021-06-01 15:07:16 +02:00
};
2021-05-30 14:33:20 +02:00
#[tokio::main]
async fn main() {
tracing_subscriber::fmt::init();
// build our application with some routes
2021-06-04 01:00:48 +02:00
let app = route(
"/:key",
get(kv_get.layer(CompressionLayer::new())).post(kv_set),
);
2021-05-30 14:33:20 +02:00
// add some middleware
let app = ServiceBuilder::new()
.timeout(Duration::from_secs(10))
.layer(TraceLayer::new_for_http())
.layer(AddExtensionLayer::new(SharedState::default()))
.service(app);
// run it with hyper
let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
tracing::debug!("listening on {}", addr);
let server = Server::bind(&addr).serve(Shared::new(app));
server.await.unwrap();
}
type SharedState = Arc<Mutex<State>>;
#[derive(Default)]
struct State {
db: HashMap<String, Bytes>,
}
2021-06-04 01:00:48 +02:00
async fn kv_get(
2021-05-30 14:33:20 +02:00
_req: Request<Body>,
2021-06-01 15:07:16 +02:00
UrlParams((key,)): UrlParams<(String,)>,
Extension(state): Extension<SharedState>,
2021-06-01 00:34:09 +02:00
) -> Result<Bytes, StatusCode> {
2021-05-30 14:33:20 +02:00
let db = &state.lock().unwrap().db;
2021-05-30 16:53:27 +02:00
if let Some(value) = db.get(&key) {
2021-05-30 14:33:20 +02:00
Ok(value.clone())
} else {
2021-06-01 00:34:09 +02:00
Err(StatusCode::NOT_FOUND)
2021-05-30 14:33:20 +02:00
}
}
2021-06-04 01:00:48 +02:00
async fn kv_set(
2021-05-30 14:33:20 +02:00
_req: Request<Body>,
2021-06-01 15:07:16 +02:00
UrlParams((key,)): UrlParams<(String,)>,
BytesMaxLength(value): BytesMaxLength<{ 1024 * 5_000 }>, // ~5mb
Extension(state): Extension<SharedState>,
2021-05-31 22:54:21 +02:00
) {
2021-06-01 15:07:16 +02:00
state.lock().unwrap().db.insert(key, value);
2021-05-30 14:33:20 +02:00
}