hive-router 0.2.0

use crate::config::coprocessor::CoprocessorConfig;
use crate::http_utils::body::read_body_stream;
use crate::telemetry::logging::targets;
use crate::telemetry::traces::spans::coprocessor::CoprocessorSpan;
use crate::telemetry::TelemetryContext;
use http::{Method as HttpMethod, Uri};
use ntex::http::HeaderMap;
use ntex::web::{self, DefaultError};
use std::ops::ControlFlow;
use std::sync::Arc;
use tracing::{debug, error, Instrument};

use crate::executor::coprocessor::client::CoprocessorClient;
use crate::executor::coprocessor::error::CoprocessorError;
use crate::executor::coprocessor::stage::Stage;
use crate::executor::coprocessor::stages::graphql::{
    GraphqlAnalysisInput, GraphqlAnalysisStage, GraphqlRequestInput, GraphqlRequestStage,
    GraphqlResponseInput, GraphqlResponseStage,
};
use crate::executor::coprocessor::stages::router::{
    RouterRequestInput, RouterRequestStage, RouterResponseInput, RouterResponseStage,
};
use crate::executor::execution::plan::FailedExecutionResult;
use crate::executor::plugins::hooks::on_graphql_params::GraphQLParams;
use crate::executor::request_context::{
    RequestContextError, RequestContextExt, RequestContextPatch, SharedRequestContext,
};
use crate::executor::response::graphql_error::GraphQLError;

pub struct CoprocessorRuntime {
    router_request: Option<StageRuntime<RouterRequestStage>>,
    router_response: Option<StageRuntime<RouterResponseStage>>,
    graphql_request: Option<StageRuntime<GraphqlRequestStage>>,
    graphql_analysis: Option<StageRuntime<GraphqlAnalysisStage>>,
    graphql_response: Option<StageRuntime<GraphqlResponseStage>>,
    body_size_limit: usize,
}

#[derive(Default)]
pub struct PerformedMutations {
    pub body: bool,
    pub headers: bool,
    pub context: bool,
}

struct StageRuntime<S: Stage> {
    client: Arc<CoprocessorClient>,
    stage: S,
    telemetry_context: Arc<TelemetryContext>,
}

pub struct MutableRequestState<'a> {
    pub method: &'a HttpMethod,
    pub uri: &'a Uri,
    pub headers: &'a mut HeaderMap,
}

impl<A: Stage> StageRuntime<A> {
    fn new(
        client: Arc<CoprocessorClient>,
        stage: A,
        telemetry_context: Arc<TelemetryContext>,
    ) -> Self {
        Self {
            client,
            stage,
            telemetry_context,
        }
    }

    async fn execute<'a>(
        &self,
        input: &mut A::Input<'a>,
        shared_context: &SharedRequestContext,
    ) -> Result<ControlFlow<web::HttpResponse, PerformedMutations>, CoprocessorError> {
        let result = self.execute_internal(input, shared_context).await;

        if result.is_err() {
            let stage_name = self.stage.stage_name();
            let metrics = &self.telemetry_context.metrics.coprocessor;
            metrics.record_error(stage_name);
        }

        result
    }

    async fn execute_internal<'a>(
        &self,
        input: &mut A::Input<'a>,
        shared_context: &SharedRequestContext,
    ) -> Result<ControlFlow<web::HttpResponse, PerformedMutations>, CoprocessorError> {
        let mut performed_mutations = PerformedMutations::default();

        // Skip remote call when condition says stage should not run
        if !self.stage.should_run(input)? {
            return Ok(ControlFlow::Continue(performed_mutations));
        }

        let stage_name = self.stage.stage_name();
        let metrics = &self.telemetry_context.metrics.coprocessor;
        metrics.record_request(stage_name);

        let start = std::time::Instant::now();
        let id = uuid::Uuid::new_v4().to_string();
        let span = CoprocessorSpan::new(stage_name, &id).span;

        async {
          // Build stage payload and call the coprocessor
          // TODO: Include `request_id` in coprocessor payloads for log correlation
          //       once Dotan's logging PR is merged.
          let request = self.stage.build_request(input, &id, shared_context)?;
            debug!(
                target: targets::COPROCESSOR,
                coprocessor_id = %id,
                coprocessor_stage = stage_name,
                "Sending coprocessor request"
            );

            let response = self.client.send(request.body).await?;

            metrics.record_duration(stage_name, start.elapsed().as_secs_f64());

            if !response.status().is_success() {
                return Err(CoprocessorError::UnexpectedStatus(response.status()));
            }

            // Parse the response
            let mut parsed = self.stage.parse_response(response.body())?;

            if parsed.body.is_some() {
                performed_mutations.body = true;
            }
            if parsed.headers.is_some() {
                performed_mutations.headers = true;
            }
            // Handle possible break decision first
            match self.stage.break_output(parsed) {
                Ok(ControlFlow::Continue(p)) => {
                    parsed = p;
                }
                Ok(ControlFlow::Break(response)) => {
                    debug!(
                        target: targets::COPROCESSOR,
                        coprocessor_id = %id,
                        coprocessor_stage = stage_name,
                        status_code = %response.status(),
                        "Coprocessor short-circuited the request"
                    );

                    return Ok(ControlFlow::Break(response));
                }
                Err(err) => {
                    return Err(err);
                }
            }

            if let Some(context_patch_json) = parsed.context.take() {
                performed_mutations.context = true;
                let context_patch = A::parse_json_body(&context_patch_json)?;
                let patch: RequestContextPatch =
                    sonic_rs::from_str(context_patch.as_ref()).map_err(RequestContextError::Json)?;
                let mut context = shared_context.read_lock()?;
                context.for_coprocessor().apply_patch(patch)?;
            }

            if performed_mutations.body || performed_mutations.headers || performed_mutations.context {
                debug!(
                    target: targets::COPROCESSOR,
                    coprocessor_id = %id,
                    coprocessor_stage = stage_name,
                    body = performed_mutations.body,
                    headers = performed_mutations.headers,
                    context = performed_mutations.context,
                    "Coprocessor mutated the request"
                );
            }

            // Apply mutations only when flow continues
            self.stage.apply_mutations(parsed, input)?;

            // Continue the pipeline.
            Ok(ControlFlow::Continue(performed_mutations))
        }
        .instrument(span)
        .await
        .inspect_err(|err| {
            error!(target: targets::COPROCESSOR, error = %err, coprocessor_id = %id, coprocessor_stage = stage_name, "Coprocessor failure");
        })
    }
}

impl CoprocessorRuntime {
    pub fn from_config(
        config: &CoprocessorConfig,
        telemetry_context: Arc<TelemetryContext>,
        body_size_limit: usize,
    ) -> Result<Self, CoprocessorError> {
        let client = Arc::new(CoprocessorClient::new(
            config.clone(),
            telemetry_context.clone(),
        )?);

        let router_request = config
            .stages
            .router
            .request
            .as_ref()
            .map(RouterRequestStage::from_config)
            .transpose()?
            .map(|adapter| StageRuntime::new(client.clone(), adapter, telemetry_context.clone()));

        let router_response = config
            .stages
            .router
            .response
            .as_ref()
            .map(RouterResponseStage::from_config)
            .transpose()?
            .map(|adapter| StageRuntime::new(client.clone(), adapter, telemetry_context.clone()));

        let graphql_request = config
            .stages
            .graphql
            .request
            .as_ref()
            .map(GraphqlRequestStage::from_config)
            .transpose()?
            .map(|adapter| StageRuntime::new(client.clone(), adapter, telemetry_context.clone()));

        let graphql_response = config
            .stages
            .graphql
            .response
            .as_ref()
            .map(GraphqlResponseStage::from_config)
            .transpose()?
            .map(|adapter| StageRuntime::new(client.clone(), adapter, telemetry_context.clone()));

        let graphql_analysis = config
            .stages
            .graphql
            .analysis
            .as_ref()
            .map(GraphqlAnalysisStage::from_config)
            .transpose()?
            .map(|adapter| StageRuntime::new(client, adapter, telemetry_context));

        Ok(Self {
            router_request,
            router_response,
            graphql_request,
            graphql_analysis,
            graphql_response,
            body_size_limit,
        })
    }

    pub async fn on_router_request(
        &self,
        mut req: web::WebRequest<DefaultError>,
    ) -> ControlFlow<web::WebResponse, web::WebRequest<DefaultError>> {
        let Some(stage) = &self.router_request else {
            return ControlFlow::Continue(req);
        };

        // We read the request body only when this stage needs to include body
        let request_body = if stage.stage.include_body() {
            let body_stream = web::types::Payload(req.take_payload());
            let new_body = match read_body_stream(&req, body_stream, self.body_size_limit).await {
                Ok(body) => body,
                // We deliberately do not map to CoprocessorError here,
                // to follow the same logic for status codes
                Err(err) => {
                    error!(target: targets::COPROCESSOR, error = %err, coprocessor_stage = stage.stage.stage_name(), "coprocessor stage failed");

                    let response =
                        build_router_stage_error_response(err.status_code(), err.error_code());
                    return ControlFlow::Break(req.into_response(response));
                }
            };

            Some(new_body)
        } else {
            None
        };

        let shared_context = match req.read_request_context() {
            Ok(context) => context,
            Err(error) => {
                let error = CoprocessorError::from(error);
                let response =
                    build_router_stage_error_response(error.status_code(), error.error_code());
                return ControlFlow::Break(req.into_response(response));
            }
        };

        let mut input = RouterRequestInput::new(req, request_body);
        match stage
            .execute(&mut input, &shared_context)
            .await
            .unwrap_or_else(|err| {
            error!(target: targets::COPROCESSOR, error = %err, coprocessor_stage = stage.stage.stage_name(), "coprocessor stage failed");
            ControlFlow::Break(build_router_stage_error_response(
                err.status_code(),
                err.error_code(),
            ))
            }) {
            ControlFlow::Continue(_) => {
                // On continue, restore the original body only when coprocessor did not replace it
                input.restore_request_body_if_unchanged();
                ControlFlow::Continue(input.request)
            }
            ControlFlow::Break(response) => {
                // On break, we return immediately and skip body restoration to avoid unnecessary work
                ControlFlow::Break(input.request.into_response(response))
            }
        }
    }

    pub async fn on_router_response(&self, response: web::WebResponse) -> web::WebResponse {
        let Some(stage) = &self.router_response else {
            return response;
        };

        let shared_context = match response.request().read_request_context() {
            Ok(context) => context,
            Err(error) => {
                let error = CoprocessorError::from(error);
                let fallback =
                    build_router_stage_error_response(error.status_code(), error.error_code());
                return response.into_response(fallback);
            }
        };

        let mut input = RouterResponseInput::new(response);

        match stage
            .execute(&mut input, &shared_context)
            .await
            .unwrap_or_else(|err| error_to_break(stage, err))
        {
            ControlFlow::Continue(_) => input.response,
            ControlFlow::Break(response) => input.response.into_response(response),
        }
    }

    pub async fn on_graphql_request(
        &self,
        request: &web::HttpRequest,
        request_headers: &mut HeaderMap,
        graphql_request: &mut GraphQLParams,
        sdl_fn: impl FnOnce() -> Arc<str>,
    ) -> Result<ControlFlow<web::HttpResponse, PerformedMutations>, CoprocessorError> {
        let Some(stage) = &self.graphql_request else {
            return Ok(ControlFlow::Continue(Default::default()));
        };

        let shared_context = request.read_request_context()?;
        let sdl = stage.stage.include_sdl().then(sdl_fn);
        let mut input =
            GraphqlRequestInput::new(request, request_headers, graphql_request, sdl.as_deref());
        stage.execute(&mut input, &shared_context).await
    }

    pub async fn on_graphql_analysis(
        &self,
        request: MutableRequestState<'_>,
        graphql_request: &GraphQLParams,
        context: &SharedRequestContext,
        sdl_fn: impl FnOnce() -> Arc<str>,
    ) -> Result<ControlFlow<web::HttpResponse, PerformedMutations>, CoprocessorError> {
        let Some(stage) = &self.graphql_analysis else {
            return Ok(ControlFlow::Continue(Default::default()));
        };

        let sdl = stage.stage.include_sdl().then(sdl_fn);
        let mut input = GraphqlAnalysisInput::new(request, graphql_request, sdl.as_deref());
        stage.execute(&mut input, context).await
    }

    pub async fn on_graphql_response(
        &self,
        response: web::HttpResponse,
        request: &web::HttpRequest,
        // We return `Option<T>` instead of `T`, because of error handling of the caller.
        // If I would return `T`, the caller would have to be wrapped with duplicated
        // error handling logic.
        // With `Option<T>`, the coprocessor runtime handles the case where sdl is not available.
        sdl_fn: impl FnOnce() -> Option<Arc<str>>,
    ) -> Result<ControlFlow<web::HttpResponse, web::HttpResponse>, CoprocessorError> {
        let Some(stage) = &self.graphql_response else {
            return Ok(ControlFlow::Continue(response));
        };

        let sdl = stage.stage.include_sdl().then(sdl_fn).flatten();
        let shared_context = request.read_request_context()?;
        let mut input = GraphqlResponseInput::new(response, request, sdl.as_deref());
        Ok(stage
            .execute(&mut input, &shared_context)
            .await?
            .map_continue(|_| input.response))
    }
}

fn error_to_break<A: Stage>(
    stage: &StageRuntime<A>,
    err: CoprocessorError,
) -> ControlFlow<web::HttpResponse, PerformedMutations> {
    // Stage error metrics and specific logs are already recorded in execute() where possible,
    // but keeping a fallback log here correctly reports the failure that causes short-circuit
    // in case it's not caught earlier, and always breaks the request.
    error!(target: targets::COPROCESSOR, error = %err, coprocessor_stage = stage.stage.stage_name(), "coprocessor stage failed");
    ControlFlow::Break(web::HttpResponse::new(err.status_code()))
}

fn build_router_stage_error_response(
    status: http::StatusCode,
    code: &'static str,
) -> web::HttpResponse {
    let body = FailedExecutionResult {
        errors: vec![GraphQLError::from_message_and_code(
            "Internal server error",
            code,
        )],
    }
    .serialize();

    web::HttpResponse::build(status)
        .header(http::header::CONTENT_TYPE, "application/json")
        .body(body)
}