use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use anyhow::{anyhow, Context, Result};
use kamaji_proto::{
decode_frame, encode_frame, DrainBudget, Error as CodecError, ErrorCode, KamajiToYubaba,
ProbeStatus, ProtocolVersion, RequestId, WorkloadEntry, WorkloadId, YubabaToKamaji,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{unix::OwnedReadHalf, unix::OwnedWriteHalf, UnixStream};
use tokio::sync::Mutex;
use tracing::{debug, info};
pub const DEFAULT_SOCKET: &str = "/run/yah/kamaji.sock";
#[derive(Debug, thiserror::Error)]
pub enum ClientError {
#[error("kamaji UDS I/O: {0}")]
Io(#[from] std::io::Error),
#[error("kamaji wire codec: {0}")]
Codec(#[from] CodecError),
#[error("kamaji closed the connection mid-request")]
PeerClosed,
#[error("kamaji error: {code:?}: {message}")]
Remote { code: ErrorCode, message: String },
#[error("kamaji replied with unexpected variant: {0}")]
Unexpected(String),
#[error("kamaji response request_id mismatch (expected {expected:?}, got {got:?})")]
RequestIdMismatch { expected: RequestId, got: RequestId },
#[error("kamaji rejected handshake (version {wanted:?})")]
HandshakeRefused { wanted: ProtocolVersion },
}
#[derive(Debug)]
struct Inner {
rd: OwnedReadHalf,
wr: OwnedWriteHalf,
buf: Vec<u8>,
}
#[derive(Debug, Clone)]
pub struct ConstableInfo {
pub version: ProtocolVersion,
pub kamaji_version: String,
}
#[derive(Debug)]
pub struct KamajiClient {
socket: PathBuf,
next_request_id: AtomicU64,
info: ConstableInfo,
inner: Mutex<Inner>,
}
impl KamajiClient {
pub async fn connect(socket: impl Into<PathBuf>) -> Result<Self, ClientError> {
let socket = socket.into();
let stream = UnixStream::connect(&socket).await?;
let (rd, wr) = stream.into_split();
let mut inner = Inner {
rd,
wr,
buf: Vec::with_capacity(4096),
};
let hello = YubabaToKamaji::Hello {
version: ProtocolVersion::CURRENT,
};
let inner_mut: &mut Inner = &mut inner;
write_frame(&mut inner_mut.wr, &hello).await?;
let reply = read_frame(&mut inner_mut.rd, &mut inner_mut.buf).await?;
let info = match reply {
KamajiToYubaba::Welcome {
version,
kamaji_version,
} => ConstableInfo {
version,
kamaji_version,
},
KamajiToYubaba::Error { code, message, .. } => {
return Err(ClientError::Remote { code, message });
}
other => return Err(ClientError::Unexpected(format!("{other:?}"))),
};
info!(
socket = %socket.display(),
kamaji_version = %info.kamaji_version,
"kamaji handshake complete"
);
Ok(Self {
socket,
next_request_id: AtomicU64::new(1),
info,
inner: Mutex::new(inner),
})
}
pub fn socket(&self) -> &Path {
&self.socket
}
pub fn info(&self) -> &ConstableInfo {
&self.info
}
fn next_request_id(&self) -> RequestId {
RequestId(self.next_request_id.fetch_add(1, Ordering::Relaxed))
}
pub async fn list(&self) -> Result<Vec<WorkloadEntry>, ClientError> {
let request_id = self.next_request_id();
let reply = self
.request(YubabaToKamaji::List { request_id }, request_id)
.await?;
match reply {
KamajiToYubaba::WorkloadList {
request_id: rid,
entries,
} => {
check_rid(request_id, rid)?;
Ok(entries)
}
other => Err(ClientError::Unexpected(format!("{other:?}"))),
}
}
pub async fn stop(&self, id: &WorkloadId) -> Result<(), ClientError> {
let request_id = self.next_request_id();
let reply = self
.request(
YubabaToKamaji::Stop {
request_id,
id: id.clone(),
},
request_id,
)
.await?;
match reply {
KamajiToYubaba::Ack {
request_id: rid, ..
} => {
check_rid(request_id, rid)?;
Ok(())
}
other => Err(ClientError::Unexpected(format!("{other:?}"))),
}
}
pub async fn drain(
&self,
id: &WorkloadId,
budget: DrainBudget,
) -> Result<(bool, Option<String>), ClientError> {
let request_id = self.next_request_id();
let reply = self
.request(
YubabaToKamaji::Drain {
request_id,
id: id.clone(),
budget,
},
request_id,
)
.await?;
match reply {
KamajiToYubaba::DrainAck {
request_id: rid,
accepted,
reason,
..
} => {
check_rid(request_id, rid)?;
Ok((accepted, reason))
}
other => Err(ClientError::Unexpected(format!("{other:?}"))),
}
}
pub async fn probe(&self, id: &WorkloadId) -> Result<ProbeStatus, ClientError> {
let request_id = self.next_request_id();
let reply = self
.request(
YubabaToKamaji::Probe {
request_id,
id: id.clone(),
},
request_id,
)
.await?;
match reply {
KamajiToYubaba::ProbeResult {
request_id: rid,
status,
..
} => {
check_rid(request_id, rid)?;
Ok(status)
}
other => Err(ClientError::Unexpected(format!("{other:?}"))),
}
}
async fn request(
&self,
req: YubabaToKamaji,
expected_rid: RequestId,
) -> Result<KamajiToYubaba, ClientError> {
let mut guard = self.inner.lock().await;
let inner: &mut Inner = &mut guard;
write_frame(&mut inner.wr, &req).await?;
let reply = read_frame(&mut inner.rd, &mut inner.buf).await?;
if let KamajiToYubaba::Error {
request_id,
code,
message,
} = reply
{
debug!(
?request_id,
?code,
%message,
"kamaji returned Error for request"
);
if let Some(rid) = request_id {
check_rid(expected_rid, rid)?;
}
return Err(ClientError::Remote { code, message });
}
Ok(reply)
}
}
fn check_rid(expected: RequestId, got: RequestId) -> Result<(), ClientError> {
if expected == got {
Ok(())
} else {
Err(ClientError::RequestIdMismatch { expected, got })
}
}
async fn write_frame(wr: &mut OwnedWriteHalf, msg: &YubabaToKamaji) -> Result<(), ClientError> {
let bytes = encode_frame(msg)?;
wr.write_all(&bytes).await?;
Ok(())
}
async fn read_frame(
rd: &mut OwnedReadHalf,
buf: &mut Vec<u8>,
) -> Result<KamajiToYubaba, ClientError> {
let mut tmp = [0u8; 4096];
loop {
match decode_frame::<KamajiToYubaba>(buf) {
Ok((msg, consumed)) => {
buf.drain(..consumed);
return Ok(msg);
}
Err(CodecError::Truncated { .. }) => {
let n = rd.read(&mut tmp).await?;
if n == 0 {
return Err(ClientError::PeerClosed);
}
buf.extend_from_slice(&tmp[..n]);
}
Err(e) => return Err(ClientError::Codec(e)),
}
}
}
use async_trait::async_trait;
use kamaji_proto::WorkloadState as ProtoWorkloadState;
fn proto_state_to_status(s: ProtoWorkloadState) -> crate::WorkloadStatus {
match s {
ProtoWorkloadState::Pending | ProtoWorkloadState::Starting => {
crate::WorkloadStatus::Pending
}
ProtoWorkloadState::Running => crate::WorkloadStatus::Running,
ProtoWorkloadState::Draining => crate::WorkloadStatus::Stopping,
ProtoWorkloadState::Exited => crate::WorkloadStatus::Stopped,
ProtoWorkloadState::Failed => crate::WorkloadStatus::Failed {
reason: "workload exited with failure".into(),
},
_ => crate::WorkloadStatus::Failed {
reason: "unknown proto WorkloadState variant".into(),
},
}
}
fn entry_to_workload_state(entry: WorkloadEntry) -> crate::WorkloadState {
crate::WorkloadState {
ident: crate::MeshIdent(entry.id.0.clone()),
container_id: entry.id.0,
status: proto_state_to_status(entry.state),
mesh_ip: None,
}
}
#[async_trait]
impl crate::Kamaji for KamajiClient {
fn backend(&self) -> crate::Backend {
crate::Backend::Native
}
async fn deploy_workload(
&self,
spec: &workload_spec::WorkloadSpec,
mesh: &crate::MeshAssignment,
) -> anyhow::Result<crate::DeployResult> {
let id = WorkloadId::new(&spec.name);
let workload_envelope = workload_spec::Workload::Container(spec.clone());
let request_id = self.next_request_id();
let reply = self
.request(
YubabaToKamaji::Deploy {
request_id,
id: id.clone(),
spec: workload_envelope,
},
request_id,
)
.await
.map_err(|e| anyhow::anyhow!("kamaji deploy_workload: {e}"))?;
match reply {
KamajiToYubaba::Ack {
request_id: rid,
kind: kamaji_proto::AckKind::Deploy,
} => {
check_rid(request_id, rid)
.map_err(|e| anyhow::anyhow!("kamaji deploy ack: {e}"))?;
Ok(crate::DeployResult {
container_id: id.0,
mesh_ip: mesh.mesh_ip,
task_pid: 0,
})
}
other => Err(anyhow::anyhow!(
"kamaji deploy_workload: unexpected reply {other:?}"
)),
}
}
async fn list_workloads(&self) -> anyhow::Result<Vec<crate::WorkloadState>> {
let entries = self
.list()
.await
.map_err(|e| anyhow::anyhow!("kamaji list_workloads: {e}"))?;
Ok(entries.into_iter().map(entry_to_workload_state).collect())
}
async fn get_workload(
&self,
ident: &crate::MeshIdent,
) -> anyhow::Result<Option<crate::WorkloadState>> {
let all = self.list_workloads().await?;
Ok(all.into_iter().find(|s| s.ident == *ident))
}
async fn stream_logs(
&self,
_ident: &crate::MeshIdent,
_opts: crate::LogOpts,
) -> anyhow::Result<crate::LogStream> {
Err(anyhow::anyhow!(
"stream_logs not yet in the sibling-Kamaji wire protocol"
))
}
async fn restart_workload(&self, ident: &crate::MeshIdent) -> anyhow::Result<()> {
let id = WorkloadId::new(&ident.0);
self.stop(&id)
.await
.map_err(|e| anyhow::anyhow!("kamaji restart_workload (via stop): {e}"))
}
async fn graceful_upgrade_workload(
&self,
spec: &workload_spec::WorkloadSpec,
mesh: &crate::MeshAssignment,
) -> anyhow::Result<crate::DeployResult> {
let id = WorkloadId::new(&spec.name);
let workload_envelope = workload_spec::Workload::Container(spec.clone());
let request_id = self.next_request_id();
let reply = self
.request(
YubabaToKamaji::GracefulUpgrade {
request_id,
id: id.clone(),
spec: workload_envelope,
},
request_id,
)
.await
.map_err(|e| anyhow::anyhow!("kamaji graceful_upgrade_workload: {e}"))?;
match reply {
KamajiToYubaba::Ack {
request_id: rid,
kind: kamaji_proto::AckKind::GracefulUpgrade,
} => {
check_rid(request_id, rid)
.map_err(|e| anyhow::anyhow!("kamaji graceful_upgrade ack: {e}"))?;
Ok(crate::DeployResult {
container_id: id.0,
mesh_ip: mesh.mesh_ip,
task_pid: 0,
})
}
other => Err(anyhow::anyhow!(
"kamaji graceful_upgrade_workload: unexpected reply {other:?}"
)),
}
}
async fn teardown_workload(&self, ident: &crate::MeshIdent) -> anyhow::Result<()> {
let id = WorkloadId::new(&ident.0);
self.stop(&id)
.await
.map_err(|e| anyhow::anyhow!("kamaji teardown_workload: {e}"))
}
async fn health(&self) -> anyhow::Result<crate::RuntimeHealth> {
Ok(crate::RuntimeHealth {
ok: true,
version: Some(self.info().kamaji_version.clone()),
detail: None,
})
}
}
pub async fn connect_with_timeout(
socket: impl Into<PathBuf>,
timeout: Duration,
) -> Result<KamajiClient> {
let socket = socket.into();
let result = tokio::time::timeout(timeout, KamajiClient::connect(socket.clone()))
.await
.map_err(|_| {
anyhow!(
"timed out after {:?} connecting to kamaji at {}",
timeout,
socket.display()
)
})?;
result.with_context(|| format!("connecting to kamaji at {}", socket.display()))
}
#[cfg(test)]
mod tests {
use super::*;
use kamaji_proto::{AckKind, WorkloadState};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::UnixListener;
async fn one_shot_server(
expect_after_hello: impl FnOnce(YubabaToKamaji) -> KamajiToYubaba + Send + 'static,
) -> (tempfile::TempDir, PathBuf) {
let tmp = tempfile::tempdir().unwrap();
let sock = tmp.path().join("kamaji.sock");
let listener = UnixListener::bind(&sock).unwrap();
let welcome = KamajiToYubaba::Welcome {
version: ProtocolVersion::CURRENT,
kamaji_version: "test-0.0.1".into(),
};
let answer = std::sync::Arc::new(std::sync::Mutex::new(Some(expect_after_hello)));
tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let mut buf = Vec::with_capacity(4096);
let mut tmpbuf = [0u8; 4096];
let hello = loop {
match decode_frame::<YubabaToKamaji>(&buf) {
Ok((m, n)) => {
buf.drain(..n);
break m;
}
Err(CodecError::Truncated { .. }) => {
let n = stream.read(&mut tmpbuf).await.unwrap();
if n == 0 {
return;
}
buf.extend_from_slice(&tmpbuf[..n]);
}
Err(e) => panic!("decode hello: {e}"),
}
};
assert!(matches!(hello, YubabaToKamaji::Hello { .. }));
let bytes = encode_frame(&welcome).unwrap();
stream.write_all(&bytes).await.unwrap();
let req = loop {
match decode_frame::<YubabaToKamaji>(&buf) {
Ok((m, n)) => {
buf.drain(..n);
break m;
}
Err(CodecError::Truncated { .. }) => {
let n = stream.read(&mut tmpbuf).await.unwrap();
if n == 0 {
return;
}
buf.extend_from_slice(&tmpbuf[..n]);
}
Err(e) => panic!("decode req: {e}"),
}
};
let answer = answer.lock().unwrap().take().expect("answer used once");
let reply = answer(req);
let bytes = encode_frame(&reply).unwrap();
stream.write_all(&bytes).await.unwrap();
});
(tmp, sock)
}
#[tokio::test]
async fn connect_handshakes_and_captures_info() {
let (_tmp, sock) = one_shot_server(|_req| KamajiToYubaba::WorkloadList {
request_id: RequestId(1),
entries: vec![],
})
.await;
let client = KamajiClient::connect(sock).await.expect("connect");
assert_eq!(client.info().kamaji_version, "test-0.0.1");
let entries = client.list().await.expect("list");
assert!(entries.is_empty());
}
#[tokio::test]
async fn list_returns_workload_entries() {
let (_tmp, sock) = one_shot_server(|req| {
let rid = match req {
YubabaToKamaji::List { request_id } => request_id,
other => panic!("expected List, got {other:?}"),
};
KamajiToYubaba::WorkloadList {
request_id: rid,
entries: vec![WorkloadEntry {
mesh_ident: None,
id: WorkloadId::new("foo"),
state: WorkloadState::Running,
pid: Some(42),
}],
}
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
let entries = client.list().await.unwrap();
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].id, WorkloadId::new("foo"));
assert_eq!(entries[0].state, WorkloadState::Running);
assert_eq!(entries[0].pid, Some(42));
}
#[tokio::test]
async fn stop_handles_ack() {
let (_tmp, sock) = one_shot_server(|req| {
let rid = match req {
YubabaToKamaji::Stop { request_id, .. } => request_id,
other => panic!("expected Stop, got {other:?}"),
};
KamajiToYubaba::Ack {
request_id: rid,
kind: AckKind::Stop,
}
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
client.stop(&WorkloadId::new("foo")).await.expect("stop");
}
#[tokio::test]
async fn drain_surfaces_ack_payload() {
let (_tmp, sock) = one_shot_server(|req| {
let (rid, id) = match req {
YubabaToKamaji::Drain {
request_id,
id,
budget,
} => {
assert_eq!(budget.flush_ms, 100);
assert_eq!(budget.checkpoint_ms, 200);
(request_id, id)
}
other => panic!("expected Drain, got {other:?}"),
};
KamajiToYubaba::DrainAck {
request_id: rid,
id,
accepted: true,
reason: Some("flushed in 50ms".into()),
}
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
let (accepted, reason) = client
.drain(
&WorkloadId::new("foo"),
DrainBudget {
flush_ms: 100,
checkpoint_ms: 200,
},
)
.await
.unwrap();
assert!(accepted);
assert_eq!(reason.as_deref(), Some("flushed in 50ms"));
}
#[tokio::test]
async fn remote_error_is_surfaced_as_client_remote() {
let (_tmp, sock) = one_shot_server(|req| {
let rid = match req {
YubabaToKamaji::Stop { request_id, .. } => request_id,
other => panic!("expected Stop, got {other:?}"),
};
KamajiToYubaba::Error {
request_id: Some(rid),
code: ErrorCode::UnknownWorkload,
message: "no such id".into(),
}
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
let err = client.stop(&WorkloadId::new("foo")).await.unwrap_err();
match err {
ClientError::Remote { code, message } => {
assert_eq!(code, ErrorCode::UnknownWorkload);
assert_eq!(message, "no such id");
}
other => panic!("expected Remote, got {other:?}"),
}
}
#[tokio::test]
async fn connect_with_timeout_fails_on_missing_socket() {
let tmp = tempfile::tempdir().unwrap();
let sock = tmp.path().join("does-not-exist.sock");
let err = connect_with_timeout(sock, Duration::from_millis(250))
.await
.unwrap_err();
let msg = format!("{err:#}");
assert!(
msg.contains("connecting to kamaji"),
"error should mention the socket path: {msg}"
);
}
use crate::Kamaji as ConstableTrait;
#[tokio::test]
async fn constable_backend_returns_native() {
let (_tmp, sock) = one_shot_server(|_| KamajiToYubaba::WorkloadList {
request_id: RequestId(1),
entries: vec![],
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
assert_eq!(client.backend(), crate::Backend::Native);
}
#[tokio::test]
async fn constable_list_workloads_maps_proto_entries() {
let (_tmp, sock) = one_shot_server(|req| {
let rid = match req {
YubabaToKamaji::List { request_id } => request_id,
other => panic!("expected List, got {other:?}"),
};
KamajiToYubaba::WorkloadList {
request_id: rid,
entries: vec![
WorkloadEntry {
mesh_ident: None,
id: WorkloadId::new("svc-a"),
state: WorkloadState::Running,
pid: Some(1000),
},
WorkloadEntry {
mesh_ident: None,
id: WorkloadId::new("svc-b"),
state: WorkloadState::Exited,
pid: None,
},
],
}
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
let states = client.list_workloads().await.unwrap();
assert_eq!(states.len(), 2);
assert_eq!(states[0].ident, crate::MeshIdent("svc-a".into()));
assert!(matches!(states[0].status, crate::WorkloadStatus::Running));
assert_eq!(states[1].ident, crate::MeshIdent("svc-b".into()));
assert!(matches!(states[1].status, crate::WorkloadStatus::Stopped));
}
#[tokio::test]
async fn constable_get_workload_finds_by_ident() {
let (_tmp, sock) = one_shot_server(|req| {
let rid = match req {
YubabaToKamaji::List { request_id } => request_id,
other => panic!("expected List, got {other:?}"),
};
KamajiToYubaba::WorkloadList {
request_id: rid,
entries: vec![WorkloadEntry {
mesh_ident: None,
id: WorkloadId::new("target"),
state: WorkloadState::Running,
pid: Some(42),
}],
}
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
let state = client
.get_workload(&crate::MeshIdent("target".into()))
.await
.unwrap();
assert!(state.is_some());
assert_eq!(state.unwrap().ident, crate::MeshIdent("target".into()));
}
#[tokio::test]
async fn constable_get_workload_returns_none_for_unknown_ident() {
let (_tmp, sock) = one_shot_server(|req| {
let rid = match req {
YubabaToKamaji::List { request_id } => request_id,
other => panic!("expected List, got {other:?}"),
};
KamajiToYubaba::WorkloadList {
request_id: rid,
entries: vec![],
}
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
let state = client
.get_workload(&crate::MeshIdent("nobody".into()))
.await
.unwrap();
assert!(state.is_none());
}
#[tokio::test]
async fn constable_teardown_workload_sends_stop() {
let (_tmp, sock) = one_shot_server(|req| {
let rid = match req {
YubabaToKamaji::Stop { request_id, id } => {
assert_eq!(id, WorkloadId::new("svc-to-stop"));
request_id
}
other => panic!("expected Stop, got {other:?}"),
};
KamajiToYubaba::Ack {
request_id: rid,
kind: AckKind::Stop,
}
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
client
.teardown_workload(&crate::MeshIdent("svc-to-stop".into()))
.await
.unwrap();
}
#[tokio::test]
async fn constable_health_returns_ok_with_kamaji_version() {
let (_tmp, sock) = one_shot_server(|_| KamajiToYubaba::WorkloadList {
request_id: RequestId(1),
entries: vec![],
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
let health = client.health().await.unwrap();
assert!(health.ok);
assert_eq!(health.version.as_deref(), Some("test-0.0.1"));
}
#[tokio::test]
async fn constable_stream_logs_returns_not_supported_err() {
let (_tmp, sock) = one_shot_server(|_| KamajiToYubaba::WorkloadList {
request_id: RequestId(1),
entries: vec![],
})
.await;
let client = KamajiClient::connect(sock).await.unwrap();
let result = client
.stream_logs(&crate::MeshIdent("any".into()), crate::LogOpts::default())
.await;
let err = match result {
Err(e) => e,
Ok(_) => panic!("expected stream_logs to return Err"),
};
assert!(err.to_string().contains("not yet in the sibling"));
}
}