132 lines
4.3 KiB
Rust
132 lines
4.3 KiB
Rust
use std::convert::Infallible;
|
|
use std::time::Duration;
|
|
|
|
use bytes::Bytes;
|
|
use http_body_util::{BodyExt, Full, Limited};
|
|
use hyper::{
|
|
Request, Response, StatusCode, body::Incoming, server::conn::http1, service::service_fn,
|
|
};
|
|
use hyper_util::rt::TokioIo;
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::value::RawValue;
|
|
use tokio::net::TcpListener;
|
|
use tokio::signal::unix::{SignalKind, signal};
|
|
use tokio::sync::watch;
|
|
use tokio::task::JoinSet;
|
|
|
|
#[derive(Deserialize, Serialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
struct Row {
|
|
id: i64,
|
|
text: String,
|
|
weight: f64,
|
|
note: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize, Serialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
struct Payload {
|
|
rows: Vec<Row>,
|
|
raw: Box<RawValue>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
struct TypedRequest {
|
|
payload: Payload,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
struct RawRequest {
|
|
source: Box<RawValue>,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct Envelope<Value> {
|
|
code: u8,
|
|
message: &'static str,
|
|
data: Value,
|
|
}
|
|
|
|
async fn handle(request: Request<Incoming>) -> Result<Response<Full<Bytes>>, Infallible> {
|
|
let typed = request.uri().path() == "/benchmark/echo/echo";
|
|
if request.method() != hyper::Method::POST
|
|
|| !typed && request.uri().path() != "/benchmark/echo/raw"
|
|
{
|
|
return Ok(Response::builder()
|
|
.status(StatusCode::NOT_FOUND)
|
|
.body(Full::new(Bytes::new()))
|
|
.unwrap());
|
|
}
|
|
let body = Limited::new(request.into_body(), 16 * 1024 * 1024)
|
|
.collect()
|
|
.await;
|
|
let encoded = body.map_err(|error| error.to_string()).and_then(|body| {
|
|
let bytes = body.to_bytes();
|
|
if typed {
|
|
serde_json::from_slice::<TypedRequest>(&bytes).and_then(|value| {
|
|
serde_json::to_vec(&Envelope {
|
|
code: 0,
|
|
message: "ok",
|
|
data: value.payload,
|
|
})
|
|
})
|
|
} else {
|
|
serde_json::from_slice::<RawRequest>(&bytes).and_then(|value| {
|
|
serde_json::to_vec(&Envelope {
|
|
code: 0,
|
|
message: "ok",
|
|
data: value.source,
|
|
})
|
|
})
|
|
}
|
|
.map_err(|error| error.to_string())
|
|
});
|
|
let (status, body) = match encoded {
|
|
Ok(body) => (StatusCode::OK, body),
|
|
Err(_) => (StatusCode::BAD_REQUEST, Vec::new()),
|
|
};
|
|
Ok(Response::builder()
|
|
.status(status)
|
|
.header("Content-Type", "application/json")
|
|
.body(Full::new(Bytes::from(body)))
|
|
.unwrap())
|
|
}
|
|
|
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
tokio::runtime::Builder::new_current_thread().enable_all().build()?.block_on(async {
|
|
let settings: serde_json::Value = serde_json::from_slice(&std::fs::read("config/setting.json")?)?;
|
|
let port = u16::try_from(settings["http"]["port"].as_u64().ok_or("missing port")?)?;
|
|
let listener = TcpListener::bind(("127.0.0.1", port)).await?;
|
|
let mut terminate = signal(SignalKind::terminate())?;
|
|
let (shutdown, stopping) = watch::channel(false);
|
|
let mut connections = JoinSet::new();
|
|
loop {
|
|
tokio::select! {
|
|
_ = terminate.recv() => break,
|
|
Some(result) = connections.join_next(), if !connections.is_empty() => { result?; },
|
|
accepted = listener.accept() => {
|
|
let (stream, _) = accepted?;
|
|
stream.set_nodelay(true)?;
|
|
let mut stopping = stopping.clone();
|
|
connections.spawn(async move {
|
|
let connection = http1::Builder::new().serve_connection(TokioIo::new(stream), service_fn(handle));
|
|
tokio::pin!(connection);
|
|
tokio::select! {
|
|
_ = &mut connection => {},
|
|
_ = stopping.changed() => {
|
|
connection.as_mut().graceful_shutdown();
|
|
let _ = tokio::time::timeout(Duration::from_secs(2), connection).await;
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
shutdown.send(true)?;
|
|
while let Some(result) = connections.join_next().await { result?; }
|
|
Ok(())
|
|
})
|
|
}
|