diff --git a/Cargo.toml b/Cargo.toml index 9d9cc15..de65d8b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ governor = "0.10.1" http-body-util = "0.1.3" log = "0.4.28" poem = { version = "3.1.12", features = ["compression"] } -reqwest = { version = "0.12.23", features = ["stream"] } +reqwest = { version = "0.12.23", features = ["stream", "json"] } reqwest-middleware = "0.4.2" reqwest-retry = "0.7.0" serde = "1.0.219" diff --git a/src/client.rs b/src/client.rs index e9cfa3a..874bad8 100644 --- a/src/client.rs +++ b/src/client.rs @@ -3,15 +3,14 @@ use reqwest_middleware::{ClientBuilder, ClientWithMiddleware}; use reqwest_retry::{RetryTransientMiddleware, policies::ExponentialBackoff}; use std::sync::LazyLock; +pub const UA: &str = concat!( + "allegedly, v", + env!("CARGO_PKG_VERSION"), + " (from @microcosm.blue; contact @bad-example.com)" +); + pub static CLIENT: LazyLock = LazyLock::new(|| { - let inner = Client::builder() - .user_agent(concat!( - "allegedly, v", - env!("CARGO_PKG_VERSION"), - " (from @microcosm.blue; contact @bad-example.com)" - )) - .build() - .unwrap(); + let inner = Client::builder().user_agent(UA).build().unwrap(); let policy = ExponentialBackoff::builder().build_with_max_retries(12); diff --git a/src/lib.rs b/src/lib.rs index 9b3dfdb..ef0ebb4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,7 +9,7 @@ mod ratelimit; mod weekly; pub use backfill::backfill; -pub use client::CLIENT; +pub use client::{CLIENT, UA}; pub use mirror::serve; pub use plc_pg::{Db, backfill_to_pg, pages_to_pg}; pub use poll::{PageBoundaryState, get_page, poll_upstream}; diff --git a/src/mirror.rs b/src/mirror.rs index f8921ab..372c404 100644 --- a/src/mirror.rs +++ b/src/mirror.rs @@ -1,4 +1,4 @@ -use crate::{GovernorMiddleware, logo}; +use crate::{GovernorMiddleware, UA, logo}; use futures::TryStreamExt; use governor::Quota; use poem::{ @@ -6,14 +6,15 @@ use poem::{ http::StatusCode, listener::TcpListener, middleware::{AddData, CatchPanic, Compression, Cors, Tracing}, - web::Data, + web::{Data, Json}, }; use reqwest::{Client, Url}; use std::{net::SocketAddr, time::Duration}; #[derive(Debug, Clone)] struct State { - client: Client, + upstream_client: Client, + wrapped_client: Client, plc: Url, upstream: Url, } @@ -60,12 +61,77 @@ Failed to reach the wrapped reference PLC server. Sorry. ) } +async fn plc_status(url: &Url, client: &Client) -> (bool, serde_json::Value) { + use serde_json::json; + + let mut url = url.clone(); + url.set_path("/_health"); + + let Ok(response) = client.get(url).send().await else { + return (false, json!({"error": "cannot reach plc server"})); + }; + + let status = response.status(); + + let Ok(text) = response.text().await else { + return (false, json!({"error": "failed to read response body"})); + }; + + let body = match serde_json::from_str(&text) { + Ok(json) => json, + Err(_) => serde_json::Value::String(text.to_string()), + }; + + if status.is_success() { + (true, body) + } else { + ( + false, + json!({ + "error": "non-ok status", + "status": status.as_str(), + "status_code": status.as_u16(), + "response": body, + }), + ) + } +} + +#[handler] +async fn health( + Data(State { + plc, + wrapped_client, + upstream, + upstream_client, + }): Data<&State>, +) -> impl IntoResponse { + let mut overall_status = StatusCode::OK; + let (ok, wrapped_status) = plc_status(plc, wrapped_client).await; + if !ok { + overall_status = StatusCode::BAD_GATEWAY; + } + let (ok, upstream_status) = plc_status(upstream, upstream_client).await; + if !ok { + overall_status = StatusCode::BAD_GATEWAY; + } + ( + overall_status, + Json(serde_json::json!({ + "server": "allegedly (mirror)", + "version": env!("CARGO_PKG_VERSION"), + "wrapped_plc": wrapped_status, + "upstream_plc": upstream_status, + })), + ) +} + #[handler] async fn proxy(req: &Request, Data(state): Data<&State>) -> Result { let mut target = state.plc.clone(); target.set_path(req.uri().path()); let upstream_res = state - .client + .upstream_client .get(target) .headers(req.headers().clone()) .send() @@ -111,19 +177,26 @@ You may wish to try upstream: {upstream} } pub async fn serve(upstream: &Url, plc: Url, bind: SocketAddr) -> std::io::Result<()> { - let wrapped_req_client = Client::builder() + let wrapped_client = Client::builder() .timeout(Duration::from_secs(3)) .build() .unwrap(); + let upstream_client = Client::builder() + .user_agent(UA) + .timeout(Duration::from_secs(6)) + .build() + .unwrap(); let state = State { - client: wrapped_req_client, + wrapped_client, + upstream_client, plc, upstream: upstream.clone(), }; let app = Route::new() .at("/", get(hello)) + .at("/_health", get(health)) .at("/:any", get(proxy).post(nope)) .with(AddData::new(state)) .with(Cors::new().allow_credentials(false))