Refactor to move caching, metrics, etc from v1 to common dir

Author: TymanWasTaken

src/api/v1/caching.rs src/api/common/caching.rssrc/api/v1/caching.rs renamed to src/api/common/caching.rs

@@ -2,20 +2,35 @@ use actix_web::{
 	HttpResponse,
 	body::{BoxBody, EitherBody, MessageBody},
 	dev::{ServiceRequest, ServiceResponse},
-	http::header::{ETAG, HeaderValue, IF_NONE_MATCH},
+	http::header::{ETAG, HeaderMap, HeaderValue, IF_NONE_MATCH},
 	middleware::Next,
-	web
+	web::{self, Bytes}
 };
 use prometheus_client::encoding::EncodeLabelSet;
 use sha2::{Digest as _, Sha256};
 
-use super::{ApiData, CacheKey, CacheValue};
+use crate::api::common::data::ApiData;
 
 #[derive(Debug, Hash, PartialEq, Eq, Clone, EncodeLabelSet)]
 pub struct CacheLabels {
 	endpoint: String
 }
 
+#[derive(Hash, PartialEq, Eq, Clone)]
+pub struct CacheKey {
+	pub path: String,
+	pub query: String
+}
+
+pub type ETagType = [u8; 32];
+
+#[derive(Clone)]
+pub struct CacheValue {
+	pub response: Bytes,
+	pub headers: HeaderMap,
+	pub etag: ETagType
+}
+
 pub async fn middleware(
 	service_request: ServiceRequest,
 	next: Next<impl MessageBody>
@@ -66,6 +81,7 @@ pub async fn middleware(
 		// Record a cache hit to the metrics
 		app_data
 			.metrics
+			.global
 			.cache_hits
 			.get_or_create(&metric_labels)
 			.inc();
@@ -104,6 +120,7 @@ pub async fn middleware(
 		// Record a cache miss to the metrics
 		app_data
 			.metrics
+			.global
 			.cache_misses
 			.get_or_create(&metric_labels)
 			.inc();

src/api/common/data.rs

@@ -0,0 +1,64 @@
+use std::{collections::HashSet, sync::Arc, time::Duration};
+
+use moka::future::Cache;
+
+use super::caching::{CacheKey, CacheValue};
+use crate::{
+	AppCommand,
+	api::common::{caching::ETagType, metrics::AppMetrics}
+};
+
+pub struct ApiData {
+	/// The maven URL prefix to expose publicly, for example https://repo.polyfrost.org/
+	pub public_maven_url: String,
+	/// The maven URL prefix to resolve artifacts internally, for example https://172.19.0.3:8080/
+	pub internal_maven_url: String,
+	/// A reqwest client to use to fetch maven data
+	pub client: Arc<reqwest::Client>,
+	/// The allowlist of paths that should be cached
+	pub cache_allowlist: HashSet<&'static str>,
+	/// The internal cache used to cache artifact responses.
+	pub cache: Cache<CacheKey, CacheValue>,
+	/// All the metrics objects used for encoding and recording metrics
+	pub metrics: AppMetrics
+}
+
+impl ApiData {
+	pub fn new(args: &AppCommand) -> Self {
+		Self {
+			public_maven_url: args.public_maven_url.to_string(),
+			internal_maven_url: args
+				.internal_maven_url
+				.as_ref()
+				.map_or(args.public_maven_url.to_string(), |u| u.to_string()),
+			client: reqwest::ClientBuilder::new()
+				.user_agent(concat!(
+					env!("CARGO_PKG_NAME"),
+					"/",
+					env!("CARGO_PKG_VERSION"),
+					" (",
+					env!("CARGO_PKG_REPOSITORY"),
+					")"
+				))
+				.build()
+				.unwrap()
+				.into(),
+			cache_allowlist: HashSet::from([
+				"/v1/artifacts/oneconfig",
+				"/v1/artifacts/{artifact:stage1|relaunch}"
+			]),
+			cache: Cache::builder()
+				.time_to_live(Duration::from_mins(2))
+				.weigher(|k: &CacheKey, v: &CacheValue| {
+					(k.path.len()
+						+ k.query.len() + const { std::mem::size_of::<ETagType>() }
+						+ v.response.len() + std::mem::size_of_val(&v.headers))
+					.try_into()
+					.unwrap_or(u32::MAX)
+				})
+				.max_capacity(/* 10 MiB */ const { 10 * 1024 * 1024 })
+				.build(),
+			metrics: AppMetrics::new()
+		}
+	}
+}

src/api/common/metrics.rs

@@ -0,0 +1,142 @@
+use actix_web::{
+	HttpResponse,
+	Responder,
+	body::MessageBody,
+	dev::{ServiceRequest, ServiceResponse},
+	get,
+	middleware::Next,
+	web::{self, ServiceConfig}
+};
+use documented::DocumentedFields;
+use prometheus_client::{
+	encoding::{EncodeLabelSet, text::encode},
+	metrics::{counter::Counter, family::Family},
+	registry::Registry
+};
+
+use crate::api::{
+	common::{caching::CacheLabels, data::ApiData},
+	v1::metrics::ApiV1Metrics
+};
+
+// TODO: Improve this macro so you can use only one macro call and it will
+// register all of them
+
+/// A macro that automatically initializes and registers a metric using inferred
+/// types and doc comments from the ApiMetrics struct
+#[macro_export]
+macro_rules! make_api_metric {
+	($registry:expr, $metrics_struct:ident, $name:ident) => {
+		make_api_metric!($registry, $metrics_struct, $name, Family)
+	};
+	($registry:expr, $metrics_struct:ident, $name:ident, $type:ident) => {
+		let $name = $type::default();
+		let name_str = stringify!($name);
+		$registry.register(
+			name_str,
+			$metrics_struct::get_field_docs(name_str)
+				.expect(&format!("No doc comment for '{}' field", name_str)),
+			$name.clone()
+		);
+	};
+}
+
+pub trait MetricsGroup {
+	fn init_metrics(registry: &mut Registry) -> Self;
+}
+
+pub struct AppMetrics {
+	/// The root metrics registry used for storing and retrieving metrics
+	pub registry: Registry,
+	/// All of the global (application-wide) metrics
+	pub global: GlobalMetrics,
+	/// All of the API v1 metrics
+	pub v1: ApiV1Metrics
+}
+
+impl AppMetrics {
+	pub fn new() -> AppMetrics {
+		let mut registry = <Registry>::default();
+
+		AppMetrics {
+			global: GlobalMetrics::init_metrics(&mut registry),
+			v1: ApiV1Metrics::init_metrics(&mut registry),
+			registry
+		}
+	}
+}
+
+#[derive(DocumentedFields)]
+pub struct GlobalMetrics {
+	/// The amount of generic API requests by endpoint and status code
+	api_requests: Family<ApiRequestLabels, Counter>,
+	/// The amount of cache hits by endpoint
+	pub cache_hits: Family<CacheLabels, Counter>,
+	/// The amount of cache misses by endpoint
+	pub cache_misses: Family<CacheLabels, Counter>
+}
+
+impl MetricsGroup for GlobalMetrics {
+	fn init_metrics(registry: &mut Registry) -> Self {
+		make_api_metric!(registry, Self, api_requests);
+		make_api_metric!(registry, Self, cache_hits);
+		make_api_metric!(registry, Self, cache_misses);
+
+		Self {
+			api_requests,
+			cache_hits,
+			cache_misses
+		}
+	}
+}
+
+pub fn configure(config: &mut ServiceConfig) { config.service(metrics_endpoint); }
+
+/// The endpoint to allow scraping metrics
+#[get("/metrics")]
+async fn metrics_endpoint(state: web::Data<ApiData>) -> impl Responder {
+	let mut body = String::new();
+	if let Err(e) = encode(&mut body, &state.metrics.registry) {
+		return HttpResponse::InternalServerError()
+			.content_type("text/plain")
+			.body(format!("Error encoding metrics: {e}"));
+	}
+
+	HttpResponse::Ok()
+		.content_type("application/openmetrics-text; version=1.0.0; charset=utf-8")
+		.body(body)
+}
+
+#[derive(Debug, Hash, PartialEq, Eq, Clone, EncodeLabelSet)]
+struct ApiRequestLabels {
+	path: String,
+	status_code: u16
+}
+
+pub async fn middleware(
+	mut service_request: ServiceRequest,
+	next: Next<impl MessageBody>
+) -> Result<ServiceResponse<impl MessageBody>, actix_web::Error> {
+	let data = service_request.extract::<web::Data<ApiData>>().await?;
+
+	let path = service_request.uri().path().to_string();
+	let response = next.call(service_request).await;
+
+	let labels = match &response {
+		Ok(r) => ApiRequestLabels {
+			path,
+			status_code: r.status().as_u16()
+		},
+		Err(e) => ApiRequestLabels {
+			path,
+			status_code: e.as_response_error().status_code().as_u16()
+		}
+	};
+	data.metrics
+		.global
+		.api_requests
+		.get_or_create(&labels)
+		.inc();
+
+	response
+}

src/api/common/mod.rs

@@ -0,0 +1,3 @@
+pub mod caching;
+pub mod data;
+pub mod metrics;

src/api/legacy/mod.rs

@@ -0,0 +1 @@
+

src/api/mod.rs

@@ -1 +1,4 @@
+pub mod common;
+
+pub mod legacy;
 pub mod v1;

src/api/v1/endpoints/artifacts.rs

@@ -11,9 +11,15 @@ use serde::{Deserialize, Serialize};
 use tokio::task::JoinSet;
 
 use crate::{
-	api::v1::{
-		ApiData,
-		responses::{ArtifactResponse, Checksum, ChecksumType, ErrorResponse, consts::*}
+	api::{
+		common::data::ApiData,
+		v1::responses::{
+			ArtifactResponse,
+			Checksum,
+			ChecksumType,
+			ErrorResponse,
+			consts::*
+		}
 	},
 	maven::{self, MavenError},
 	types::gradle_module_metadata::{
@@ -27,14 +33,12 @@ use crate::{
 
 const ONECONFIG_GROUP: &str = "org.polyfrost.oneconfig";
 
-pub fn configure() -> impl FnOnce(&mut ServiceConfig) {
-	|config| {
-		config.service(
-			web::scope("/artifacts")
-				.service(oneconfig)
-				.service(platform_agnostic_artifacts)
-		);
-	}
+pub fn configure(config: &mut ServiceConfig) {
+	config.service(
+		web::scope("/artifacts")
+			.service(oneconfig)
+			.service(platform_agnostic_artifacts)
+	);
 }
 
 #[derive(Serialize, Deserialize, Debug, Hash, PartialEq, Eq, Clone)]
@@ -190,10 +194,6 @@ async fn oneconfig(
 	};
 
 	let mut join_set: JoinSet<Result<ArtifactResponse, anyhow::Error>> = JoinSet::new();
-	let internal_maven_url = state
-		.internal_maven_url
-		.clone()
-		.unwrap_or(state.public_maven_url.clone());
 
 	for variant in dependency.variants {
 		let Variant::OneConfigModulesApiElements { dependencies } = variant else {
@@ -206,7 +206,7 @@ async fn oneconfig(
 			}
 
 			let internal_dep_url =
-				maven::get_dep_url(&internal_maven_url, repository, &dep);
+				maven::get_dep_url(&state.internal_maven_url, repository, &dep);
 			let dep_url = maven::get_dep_url(&state.public_maven_url, repository, &dep);
 
 			let client = state.client.clone();
@@ -295,14 +295,7 @@ async fn platform_agnostic_artifacts(
 
 	let checksum = match maven::fetch_checksum(
 		&state.client,
-		&maven::get_dep_url(
-			&state
-				.internal_maven_url
-				.clone()
-				.unwrap_or(state.public_maven_url.clone()),
-			repository,
-			&dep
-		)
+		&maven::get_dep_url(&state.internal_maven_url, repository, &dep)
 	)
 	.await
 	{