From 21e94d965451ae91bed65a15d8e6f305f5a4098e Mon Sep 17 00:00:00 2001 From: Lenny Burdette Date: Mon, 13 Mar 2023 11:33:43 -0400 Subject: [PATCH] cache-control example uses cache-control headers from subgraphs to determine an overall cache-control policy partially addresses #326 --- Cargo.lock | 12 ++ Cargo.toml | 3 +- apollo-router/src/plugins/external_tests.rs | 2 + apollo-router/src/plugins/rhai/mod.rs | 1 + .../src/plugins/traffic_shaping/cache.rs | 1 + apollo-router/src/services/subgraph.rs | 8 +- apollo-router/src/services/supergraph.rs | 26 +-- apollo-router/src/test_harness.rs | 1 + .../tests/fixtures/request_response_test.rhai | 20 +-- examples/cache-control/rhai/Cargo.toml | 13 ++ examples/cache-control/rhai/README.md | 9 ++ examples/cache-control/rhai/router.yaml | 3 + .../cache-control/rhai/src/cache_control.rhai | 80 ++++++++++ examples/cache-control/rhai/src/main.rs | 151 ++++++++++++++++++ 14 files changed, 306 insertions(+), 24 deletions(-) create mode 100644 examples/cache-control/rhai/Cargo.toml create mode 100644 examples/cache-control/rhai/README.md create mode 100644 examples/cache-control/rhai/router.yaml create mode 100644 examples/cache-control/rhai/src/cache_control.rhai create mode 100644 examples/cache-control/rhai/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 72a364de227..d1314ea22e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -755,6 +755,18 @@ dependencies = [ "either", ] +[[package]] +name = "cache-control" +version = "0.1.0" +dependencies = [ + "anyhow", + "apollo-router", + "http", + "serde_json", + "tokio", + "tower", +] + [[package]] name = "cargo-scaffold" version = "0.8.7" diff --git a/Cargo.toml b/Cargo.toml index bd93f1b16e2..f2627035411 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ members = [ "apollo-router-scaffold", "examples/add-timestamp-header/rhai", "examples/async-auth/rust", + "examples/cache-control/rhai", "examples/context/rust", "examples/cookies-to-headers/rhai", "examples/data-response-mutate/rhai", @@ -35,7 +36,7 @@ strip = "debuginfo" incremental = false # If building a dhat feature, you must use this profile -# e.g. heap allocation tracing: cargo build --profile release-dhat --features dhat-heap +# e.g. heap allocation tracing: cargo build --profile release-dhat --features dhat-heap # e.g. heap and ad-hoc allocation tracing: cargo build --profile release-dhat --features dhat-heap,dhat-ad-hoc [profile.release-dhat] inherits = "release" diff --git a/apollo-router/src/plugins/external_tests.rs b/apollo-router/src/plugins/external_tests.rs index 0e081791cc6..d7aa37ce07b 100644 --- a/apollo-router/src/plugins/external_tests.rs +++ b/apollo-router/src/plugins/external_tests.rs @@ -377,6 +377,7 @@ mod tests { .errors(Vec::new()) .extensions(crate::json_ext::Object::new()) .context(req.context) + .headers(HeaderMap::new()) .build()) }); @@ -532,6 +533,7 @@ mod tests { .errors(Vec::new()) .extensions(crate::json_ext::Object::new()) .context(req.context) + .headers(HeaderMap::new()) .build()) }); diff --git a/apollo-router/src/plugins/rhai/mod.rs b/apollo-router/src/plugins/rhai/mod.rs index dc951e76d83..64b1f6a25f0 100644 --- a/apollo-router/src/plugins/rhai/mod.rs +++ b/apollo-router/src/plugins/rhai/mod.rs @@ -641,6 +641,7 @@ macro_rules! gen_map_request { .and_data(body.data) .and_label(body.label) .and_path(body.path) + .headers(HeaderMap::new()) .build() } else { $base::Response::error_builder() diff --git a/apollo-router/src/plugins/traffic_shaping/cache.rs b/apollo-router/src/plugins/traffic_shaping/cache.rs index f7e6eee4cb1..eab910e3532 100644 --- a/apollo-router/src/plugins/traffic_shaping/cache.rs +++ b/apollo-router/src/plugins/traffic_shaping/cache.rs @@ -179,6 +179,7 @@ where .data(data) .extensions(Object::new()) .context(request.context) + .headers(http::HeaderMap::new()) .build()) } } diff --git a/apollo-router/src/services/subgraph.rs b/apollo-router/src/services/subgraph.rs index 1b9ea53a7cc..7e00b220878 100644 --- a/apollo-router/src/services/subgraph.rs +++ b/apollo-router/src/services/subgraph.rs @@ -138,6 +138,7 @@ impl Response { extensions: Object, status_code: Option, context: Context, + headers: http::HeaderMap, ) -> Response { // Build a response let res = graphql::Response::builder() @@ -149,11 +150,13 @@ impl Response { .build(); // Build an http Response - let response = http::Response::builder() + let mut response = http::Response::builder() .status(status_code.unwrap_or(StatusCode::OK)) .body(res) .expect("Response is serializable; qed"); + *response.headers_mut() = headers; + Self { response, context } } @@ -172,6 +175,7 @@ impl Response { extensions: JsonMap, status_code: Option, context: Option, + headers: Option>, ) -> Response { Response::new( label, @@ -181,6 +185,7 @@ impl Response { extensions, status_code, context.unwrap_or_default(), + headers.unwrap_or_default(), ) } @@ -201,6 +206,7 @@ impl Response { Default::default(), status_code, context, + Default::default(), )) } } diff --git a/apollo-router/src/services/supergraph.rs b/apollo-router/src/services/supergraph.rs index f28f5094a13..cc4560fc98f 100644 --- a/apollo-router/src/services/supergraph.rs +++ b/apollo-router/src/services/supergraph.rs @@ -131,29 +131,31 @@ impl Request { /// Create a request with an example query, for tests #[builder(visibility = "pub")] fn canned_new( + query: Option, operation_name: Option, // Skip the `Object` type alias in order to use buildstructor’s map special-casing extensions: JsonMap, context: Option, headers: MultiMap, ) -> Result { - let query = " - query TopProducts($first: Int) { - topProducts(first: $first) { - upc - name - reviews { - id - product { name } - author { id name } - } - } + let default_query = " + query TopProducts($first: Int) { + topProducts(first: $first) { + upc + name + reviews { + id + product { name } + author { id name } + } + } } "; + let query = query.unwrap_or(default_query.to_string()); let mut variables = JsonMap::new(); variables.insert("first", 2_usize.into()); Self::fake_new( - Some(query.to_owned()), + Some(query), operation_name, variables, extensions, diff --git a/apollo-router/src/test_harness.rs b/apollo-router/src/test_harness.rs index e49cbe9e17d..93a82446dcd 100644 --- a/apollo-router/src/test_harness.rs +++ b/apollo-router/src/test_harness.rs @@ -214,6 +214,7 @@ impl<'a> TestHarness<'a> { let empty_response = subgraph::Response::builder() .extensions(crate::json_ext::Object::new()) .context(request.context) + .headers(http::HeaderMap::new()) .build(); std::future::ready(Ok(empty_response)) }) diff --git a/apollo-router/tests/fixtures/request_response_test.rhai b/apollo-router/tests/fixtures/request_response_test.rhai index 87a807e5dc4..1b093272898 100644 --- a/apollo-router/tests/fixtures/request_response_test.rhai +++ b/apollo-router/tests/fixtures/request_response_test.rhai @@ -41,16 +41,16 @@ fn process_supergraph_request(request) { } let expected_query = ` - query TopProducts($first: Int) { - topProducts(first: $first) { - upc - name - reviews { - id - product { name } - author { id name } - } - } + query TopProducts($first: Int) { + topProducts(first: $first) { + upc + name + reviews { + id + product { name } + author { id name } + } + } } `; if request.body.query != expected_query { diff --git a/examples/cache-control/rhai/Cargo.toml b/examples/cache-control/rhai/Cargo.toml new file mode 100644 index 00000000000..70039854a65 --- /dev/null +++ b/examples/cache-control/rhai/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "cache-control" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[dependencies] +anyhow = "1" +apollo-router = { path = "../../../apollo-router" } +http = "0.2" +serde_json = "1" +tokio = { version = "1", features = ["full"] } +tower = { version = "0.4", features = ["full"] } diff --git a/examples/cache-control/rhai/README.md b/examples/cache-control/rhai/README.md new file mode 100644 index 00000000000..980126d1f5b --- /dev/null +++ b/examples/cache-control/rhai/README.md @@ -0,0 +1,9 @@ +# Rhai script + +Demonstrates header and context manipulation via Rhai script. + +Usage: + +```bash +cargo run -- -s ../../graphql/supergraph.graphql -c ./router.yaml +``` diff --git a/examples/cache-control/rhai/router.yaml b/examples/cache-control/rhai/router.yaml new file mode 100644 index 00000000000..f1e78d88312 --- /dev/null +++ b/examples/cache-control/rhai/router.yaml @@ -0,0 +1,3 @@ +rhai: + scripts: src + main: cache_control.rhai diff --git a/examples/cache-control/rhai/src/cache_control.rhai b/examples/cache-control/rhai/src/cache_control.rhai new file mode 100644 index 00000000000..9cee2d14edd --- /dev/null +++ b/examples/cache-control/rhai/src/cache_control.rhai @@ -0,0 +1,80 @@ +fn subgraph_service(service, subgraph) { + // collect the max-age and scope values from cache-control headers and store + // on the context for use in supergraph_service + service.map_response(|response| { + let cache_control = response.headers.values("cache-control").get(0); + + // if a subgraph response is uncacheable, the whole response is uncacheable + if cache_control == () { + response.context.cache_control_uncacheable = true; + return; + } + + let max_age = get_max_age(cache_control); + + // use the smallest max age + response.context.upsert("cache_control_max_age", |current| { + if current == () { + max_age + } else if max_age < current { + max_age + } else { + current + } + }); + + let scope = if cache_control.contains("public") { + "public" + } else { + "private" + }; + + // if the scope is ever private, it cannot become public + response.context.upsert("cache_control_scope", |current| { + if current == "private" || scope == "private" { + "private" + } else { + scope + } + }); + }); +} + +fn supergraph_service(service) { + // attach the cache-control header if enough data is available + service.map_response(|response| { + let uncacheable = response.context.cache_control_uncacheable; + let max_age = response.context.cache_control_max_age; + let scope = response.context.cache_control_scope; + + if uncacheable != true && max_age != () && scope != () { + response.headers["cache-control"] = `max-age=${max_age}, ${scope}`; + } + }); +} + +// find the the max-age= part and parse the value into an integer +fn get_max_age(str) { + let max_age = 0; + + for part in str.split(",") { + part.remove(" "); + + if part.starts_with("max-age=") { + let num = part.split("=").get(1); + + if num == () || num == "" { + break; + } + + try { + max_age = num.parse_int(); + } catch (err) { + log_error(`error parsing max-age from "${str}": ${err}`); + } + break; + } + } + + max_age +} diff --git a/examples/cache-control/rhai/src/main.rs b/examples/cache-control/rhai/src/main.rs new file mode 100644 index 00000000000..14d1d53e0bd --- /dev/null +++ b/examples/cache-control/rhai/src/main.rs @@ -0,0 +1,151 @@ +use anyhow::Result; + +// `cargo run -- -s ../../graphql/supergraph.graphql -c ./router.yaml` +fn main() -> Result<()> { + apollo_router::main() +} + +#[cfg(test)] +mod tests { + use apollo_router::graphql; + use apollo_router::plugin::test; + use apollo_router::services::subgraph; + use apollo_router::services::supergraph; + use http::HeaderMap; + use http::StatusCode; + use tower::util::ServiceExt; + + async fn cache_control_header( + header_one: Option, + header_two: Option, + ) -> Option { + let mut mock_service1 = test::MockSubgraphService::new(); + let mut mock_service2 = test::MockSubgraphService::new(); + + mock_service1.expect_clone().return_once(|| { + let mut mock_service = test::MockSubgraphService::new(); + mock_service + .expect_call() + .once() + .returning(move |req: subgraph::Request| { + let mut headers = HeaderMap::new(); + if let Some(value) = &header_one { + headers.insert("cache-control", value.parse().unwrap()); + } + + Ok(subgraph::Response::fake_builder() + .headers(headers) + .context(req.context) + .build()) + }); + mock_service + }); + + mock_service2.expect_clone().return_once(move || { + let mut mock_service = test::MockSubgraphService::new(); + mock_service + .expect_call() + .once() + .returning(move |req: subgraph::Request| { + let mut headers = HeaderMap::new(); + if let Some(value) = &header_two { + headers.insert("cache-control", value.parse().unwrap()); + } + + Ok(subgraph::Response::fake_builder() + .headers(headers) + .context(req.context) + .build()) + }); + mock_service + }); + + let config = serde_json::json!({ + "rhai": { + "scripts": "src", + "main": "cache_control.rhai", + } + }); + + let test_harness = apollo_router::TestHarness::builder() + .configuration_json(config) + .unwrap() + .subgraph_hook(move |name, _| match name { + "accounts" => mock_service1.clone().boxed(), + _ => mock_service2.clone().boxed(), + }) + // .log_level("DEBUG") + .build_router() + .await + .unwrap(); + + let query = " + query TopProducts { + me { id } + topProducts { name } + } + "; + + let request = supergraph::Request::fake_builder() + .query(query) + .build() + .expect("a valid SupergraphRequest"); + + let mut service_response = test_harness + .oneshot(request.try_into().unwrap()) + .await + .unwrap(); + + let response: graphql::Response = serde_json::from_slice( + service_response + .next_response() + .await + .unwrap() + .unwrap() + .to_vec() + .as_slice(), + ) + .unwrap(); + assert_eq!(response.errors, []); + + assert_eq!(StatusCode::OK, service_response.response.status()); + + service_response + .response + .headers() + .get("cache-control") + .map(|v| v.to_str().expect("can parse header value").to_string()) + } + + #[tokio::test] + async fn test_subgraph_cache_control() { + assert_eq!( + cache_control_header( + Some("max-age=100, private".to_string()), + Some("max-age=50, public".to_string()) + ) + .await, + Some("max-age=50, private".to_owned()) + ); + } + + #[tokio::test] + async fn test_subgraph_cache_control_public() { + assert_eq!( + cache_control_header( + Some("max-age=100, public".to_string()), + Some("max-age=50, public".to_string()) + ) + .await, + Some("max-age=50, public".to_owned()) + ); + } + + #[tokio::test] + async fn test_subgraph_cache_control_missing() { + assert_eq!( + cache_control_header(Some("max-age=100, private".to_string()), None).await, + None + ); + } +}