krata: utilize gRPC for control service

This commit is contained in:
Alex Zenla
2024-03-06 12:05:01 +00:00
parent 31cf3044a4
commit 3628422168
24 changed files with 532 additions and 1159 deletions

View File

@ -32,10 +32,8 @@ bytes = { workspace = true }
tokio-stream = { workspace = true }
async-trait = { workspace = true }
signal-hook = { workspace = true }
[dependencies.tokio-listener]
workspace = true
features = ["clap"]
async-stream = { workspace = true }
tonic = { workspace = true, features = ["tls"]}
[dependencies.krata]
path = "../shared"
@ -62,7 +60,3 @@ path = "src/lib.rs"
[[bin]]
name = "kratad"
path = "bin/daemon.rs"
[[example]]
name = "kratad-dial"
path = "examples/dial.rs"

View File

@ -1,15 +1,17 @@
use std::sync::{atomic::AtomicBool, Arc};
use anyhow::{anyhow, Result};
use anyhow::Result;
use clap::Parser;
use env_logger::Env;
use krata::dial::ControlDialAddress;
use kratad::{runtime::Runtime, Daemon};
use tokio_listener::ListenerAddressLFlag;
use std::{
str::FromStr,
sync::{atomic::AtomicBool, Arc},
};
#[derive(Parser)]
struct Args {
#[clap(flatten)]
listener: ListenerAddressLFlag,
#[arg(short, long, default_value = "unix:///var/lib/krata/daemon.socket")]
listen: String,
#[arg(short, long, default_value = "/var/lib/krata")]
store: String,
}
@ -20,12 +22,10 @@ async fn main() -> Result<()> {
mask_sighup()?;
let args = Args::parse();
let Some(listener) = args.listener.bind().await else {
return Err(anyhow!("no listener specified"));
};
let addr = ControlDialAddress::from_str(&args.listen)?;
let runtime = Runtime::new(args.store.clone()).await?;
let mut daemon = Daemon::new(runtime).await?;
daemon.listen(listener?).await?;
let mut daemon = Daemon::new(args.store.clone(), runtime).await?;
daemon.listen(addr).await?;
Ok(())
}

View File

@ -1,28 +0,0 @@
use anyhow::Result;
use krata::control::{ListRequest, Message, Request, RequestBox};
use tokio::{
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
net::TcpStream,
};
use tokio_stream::{wrappers::LinesStream, StreamExt};
#[tokio::main]
async fn main() -> Result<()> {
let mut stream = TcpStream::connect("127.0.0.1:4050").await?;
let (read, mut write) = stream.split();
let mut read = LinesStream::new(BufReader::new(read).lines());
let send = Message::Request(RequestBox {
id: 1,
request: Request::List(ListRequest {}),
});
let mut line = serde_json::to_string(&send)?;
line.push('\n');
write.write_all(line.as_bytes()).await?;
println!("sent: {:?}", send);
while let Some(line) = read.try_next().await? {
let message: Message = serde_json::from_str(&line)?;
println!("received: {:?}", message);
}
Ok(())
}

172
daemon/src/control.rs Normal file
View File

@ -0,0 +1,172 @@
use std::{io, pin::Pin};
use async_stream::try_stream;
use futures::Stream;
use krata::control::{
control_service_server::ControlService, ConsoleDataReply, ConsoleDataRequest,
DestroyGuestReply, DestroyGuestRequest, GuestInfo, LaunchGuestReply, LaunchGuestRequest,
ListGuestsReply, ListGuestsRequest,
};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
select,
};
use tokio_stream::StreamExt;
use tonic::{Request, Response, Status, Streaming};
use crate::runtime::{launch::GuestLaunchRequest, Runtime};
pub struct ApiError {
message: String,
}
impl From<anyhow::Error> for ApiError {
fn from(value: anyhow::Error) -> Self {
ApiError {
message: value.to_string(),
}
}
}
impl From<ApiError> for Status {
fn from(value: ApiError) -> Self {
Status::unknown(value.message)
}
}
#[derive(Clone)]
pub struct RuntimeControlService {
runtime: Runtime,
}
impl RuntimeControlService {
pub fn new(runtime: Runtime) -> Self {
Self { runtime }
}
}
enum ConsoleDataSelect {
Read(io::Result<usize>),
Write(Option<Result<ConsoleDataRequest, tonic::Status>>),
}
#[tonic::async_trait]
impl ControlService for RuntimeControlService {
type ConsoleDataStream =
Pin<Box<dyn Stream<Item = Result<ConsoleDataReply, Status>> + Send + 'static>>;
async fn launch_guest(
&self,
request: Request<LaunchGuestRequest>,
) -> Result<Response<LaunchGuestReply>, Status> {
let request = request.into_inner();
let guest: GuestInfo = self
.runtime
.launch(GuestLaunchRequest {
image: &request.image,
vcpus: request.vcpus,
mem: request.mem,
env: empty_vec_optional(request.env),
run: empty_vec_optional(request.run),
debug: false,
})
.await
.map_err(ApiError::from)?
.into();
Ok(Response::new(LaunchGuestReply { guest: Some(guest) }))
}
async fn destroy_guest(
&self,
request: Request<DestroyGuestRequest>,
) -> Result<Response<DestroyGuestReply>, Status> {
let request = request.into_inner();
self.runtime
.destroy(&request.guest_id)
.await
.map_err(ApiError::from)?;
Ok(Response::new(DestroyGuestReply {}))
}
async fn list_guests(
&self,
request: Request<ListGuestsRequest>,
) -> Result<Response<ListGuestsReply>, Status> {
let _ = request.into_inner();
let guests = self.runtime.list().await.map_err(ApiError::from)?;
let guests = guests
.into_iter()
.map(GuestInfo::from)
.collect::<Vec<GuestInfo>>();
Ok(Response::new(ListGuestsReply { guests }))
}
async fn console_data(
&self,
request: Request<Streaming<ConsoleDataRequest>>,
) -> Result<Response<Self::ConsoleDataStream>, Status> {
let mut input = request.into_inner();
let Some(request) = input.next().await else {
return Err(ApiError {
message: "expected to have at least one request".to_string(),
}
.into());
};
let request = request?;
let mut console = self
.runtime
.console(&request.guest)
.await
.map_err(ApiError::from)?;
let output = try_stream! {
let mut buffer: Vec<u8> = vec![0u8; 256];
loop {
let what = select! {
x = console.read_handle.read(&mut buffer) => ConsoleDataSelect::Read(x),
x = input.next() => ConsoleDataSelect::Write(x),
};
match what {
ConsoleDataSelect::Read(result) => {
let size = result?;
let data = buffer[0..size].to_vec();
yield ConsoleDataReply { data, };
},
ConsoleDataSelect::Write(Some(request)) => {
let request = request?;
if !request.data.is_empty() {
console.write_handle.write_all(&request.data).await?;
}
},
ConsoleDataSelect::Write(None) => {
break;
}
}
}
};
Ok(Response::new(Box::pin(output) as Self::ConsoleDataStream))
}
}
impl From<crate::runtime::GuestInfo> for GuestInfo {
fn from(value: crate::runtime::GuestInfo) -> Self {
GuestInfo {
id: value.uuid.to_string(),
image: value.image,
ipv4: value.ipv4.map(|x| x.ip().to_string()).unwrap_or_default(),
ipv6: value.ipv6.map(|x| x.ip().to_string()).unwrap_or_default(),
}
}
}
fn empty_vec_optional<T>(value: Vec<T>) -> Option<Vec<T>> {
if value.is_empty() {
None
} else {
Some(value)
}
}

View File

@ -1,91 +0,0 @@
use anyhow::{anyhow, Result};
use krata::control::{ConsoleStreamResponse, ConsoleStreamUpdate, Request, Response, StreamUpdate};
use log::warn;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
select,
};
use crate::{
listen::DaemonRequestHandler,
runtime::{console::XenConsole, Runtime},
};
use krata::stream::{ConnectionStreams, StreamContext};
pub struct ConsoleStreamRequestHandler {}
impl Default for ConsoleStreamRequestHandler {
fn default() -> Self {
Self::new()
}
}
impl ConsoleStreamRequestHandler {
pub fn new() -> Self {
Self {}
}
async fn link_console_stream(mut stream: StreamContext, mut console: XenConsole) -> Result<()> {
loop {
let mut buffer = vec![0u8; 256];
select! {
x = console.read_handle.read(&mut buffer) => match x {
Ok(size) => {
let data = buffer[0..size].to_vec();
let update = StreamUpdate::ConsoleStream(ConsoleStreamUpdate {
data,
});
stream.send(update).await?;
},
Err(error) => {
return Err(error.into());
}
},
x = stream.receiver.recv() => match x {
Some(StreamUpdate::ConsoleStream(update)) => {
console.write_handle.write_all(&update.data).await?;
}
None => {
break;
}
}
};
}
Ok(())
}
}
#[async_trait::async_trait]
impl DaemonRequestHandler for ConsoleStreamRequestHandler {
fn accepts(&self, request: &Request) -> bool {
matches!(request, Request::ConsoleStream(_))
}
async fn handle(
&self,
streams: ConnectionStreams,
runtime: Runtime,
request: Request,
) -> Result<Response> {
let console_stream = match request {
Request::ConsoleStream(stream) => stream,
_ => return Err(anyhow!("unknown request")),
};
let console = runtime.console(&console_stream.guest).await?;
let stream = streams.open().await?;
let id = stream.id;
tokio::task::spawn(async move {
if let Err(error) =
ConsoleStreamRequestHandler::link_console_stream(stream, console).await
{
warn!("failed to process console stream: {}", error);
}
});
Ok(Response::ConsoleStream(ConsoleStreamResponse {
stream: id,
}))
}
}

View File

@ -1,44 +0,0 @@
use anyhow::{anyhow, Result};
use krata::{
control::{DestroyResponse, Request, Response},
stream::ConnectionStreams,
};
use crate::{listen::DaemonRequestHandler, runtime::Runtime};
pub struct DestroyRequestHandler {}
impl Default for DestroyRequestHandler {
fn default() -> Self {
Self::new()
}
}
impl DestroyRequestHandler {
pub fn new() -> Self {
Self {}
}
}
#[async_trait::async_trait]
impl DaemonRequestHandler for DestroyRequestHandler {
fn accepts(&self, request: &Request) -> bool {
matches!(request, Request::Destroy(_))
}
async fn handle(
&self,
_: ConnectionStreams,
runtime: Runtime,
request: Request,
) -> Result<Response> {
let destroy = match request {
Request::Destroy(destroy) => destroy,
_ => return Err(anyhow!("unknown request")),
};
let guest = runtime.destroy(&destroy.guest).await?;
Ok(Response::Destroy(DestroyResponse {
guest: guest.to_string(),
}))
}
}

View File

@ -1,55 +0,0 @@
use anyhow::{anyhow, Result};
use krata::{
control::{GuestInfo, LaunchResponse, Request, Response},
stream::ConnectionStreams,
};
use crate::{
listen::DaemonRequestHandler,
runtime::{launch::GuestLaunchRequest, Runtime},
};
pub struct LaunchRequestHandler {}
impl Default for LaunchRequestHandler {
fn default() -> Self {
Self::new()
}
}
impl LaunchRequestHandler {
pub fn new() -> Self {
Self {}
}
}
#[async_trait::async_trait]
impl DaemonRequestHandler for LaunchRequestHandler {
fn accepts(&self, request: &Request) -> bool {
matches!(request, Request::Launch(_))
}
async fn handle(
&self,
_: ConnectionStreams,
runtime: Runtime,
request: Request,
) -> Result<Response> {
let launch = match request {
Request::Launch(launch) => launch,
_ => return Err(anyhow!("unknown request")),
};
let guest: GuestInfo = runtime
.launch(GuestLaunchRequest {
image: &launch.image,
vcpus: launch.vcpus,
mem: launch.mem,
env: launch.env,
run: launch.run,
debug: false,
})
.await?
.into();
Ok(Response::Launch(LaunchResponse { guest }))
}
}

View File

@ -1,37 +0,0 @@
use anyhow::Result;
use krata::{
control::{GuestInfo, ListResponse, Request, Response},
stream::ConnectionStreams,
};
use crate::{listen::DaemonRequestHandler, runtime::Runtime};
pub struct ListRequestHandler {}
impl Default for ListRequestHandler {
fn default() -> Self {
Self::new()
}
}
impl ListRequestHandler {
pub fn new() -> Self {
Self {}
}
}
#[async_trait::async_trait]
impl DaemonRequestHandler for ListRequestHandler {
fn accepts(&self, request: &Request) -> bool {
matches!(request, Request::List(_))
}
async fn handle(&self, _: ConnectionStreams, runtime: Runtime, _: Request) -> Result<Response> {
let guests = runtime.list().await?;
let guests = guests
.into_iter()
.map(GuestInfo::from)
.collect::<Vec<GuestInfo>>();
Ok(Response::List(ListResponse { guests }))
}
}

View File

@ -1,15 +0,0 @@
pub mod console;
pub mod destroy;
pub mod launch;
pub mod list;
impl From<crate::runtime::GuestInfo> for krata::control::GuestInfo {
fn from(value: crate::runtime::GuestInfo) -> Self {
krata::control::GuestInfo {
id: value.uuid.to_string(),
image: value.image.clone(),
ipv4: value.ipv4.map(|x| x.ip().to_string()),
ipv6: value.ipv6.map(|x| x.ip().to_string()),
}
}
}

View File

@ -1,37 +1,74 @@
use anyhow::Result;
use handlers::{
console::ConsoleStreamRequestHandler, destroy::DestroyRequestHandler,
launch::LaunchRequestHandler, list::ListRequestHandler,
};
use listen::{DaemonListener, DaemonRequestHandlers};
use runtime::Runtime;
use tokio_listener::Listener;
use std::{net::SocketAddr, path::PathBuf, str::FromStr};
pub mod handlers;
pub mod listen;
use anyhow::Result;
use control::RuntimeControlService;
use krata::{control::control_service_server::ControlServiceServer, dial::ControlDialAddress};
use log::info;
use runtime::Runtime;
use tokio::net::UnixListener;
use tokio_stream::wrappers::UnixListenerStream;
use tonic::transport::{Identity, Server, ServerTlsConfig};
pub mod control;
pub mod runtime;
pub struct Daemon {
store: String,
runtime: Runtime,
}
impl Daemon {
pub async fn new(runtime: Runtime) -> Result<Self> {
Ok(Self { runtime })
pub async fn new(store: String, runtime: Runtime) -> Result<Self> {
Ok(Self { store, runtime })
}
pub async fn listen(&mut self, listener: Listener) -> Result<()> {
let handlers = DaemonRequestHandlers::new(
self.runtime.clone(),
vec![
Box::new(LaunchRequestHandler::new()),
Box::new(DestroyRequestHandler::new()),
Box::new(ConsoleStreamRequestHandler::new()),
Box::new(ListRequestHandler::new()),
],
);
let mut listener = DaemonListener::new(listener, handlers);
listener.handle().await?;
pub async fn listen(&mut self, addr: ControlDialAddress) -> Result<()> {
let control_service = RuntimeControlService::new(self.runtime.clone());
let mut server = Server::builder();
if let ControlDialAddress::Tls {
host: _,
port: _,
insecure,
} = &addr
{
let mut tls_config = ServerTlsConfig::new();
if !insecure {
let certificate_path = format!("{}/tls/daemon.pem", self.store);
let key_path = format!("{}/tls/daemon.key", self.store);
tls_config = tls_config.identity(Identity::from_pem(certificate_path, key_path));
}
server = server.tls_config(tls_config)?;
}
let server = server.add_service(ControlServiceServer::new(control_service));
info!("listening on address {}", addr);
match addr {
ControlDialAddress::UnixSocket { path } => {
let path = PathBuf::from(path);
if path.exists() {
tokio::fs::remove_file(&path).await?;
}
let listener = UnixListener::bind(path)?;
let stream = UnixListenerStream::new(listener);
server.serve_with_incoming(stream).await?;
}
ControlDialAddress::Tcp { host, port } => {
let address = format!("{}:{}", host, port);
server.serve(SocketAddr::from_str(&address)?).await?;
}
ControlDialAddress::Tls {
host,
port,
insecure: _,
} => {
let address = format!("{}:{}", host, port);
server.serve(SocketAddr::from_str(&address)?).await?;
}
}
Ok(())
}
}

View File

@ -1,228 +0,0 @@
use std::collections::HashMap;
use std::sync::Arc;
use anyhow::{anyhow, Result};
use krata::control::{ErrorResponse, Message, Request, RequestBox, Response, ResponseBox};
use log::trace;
use log::warn;
use tokio::sync::Mutex;
use tokio::{
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
select,
sync::mpsc::{channel, Receiver, Sender},
};
use tokio_listener::{Connection, Listener, SomeSocketAddrClonable};
use tokio_stream::{wrappers::LinesStream, StreamExt};
use crate::runtime::Runtime;
use krata::stream::ConnectionStreams;
const QUEUE_MAX_LEN: usize = 100;
#[async_trait::async_trait]
pub trait DaemonRequestHandler: Send + Sync {
fn accepts(&self, request: &Request) -> bool;
async fn handle(
&self,
streams: ConnectionStreams,
runtime: Runtime,
request: Request,
) -> Result<Response>;
}
#[derive(Clone)]
pub struct DaemonRequestHandlers {
runtime: Runtime,
handlers: Arc<Vec<Box<dyn DaemonRequestHandler>>>,
}
impl DaemonRequestHandlers {
pub fn new(runtime: Runtime, handlers: Vec<Box<dyn DaemonRequestHandler>>) -> Self {
DaemonRequestHandlers {
runtime,
handlers: Arc::new(handlers),
}
}
async fn dispatch(&self, streams: ConnectionStreams, request: Request) -> Result<Response> {
for handler in self.handlers.iter() {
if handler.accepts(&request) {
return handler.handle(streams, self.runtime.clone(), request).await;
}
}
Err(anyhow!("daemon cannot handle that request"))
}
}
pub struct DaemonListener {
listener: Listener,
handlers: DaemonRequestHandlers,
connections: Arc<Mutex<HashMap<u64, DaemonConnection>>>,
next: Arc<Mutex<u64>>,
}
impl DaemonListener {
pub fn new(listener: Listener, handlers: DaemonRequestHandlers) -> DaemonListener {
DaemonListener {
listener,
handlers,
connections: Arc::new(Mutex::new(HashMap::new())),
next: Arc::new(Mutex::new(0)),
}
}
pub async fn handle(&mut self) -> Result<()> {
loop {
let (connection, addr) = self.listener.accept().await?;
let connection =
DaemonConnection::new(connection, addr.clonable(), self.handlers.clone()).await?;
let id = {
let mut next = self.next.lock().await;
let id = *next;
*next = id + 1;
id
};
trace!("new connection from {}", connection.addr);
let tx_channel = connection.tx_sender.clone();
let addr = connection.addr.clone();
self.connections.lock().await.insert(id, connection);
let connections_for_close = self.connections.clone();
tokio::task::spawn(async move {
tx_channel.closed().await;
trace!("connection from {} closed", addr);
connections_for_close.lock().await.remove(&id);
});
}
}
}
#[derive(Clone)]
pub struct DaemonConnection {
tx_sender: Sender<Message>,
addr: SomeSocketAddrClonable,
handlers: DaemonRequestHandlers,
streams: ConnectionStreams,
}
impl DaemonConnection {
pub async fn new(
connection: Connection,
addr: SomeSocketAddrClonable,
handlers: DaemonRequestHandlers,
) -> Result<Self> {
let (tx_sender, tx_receiver) = channel::<Message>(QUEUE_MAX_LEN);
let streams_tx_sender = tx_sender.clone();
let instance = DaemonConnection {
tx_sender,
addr,
handlers,
streams: ConnectionStreams::new(streams_tx_sender),
};
{
let mut instance = instance.clone();
tokio::task::spawn(async move {
if let Err(error) = instance.process(tx_receiver, connection).await {
warn!(
"failed to process daemon connection for {}: {}",
instance.addr, error
);
}
});
}
Ok(instance)
}
async fn process(
&mut self,
mut tx_receiver: Receiver<Message>,
connection: Connection,
) -> Result<()> {
let (read, mut write) = tokio::io::split(connection);
let mut read = LinesStream::new(BufReader::new(read).lines());
loop {
select! {
x = read.next() => match x {
Some(Ok(line)) => {
let message: Message = serde_json::from_str(&line)?;
trace!("received message '{}' from {}", serde_json::to_string(&message)?, self.addr);
let mut context = self.clone();
tokio::task::spawn(async move {
if let Err(error) = context.handle_message(&message).await {
let line = serde_json::to_string(&message).unwrap_or("<invalid>".to_string());
warn!("failed to handle message '{}' from {}: {}", line, context.addr, error);
}
});
},
Some(Err(error)) => {
return Err(error.into());
},
None => {
break;
}
},
x = tx_receiver.recv() => match x {
Some(message) => {
if let Message::StreamUpdated(ref update) = message {
self.streams.outgoing(update).await?;
}
let mut line = serde_json::to_string(&message)?;
trace!("sending message '{}' to {}", line, self.addr);
line.push('\n');
write.write_all(line.as_bytes()).await?;
},
None => {
break;
}
}
};
}
Ok(())
}
async fn handle_message(&mut self, message: &Message) -> Result<()> {
match message {
Message::Request(req) => {
self.handle_request(req.clone()).await?;
}
Message::Response(_) => {
return Err(anyhow!(
"received a response message from client {}, but this is the daemon",
self.addr
));
}
Message::StreamUpdated(updated) => {
self.streams.incoming(updated.clone()).await?;
}
}
Ok(())
}
async fn handle_request(&mut self, req: RequestBox) -> Result<()> {
let id = req.id;
let response = self
.handlers
.dispatch(self.streams.clone(), req.request)
.await
.map_err(|error| {
Response::Error(ErrorResponse {
message: error.to_string(),
})
});
let response = if let Err(response) = response {
response
} else {
response.unwrap()
};
let resp = ResponseBox { id, response };
self.tx_sender.send(Message::Response(resp)).await?;
Ok(())
}
}