lambda improvements (#6171)

* make lambda support optional

* map lambda error to actual variant

* retry rate limited lambdas

Author: trinity-1686a

quickwit/quickwit-cli/Cargo.toml

@@ -101,6 +101,7 @@ release-feature-set = [
   "quickwit-indexing/pulsar",
   "quickwit-indexing/sqs",
   "quickwit-indexing/vrl",
+  "quickwit-serve/lambda",
   "quickwit-storage/azure",
   "quickwit-storage/gcs",
   "quickwit-metastore/postgres",
@@ -114,6 +115,7 @@ release-feature-vendored-set = [
   "quickwit-indexing/sqs",
   "quickwit-indexing/vrl",
   "quickwit-indexing/vendored-kafka",
+  "quickwit-serve/lambda",
   "quickwit-storage/azure",
   "quickwit-storage/gcs",
   "quickwit-metastore/postgres",
@@ -126,6 +128,7 @@ release-macos-feature-vendored-set = [
   "quickwit-indexing/sqs",
   "quickwit-indexing/vrl",
   "quickwit-indexing/vendored-kafka-macos",
+  "quickwit-serve/lambda",
   "quickwit-storage/azure",
   "quickwit-storage/gcs",
   "quickwit-metastore/postgres",

quickwit/quickwit-lambda-client/src/deploy.rs

@@ -26,7 +26,7 @@
 //! - Old versions are garbage collected (keep current + top 5 most recent)
 
 use std::collections::HashMap;
-use std::sync::{Arc, OnceLock};
+use std::sync::OnceLock;
 
 use anyhow::{Context, anyhow};
 use aws_sdk_lambda::Client as LambdaClient;
@@ -108,7 +108,7 @@ fn version_description(deploy_config_opt: Option<&LambdaDeployConfig>) -> String
 /// ensuring the deployed Lambda matches the embedded binary.
 pub async fn try_get_or_deploy_invoker(
     lambda_config: &LambdaConfig,
-) -> anyhow::Result<Arc<dyn LambdaLeafSearchInvoker>> {
+) -> anyhow::Result<impl LambdaLeafSearchInvoker> {
     let aws_config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
     let client = LambdaClient::new(&aws_config);
     let function_name = &lambda_config.function_name;

quickwit/quickwit-lambda-client/src/invoker.rs

@@ -12,30 +12,65 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::sync::Arc;
-
 use anyhow::Context as _;
 use async_trait::async_trait;
 use aws_sdk_lambda::Client as LambdaClient;
+use aws_sdk_lambda::error::{DisplayErrorContext, SdkError};
+use aws_sdk_lambda::operation::invoke::InvokeError;
 use aws_sdk_lambda::primitives::Blob;
 use aws_sdk_lambda::types::InvocationType;
 use base64::prelude::*;
 use prost::Message;
+use quickwit_common::retry::{RetryParams, retry};
 use quickwit_lambda_server::{LambdaSearchRequestPayload, LambdaSearchResponsePayload};
 use quickwit_proto::search::{LambdaSearchResponses, LambdaSingleSplitResult, LeafSearchRequest};
 use quickwit_search::{LambdaLeafSearchInvoker, SearchError};
 use tracing::{debug, info, instrument};
 
 use crate::metrics::LAMBDA_METRICS;
 
+/// Maps an AWS Lambda SDK invocation error to a `SearchError`.
+fn invoke_error_to_search_error(error: SdkError<InvokeError>) -> SearchError {
+    if let SdkError::ServiceError(ref service_error) = error
+        && matches!(
+            service_error.err(),
+            InvokeError::TooManyRequestsException(_)
+                | InvokeError::EniLimitReachedException(_)
+                | InvokeError::SubnetIpAddressLimitReachedException(_)
+                | InvokeError::Ec2ThrottledException(_)
+                | InvokeError::ResourceConflictException(_)
+        )
+    {
+        return SearchError::TooManyRequests;
+    }
+
+    let is_timeout = match &error {
+        SdkError::TimeoutError(_) => true,
+        SdkError::DispatchFailure(failure) => failure.is_timeout(),
+        SdkError::ServiceError(service_error) => matches!(
+            service_error.err(),
+            InvokeError::EfsMountTimeoutException(_) | InvokeError::SnapStartTimeoutException(_)
+        ),
+        _ => false,
+    };
+
+    let error_msg = format!("lambda invocation failed: {}", DisplayErrorContext(&error));
+
+    if is_timeout {
+        SearchError::Timeout(error_msg)
+    } else {
+        SearchError::Internal(error_msg)
+    }
+}
+
 /// Create a Lambda invoker for a specific version.
 ///
 /// The version number is used as the qualifier when invoking, ensuring we call
 /// the exact published version (not $LATEST).
 pub(crate) async fn create_lambda_invoker_for_version(
     function_name: String,
     version: String,
-) -> anyhow::Result<Arc<dyn LambdaLeafSearchInvoker>> {
+) -> anyhow::Result<AwsLambdaInvoker> {
     let aws_config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
     let client = LambdaClient::new(&aws_config);
     let invoker = AwsLambdaInvoker {
@@ -44,11 +79,11 @@ pub(crate) async fn create_lambda_invoker_for_version(
         version,
     };
     invoker.validate().await?;
-    Ok(Arc::new(invoker))
+    Ok(invoker)
 }
 
 /// AWS Lambda implementation of RemoteFunctionInvoker.
-struct AwsLambdaInvoker {
+pub(crate) struct AwsLambdaInvoker {
     client: LambdaClient,
     function_name: String,
     /// The version number to invoke (e.g., "7", "12").
@@ -79,6 +114,12 @@ impl AwsLambdaInvoker {
     }
 }
 
+const LAMBDA_RETRY_PARAMS: RetryParams = RetryParams {
+    base_delay: std::time::Duration::from_secs(1),
+    max_delay: std::time::Duration::from_secs(10),
+    max_attempts: 3,
+};
+
 #[async_trait]
 impl LambdaLeafSearchInvoker for AwsLambdaInvoker {
     #[instrument(skip(self, request), fields(function_name = %self.function_name, version = %self.version))]
@@ -88,7 +129,10 @@ impl LambdaLeafSearchInvoker for AwsLambdaInvoker {
     ) -> Result<Vec<LambdaSingleSplitResult>, SearchError> {
         let start = std::time::Instant::now();
 
-        let result = self.invoke_leaf_search_inner(request).await;
+        let result = retry(&LAMBDA_RETRY_PARAMS, || {
+            self.invoke_leaf_search_inner(request.clone())
+        })
+        .await;
 
         let elapsed = start.elapsed().as_secs_f64();
         let status = if result.is_ok() { "success" } else { "error" };
@@ -141,7 +185,7 @@ impl AwsLambdaInvoker {
         let response = invoke_builder
             .send()
             .await
-            .map_err(|e| SearchError::Internal(format!("Lambda invocation error: {}", e)))?;
+            .map_err(invoke_error_to_search_error)?;
 
         // Check for function error
         if let Some(error) = response.function_error() {

quickwit/quickwit-lambda-server/src/context.rs

@@ -33,7 +33,8 @@ impl LambdaSearcherContext {
         info!("initializing lambda searcher context");
 
         let searcher_config = try_searcher_config_from_env()?;
-        let searcher_context = Arc::new(SearcherContext::new(searcher_config, None, None));
+        let searcher_context =
+            Arc::new(SearcherContext::new_without_invoker(searcher_config, None));
         let storage_resolver = StorageResolver::configured(&Default::default());
 
         Ok(Self {

quickwit/quickwit-search/src/error.rs

@@ -14,6 +14,7 @@
 
 use itertools::Itertools;
 use quickwit_common::rate_limited_error;
+use quickwit_common::retry::Retryable;
 use quickwit_doc_mapper::QueryParserError;
 use quickwit_proto::error::grpc_error_to_grpc_status;
 use quickwit_proto::metastore::{EntityKind, MetastoreError};
@@ -175,6 +176,12 @@ impl From<MetastoreError> for SearchError {
     }
 }
 
+impl Retryable for SearchError {
+    fn is_retryable(&self) -> bool {
+        matches!(self, SearchError::TooManyRequests | SearchError::Timeout(_))
+    }
+}
+
 impl From<JoinError> for SearchError {
     fn from(join_error: JoinError) -> SearchError {
         SearchError::Internal(format!("spawned task in root join failed: {join_error}"))

quickwit/quickwit-search/src/invoker.rs

@@ -36,3 +36,27 @@ pub trait LambdaLeafSearchInvoker: Send + Sync + 'static {
         request: LeafSearchRequest,
     ) -> Result<Vec<LambdaSingleSplitResult>, SearchError>;
 }
+
+#[async_trait]
+impl<T> LambdaLeafSearchInvoker for Box<T>
+where T: LambdaLeafSearchInvoker + ?Sized
+{
+    async fn invoke_leaf_search(
+        &self,
+        request: LeafSearchRequest,
+    ) -> Result<Vec<LambdaSingleSplitResult>, SearchError> {
+        (**self).invoke_leaf_search(request).await
+    }
+}
+
+#[async_trait]
+impl<T> LambdaLeafSearchInvoker for std::sync::Arc<T>
+where T: LambdaLeafSearchInvoker + ?Sized
+{
+    async fn invoke_leaf_search(
+        &self,
+        request: LeafSearchRequest,
+    ) -> Result<Vec<LambdaSingleSplitResult>, SearchError> {
+        (**self).invoke_leaf_search(request).await
+    }
+}

quickwit/quickwit-search/src/lib.rs

@@ -287,7 +287,7 @@ pub async fn single_node_search(
     let search_job_placer = SearchJobPlacer::new(searcher_pool.clone());
     let cluster_client = ClusterClient::new(search_job_placer);
     let searcher_config = SearcherConfig::default();
-    let searcher_context = Arc::new(SearcherContext::new(searcher_config, None, None));
+    let searcher_context = Arc::new(SearcherContext::new_without_invoker(searcher_config, None));
     let search_service = Arc::new(SearchServiceImpl::new(
         metastore.clone(),
         storage_resolver,

quickwit/quickwit-search/src/service.rs

@@ -436,14 +436,26 @@ impl SearcherContext {
     #[cfg(test)]
     pub fn for_test() -> SearcherContext {
         let searcher_config = SearcherConfig::default();
-        SearcherContext::new(searcher_config, None, None)
+        SearcherContext::new_without_invoker(searcher_config, None)
+    }
+
+    /// Creates a new searcher context without a lambda invoker.
+    pub fn new_without_invoker(
+        searcher_config: SearcherConfig,
+        split_cache_opt: Option<Arc<SplitCache>>,
+    ) -> Self {
+        Self::new(
+            searcher_config,
+            split_cache_opt,
+            None::<Box<dyn LambdaLeafSearchInvoker>>,
+        )
     }
 
     /// Creates a new searcher context, given a searcher config, and an optional `SplitCache`.
     pub fn new(
         searcher_config: SearcherConfig,
         split_cache_opt: Option<Arc<SplitCache>>,
-        lambda_invoker: Option<Arc<dyn LambdaLeafSearchInvoker>>,
+        lambda_invoker: Option<impl LambdaLeafSearchInvoker + 'static>,
     ) -> Self {
         let global_split_footer_cache = MemorySizedCache::from_config(
             &searcher_config.split_footer_cache,
@@ -463,6 +475,9 @@ impl SearcherContext {
             Some(searcher_config.aggregation_bucket_limit),
         );
 
+        let lambda_invoker =
+            lambda_invoker.map(|invoker| Arc::new(invoker) as Arc<dyn LambdaLeafSearchInvoker>);
+
         Self {
             searcher_config,
             fast_fields_cache: storage_long_term_cache,

quickwit/quickwit-search/src/tests.rs

@@ -1028,8 +1028,10 @@ async fn test_search_util(test_sandbox: &TestSandbox, query: &str) -> Vec<u32> {
         max_hits: 100,
         ..Default::default()
     });
-    let searcher_context: Arc<SearcherContext> =
-        Arc::new(SearcherContext::new(SearcherConfig::default(), None, None));
+    let searcher_context: Arc<SearcherContext> = Arc::new(SearcherContext::new_without_invoker(
+        SearcherConfig::default(),
+        None,
+    ));
 
     let search_response = single_doc_mapping_leaf_search(
         searcher_context,
@@ -1666,7 +1668,10 @@ async fn test_single_node_list_terms() -> anyhow::Result<()> {
         .into_iter()
         .map(|split| extract_split_and_footer_offsets(&split.split_metadata))
         .collect();
-    let searcher_context = Arc::new(SearcherContext::new(SearcherConfig::default(), None, None));
+    let searcher_context = Arc::new(SearcherContext::new_without_invoker(
+        SearcherConfig::default(),
+        None,
+    ));
 
     {
         let request = ListTermsRequest {

quickwit/quickwit-serve/Cargo.toml

@@ -72,7 +72,7 @@ quickwit-opentelemetry = { workspace = true }
 quickwit-proto = { workspace = true }
 quickwit-query = { workspace = true }
 quickwit-search = { workspace = true }
-quickwit-lambda-client = { workspace = true }
+quickwit-lambda-client = { workspace = true, optional = true }
 quickwit-storage = { workspace = true }
 quickwit-telemetry = { workspace = true }
 
@@ -115,3 +115,6 @@ sqs-for-tests = [
   "quickwit-indexing/sqs",
   "quickwit-indexing/sqs-test-helpers"
 ]
+lambda = [
+  "quickwit-lambda-client"
+]