diff --git a/src/handler/http/request/logs/ingest.rs b/src/handler/http/request/logs/ingest.rs index e6e8f88899..ca19957244 100644 --- a/src/handler/http/request/logs/ingest.rs +++ b/src/handler/http/request/logs/ingest.rs @@ -15,7 +15,12 @@ use std::io::Error; -use actix_web::{HttpRequest, HttpResponse, http, http::header, post, web}; +use actix_http::header::HeaderMap; +use actix_web::{ + HttpRequest, HttpResponse, http, + http::header::{HeaderName, HeaderValue}, + post, web, +}; use config::meta::otlp::OtlpRequestType; #[cfg(feature = "cloud")] use config::meta::stream::StreamType; @@ -42,6 +47,26 @@ use crate::{ service::logs::{self, otlp::handle_request}, }; +const HEADER_O2_PROCESS_TIME: HeaderName = HeaderName::from_static("o2_process_time"); + +/// Returns the current time, accounting for `ZO_ACTIX_SLOW_LOG_THRESHOLD` env +fn get_process_time() -> i64 { + if config::get_config().limit.http_slow_log_threshold > 0 { + config::utils::time::now_micros() + } else { + 0 + } +} + +fn insert_process_time_header(time: i64, headers: &mut HeaderMap) { + if time > 0 { + headers.insert( + HEADER_O2_PROCESS_TIME, + HeaderValue::from_str(&time.to_string()).unwrap(), + ); + } +} + /// _bulk ES compatible ingestion API #[utoipa::path( context_path = "/api", @@ -72,30 +97,22 @@ pub async fn bulk( org_id: web::Path, body: web::Bytes, Headers(user_email): Headers, - _in_req: HttpRequest, ) -> Result { let org_id = org_id.into_inner(); let user_email = &user_email.user_id; #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Logs, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Logs, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } // log start processing time - let process_time = if config::get_config().limit.http_slow_log_threshold > 0 { - config::utils::time::now_micros() - } else { - 0 - }; + let process_time = get_process_time(); let mut resp = match logs::bulk::ingest( **thread_id, @@ -123,12 +140,7 @@ pub async fn bulk( } }; - if process_time > 0 { - resp.headers_mut().insert( - header::HeaderName::from_static("o2_process_time"), - header::HeaderValue::from_str(&process_time.to_string()).unwrap(), - ); - } + insert_process_time_header(process_time, resp.headers_mut()); Ok(resp) } @@ -164,30 +176,22 @@ pub async fn multi( path: web::Path<(String, String)>, body: web::Bytes, Headers(user_email): Headers, - _in_req: HttpRequest, ) -> Result { let (org_id, stream_name) = path.into_inner(); let user_email = &user_email.user_id; #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Logs, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Logs, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } // log start processing time - let process_time = if config::get_config().limit.http_slow_log_threshold > 0 { - config::utils::time::now_micros() - } else { - 0 - }; + let process_time = get_process_time(); let mut resp = match logs::ingest::ingest( **thread_id, @@ -221,12 +225,7 @@ pub async fn multi( } }; - if process_time > 0 { - resp.headers_mut().insert( - header::HeaderName::from_static("o2_process_time"), - header::HeaderValue::from_str(&process_time.to_string()).unwrap(), - ); - } + insert_process_time_header(process_time, resp.headers_mut()); Ok(resp) } @@ -259,30 +258,22 @@ pub async fn json( path: web::Path<(String, String)>, body: web::Bytes, Headers(user_email): Headers, - _in_req: HttpRequest, ) -> Result { let (org_id, stream_name) = path.into_inner(); let user_email = &user_email.user_id; #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Logs, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Logs, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } // log start processing time - let process_time = if config::get_config().limit.http_slow_log_threshold > 0 { - config::utils::time::now_micros() - } else { - 0 - }; + let process_time = get_process_time(); let mut resp = match logs::ingest::ingest( **thread_id, @@ -310,18 +301,12 @@ pub async fn json( e, )) } else { - HttpResponse::BadRequest() - .json(MetaHttpResponse::error(http::StatusCode::BAD_REQUEST, e)) + MetaHttpResponse::bad_request(e) } } }; - if process_time > 0 { - resp.headers_mut().insert( - header::HeaderName::from_static("o2_process_time"), - header::HeaderValue::from_str(&process_time.to_string()).unwrap(), - ); - } + insert_process_time_header(process_time, resp.headers_mut()); Ok(resp) } @@ -357,23 +342,19 @@ pub async fn handle_kinesis_request( path: web::Path<(String, String)>, post_data: web::Json, Headers(user_email): Headers, - _in_req: HttpRequest, ) -> Result { let (org_id, stream_name) = path.into_inner(); let user_email = &user_email.user_id; let request_id = post_data.request_id.clone(); #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Logs, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Logs, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } let request_time = post_data @@ -425,22 +406,18 @@ pub async fn handle_gcp_request( path: web::Path<(String, String)>, post_data: web::Json, Headers(user_email): Headers, - _in_req: HttpRequest, ) -> Result { let (org_id, stream_name) = path.into_inner(); let user_email = &user_email.user_id; #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Logs, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Logs, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } Ok( @@ -510,16 +487,13 @@ pub async fn otlp_logs_write( .map(|header| header.to_str().unwrap()); #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Logs, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Logs, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } let (request, request_type) = match content_type { @@ -546,10 +520,7 @@ pub async fn otlp_logs_write( } } _ => { - return Ok(HttpResponse::BadRequest().json(MetaHttpResponse::error( - http::StatusCode::BAD_REQUEST, - "Bad Request", - ))); + return Ok(MetaHttpResponse::bad_request("Bad Request")); } }; @@ -616,30 +587,22 @@ pub async fn hec( org_id: web::Path, body: web::Bytes, Headers(user_email): Headers, - _in_req: HttpRequest, ) -> Result { let org_id = org_id.into_inner(); let user_email = &user_email.user_id; #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Logs, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Logs, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } // log start processing time - let process_time = if config::get_config().limit.http_slow_log_threshold > 0 { - config::utils::time::now_micros() - } else { - 0 - }; + let process_time = get_process_time(); let mut resp = match logs::hec::ingest(**thread_id, &org_id, body, user_email).await { Ok(v) => { @@ -663,12 +626,7 @@ pub async fn hec( } }; - if process_time > 0 { - resp.headers_mut().insert( - header::HeaderName::from_static("o2_process_time"), - header::HeaderValue::from_str(&process_time.to_string()).unwrap(), - ); - } + insert_process_time_header(process_time, resp.headers_mut()); Ok(resp) } diff --git a/src/handler/http/request/metrics/ingest.rs b/src/handler/http/request/metrics/ingest.rs index 6027ff9fc5..d50a607375 100644 --- a/src/handler/http/request/metrics/ingest.rs +++ b/src/handler/http/request/metrics/ingest.rs @@ -74,16 +74,13 @@ pub async fn json( let user = IngestUser::from_user_email(&user_email.user_id); #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Metrics, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Metrics, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } let mut resp = match metrics::json::ingest(&org_id, body, user).await { @@ -142,16 +139,13 @@ pub async fn otlp_metrics_write( let user = IngestUser::from_user_email(&user_email.user_id); #[cfg(feature = "cloud")] - match check_ingestion_allowed(&org_id, StreamType::Metrics, None).await { - Ok(_) => {} - Err(e) => { - return Ok( - HttpResponse::TooManyRequests().json(MetaHttpResponse::error( - http::StatusCode::TOO_MANY_REQUESTS, - e, - )), - ); - } + if let Err(e) = check_ingestion_allowed(&org_id, StreamType::Metrics, None).await { + return Ok( + HttpResponse::TooManyRequests().json(MetaHttpResponse::error( + http::StatusCode::TOO_MANY_REQUESTS, + e, + )), + ); } let content_type = req.headers().get("Content-Type").unwrap().to_str().unwrap();