pub mod actions;
pub mod exposition;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;
use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper::header::{HeaderValue, ACCEPT, ALLOW, CONTENT_TYPE, RETRY_AFTER, USER_AGENT};
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Method, Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use serde_json::{json, Value};
use tokio::sync::mpsc;
use tracing::{debug, error, info};
use crate::llm::action_helper::call_llm;
use crate::llm::ollama_client::OllamaClient;
use crate::protocol::Event;
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use crate::state::ServerId;
use crate::{console_debug, console_error, console_info};
use exposition::MetricFamilies;
pub const MAX_REQUEST_BODY_BYTES: usize = 64 * 1024;
const FIRST_BYTE_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const IDLE_BETWEEN_REQUESTS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
const CONNECTION_CAP_REFUSAL: &[u8] = b"HTTP/1.1 503 Service Unavailable\r\n\
Content-Length: 0\r\nRetry-After: 5\r\nConnection: close\r\n\r\n";
const INDEX_HTML: &str = "<html>\n<head><title>NetGet Prometheus Exporter</title></head>\n\
<body>\n<h1>NetGet Prometheus Exporter</h1>\n<p><a href=\"/metrics\">Metrics</a></p>\n</body>\n\
</html>\n";
pub struct PrometheusServer;
impl PrometheusServer {
pub async fn spawn_with_llm_actions(
listen_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: ServerId,
) -> anyhow::Result<SocketAddr> {
let listener =
crate::server::socket_helpers::create_reusable_tcp_listener(listen_addr).await?;
let local_addr = listener.local_addr()?;
console_info!(
status_tx,
"Prometheus exporter listening on http://{}/metrics",
local_addr
);
let protocol = Arc::new(actions::PrometheusProtocol::new());
let task_registrar = app_state.clone();
let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
let accept_handle = tokio::spawn(async move {
loop {
match crate::server::accept_bounded::accept_bounded(
&listener,
&limiter,
CONNECTION_CAP_REFUSAL,
"Prometheus",
Some(&status_tx),
)
.await
{
Ok((stream, remote_addr, permit)) => {
let connection_id =
ConnectionId::new(app_state.get_next_unified_id().await);
let local_addr_conn = stream.local_addr().unwrap_or(local_addr);
info!(
"Prometheus connection {} from {}",
connection_id, remote_addr
);
use crate::state::server::{
ConnectionState as ServerConnectionState, ConnectionStatus,
ProtocolConnectionInfo,
};
let now = crate::utils::clock::Instant::now();
app_state
.add_connection_to_server(
server_id,
ServerConnectionState {
id: connection_id,
remote_addr,
local_addr: local_addr_conn,
bytes_sent: 0,
bytes_received: 0,
packets_sent: 0,
packets_received: 0,
last_activity: now,
status: ConnectionStatus::Active,
status_changed_at: now,
protocol_info: ProtocolConnectionInfo::empty(),
},
)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
let ctx = RequestContext {
llm_client: llm_client.clone(),
app_state: app_state.clone(),
status_tx: status_tx.clone(),
protocol: protocol.clone(),
server_id,
connection_id,
};
let app_state_for_close = app_state.clone();
let status_for_close = status_tx.clone();
let task_owner = app_state.clone();
task_owner
.spawn_server_task(server_id, async move {
let _permit = permit;
match tokio::time::timeout(
FIRST_BYTE_READ_TIMEOUT,
stream.peek(&mut [0u8; 1]),
)
.await
{
Ok(Ok(n)) if n > 0 => {
serve_connection(TokioIo::new(stream), ctx).await;
}
Ok(_) => {}
Err(_) => {
debug!(
"Prometheus peer {} sent nothing for {}s; closing \
before the request line",
remote_addr,
FIRST_BYTE_READ_TIMEOUT.as_secs()
);
}
}
app_state_for_close
.close_connection_on_server(server_id, connection_id)
.await;
let _ = status_for_close.send(format!(
"[INFO] Prometheus connection {connection_id} closed"
));
let _ = status_for_close.send("__UPDATE_UI__".to_string());
})
.await;
}
Err(e) => {
console_error!(status_tx, "Failed to accept Prometheus connection: {}", e);
break;
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(local_addr)
}
}
#[derive(Clone)]
struct RequestContext {
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
protocol: Arc<actions::PrometheusProtocol>,
server_id: ServerId,
connection_id: ConnectionId,
}
async fn serve_connection<T>(io: TokioIo<T>, ctx: RequestContext)
where
T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
let activity = Arc::new(crate::server::accept_bounded::ConnectionActivity::new());
let activity_for_service = Arc::clone(&activity);
let service = service_fn(move |req: Request<Incoming>| {
let ctx = ctx.clone();
let activity = Arc::clone(&activity_for_service);
async move {
let _busy = activity.busy();
Ok::<_, Infallible>(handle_request(req, ctx).await)
}
});
let conn = http1::Builder::new().serve_connection(io, service);
tokio::pin!(conn);
tokio::select! {
result = &mut conn => {
if let Err(err) = result {
debug!("Prometheus connection ended: {:?}", err);
}
}
_ = crate::server::accept_bounded::watch_idle(
Arc::clone(&activity),
IDLE_BETWEEN_REQUESTS_TIMEOUT,
) => {
debug!(
"Prometheus connection idle for {}s; closing",
IDLE_BETWEEN_REQUESTS_TIMEOUT.as_secs()
);
}
}
}
fn text_response(
status: StatusCode,
content_type: &'static str,
body: String,
) -> Response<Full<Bytes>> {
let mut response = Response::new(Full::new(Bytes::from(body)));
*response.status_mut() = status;
response
.headers_mut()
.insert(CONTENT_TYPE, HeaderValue::from_static(content_type));
response
}
fn error_response(status: StatusCode, message: &'static str) -> Response<Full<Bytes>> {
text_response(status, "text/plain; charset=utf-8", format!("{message}\n"))
}
async fn handle_request(req: Request<Incoming>, ctx: RequestContext) -> Response<Full<Bytes>> {
let (parts, body) = req.into_parts();
let method = parts.method.clone();
let path = parts.uri.path().to_string();
let header = |name| {
parts
.headers
.get(name)
.and_then(|v: &HeaderValue| v.to_str().ok())
.map(str::to_string)
};
let accept = header(ACCEPT);
let user_agent = header(USER_AGENT).unwrap_or_default();
let scrape_timeout = header(hyper::header::HeaderName::from_static(
"x-prometheus-scrape-timeout-seconds",
));
debug!("Prometheus {} {}", method, path);
let _ = ctx
.status_tx
.send(format!("[DEBUG] Prometheus {method} {path}"));
let body_len = match http_body_util::Limited::new(body, MAX_REQUEST_BODY_BYTES)
.collect()
.await
{
Ok(collected) => collected.to_bytes().len(),
Err(e) => {
error!(
"Prometheus {} {}: decision=fail_closed_body_rejected (limit {} bytes): {}",
method, path, MAX_REQUEST_BODY_BYTES, e
);
return error_response(
StatusCode::PAYLOAD_TOO_LARGE,
"netget: request body is larger than this exporter accepts",
);
}
};
ctx.app_state
.update_connection_stats(
ctx.server_id,
ctx.connection_id,
Some(body_len as u64),
None,
Some(1),
None,
)
.await;
match (path.as_str(), &method) {
("/", &Method::GET) | ("/", &Method::HEAD) => {
return text_response(
StatusCode::OK,
"text/html; charset=utf-8",
INDEX_HTML.to_string(),
);
}
("/metrics", &Method::GET) => {}
("/metrics", &Method::HEAD) => {
let format = exposition::negotiate(accept.as_deref());
return text_response(StatusCode::OK, format.content_type(), String::new());
}
("/metrics", _) | ("/", _) => {
let mut response = error_response(
StatusCode::METHOD_NOT_ALLOWED,
"netget: only GET and HEAD are served here",
);
response
.headers_mut()
.insert(ALLOW, HeaderValue::from_static("GET, HEAD"));
return response;
}
_ => {
return error_response(StatusCode::NOT_FOUND, "404 page not found");
}
}
let format = exposition::negotiate(accept.as_deref());
let mut data = json!({
"path": path,
"accept": accept.clone().unwrap_or_default(),
"user_agent": user_agent,
"format": format.as_str(),
});
if let Some(timeout) = scrape_timeout
.as_deref()
.and_then(|t| t.parse::<f64>().ok())
{
data["scrape_timeout_seconds"] = json!(timeout);
}
console_debug!(
ctx.status_tx,
"Calling LLM for Prometheus scrape of {}",
path
);
let event = Event::new(&actions::PROMETHEUS_SCRAPE_EVENT, data);
let llm_result = call_llm(
&ctx.llm_client,
&ctx.app_state,
ctx.server_id,
Some(ctx.connection_id),
&event,
ctx.protocol.as_ref(),
)
.await;
match llm_result {
Ok(execution) => {
let failure_summary = execution.failure_summary();
for result in execution.protocol_results {
let crate::llm::ActionResult::Custom { name, data } = result else {
continue;
};
match name.as_str() {
"send_metrics" => {
let Some(families) = data
.get("metrics")
.and_then(|m| MetricFamilies::parse(m).ok())
else {
error!(
"Prometheus {}: decision=fail_closed_invalid_exposition \
(executor output did not re-validate)",
path
);
return error_response(
StatusCode::INTERNAL_SERVER_ERROR,
"netget: the handler's metrics were refused as invalid",
);
};
let body = families.render(format);
info!(
"Prometheus {}: decision=model_answer format={} families={} samples={}",
path,
format.as_str(),
families.family_count(),
families.sample_count()
);
ctx.app_state
.update_connection_stats(
ctx.server_id,
ctx.connection_id,
None,
Some(body.len() as u64),
None,
Some(1),
)
.await;
return text_response(StatusCode::OK, format.content_type(), body);
}
"send_scrape_error" => {
let status = data
.get("status")
.and_then(Value::as_u64)
.and_then(|s| u16::try_from(s).ok())
.and_then(|s| StatusCode::from_u16(s).ok())
.filter(|s| s.is_client_error() || s.is_server_error())
.unwrap_or(StatusCode::SERVICE_UNAVAILABLE);
let message = data
.get("message")
.and_then(Value::as_str)
.unwrap_or("scrape refused")
.to_string();
info!(
"Prometheus {}: decision=model_reject status={}",
path,
status.as_u16()
);
return text_response(
status,
"text/plain; charset=utf-8",
format!("{}\n", crate::utils::sanitize::line_field(&message)),
);
}
other => debug!("Prometheus: ignoring non-exporter action '{}'", other),
}
}
match failure_summary {
None => error!(
"Prometheus {}: decision=fail_closed_no_action (no send_metrics)",
path
),
Some(summary) => error!(
"Prometheus {}: decision=fail_closed_invalid_exposition ({})",
path, summary
),
}
console_error!(
ctx.status_tx,
"Prometheus: the handler returned no usable metrics for {}",
path
);
error_response(
StatusCode::INTERNAL_SERVER_ERROR,
"netget: the handler returned no usable metrics",
)
}
Err(e) => {
let failure = crate::utils::WireFailure::classify(&e);
error!(
"Prometheus {}: decision=fail_closed_llm_error category={:?} error={}",
path, failure, e
);
console_error!(ctx.status_tx, "Prometheus LLM call failed: {}", e);
match failure {
crate::utils::WireFailure::Overloaded => {
let mut response =
error_response(StatusCode::SERVICE_UNAVAILABLE, failure.prefixed_text());
response
.headers_mut()
.insert(RETRY_AFTER, HeaderValue::from_static("5"));
response
}
crate::utils::WireFailure::Unavailable => {
error_response(StatusCode::INTERNAL_SERVER_ERROR, failure.prefixed_text())
}
}
}
}
}