improve handling of deferred response errors in rhai scripts (#2945)

Currently, if a rhai script errors in rhai, the rhai plugin ignores the
error and returns None in the stream of results.

This has two unfortunate aspects:
 - the error is not propagated to the client
 - the stream is terminated (silently)

The fix captures the error and propagates the response to the client.

Fixes #2935 #2936 

<!-- start metadata -->

**Checklist**

Complete the checklist (and note appropriate exceptions) before a final
PR is raised.

- [x] Changes are compatible[^1]
- [x] Documentation[^2] completed
- [x] Performance impact assessed and acceptable
- Tests added and passing[^3]
    - [x] Unit Tests
    ~- [ ] Integration Tests~
    - [x] Manual Tests

**Exceptions**

*Note any exceptions here*

**Notes**

[^1]. It may be appropriate to bring upcoming changes to the attention
of other (impacted) groups. Please endeavour to do this before seeking
PR approval. The mechanism for doing this will vary considerably, so use
your judgement as to how and when to do this.
[^2]. Configuration is an important part of many changes. Where
applicable please try to document configuration examples.
[^3]. Tick whichever testing boxes are applicable. If you are adding
Manual Tests:
- please document the manual testing (extensively) in the Exceptions.
- please raise a separate issue to automate the test and label it (or
ask for it to be labeled) as `manual test`

Author: garypen

.changesets/feat_garypen_2935_rhai_defer_error.md

@@ -0,0 +1,38 @@
+### Improve handling of deferred response errors from rhai scripts ([Issue #2935](https://github.com/apollographql/router/issues/2935)) ([Issue #2936](https://github.com/apollographql/router/issues/2936))
+
+
+Whilst processing a deferred response; if a rhai script errors in rhai the router ignores the error and returns None in the stream of results. This has two unfortunate aspects:
+
+ - the error is not propagated to the client
+ - the stream is terminated (silently)
+
+The fix captures the error and propagates the response to the client.
+
+This fix also adds support for the `is_primary()` method which may be invoked on both supergraph_service() and execution_service() responses. It may be used to avoid implementing exception handling for header interactions and to determine if a response `is_primary()` (i.e.: first) or not.
+
+e.g.:
+
+```
+    if response.is_primary() {
+        print(`all response headers: ${response.headers}`);
+    } else {
+        print(`don't try to access headers`);
+    }
+```
+
+vs
+
+```
+    try {
+        print(`all response headers: ${response.headers}`);
+    }
+    catch(err) {
+        if err == "cannot access headers on a deferred response" {
+            print(`don't try to access headers`);
+        }
+    }
+```
+Note: This is a minimal example for purposes of illustration. A real exception handler would handle all error conditions.
+
+
+By [@garypen](https://github.com/garypen) in https://github.com/apollographql/router/pull/2945

apollo-router/src/plugins/rhai/mod.rs

@@ -85,6 +85,9 @@ type SharedMut<T> = rhai::Shared<Mutex<Option<T>>>;
 
 pub(crate) const RHAI_SPAN_NAME: &str = "rhai_plugin";
 
+const CANNOT_ACCESS_HEADERS_ON_A_DEFERRED_RESPONSE: &str =
+    "cannot access headers on a deferred response";
+
 impl<T> OptionDance<T> for SharedMut<T> {
     fn with_mut<R>(&self, f: impl FnOnce(&mut T) -> R) -> R {
         let mut guard = self.lock().expect("poisoned mutex");
@@ -214,11 +217,25 @@ mod router_plugin {
         Ok(obj.with_mut(|response| response.response.headers().clone()))
     }
 
+    #[rhai_fn(name = "is_primary", pure)]
+    pub(crate) fn supergraph_response_is_primary(
+        _obj: &mut SharedMut<supergraph::Response>,
+    ) -> bool {
+        true
+    }
+
     #[rhai_fn(get = "headers", pure, return_raw)]
     pub(crate) fn get_originating_headers_router_deferred_response(
         _obj: &mut SharedMut<supergraph::DeferredResponse>,
     ) -> Result<HeaderMap, Box<EvalAltResult>> {
-        Err("cannot access headers on a deferred response".into())
+        Err(CANNOT_ACCESS_HEADERS_ON_A_DEFERRED_RESPONSE.into())
+    }
+
+    #[rhai_fn(name = "is_primary", pure)]
+    pub(crate) fn supergraph_deferred_response_is_primary(
+        _obj: &mut SharedMut<supergraph::DeferredResponse>,
+    ) -> bool {
+        false
     }
 
     #[rhai_fn(get = "headers", pure, return_raw)]
@@ -228,11 +245,23 @@ mod router_plugin {
         Ok(obj.with_mut(|response| response.response.headers().clone()))
     }
 
+    #[rhai_fn(name = "is_primary", pure)]
+    pub(crate) fn execution_response_is_primary(_obj: &mut SharedMut<execution::Response>) -> bool {
+        true
+    }
+
     #[rhai_fn(get = "headers", pure, return_raw)]
     pub(crate) fn get_originating_headers_execution_deferred_response(
         _obj: &mut SharedMut<execution::DeferredResponse>,
     ) -> Result<HeaderMap, Box<EvalAltResult>> {
-        Err("cannot access headers on a deferred response".into())
+        Err(CANNOT_ACCESS_HEADERS_ON_A_DEFERRED_RESPONSE.into())
+    }
+
+    #[rhai_fn(name = "is_primary", pure)]
+    pub(crate) fn execution_deferred_response_is_primary(
+        _obj: &mut SharedMut<execution::DeferredResponse>,
+    ) -> bool {
+        false
     }
 
     #[rhai_fn(get = "headers", pure, return_raw)]
@@ -291,7 +320,7 @@ mod router_plugin {
         _obj: &mut SharedMut<supergraph::DeferredResponse>,
         _headers: HeaderMap,
     ) -> Result<(), Box<EvalAltResult>> {
-        Err("cannot access headers on a deferred response".into())
+        Err(CANNOT_ACCESS_HEADERS_ON_A_DEFERRED_RESPONSE.into())
     }
 
     #[rhai_fn(set = "headers", return_raw)]
@@ -308,7 +337,7 @@ mod router_plugin {
         _obj: &mut SharedMut<execution::DeferredResponse>,
         _headers: HeaderMap,
     ) -> Result<(), Box<EvalAltResult>> {
-        Err("cannot access headers on a deferred response".into())
+        Err(CANNOT_ACCESS_HEADERS_ON_A_DEFERRED_RESPONSE.into())
     }
 
     #[rhai_fn(set = "headers", return_raw)]
@@ -907,7 +936,16 @@ macro_rules! gen_map_deferred_response {
                             );
                             if let Err(error) = result {
                                 tracing::error!("map_response callback failed: {error}");
-                                return None;
+                                let error_details = process_error(error);
+                                let mut guard = shared_response.lock().unwrap();
+                                let response_opt = guard.take();
+                                let $rhai_deferred_response { mut response, .. } = response_opt.unwrap();
+                                let error = Error {
+                                    message: error_details.message.unwrap_or_default(),
+                                    ..Default::default()
+                                };
+                                response.errors = vec![error];
+                                return Some(response);
                             }
 
                             let mut guard = shared_response.lock().unwrap();

apollo-router/src/plugins/rhai/tests.rs

@@ -407,6 +407,14 @@ async fn it_can_process_supergraph_response() {
     gen_response_test!(RhaiSupergraphResponse, "process_supergraph_response");
 }
 
+#[tokio::test]
+async fn it_can_process_supergraph_response_is_primary() {
+    gen_response_test!(
+        RhaiSupergraphResponse,
+        "process_supergraph_response_is_primary"
+    );
+}
+
 #[tokio::test]
 async fn it_can_process_supergraph_deferred_response() {
     gen_response_test!(
@@ -415,16 +423,40 @@ async fn it_can_process_supergraph_deferred_response() {
     );
 }
 
+#[tokio::test]
+async fn it_can_process_supergraph_deferred_response_is_not_primary() {
+    gen_response_test!(
+        RhaiSupergraphDeferredResponse,
+        "process_supergraph_deferred_response_is_not_primary"
+    );
+}
+
 #[tokio::test]
 async fn it_can_process_execution_response() {
     gen_response_test!(RhaiExecutionResponse, "process_execution_response");
 }
 
+#[tokio::test]
+async fn it_can_process_execution_response_is_primary() {
+    gen_response_test!(
+        RhaiExecutionResponse,
+        "process_execution_response_is_primary"
+    );
+}
+
 #[tokio::test]
 async fn it_can_process_execution_deferred_response() {
     gen_response_test!(RhaiExecutionDeferredResponse, "process_execution_response");
 }
 
+#[tokio::test]
+async fn it_can_process_execution_deferred_response_is_not_primary() {
+    gen_response_test!(
+        RhaiExecutionDeferredResponse,
+        "process_execution_deferred_response_is_not_primary"
+    );
+}
+
 #[tokio::test]
 async fn it_can_process_subgraph_response() {
     let dyn_plugin: Box<dyn DynPlugin> = crate::plugin::plugins()

apollo-router/tests/fixtures/request_response_test.rhai

@@ -161,3 +161,40 @@ fn process_subgraph_response_om_missing_message(response) {
         status: 400,
     };
 }
+
+fn process_supergraph_response_is_primary(response) {
+    if !response.is_primary() {
+        throw #{
+            status: 500,
+            message: "should be primary"
+        };
+    }
+}
+
+fn process_execution_response_is_primary(response) {
+    if !response.is_primary() {
+        throw #{
+            status: 500,
+            message: "should be primary"
+        };
+    }
+}
+
+fn process_supergraph_deferred_response_is_not_primary(response) {
+    if response.is_primary() {
+        throw #{
+            status: 500,
+            message: "should not be primary"
+        };
+    }
+}
+
+fn process_execution_deferred_response_is_not_primary(response) {
+    if response.is_primary() {
+        throw #{
+            status: 500,
+            message: "should not be primary"
+        };
+    }
+}
+

docs/source/customizations/rhai-api.mdx

@@ -435,6 +435,16 @@ The following fields are identical in behavior to their `request` counterparts:
 * [`headers`](#requestheaders)
 * [`body.extensions`](#requestbodyextensions)
 
+Note: Be particularly careful when interacting with headers in a response context. For supergraph_service() and execution_service(), response headers only exist for the first response in a deferred response stream. You can handle this by making use of the `is_primary()` function which will return true if a response is the first (or primary) response. If you do try to access the headers in a non-primary response, then you'll raise an exception which can be handled like any other rhai exception, but is not so convenient as using the `is_primary()` method.
+
+```rhai
+    if response.is_primary() {
+        print(`all response headers: ${response.headers}`);
+    } else {
+        print(`don't try to access headers`);
+    }
+```
+
 Other fields are described below.
 
 ### `response.body.label`