garage/src/rpc_server.rs

70 lines
2.1 KiB
Rust
Raw Normal View History

2020-04-06 17:55:39 +00:00
use std::net::SocketAddr;
use std::sync::Arc;
use bytes::IntoBuf;
use hyper::service::{make_service_fn, service_fn};
2020-04-06 17:55:39 +00:00
use hyper::server::conn::AddrStream;
use hyper::{Body, Method, Request, Response, Server, StatusCode};
use futures::future::Future;
use crate::error::Error;
use crate::proto::Message;
2020-04-06 17:55:39 +00:00
use crate::membership::System;
fn err_to_msg(x: Result<Message, Error>) -> Message {
match x {
Err(e) => Message::Error(format!("{}", e)),
Ok(msg) => msg,
}
}
async fn handler(sys: Arc<System>, req: Request<Body>, addr: SocketAddr) -> Result<Response<Body>, Error> {
if req.method() != &Method::POST {
let mut bad_request = Response::default();
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
return Ok(bad_request);
}
let whole_body = hyper::body::to_bytes(req.into_body()).await?;
2020-04-06 17:55:39 +00:00
let msg = rmp_serde::decode::from_read::<_, Message>(whole_body.into_buf())?;
eprintln!("RPC from {}: {:?}", addr, msg);
let resp = err_to_msg(match &msg {
2020-04-06 19:02:15 +00:00
Message::Ping(ping) => sys.handle_ping(&addr, ping).await,
2020-04-06 20:27:51 +00:00
Message::PullStatus => sys.handle_pull_status().await,
Message::PullConfig => sys.handle_pull_config().await,
Message::AdvertiseNodesUp(adv) => sys.handle_advertise_nodes_up(adv).await,
Message::AdvertiseConfig(adv) => sys.handle_advertise_config(adv).await,
2020-04-06 17:55:39 +00:00
_ => Ok(Message::Error(format!("Unexpected message: {:?}", msg))),
});
Ok(Response::new(Body::from(
rmp_serde::encode::to_vec_named(&resp)?
)))
}
2020-04-06 17:55:39 +00:00
pub async fn run_rpc_server(sys: Arc<System>, shutdown_signal: impl Future<Output=()>) -> Result<(), hyper::Error> {
let bind_addr = ([0, 0, 0, 0], sys.config.rpc_port).into();
2020-04-06 17:55:39 +00:00
let service = make_service_fn(|conn: &AddrStream| {
let client_addr = conn.remote_addr();
let sys = sys.clone();
async move {
Ok::<_, Error>(service_fn(move |req: Request<Body>| {
let sys = sys.clone();
handler(sys, req, client_addr)
}))
}
});
2020-04-06 17:55:39 +00:00
let server = Server::bind(&bind_addr).serve(service) ;
let graceful = server.with_graceful_shutdown(shutdown_signal);
2020-04-06 17:55:39 +00:00
println!("RPC server listening on http://{}", bind_addr);
graceful.await
}