init nix-racer substituter proxy
This commit is contained in:
@@ -0,0 +1,212 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
extract::{Path, State},
|
||||
http::{StatusCode, header},
|
||||
response::{IntoResponse, Redirect},
|
||||
routing::get,
|
||||
};
|
||||
use reqwest::Client;
|
||||
use tokio::task::JoinSet;
|
||||
use tracing_subscriber::{EnvFilter, layer::SubscriberExt as _, util::SubscriberInitExt as _};
|
||||
|
||||
mod config;
|
||||
use config::*;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct AppState {
|
||||
client: Client,
|
||||
config: Arc<AppConfig>,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
tracing_subscriber::registry()
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.with(
|
||||
EnvFilter::builder()
|
||||
.with_env_var("NIX_RACER_LOG")
|
||||
.with_default_directive("nix_racer=debug".parse().unwrap())
|
||||
.from_env()
|
||||
.unwrap(),
|
||||
)
|
||||
.init();
|
||||
|
||||
let config: Arc<AppConfig> = match std::fs::read_to_string("/etc/nix/nix-racer.toml") {
|
||||
Ok(file) => toml::from_str::<AppConfig>(&file).unwrap_or_else(|x| {
|
||||
tracing::warn!("Failed to parse config file ({x}), fallback to default configuration");
|
||||
AppConfig::default()
|
||||
}),
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
|
||||
tracing::info!("Config file not found, fallback to default configuration");
|
||||
AppConfig::default()
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!("Failed to read config file ({err}), fallback to default configuration");
|
||||
AppConfig::default()
|
||||
}
|
||||
}
|
||||
.into();
|
||||
|
||||
if config.substituters.is_empty() {
|
||||
tracing::error!("No substituters found");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
let listen = config.listen;
|
||||
|
||||
let client = Client::builder()
|
||||
.timeout(Duration::from_secs(5))
|
||||
.pool_idle_timeout(Duration::from_secs(90))
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let state = AppState { client, config };
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(listen).await.unwrap();
|
||||
|
||||
let app = Router::new()
|
||||
.route("/nix-cache-info", get(nix_cache_info_handler))
|
||||
.route("/{*path}", get(proxy_handler))
|
||||
.with_state(state);
|
||||
|
||||
tracing::info!("Smart Nix Proxy listening on http://{listen}");
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
}
|
||||
|
||||
async fn nix_cache_info_handler() -> impl IntoResponse {
|
||||
let info = "StoreDir: /nix/store\nWantMassQuery: 1\nPriority: 30\n";
|
||||
([(header::CONTENT_TYPE, "text/x-nix-cache-info")], info)
|
||||
}
|
||||
|
||||
async fn proxy_handler(
|
||||
Path(path): Path<String>,
|
||||
State(state): State<AppState>,
|
||||
) -> impl IntoResponse {
|
||||
if !path.ends_with(".narinfo") {
|
||||
let redirect_url = format!("{}/{}", state.config.substituters[0].url, path);
|
||||
return Redirect::temporary(&redirect_url).into_response();
|
||||
}
|
||||
|
||||
let start_time = Instant::now();
|
||||
let mut set = JoinSet::new();
|
||||
|
||||
let mut penalties = HashMap::new();
|
||||
|
||||
for upstream in state.config.substituters.iter() {
|
||||
let url = format!("{}/{}", upstream.url, path);
|
||||
let client = state.client.clone();
|
||||
let base_url = upstream.url;
|
||||
|
||||
let handle = set.spawn(async move {
|
||||
let max_retries = 2;
|
||||
let mut delay = Duration::from_millis(50);
|
||||
|
||||
for attempt in 0..=max_retries {
|
||||
match client.get(&url).send().await {
|
||||
Ok(resp) if resp.status().is_success() => {
|
||||
let bytes = resp.bytes().await.map_err(|_| "Body error")?;
|
||||
|
||||
let text = String::from_utf8_lossy(&bytes);
|
||||
let mut new_text = String::with_capacity(text.len() + 64);
|
||||
for line in text.lines() {
|
||||
if let Some(rel_url) = line.strip_prefix("URL: ") {
|
||||
if rel_url.starts_with("http://") || rel_url.starts_with("https://")
|
||||
{
|
||||
new_text.push_str(line);
|
||||
} else {
|
||||
new_text.push_str(&format!("URL: {}/{}", base_url, rel_url));
|
||||
}
|
||||
} else {
|
||||
new_text.push_str(line);
|
||||
}
|
||||
new_text.push('\n');
|
||||
}
|
||||
return Ok((new_text.into_bytes(), base_url));
|
||||
}
|
||||
Ok(resp) if resp.status() == StatusCode::NOT_FOUND => return Err("404"),
|
||||
Err(_) => {
|
||||
if attempt == max_retries {
|
||||
return Err("Max retries");
|
||||
}
|
||||
tokio::time::sleep(delay).await;
|
||||
delay *= 2;
|
||||
}
|
||||
_ => return Err("Other HTTP Error"),
|
||||
}
|
||||
}
|
||||
Err("Unreachable")
|
||||
});
|
||||
|
||||
penalties.insert(handle.id(), upstream.penalty);
|
||||
}
|
||||
|
||||
struct ProxyResult {
|
||||
score: Duration,
|
||||
body: Vec<u8>,
|
||||
url: &'static str,
|
||||
}
|
||||
let mut best_result: Option<ProxyResult> = None;
|
||||
|
||||
loop {
|
||||
let min_active_penalty = penalties.values().min().copied();
|
||||
|
||||
if let Some(ProxyResult { score, .. }) = best_result {
|
||||
if let Some(min_p) = min_active_penalty {
|
||||
if start_time.elapsed() + Duration::from_millis(u64::from(min_p)) >= score {
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
} else if min_active_penalty.is_none() {
|
||||
break;
|
||||
}
|
||||
|
||||
let timeout_dur = if let Some(ProxyResult { score, .. }) = best_result {
|
||||
score.saturating_sub(start_time.elapsed())
|
||||
} else {
|
||||
Duration::from_secs(86400)
|
||||
};
|
||||
|
||||
if timeout_dur.is_zero() && best_result.is_some() {
|
||||
break;
|
||||
}
|
||||
|
||||
let res = tokio::time::timeout(timeout_dur, set.join_next_with_id()).await;
|
||||
|
||||
match res {
|
||||
Ok(Some(Ok((task_id, task_res)))) => {
|
||||
let penalty = penalties.remove(&task_id).unwrap();
|
||||
if let Ok((body, url)) = task_res {
|
||||
let score = start_time.elapsed() + Duration::from_millis(u64::from(penalty));
|
||||
if best_result.as_ref().is_none_or(|best| score < best.score) {
|
||||
best_result = Some(ProxyResult { score, body, url })
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Some(Err(join_err))) => {
|
||||
penalties.remove(&join_err.id());
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ProxyResult { score, body, url }) = best_result {
|
||||
tracing::debug!("Winner for {}: {} (Score: {:?})", path, url, score,);
|
||||
(
|
||||
StatusCode::OK,
|
||||
[(header::CONTENT_TYPE, "text/x-nix-narinfo")],
|
||||
body,
|
||||
)
|
||||
.into_response()
|
||||
} else {
|
||||
(StatusCode::NOT_FOUND, "Not Found").into_response()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user