use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
use tokio::sync::Semaphore;
use tropel_metrics::collector::{MetricsCollector, MetricsSnapshot, SeriesSnapshot};
use tropel_scheduler::VUScheduler;
use tropel_sdk::{Result, TropelError};
const MAX_BODY_SIZE: usize = 64 * 1024;
const MAX_HEADER_LINE_LEN: usize = 8 * 1024;
const CONN_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_CONNS: usize = 8;
const ACCEPT_BACKOFF: Duration = Duration::from_millis(100);
pub struct ControlApiState {
pub scheduler: Arc<VUScheduler>,
pub metrics: Arc<MetricsCollector>,
pub setup_data: Arc<std::sync::Mutex<Option<Vec<u8>>>>,
pub scenario_name: String,
}
pub async fn serve_control_api(port: u16, state: ControlApiState) -> Result<()> {
let addr = format!("127.0.0.1:{}", port);
let listener = match TcpListener::bind(&addr).await {
Ok(l) => l,
Err(e) => {
tracing::error!("control API: failed to bind {}: {}", addr, e);
return Err(tropel_sdk::TropelError::Config(format!(
"control API: failed to bind {}: {}",
addr, e
)));
}
};
tracing::info!("Control API listening on http://{addr}");
let conn_permits = Arc::new(Semaphore::new(MAX_CONNS));
let state = Arc::new(state);
loop {
let (stream, _peer) = match listener.accept().await {
Ok(x) => x,
Err(e) => {
tracing::debug!("control API: accept error: {}; backing off", e);
tokio::time::sleep(ACCEPT_BACKOFF).await;
continue;
}
};
let permit = match conn_permits.clone().try_acquire_owned() {
Ok(p) => p,
Err(_) => {
let mut out = stream;
let _ = out
.write_all(
b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
)
.await;
continue;
}
};
let state = state.clone();
tokio::spawn(async move {
let _permit = permit;
if let Err(e) = handle_conn(stream, &state).await {
tracing::debug!("control API: connection error: {}", e);
}
});
}
}
async fn handle_conn<S>(stream: S, state: &Arc<ControlApiState>) -> Result<()>
where
S: AsyncRead + AsyncWrite + Unpin,
{
match tokio::time::timeout(CONN_TIMEOUT, serve_request(stream, state)).await {
Ok(r) => r,
Err(_elapsed) => {
tracing::debug!("control API: connection timed out");
Ok(())
}
}
}
async fn serve_request<S>(stream: S, state: &Arc<ControlApiState>) -> Result<()>
where
S: AsyncRead + AsyncWrite + Unpin,
{
let mut reader = BufReader::new(stream);
let mut request_line = String::new();
let mut limited = (&mut reader).take(MAX_HEADER_LINE_LEN as u64 + 1);
if limited.read_line(&mut request_line).await? == 0 {
return Ok(());
}
if request_line.len() > MAX_HEADER_LINE_LEN {
return Err(TropelError::Http(format!(
"control API: request line too long ({} > {})",
request_line.len(),
MAX_HEADER_LINE_LEN
)));
}
let request_line = request_line.trim_end().to_string();
let mut parts = request_line.split_whitespace();
let method = parts.next().unwrap_or("").to_string();
let path = parts.next().unwrap_or("").to_string();
let mut content_length: usize = 0;
loop {
let mut line = String::new();
let mut limited = (&mut reader).take(MAX_HEADER_LINE_LEN as u64 + 1);
if limited.read_line(&mut line).await? == 0 {
break;
}
if line.len() > MAX_HEADER_LINE_LEN {
return Err(TropelError::Http(format!(
"control API: header line too long ({} > {})",
line.len(),
MAX_HEADER_LINE_LEN
)));
}
let line = line.trim_end();
if line.is_empty() {
break;
}
if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") {
content_length = v.trim().parse().unwrap_or(0);
}
}
if content_length > MAX_BODY_SIZE {
let mut out = reader.into_inner();
out.write_all(
b"HTTP/1.1 413 Payload Too Large\r\nContent-Type: application/json\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
)
.await?;
out.flush().await?;
return Ok(());
}
let mut body = Vec::new();
if content_length > 0 {
body.resize(content_length, 0);
reader.read_exact(&mut body).await?;
}
let (status, response_body) = route(&method, &path, &body, state).await;
let response = format!(
"HTTP/1.1 {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
status,
response_body.len(),
response_body
);
let mut out = reader.into_inner();
out.write_all(response.as_bytes()).await?;
out.flush().await?;
Ok(())
}
async fn route(
method: &str,
path: &str,
body: &[u8],
state: &Arc<ControlApiState>,
) -> (String, String) {
let sched = &state.scheduler;
match (method, path) {
("GET", "/v1/status") => ("200 OK".to_string(), status_json(sched)),
("PATCH", "/v1/status") => match parse_status_body(body) {
Some(patch) => {
if patch.vus.is_some() || patch.max.is_some() {
let vus = patch.vus.unwrap_or_else(|| sched.control_target());
let max = patch.max.unwrap_or_else(|| sched.control_max());
sched.set_control_target(vus, max);
tracing::info!("Control API: set VUs target={} max={}", vus, max);
}
if let Some(paused) = patch.paused {
sched.set_paused(paused);
tracing::info!("Control API: paused={}", paused);
}
if patch.stopped == Some(true) {
sched.request_stop();
tracing::info!("Control API: stop requested");
}
("200 OK".to_string(), status_json(sched))
}
None => (
"400 Bad Request".to_string(),
"{\"error\":\"expected {\\\"vus\\\":N,\\\"max\\\":M,\\\"paused\\\":bool}\"}"
.to_string(),
),
},
("PATCH", "/v1/stop") => {
match parse_status_body(body) {
Some(patch) if patch.stopped == Some(true) || patch.stopped.is_none() => {
if patch.stopped == Some(true) {
sched.request_stop();
tracing::info!("Control API: stop requested (PATCH /v1/stop)");
}
("200 OK".to_string(), status_json(sched))
}
_ => (
"400 Bad Request".to_string(),
"{\"error\":\"expected {\\\"data\\\":{\\\"attributes\\\":{\\\"stopped\\\":true}}}\"}"
.to_string(),
),
}
}
("POST", "/v1/stop") => {
sched.request_stop();
("200 OK".to_string(), status_json(sched))
}
("GET", "/v1/metrics") => {
let snap = state.metrics.snapshot().await;
("200 OK".to_string(), metrics_json(&snap))
}
("GET", "/v1/groups") => {
let snap = state.metrics.snapshot().await;
("200 OK".to_string(), groups_json(&snap, &state.scenario_name))
}
("GET", "/v1/setup") => {
let data = state.setup_data.lock().unwrap().clone();
("200 OK".to_string(), setup_json(data.as_deref()))
}
("PUT", "/v1/setup") => {
let parsed: Option<serde_json::Value> = if body.is_empty() {
None
} else {
match serde_json::from_slice(body) {
Ok(v) => Some(v),
Err(e) => {
return (
"400 Bad Request".to_string(),
format!(r#"{{"error":"invalid setup data: {e}"}}"#),
)
}
}
};
*state.setup_data.lock().unwrap() = parsed.map(|v| v.to_string().into_bytes());
let data = state.setup_data.lock().unwrap().clone();
("200 OK".to_string(), setup_json(data.as_deref()))
}
("POST", "/v1/setup") => {
(
"405 Method Not Allowed".to_string(),
r#"{"error":"setup() runs once at engine start; POST re-run not supported"}"#
.to_string(),
)
}
("POST", "/v1/teardown") => {
(
"405 Method Not Allowed".to_string(),
r#"{"error":"teardown() runs once at engine stop; POST re-run not supported"}"#
.to_string(),
)
}
_ => (
"404 Not Found".to_string(),
r#"{"error":"not found"}"#.to_string(),
),
}
}
fn status_json(sched: &Arc<VUScheduler>) -> String {
let vus = sched.control_target();
let max = sched.control_max();
let paused = sched.is_paused();
let stopped = sched.is_stop_requested();
let running = !stopped;
let tainted = sched.is_tainted();
format!(
r#"{{"data":{{"type":"status","id":"default","attributes":{{"vus":{},"vus-max":{},"max":{},"paused":{},"running":{},"stopped":{},"tainted":{}}}}}}}"#,
vus, max, max, paused, running, stopped, tainted
)
}
fn metrics_json(snap: &MetricsSnapshot) -> String {
let mut entries: Vec<serde_json::Value> = Vec::new();
let mut by_metric: std::collections::BTreeMap<&str, &SeriesSnapshot> = Default::default();
for s in &snap.series {
by_metric.insert(&s.metric, s);
}
for (name, s) in by_metric {
let sample = serde_json::json!({
"value": s.last,
});
entries.push(serde_json::json!({
"type": "metrics",
"id": name,
"attributes": {
"type": metric_type_name(s.metric_type),
"contains": "default",
"tainted": false,
"sample": sample,
},
}));
}
serde_json::json!({ "data": entries }).to_string()
}
fn metric_type_name(t: tropel_metrics::collector::MetricType) -> &'static str {
use tropel_metrics::collector::MetricType;
match t {
MetricType::Counter => "counter",
MetricType::Gauge => "gauge",
MetricType::Rate => "rate",
MetricType::Trend => "trend",
}
}
fn groups_json(_snap: &MetricsSnapshot, scenario_name: &str) -> String {
serde_json::json!({
"data": [{
"type": "groups",
"id": "0",
"attributes": {
"path": "",
"name": scenario_name,
"checks": [],
},
"relationships": {
"groups": { "data": [] },
"parent": { "data": null },
},
}]
})
.to_string()
}
fn setup_json(data: Option<&[u8]>) -> String {
let value = match data {
Some(bytes) => serde_json::from_slice(bytes).unwrap_or(serde_json::Value::Null),
None => serde_json::Value::Null,
};
serde_json::json!({ "data": { "data": value } }).to_string()
}
#[derive(Debug, Clone, Copy, PartialEq)]
struct StatusPatch {
vus: Option<u32>,
max: Option<u32>,
paused: Option<bool>,
stopped: Option<bool>,
}
fn parse_status_body(body: &[u8]) -> Option<StatusPatch> {
let text = std::str::from_utf8(body).ok()?;
let json: serde_json::Value = serde_json::from_str(text).ok()?;
let attrs = json
.get("data")
.and_then(|d| d.get("attributes"))
.or(Some(&json))?;
let vus = attrs.get("vus").and_then(|v| v.as_u64()).map(|v| v as u32);
let max = attrs
.get("vus-max")
.and_then(|v| v.as_u64())
.map(|v| v as u32)
.or_else(|| attrs.get("max").and_then(|v| v.as_u64()).map(|v| v as u32));
let paused = attrs.get("paused").and_then(|v| v.as_bool());
let stopped = attrs.get("stopped").and_then(|v| v.as_bool());
if vus.is_none() && max.is_none() && paused.is_none() && stopped.is_none() {
return None;
}
Some(StatusPatch {
vus,
max,
paused,
stopped,
})
}
#[cfg(test)]
mod tests {
use super::*;
use tropel_metrics::collector::MetricsCollector;
fn test_state(sched: Arc<VUScheduler>) -> Arc<ControlApiState> {
Arc::new(ControlApiState {
scheduler: sched,
metrics: Arc::new(MetricsCollector::new()),
setup_data: Arc::new(std::sync::Mutex::new(None)),
scenario_name: "s".to_string(),
})
}
#[test]
fn parses_flat_body() {
assert_eq!(
parse_status_body(br#"{"vus":5,"max":20}"#),
Some(StatusPatch {
vus: Some(5),
max: Some(20),
paused: None,
stopped: None,
})
);
}
#[test]
fn parses_k6_envelope() {
assert_eq!(
parse_status_body(br#"{"data":{"attributes":{"vus":3,"max":9}}}"#),
Some(StatusPatch {
vus: Some(3),
max: Some(9),
paused: None,
stopped: None,
})
);
}
#[test]
fn parses_paused_only() {
assert_eq!(
parse_status_body(br#"{"paused":true}"#),
Some(StatusPatch {
vus: None,
max: None,
paused: Some(true),
stopped: None,
})
);
}
#[test]
fn partial_patch_with_only_vus_is_valid() {
assert_eq!(
parse_status_body(br#"{"vus":5}"#),
Some(StatusPatch {
vus: Some(5),
max: None,
paused: None,
stopped: None,
})
);
}
#[test]
fn rejects_garbage_and_unknown_only() {
assert_eq!(parse_status_body(br#"{"foo":1}"#), None); assert_eq!(parse_status_body(b"garbage"), None);
assert_eq!(parse_status_body(b"{}"), None);
}
#[test]
fn parses_stopped_only() {
assert_eq!(
parse_status_body(br#"{"stopped":true}"#),
Some(StatusPatch {
vus: None,
max: None,
paused: None,
stopped: Some(true),
})
);
}
#[test]
fn route_stopped_true_requests_stop() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let state = test_state(sched.clone());
assert!(!sched.is_stop_requested());
let rt = tokio::runtime::Runtime::new().unwrap();
let (status, body) =
rt.block_on(route("PATCH", "/v1/status", br#"{"stopped":true}"#, &state));
assert_eq!(status, "200 OK");
assert!(
sched.is_stop_requested(),
"PATCH stopped:true must request a stop"
);
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(v["data"]["attributes"]["stopped"], true);
}
#[test]
fn route_stopped_false_is_noop() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let state = test_state(sched.clone());
let rt = tokio::runtime::Runtime::new().unwrap();
let (status, _) = rt.block_on(route(
"PATCH",
"/v1/status",
br#"{"stopped":false}"#,
&state,
));
assert_eq!(status, "200 OK");
assert!(!sched.is_stop_requested());
}
#[tokio::test]
async fn oversized_body_rejected_before_alloc() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let (mut client, server) = tokio::io::duplex(4096);
let state = test_state(sched);
let server_task = tokio::spawn(async move { serve_request(server, &state).await });
client
.write_all(b"PATCH /v1/status HTTP/1.1\r\nContent-Length: 68719476736\r\n\r\n")
.await
.unwrap();
let mut resp = Vec::new();
client.read_to_end(&mut resp).await.unwrap();
let text = String::from_utf8_lossy(&resp);
assert!(
text.contains("413 Payload Too Large"),
"expected 413, got: {}",
text
);
server_task.await.unwrap().unwrap();
}
#[tokio::test(start_paused = true)]
async fn stalled_client_is_timed_out() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let (mut client, server) = tokio::io::duplex(4096);
client.write_all(b"PATCH /v1/status HTT").await.unwrap();
let state = test_state(sched);
let server_task = tokio::spawn(async move { handle_conn(server, &state).await });
tokio::time::advance(Duration::from_secs(11)).await;
let result = server_task.await.unwrap();
assert!(
result.is_ok(),
"stalled client must be cut off by the read timeout"
);
}
#[test]
fn status_json_is_k6_shape() {
let sched = VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
);
let sched = Arc::new(sched);
sched.set_control_target(4, 10);
let body = status_json(&sched);
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
let attrs = &v["data"]["attributes"];
assert_eq!(attrs["type"], serde_json::Value::Null); assert_eq!(v["data"]["type"], "status");
assert_eq!(v["data"]["id"], "default");
assert_eq!(attrs["vus"], 4);
assert_eq!(attrs["vus-max"], 10);
assert_eq!(attrs["max"], 10);
assert_eq!(attrs["paused"], false);
assert_eq!(attrs["running"], true);
assert_eq!(attrs["stopped"], false);
assert_eq!(attrs["tainted"], false);
sched.set_tainted();
let body = status_json(&sched);
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(v["data"]["attributes"]["tainted"], true);
}
#[test]
fn patch_stop_accepts_k6_envelope() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let state = test_state(sched.clone());
assert!(!sched.is_stop_requested());
let rt = tokio::runtime::Runtime::new().unwrap();
let (status, body) = rt.block_on(route(
"PATCH",
"/v1/stop",
br#"{"data":{"attributes":{"stopped":true}}}"#,
&state,
));
assert_eq!(status, "200 OK");
assert!(
sched.is_stop_requested(),
"PATCH /v1/stop with the k6 envelope must stop the run"
);
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(v["data"]["attributes"]["stopped"], true);
}
#[tokio::test]
async fn get_metrics_returns_k6_envelope() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let state = test_state(sched);
let (status, body) = route("GET", "/v1/metrics", b"", &state).await;
assert_eq!(status, "200 OK");
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(
v["data"],
serde_json::Value::Array(vec![]),
"no metrics yet"
);
}
#[tokio::test]
async fn get_groups_returns_k6_envelope() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let state = test_state(sched);
let (status, body) = route("GET", "/v1/groups", b"", &state).await;
assert_eq!(status, "200 OK");
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
let group = &v["data"][0];
assert_eq!(group["type"], "groups");
assert_eq!(group["id"], "0");
assert_eq!(group["attributes"]["path"], "");
assert_eq!(group["attributes"]["name"], "s");
assert_eq!(
group["relationships"]["groups"]["data"],
serde_json::Value::Array(vec![])
);
assert_eq!(
group["relationships"]["parent"]["data"],
serde_json::Value::Null
);
}
#[tokio::test]
async fn setup_get_put_roundtrip() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let state = test_state(sched);
let (status, body) = route("GET", "/v1/setup", b"", &state).await;
assert_eq!(status, "200 OK");
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(v["data"]["data"], serde_json::Value::Null);
let (status, body) = route("PUT", "/v1/setup", br#"{"token":"abc"}"#, &state).await;
assert_eq!(status, "200 OK");
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(v["data"]["data"]["token"], "abc");
let (status, body) = route("GET", "/v1/setup", b"", &state).await;
assert_eq!(status, "200 OK");
let v: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(v["data"]["data"]["token"], "abc");
}
#[tokio::test]
async fn setup_teardown_post_reexecution_rejected() {
let sched = Arc::new(VUScheduler::new(
&tropel_core::config::ExecutionConfig::ExternallyControlled {
vus: 2,
max_vus: 10,
duration: None,
graceful_stop: None,
think_time: Default::default(),
},
));
let state = test_state(sched);
let (status, body) = route("POST", "/v1/setup", b"", &state).await;
assert_eq!(status, "405 Method Not Allowed");
assert!(body.contains("setup() runs once"));
let (status, body) = route("POST", "/v1/teardown", b"", &state).await;
assert_eq!(status, "405 Method Not Allowed");
assert!(body.contains("teardown() runs once"));
}
}