Fix request id span

Author: rdettai

distribution/lambda/cdk/cli.py

@@ -72,7 +72,7 @@ def from_lambda_response(lambda_resp: dict) -> "LambdaResult":
             log_tail=base64.b64decode(lambda_resp["LogResult"]).decode(),
             payload=payload,
             raw_size_bytes=len(payload),
-            status_code=200,
+            status_code=0,
         )
 
     @staticmethod
@@ -119,8 +119,6 @@ def _format_lambda_output(
     if lambda_result.function_error != "":
         print("\n## FUNCTION ERROR:")
         print(lambda_result.function_error)
-    print("\n## LOG TAIL:")
-    print(lambda_result.log_tail)
     print("\n## RAW RESPONSE SIZE (BYTES):")
     if len(lambda_result.payload) == 0:
         ratio = "empty payload"
@@ -213,12 +211,14 @@ def _invoke_searcher(
         LogType="Tail",
         Payload=json.dumps(
             {
-                "path": f"/api/v1/{index_id}/search",
                 "resource": f"/api/v1/{index_id}/search",
+                "path": f"/api/v1/{index_id}/search",
                 "httpMethod": "POST",
                 "headers": {
                     "Content-Type": "application/json",
-                    "Content-Length": f"{len(payload)}",
+                },
+                "requestContext": {
+                    "httpMethod": "POST",
                 },
                 "body": payload,
                 "isBase64Encoded": False,
@@ -268,7 +268,6 @@ def get_logs(
                     last_event_id = event["eventId"]
                     yield event["message"]
                     if event["message"].startswith("REPORT"):
-                        print(event["message"])
                         lower_time_bound = int(event["timestamp"])
                         last_event_id = "REPORT"
                         break
@@ -296,6 +295,7 @@ def download_logs_to_file(request_id: str, function_name: str, invoke_start: flo
                 int(invoke_start * 1000),
             ):
                 f.write(log)
+            print(f"Logs written to lambda.{request_id}.log")
     except Exception as e:
         print(f"Failed to download logs: {e}")
 

quickwit/Cargo.lock

@@ -39,7 +39,6 @@ rand = { workspace = true }
 reqwest = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
-thiserror = { workspace = true }
 time = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }

quickwit/quickwit-lambda/Cargo.toml

@@ -18,12 +18,12 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 use quickwit_lambda::logger;
-use quickwit_lambda::searcher::{searcher_api, warp_lambda};
+use quickwit_lambda::searcher::{setup_searcher_api, warp_lambda};
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     logger::setup_lambda_tracer(tracing::Level::INFO)?;
-    let routes = searcher_api().await?;
+    let routes = setup_searcher_api().await?;
     let warp_service = warp::service(routes);
     warp_lambda::run(warp_service)
         .await

quickwit/quickwit-lambda/src/bin/searcher.rs

@@ -32,16 +32,15 @@ use quickwit_search::{
 use quickwit_serve::lambda_search_api::*;
 use quickwit_storage::StorageResolver;
 use quickwit_telemetry::payload::{QuickwitFeature, QuickwitTelemetryInfo, TelemetryEvent};
-use tracing::{error, info, info_span};
+use tracing::{error, info};
 use warp::filters::path::FullPath;
 use warp::reject::Rejection;
 use warp::Filter;
 
-use super::LAMBDA_REQUEST_ID_HEADER;
 use crate::searcher::environment::CONFIGURATION_TEMPLATE;
 use crate::utils::load_node_config;
 
-async fn get_search_service(
+async fn create_local_search_service(
     searcher_config: SearcherConfig,
     metastore: MetastoreServiceClient,
     storage_resolver: StorageResolver,
@@ -88,18 +87,31 @@ fn es_compat_api(
         .or(es_compat_cat_indices_handler(metastore.clone()))
 }
 
-pub async fn searcher_api(
+fn v1_searcher_api(
+    search_service: Arc<dyn SearchService>,
+    metastore: MetastoreServiceClient,
+) -> impl Filter<Extract = (impl warp::Reply,), Error = Rejection> + Clone {
+    warp::path!("api" / "v1" / ..)
+        .and(native_api(search_service.clone()).or(es_compat_api(search_service, metastore)))
+        .with(warp::filters::compression::gzip())
+        .recover(|rejection| {
+            error!(?rejection, "request rejected");
+            recover_fn(rejection)
+        })
+}
+
+pub async fn setup_searcher_api(
 ) -> anyhow::Result<impl Filter<Extract = (impl warp::Reply,), Error = Rejection> + Clone> {
     let (node_config, storage_resolver, metastore) =
         load_node_config(CONFIGURATION_TEMPLATE).await?;
 
-    let services: HashSet<String> =
-        HashSet::from_iter([QuickwitService::Searcher.as_str().to_string()]);
-    let telemetry_info =
-        QuickwitTelemetryInfo::new(services, HashSet::from_iter([QuickwitFeature::AwsLambda]));
+    let telemetry_info = QuickwitTelemetryInfo::new(
+        HashSet::from_iter([QuickwitService::Searcher.as_str().to_string()]),
+        HashSet::from_iter([QuickwitFeature::AwsLambda]),
+    );
     let _telemetry_handle_opt = quickwit_telemetry::start_telemetry_loop(telemetry_info);
 
-    let search_service = get_search_service(
+    let search_service = create_local_search_service(
         node_config.searcher_config,
         metastore.clone(),
         storage_resolver,
@@ -122,26 +134,11 @@ pub async fn searcher_api(
     let after_hook = warp::log::custom(|info| {
         info!(status = info.status().as_str(), "request completed");
     });
-    let api = warp::any().and(before_hook).and(
-        warp::path!("api" / "v1" / ..)
-            .and(native_api(search_service.clone()).or(es_compat_api(search_service, metastore)))
-            .with(warp::filters::compression::gzip())
-            .recover(|rejection| {
-                error!(?rejection, "request rejected");
-                recover_fn(rejection)
-            })
-            .with(after_hook)
-            .with(warp::trace(|info| {
-                info_span!(
-                    "request",
-                    request_id = info
-                        .request_headers()
-                        .get(LAMBDA_REQUEST_ID_HEADER)
-                        .and_then(|v| v.to_str().ok())
-                        .unwrap_or(&""),
-                )
-            })),
-    );
+
+    let api = warp::any()
+        .and(before_hook)
+        .and(v1_searcher_api(search_service, metastore))
+        .with(after_hook);
 
     Ok(api)
 }

quickwit/quickwit-lambda/src/searcher/api.rs

@@ -21,6 +21,4 @@ mod api;
 mod environment;
 pub mod warp_lambda;
 
-pub use api::searcher_api;
-
-const LAMBDA_REQUEST_ID_HEADER: &str = "lambda-request-id";
+pub use api::setup_searcher_api;

quickwit/quickwit-lambda/src/searcher/mod.rs

@@ -31,31 +31,19 @@ use anyhow::anyhow;
 use http::header::Entry;
 use lambda_http::http::response::Parts;
 use lambda_http::http::HeaderValue;
-use lambda_http::request::RequestContext;
 use lambda_http::{
     lambda_runtime, Adapter, Body as LambdaBody, Error as LambdaError, Request, RequestExt,
     Response, Service,
 };
 use mime_guess::{mime, Mime};
 use once_cell::sync::Lazy;
+use tracing::{info_span, Instrument};
 use warp::hyper::Body as WarpBody;
 pub use {lambda_http, warp};
 
-use super::LAMBDA_REQUEST_ID_HEADER;
-
 pub type WarpRequest = warp::http::Request<warp::hyper::Body>;
 pub type WarpResponse = warp::http::Response<warp::hyper::Body>;
 
-#[derive(thiserror::Error, Debug)]
-pub enum WarpAdapterError {
-    #[error("This may never occur, it's infallible!!")]
-    Infallible(#[from] std::convert::Infallible),
-    #[error("Warp error: `{0:#?}`")]
-    HyperError(#[from] warp::hyper::Error),
-    #[error("Unexpected error: `{0:#?}`")]
-    Unexpected(#[from] LambdaError),
-}
-
 pub async fn run<'a, S>(service: S) -> Result<(), LambdaError>
 where
     S: Service<WarpRequest, Response = WarpResponse, Error = Infallible> + Send + 'a,
@@ -141,18 +129,7 @@ where
 
     fn call(&mut self, req: Request) -> Self::Future {
         let query_params = req.query_string_parameters();
-        let request_id_opt = match req.request_context() {
-            RequestContext::ApiGatewayV2(ctx) => ctx.request_id.clone(),
-            RequestContext::ApiGatewayV1(ctx) => ctx.request_id.clone(),
-            RequestContext::Alb(_) => None,
-            RequestContext::WebSocket(ctx) => ctx.request_id.clone(),
-        };
-        let request_id_header_opt = match request_id_opt.as_ref().map(|a| HeaderValue::from_str(a))
-        {
-            Some(Ok(rid)) => Some(rid),
-            Some(Err(_)) => None,
-            None => None,
-        };
+        let request_id = req.lambda_context().request_id.clone();
         let (mut parts, body) = req.into_parts();
         let (content_len, body) = match body {
             LambdaBody::Empty => (0, WarpBody::empty()),
@@ -175,9 +152,6 @@ where
         if let Entry::Vacant(v) = parts.headers.entry("Content-Length") {
             v.insert(content_len.into());
         }
-        if let Some(rid) = &request_id_header_opt {
-            parts.headers.insert(LAMBDA_REQUEST_ID_HEADER, rid.clone());
-        }
 
         parts.uri = warp::hyper::Uri::from_str(uri.as_str())
             .map_err(|e| e)
@@ -190,14 +164,12 @@ where
         // Create lambda future
         let fut = async move {
             let warp_response = warp_fut.await?;
-            let (mut parts, res_body): (_, _) = warp_response.into_parts();
+            let (parts, res_body): (_, _) = warp_response.into_parts();
             let body = warp_body_as_lambda_body(res_body, &parts).await?;
-            if let Some(rid) = request_id_header_opt {
-                parts.headers.insert(LAMBDA_REQUEST_ID_HEADER, rid);
-            }
             let lambda_response = Response::from_parts(parts, body);
             Ok::<Self::Response, Self::Error>(lambda_response)
-        };
+        }
+        .instrument(info_span!("searcher request", request_id));
         Box::pin(fut)
     }
 }