apollo-router 3.0.0-alpha.2

//! Tests for this functionality are still mostly in the `crate::services::supergraph::tests` module.

use std::pin::Pin;
use std::sync::Arc;
use std::sync::Weak;
use std::sync::atomic::Ordering;
use std::time::Instant;

use futures::FutureExt;
use futures::Stream;
use futures::future::BoxFuture;
use futures::stream::StreamExt;
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::SendError;
use tokio_stream::wrappers::ReceiverStream;
use tower::Layer;
use tower::Service;
use tower::ServiceExt as _;
use tracing::Span;
use tracing::field;
use tracing_futures::Instrument;

use crate::Configuration;
use crate::Context;
use crate::Notify;
use crate::apollo_studio_interop::UsageReporting;
use crate::context::OPERATION_NAME;
use crate::graphql;
use crate::graphql::Response;
use crate::plugins::authentication::APOLLO_AUTHENTICATION_JWT_CLAIMS;
use crate::plugins::subscription::SUBSCRIPTION_ERROR_EXTENSION_KEY;
use crate::plugins::subscription::SubscriptionConfig;
use crate::plugins::telemetry::tracing::apollo_telemetry::APOLLO_PRIVATE_DURATION_NS;
use crate::query_planner::subscription::OPENED_SUBSCRIPTIONS;
use crate::query_planner::subscription::SUBSCRIPTION_EVENT_SPAN_NAME;
use crate::query_planner::subscription::SubscriptionHandle;
use crate::services::ExecutionRequest;
use crate::services::execution;
use crate::services::execution::QueryPlan;
use crate::services::subgraph::BoxGqlStream;
use crate::spec::Schema;

pub(crate) const SUBSCRIPTION_CONFIG_RELOAD_EXTENSION_CODE: &str = "SUBSCRIPTION_CONFIG_RELOAD";
pub(crate) const SUBSCRIPTION_SCHEMA_RELOAD_EXTENSION_CODE: &str = "SUBSCRIPTION_SCHEMA_RELOAD";
const SUBSCRIPTION_JWT_EXPIRED_EXTENSION_CODE: &str = "SUBSCRIPTION_JWT_EXPIRED";
const SUBSCRIPTION_EXECUTION_ERROR_EXTENSION_CODE: &str = "SUBSCRIPTION_EXECUTION_ERROR";
pub(crate) const SUBSCRIPTION_MAX_LIFETIME_EXTENSION_CODE: &str =
    "SUBSCRIPTION_MAX_LIFETIME_EXCEEDED";

/// The execution side of the subscriptions implementation starts up a side-channel task used to
/// handle messages received from the subgraph that we subscribed to.
pub(crate) struct SubscriptionExecutionLayer {
    notify: Notify<String, graphql::Response>,
}

impl SubscriptionExecutionLayer {
    pub(crate) fn new(notify: Notify<String, graphql::Response>) -> Self {
        Self { notify }
    }
}

impl<S> Layer<S> for SubscriptionExecutionLayer {
    type Service = SubscriptionExecutionService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        SubscriptionExecutionService {
            inner,
            notify: self.notify.clone(),
        }
    }
}

#[derive(Clone)]
pub(crate) struct SubscriptionExecutionService<S> {
    inner: S,
    notify: Notify<String, graphql::Response>,
}

impl<S> Service<ExecutionRequest> for SubscriptionExecutionService<S>
where
    S: Service<ExecutionRequest, Response = execution::Response, Error = tower::BoxError>
        + Clone
        + Send
        + 'static,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut req: ExecutionRequest) -> Self::Future {
        let inner = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, inner);
        let notify = self.notify.clone();

        Box::pin(async move {
            if req.query_plan.is_subscription() {
                let context = req.context.clone();
                let (subs_tx, subs_rx) = mpsc::channel(1);
                let query_plan = req.query_plan.clone();
                let execution_service_cloned = inner.clone();
                let supergraph_http_request = req.supergraph_request.clone();

                // Subscribe to config and schema changes immediately, so we don't
                // miss out on any events happening in-between here and when we
                // actually start responding to changes (which happens asynchronously
                // inside the subscription task).
                let configuration_updated_rx = Box::pin(notify.subscribe_configuration());
                let schema_updated_rx = Box::pin(notify.subscribe_schema());

                // Spawn the side-channel task for subscription.
                tokio::spawn(async move {
                    subscription_task(
                        execution_service_cloned,
                        context,
                        query_plan,
                        subs_rx,
                        supergraph_http_request,
                        configuration_updated_rx,
                        schema_updated_rx,
                    )
                    .await;
                });
                req.subscription_tx = subs_tx.into();
            }

            inner.call(req).await
        })
    }
}

// WARN: you might be tempted to change this to pub(crate) because its fields are pub(crate) and
// there's no actual use in having it marked pub _apart_ from it being used in the
// execution::Request builder and, more importantly, the fake_builder we get from BuildStructor.
// Customers use that in unit tests for their plugins, and they'll fail to compile if we mark this
// pub(crate)
pub struct SubscriptionTaskParams {
    pub(crate) client_sender: tokio::sync::mpsc::Sender<Response>,
    pub(crate) subscription_handle: SubscriptionHandle,
    pub(crate) subscription_config: SubscriptionConfig,
    pub(crate) stream_rx: ReceiverStream<BoxGqlStream>,
}

fn subscription_fatal_error(message: impl Into<String>, extension_code: &str) -> Response {
    Response::builder()
        .subscribed(false)
        .extension(SUBSCRIPTION_ERROR_EXTENSION_KEY, true)
        .error(
            graphql::Error::builder()
                .message(message)
                .extension_code(extension_code)
                .build(),
        )
        .build()
}

/// The execution-side end of the subscriptions side-channel, which propagates messages from the
/// subgraph to the client.
///
/// This end receives the messages from the subgraph, executes query plans to resolve federated
/// data in those messages, and sends the response back on a channel that is part of the eventual
/// response.
async fn subscription_task(
    execution_service: impl Service<
        ExecutionRequest,
        Response = execution::Response,
        Error = tower::BoxError,
    > + Clone,
    context: Context,
    query_plan: Arc<QueryPlan>,
    mut rx: mpsc::Receiver<SubscriptionTaskParams>,
    supergraph_http_request: http::Request<graphql::Request>,
    mut configuration_updated_rx: Pin<Box<dyn Stream<Item = Weak<Configuration>> + Send>>,
    mut schema_updated_rx: Pin<Box<dyn Stream<Item = Arc<Schema>> + Send>>,
) {
    let sub_params = match rx.recv().await {
        Some(sub_params) => sub_params,
        None => {
            return;
        }
    };
    let subscription_config = sub_params.subscription_config;
    let subscription_handle = sub_params.subscription_handle;
    let mut receiver = sub_params.stream_rx;
    let sender = sub_params.client_sender;

    // Get the rest of the query_plan to execute for subscription events. If `rest` is `None`, we
    // bubble that up to the whole query plan instead of having an empty node: this lets us avoid
    // doing excessive calls into the execution service when no fetches are required.
    let query_plan = match query_plan.root.as_deref() {
        Some(crate::query_planner::PlanNode::Subscription { rest, .. }) => rest.clone().map(|r| {
            Arc::new(QueryPlan {
                usage_reporting: query_plan.usage_reporting.clone(),
                root: Some(Arc::new(*r)),
                formatted_query_plan: query_plan.formatted_query_plan.clone(),
                query: query_plan.query.clone(),
                estimated_size: Default::default(),
            })
        }),
        _ => {
            let _ = sender
                .send(subscription_fatal_error(
                    "cannot execute the subscription event",
                    SUBSCRIPTION_EXECUTION_ERROR_EXTENSION_CODE,
                ))
                .await;
            return;
        }
    };

    let limit_is_set = subscription_config.max_opened_subscriptions.is_some();
    let mut subscription_handle = subscription_handle.clone();
    let operation_signature = context
        .extensions()
        .with_lock(|lock| {
            lock.get::<Arc<UsageReporting>>()
                .map(|usage_reporting| usage_reporting.get_stats_report_key().clone())
        })
        .unwrap_or_default();

    let operation_name = context
        .get::<_, String>(OPERATION_NAME)
        .ok()
        .flatten()
        .unwrap_or_default();

    let mut receiver = match receiver.next().await {
        Some(receiver) => receiver,
        None => {
            tracing::trace!("receiver channel closed");
            return;
        }
    };

    if limit_is_set {
        OPENED_SUBSCRIPTIONS.fetch_add(1, Ordering::Relaxed);
    }

    // Note: `configuration_updated_rx` and `schema_updated_rx` are now
    // subscribed in `SubscriptionExecutionService::call` (before the spawn)
    // and passed in as parameters. See the comment in `call` for the race
    // this fix addresses.

    let mut timeout = if context
        .get_json_value(APOLLO_AUTHENTICATION_JWT_CLAIMS)
        .is_some()
    {
        let expires_in = crate::plugins::authentication::jwks::jwt_expires_in(&context);
        tokio::time::sleep(expires_in).boxed()
    } else {
        futures::future::pending().boxed()
    };

    let mut max_lifetime_timeout = if let Some(max_lifetime) = subscription_config.max_lifetime {
        tokio::time::sleep(max_lifetime).boxed()
    } else {
        futures::future::pending().boxed()
    };

    loop {
        tokio::select! {
            // We prefer to specify the order of checks within the select
            biased;
            _ = subscription_handle.closed_signal.recv() => {
                break;
            }
            _ = &mut timeout => {
                let _ = sender.send(subscription_fatal_error("subscription closed because the JWT has expired", SUBSCRIPTION_JWT_EXPIRED_EXTENSION_CODE)).await;
                break;
            },
            _ = &mut max_lifetime_timeout => {
                let _ = sender.send(subscription_fatal_error("subscription closed because the maximum lifetime has been reached", SUBSCRIPTION_MAX_LIFETIME_EXTENSION_CODE)).await;
                break;
            },
            message = receiver.next() => {
                match message {
                    Some(mut val) => {
                        val.created_at = Some(Instant::now());
                        let res = dispatch_subscription_event(
                            &supergraph_http_request,
                            execution_service.clone(),
                            query_plan.as_ref(),
                            context.clone(),
                            val,
                            sender.clone(),
                        )
                            .instrument(tracing::info_span!(SUBSCRIPTION_EVENT_SPAN_NAME,
                                graphql.operation.name = %operation_name,
                                otel.kind = "INTERNAL",
                                apollo_private.operation_signature = %operation_signature,
                                apollo_private.duration_ns = field::Empty,)
                            )
                            .await;
                        if let Err(err) = res {
                            tracing::error!("cannot send the subscription to the client: {err:?}");
                            break;
                        }
                    }
                    None => break,
                }
            }
            Some(_new_configuration) = configuration_updated_rx.next() => {
                let _ = sender
                    .send(subscription_fatal_error("subscription has been closed due to a configuration reload", SUBSCRIPTION_CONFIG_RELOAD_EXTENSION_CODE))
                    .await;
            }
            Some(_new_schema) = schema_updated_rx.next() => {
                let _ = sender
                    .send(subscription_fatal_error("subscription has been closed due to a schema reload", SUBSCRIPTION_SCHEMA_RELOAD_EXTENSION_CODE))
                    .await;

                break;
            }
        }
    }
    drop(sender);
    tracing::trace!("Leaving the task for subscription");
    if limit_is_set {
        OPENED_SUBSCRIPTIONS.fetch_sub(1, Ordering::Relaxed);
    }
}

async fn dispatch_subscription_event(
    supergraph_http_request: &http::Request<graphql::Request>,
    execution_service: impl Service<
        ExecutionRequest,
        Response = execution::Response,
        Error = tower::BoxError,
    > + Clone,
    query_plan: Option<&Arc<QueryPlan>>,
    context: Context,
    mut val: graphql::Response,
    sender: mpsc::Sender<Response>,
) -> Result<(), SendError<Response>> {
    let start = Instant::now();
    let span = Span::current();
    let res = match query_plan {
        Some(query_plan) => {
            let cloned_supergraph_req = supergraph_http_request.clone();
            let execution_request = ExecutionRequest::internal_builder()
                .supergraph_request(cloned_supergraph_req)
                .query_plan(query_plan.clone())
                .context(context)
                .source_stream_value(val.data.take().unwrap_or_default())
                .build()
                .await;

            let execution_response = execution_service.oneshot(execution_request).await;
            let next_response = match execution_response {
                Ok(mut execution_response) => execution_response.next_response().await,
                Err(err) => {
                    tracing::error!("cannot execute the subscription event: {err:?}");
                    let _ = sender
                        .send(subscription_fatal_error(
                            "cannot execute the subscription event",
                            SUBSCRIPTION_EXECUTION_ERROR_EXTENSION_CODE,
                        ))
                        .await;
                    return Ok(());
                }
            };

            if let Some(mut next_response) = next_response {
                next_response.created_at = val.created_at;
                next_response.subscribed = val.subscribed;
                val.errors.append(&mut next_response.errors);
                next_response.errors = val.errors;

                sender.send(next_response).await
            } else {
                Ok(())
            }
        }
        None => sender.send(val).await,
    };
    span.record(
        APOLLO_PRIVATE_DURATION_NS,
        start.elapsed().as_nanos() as i64,
    );

    res
}