diff --git a/Cargo.lock b/Cargo.lock index 2fc5212..60dd5b2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -123,6 +123,16 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7d5a26814d8dcb93b0e5a0ff3c6d80a8843bafb21b39e8e18a6f05471870e110" +[[package]] +name = "ariadne" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44055e597c674aef7cb903b2b9f6e4cba1277ed0d2d61dae7cd52d7ffa81f8e2" +dependencies = [ + "unicode-width", + "yansi", +] + [[package]] name = "assert_cmd" version = "2.0.16" @@ -192,6 +202,12 @@ version = "0.21.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" +[[package]] +name = "beef" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a8241f3ebb85c056b509d4327ad0358fbbba6ffb340bf388f26350aeda225b1" + [[package]] name = "bitflags" version = "1.3.2" @@ -213,6 +229,44 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bluejay-core" +version = "0.1.0" +source = "git+https://github.com/Shopify/bluejay.git?rev=c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf#c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf" +dependencies = [ + "enum-as-inner", + "itertools 0.13.0", + "paste", + "serde_json", + "strum", +] + +[[package]] +name = "bluejay-parser" +version = "0.1.0" +source = "git+https://github.com/Shopify/bluejay.git?rev=c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf#c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf" +dependencies = [ + "ariadne", + "bluejay-core", + "enum-as-inner", + "itertools 0.13.0", + "logos", + "strum", +] + +[[package]] +name = "bluejay-validator" +version = "0.1.0" +source = "git+https://github.com/Shopify/bluejay.git?rev=c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf#c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf" +dependencies = [ + "bluejay-core", + "bluejay-parser", + "itertools 0.13.0", + "paste", + "seq-macro", + "serde_json", +] + [[package]] name = "bstr" version = "1.10.0" @@ -524,7 +578,7 @@ dependencies = [ "cranelift-codegen", "cranelift-entity", "cranelift-frontend", - "itertools", + "itertools 0.12.1", "log", "smallvec", "wasmparser", @@ -682,6 +736,18 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "enum-as-inner" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ffccbb6966c05b32ef8fbac435df276c4ae4d3dc55a8cd0eb9745e6c12f546a" +dependencies = [ + "heck 0.4.1", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "equivalent" version = "1.0.1" @@ -741,6 +807,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "form_urlencoded" version = "1.2.1" @@ -768,6 +840,9 @@ dependencies = [ "anyhow", "assert_cmd", "assert_fs", + "bluejay-core", + "bluejay-parser", + "bluejay-validator", "clap", "colored", "deterministic-wasi-ctx", @@ -1079,6 +1154,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.11" @@ -1177,6 +1261,39 @@ dependencies = [ "serde_json", ] +[[package]] +name = "logos" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff1ceb190eb9bdeecdd8f1ad6a71d6d632a50905948771718741b5461fb01e13" +dependencies = [ + "logos-derive", +] + +[[package]] +name = "logos-codegen" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90be66cb7bd40cb5cc2e9cfaf2d1133b04a3d93b72344267715010a466e0915a" +dependencies = [ + "beef", + "fnv", + "lazy_static", + "proc-macro2", + "quote", + "regex-syntax", + "syn", +] + +[[package]] +name = "logos-derive" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45154231e8e96586b39494029e58f12f8ffcb5ecf80333a603a13aa205ea8cbd" +dependencies = [ + "logos-codegen", +] + [[package]] name = "mach2" version = "0.4.2" @@ -1598,6 +1715,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rustversion" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "955d28af4278de8121b7ebeb796b6a45735dc01436d898801014aced2773a3d6" + [[package]] name = "ryu" version = "1.0.18" @@ -1619,6 +1742,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61697e0a1c7e512e84a621326239844a24d8207b4669b41bc18b32ea5cbf988b" +[[package]] +name = "seq-macro" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" + [[package]] name = "serde" version = "1.0.206" @@ -1732,6 +1861,28 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "rustversion", + "syn", +] + [[package]] name = "syn" version = "2.0.72" @@ -2730,6 +2881,12 @@ dependencies = [ "wast 35.0.2", ] +[[package]] +name = "yansi" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" + [[package]] name = "zerocopy" version = "0.6.6" diff --git a/Cargo.toml b/Cargo.toml index b62da16..72ea586 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,9 @@ rust-embed = "8.5.0" rmp-serde = "1.3" is-terminal = "0.4.12" wasmprof = "0.7.0" +bluejay-parser = { git = "https://github.com/Shopify/bluejay.git", rev = "c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf", features = ["format-errors"] } +bluejay-core = { git = "https://github.com/Shopify/bluejay.git", rev = "c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf" } +bluejay-validator = { git = "https://github.com/Shopify/bluejay.git", rev = "c7e7c2bfb73c7b4869aa8569c15cd3c4eb48b8bf" } [dev-dependencies] assert_cmd = "2.0" diff --git a/src/bluejay_schema_analyzer.rs b/src/bluejay_schema_analyzer.rs new file mode 100644 index 0000000..54804bf --- /dev/null +++ b/src/bluejay_schema_analyzer.rs @@ -0,0 +1,281 @@ +use crate::scale_limits_analyzer::ScaleLimitsAnalyzer; +use anyhow::{anyhow, Result}; +use bluejay_parser::{ + ast::{ + definition::{DefinitionDocument, SchemaDefinition}, + executable::ExecutableDocument, + Parse, + }, + Error, +}; + +pub struct BluejaySchemaAnalyzer; + +impl BluejaySchemaAnalyzer { + pub fn analyze_schema_definition( + schema_string: &str, + schema_path: Option<&str>, + query: &str, + query_path: Option<&str>, + input: &serde_json::Value, + ) -> Result { + let document_definition = DefinitionDocument::parse(schema_string) + .map_err(|errors| anyhow!(Error::format_errors(schema_string, schema_path, errors)))?; + + let schema_definition = SchemaDefinition::try_from(&document_definition) + .map_err(|errors| anyhow!(Error::format_errors(schema_string, schema_path, errors)))?; + + let executable_document = ExecutableDocument::parse(query) + .map_err(|errors| anyhow!(Error::format_errors(query, query_path, errors)))?; + + let cache = + bluejay_validator::executable::Cache::new(&executable_document, &schema_definition); + + ScaleLimitsAnalyzer::analyze( + &executable_document, + &schema_definition, + None, + &Default::default(), + &cache, + input, + ) + .map_err(|e| anyhow!("Unable to analyze scale limits: {}", e.message())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn test_analyze_schema_definition() { + let schema_string = r#" + directive @scaleLimits(rate: Float!) on FIELD_DEFINITION + type Query { + field: String @scaleLimits(rate: 0.005) + } + "#; + let query = "{ field }"; + let input_json = json!({ + "field": "value" + }); + + let result = BluejaySchemaAnalyzer::analyze_schema_definition( + schema_string, + Some("schema.graphql"), + query, + Some("query.graphql"), + &input_json, + ); + assert!( + result.is_ok(), + "Expected successful analysis but got an error: {:?}", + result + ); + + let scale_factor = result.unwrap(); + let expected_scale_factor = 1.0; + assert_eq!( + scale_factor, expected_scale_factor, + "The scale factor did not match the expected value" + ); + } + + #[test] + fn test_analyze_schema_with_array_length_scaling() { + let schema_string = r#" + directive @scaleLimits(rate: Float!) on FIELD_DEFINITION + type Query { + cartLines: [String] @scaleLimits(rate: 0.005) + } + "#; + let query = "{ cartLines }"; + let input_json = json!({ + "cartLines": vec!["moeowomeow"; 500] + }); + + let result = BluejaySchemaAnalyzer::analyze_schema_definition( + schema_string, + Some("schema.graphql"), + query, + Some("query.graphql"), + &input_json, + ); + assert!( + result.is_ok(), + "Expected successful analysis but got an error: {:?}", + result + ); + + let scale_factor = result.unwrap(); + let expected_scale_factor = 2.5; // Adjust this based on how your scale limits are defined + assert_eq!( + scale_factor, expected_scale_factor, + "The scale factor did not match the expected value for array length scaling" + ); + } + + #[test] + fn test_analyze_schema_with_array_length_scaling_to_max_scale_factor() { + let schema_string = r#" + directive @scaleLimits(rate: Float!) on FIELD_DEFINITION + type Query { + cartLines: [String] @scaleLimits(rate: 0.005) + } + "#; + let query = "{ cartLines }"; + let input_json = json!({ + "cartLines": vec!["item"; 1000000] // value that would scale well beyond the max + }); + + let result = BluejaySchemaAnalyzer::analyze_schema_definition( + schema_string, + Some("schema.graphql"), + query, + Some("query.graphql"), + &input_json, + ); + assert!( + result.is_ok(), + "Expected successful analysis but got an error: {:?}", + result + ); + + let scale_factor = result.unwrap(); + let expected_scale_factor = 10.0; + assert_eq!( + scale_factor, expected_scale_factor, + "The scale factor did not match the expected value for array length scaling" + ); + } + + #[test] + fn test_invalid_schema() { + let invalid_schema_string = r#" + directive @scaleLimits(rate: Float!) on FIELD_DEFINITION + type Query { + field: String @scaleLimits(rate: "invalid") // Invalid rate type + } + "#; + let valid_query = "query { field }"; + let input_json = json!({ + "field": "value" + }); + + let result = BluejaySchemaAnalyzer::analyze_schema_definition( + invalid_schema_string, + Some("invalid_schema.graphql"), + valid_query, + Some("query.graphql"), + &input_json, + ); + + assert!( + result.is_err(), + "Expected an error due to invalid schema and query, but got success: {:?}", + result + ); + } + + #[test] + fn test_invalid_query() { + let schema_string = r#" + directive @scaleLimits(rate: Float!) on FIELD_DEFINITION + type Query { + field: String @scaleLimits(rate: 0.005) + } + "#; + let invalid_query = "query { field "; + let input_json = json!({ + "field": "value" + }); + + let result = BluejaySchemaAnalyzer::analyze_schema_definition( + schema_string, + Some("schema.graphql"), + invalid_query, + Some("invalid_query.graphql"), + &input_json, + ); + + assert!( + result.is_err(), + "Expected an error due to invalid schema and query, but got success: {:?}", + result + ); + } + + #[test] + fn test_no_double_counting_for_duplicate_fields_with_array() { + let schema_string = r#" + directive @scaleLimits(rate: Float!) on FIELD_DEFINITION + type Query { + field: [String] @scaleLimits(rate: 0.005) + } + "#; + let query = "{ field field }"; + let input_json = json!({ + "field": vec!["value"; 200] + }); + + let result = BluejaySchemaAnalyzer::analyze_schema_definition( + schema_string, + Some("schema.graphql"), + query, + Some("query.graphql"), + &input_json, + ); + assert!( + result.is_ok(), + "Expected successful analysis but got an error: {:?}", + result + ); + + let scale_factor = result.unwrap(); + let expected_scale_factor = 1.0; + assert_eq!( + scale_factor, expected_scale_factor, + "The scale factor did not match the expected value, indicating potential double counting" + ); + } + + #[test] + fn test_no_double_counting_for_duplicate_fields_with_nested_array() { + let schema_string = r#" + directive @scaleLimits(rate: Float!) on FIELD_DEFINITION + type Query { + field: [MyObject] + } + + type MyObject { + field: [String] @scaleLimits(rate: 0.005) + } + "#; + let query = "{ field { field } }"; + let nested_field = json!({ "field": vec!["value"; 200] }); + let input_json = json!({ + "field": vec![nested_field; 2] + }); + + let result = BluejaySchemaAnalyzer::analyze_schema_definition( + schema_string, + Some("schema.graphql"), + query, + Some("query.graphql"), + &input_json, + ); + assert!( + result.is_ok(), + "Expected successful analysis but got an error: {:?}", + result + ); + + let scale_factor = result.unwrap(); + let expected_scale_factor = 2.0; + assert_eq!( + scale_factor, expected_scale_factor, + "The scale factor did not match the expected value, indicating potential double counting" + ); + } +} diff --git a/src/engine.rs b/src/engine.rs index bec94cc..dda4685 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -53,6 +53,7 @@ pub struct FunctionRunParams<'a> { pub input: Vec, pub export: &'a str, pub profile_opts: Option<&'a ProfileOpts>, + pub scale_factor: f64, } const STARTING_FUEL: u64 = u64::MAX; @@ -114,6 +115,7 @@ pub fn run(params: FunctionRunParams) -> Result { input, export, profile_opts, + scale_factor, } = params; let engine = Engine::new( @@ -231,6 +233,7 @@ pub fn run(params: FunctionRunParams) -> Result { input: function_run_input, output, profile: profile_data, + scale_factor, }; Ok(function_run_result) diff --git a/src/function_run_result.rs b/src/function_run_result.rs index f9ce1af..78ab01b 100644 --- a/src/function_run_result.rs +++ b/src/function_run_result.rs @@ -28,15 +28,46 @@ pub struct FunctionRunResult { pub output: FunctionOutput, #[serde(skip)] pub profile: Option, + #[serde(skip)] + pub scale_factor: f64, +} + +const DEFAULT_INSTRUCTIONS_LIMIT: u64 = 11_000_000; +const DEFAULT_INPUT_SIZE_LIMIT: u64 = 64_000; +const DEFAULT_OUTPUT_SIZE_LIMIT: u64 = 20_000; + +pub fn get_json_size_as_bytes(value: &serde_json::Value) -> usize { + serde_json::to_vec(value).map(|v| v.len()).unwrap_or(0) } impl FunctionRunResult { pub fn to_json(&self) -> String { serde_json::to_string_pretty(&self).unwrap_or_else(|error| error.to_string()) } + + pub fn input_size(&self) -> usize { + get_json_size_as_bytes(&self.input) + } + + pub fn output_size(&self) -> usize { + match &self.output { + FunctionOutput::JsonOutput(value) => get_json_size_as_bytes(value), + FunctionOutput::InvalidJsonOutput(_value) => 0, + } + } } -fn humanize_instructions(instructions: u64) -> String { +fn humanize_size(size_bytes: u64, size_limit: u64, title: &str) -> String { + let size_kb = size_bytes as f64 / 1024.0; // Convert bytes to kilobytes + + if size_bytes > size_limit { + format!("{}: {:.2}KB", title, size_kb).red().to_string() + } else { + format!("{}: {:.2}KB", title, size_kb) + } +} + +fn humanize_instructions(instructions: u64, scale_factor: f64) -> String { let instructions_humanized = match instructions { 0..=999 => instructions.to_string(), 1000..=999_999 => format!("{}K", instructions as f64 / 1000.0), @@ -44,11 +75,14 @@ fn humanize_instructions(instructions: u64) -> String { 1_000_000_000..=u64::MAX => format!("{}B", instructions as f64 / 1_000_000_000.0), }; - match instructions { - 0..=11_000_000 => format!("Instructions: {instructions_humanized}"), - 11_000_001..=u64::MAX => format!("Instructions: {instructions_humanized}") + let instructions_limit = DEFAULT_INSTRUCTIONS_LIMIT as f64 * scale_factor; + + if instructions > instructions_limit as u64 { + format!("Instructions: {}", instructions_humanized) .red() - .to_string(), + .to_string() + } else { + format!("Instructions: {}", instructions_humanized) } } @@ -107,15 +141,65 @@ impl fmt::Display for FunctionRunResult { } } + writeln!( + formatter, + "\n{}\n\n", + " Resource Limits " + .black() + .on_bright_magenta() + )?; + + writeln!( + formatter, + "Input Size: {:.2}KB", + (DEFAULT_INPUT_SIZE_LIMIT as f64) * self.scale_factor / 1024.0 + )?; + writeln!( + formatter, + "Output Size: {:.2}KB", + (DEFAULT_OUTPUT_SIZE_LIMIT as f64) * self.scale_factor / 1024.0 + )?; + writeln!( + formatter, + "Instructions: {:.2}M", + (DEFAULT_INSTRUCTIONS_LIMIT as f64) * self.scale_factor / 1_000_000.0 + )?; + let title = " Benchmark Results " .black() .on_truecolor(150, 191, 72); + let input_size_limit = self.scale_factor * DEFAULT_INPUT_SIZE_LIMIT as f64; + let output_size_limit = self.scale_factor * DEFAULT_OUTPUT_SIZE_LIMIT as f64; + write!(formatter, "\n\n{title}\n\n")?; writeln!(formatter, "Name: {}", self.name)?; writeln!(formatter, "Linear Memory Usage: {}KB", self.memory_usage)?; - writeln!(formatter, "{}", humanize_instructions(self.instructions))?; - writeln!(formatter, "Size: {}KB\n", self.size)?; + writeln!( + formatter, + "{}", + humanize_instructions(self.instructions, self.scale_factor) + )?; + writeln!( + formatter, + "{}", + humanize_size( + self.input_size() as u64, + input_size_limit as u64, + "Input Size", + ) + )?; + writeln!( + formatter, + "{}", + humanize_size( + self.output_size() as u64, + output_size_limit as u64, + "Output Size" + ) + )?; + + writeln!(formatter, "Module Size: {}KB\n", self.size)?; Ok(()) } @@ -145,11 +229,15 @@ mod tests { "test": "test" })), profile: None, + scale_factor: 1.0, }; let predicate = predicates::str::contains("Instructions: 1.001K") .and(predicates::str::contains("Linear Memory Usage: 1000KB")) - .and(predicates::str::contains(expected_input_display)); + .and(predicates::str::contains(expected_input_display)) + .and(predicates::str::contains("Input Size: 62.50KB")) + .and(predicates::str::contains("Output Size: 19.53KB")) + .and(predicates::str::contains("Instructions: 11.00M")); assert!(predicate.eval(&function_run_result.to_string())); Ok(()) @@ -172,6 +260,7 @@ mod tests { "test": "test" })), profile: None, + scale_factor: 1.0, }; let predicate = predicates::str::contains("Instructions: 1") @@ -198,6 +287,7 @@ mod tests { "test": "test" })), profile: None, + scale_factor: 1.0, }; let predicate = predicates::str::contains("Instructions: 999") diff --git a/src/lib.rs b/src/lib.rs index d54cd93..4650e09 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +pub mod bluejay_schema_analyzer; +pub mod scale_limits_analyzer; pub mod engine; pub mod function_run_result; pub mod logs; diff --git a/src/main.rs b/src/main.rs index 425af2d..d7c16fc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,12 +5,17 @@ use std::{ }; use anyhow::{anyhow, Result}; + use clap::{Parser, ValueEnum}; -use function_runner::engine::{run, FunctionRunParams, ProfileOpts}; +use function_runner::{ + bluejay_schema_analyzer::BluejaySchemaAnalyzer, + engine::{run, FunctionRunParams, ProfileOpts}, +}; use is_terminal::IsTerminal; const PROFILE_DEFAULT_INTERVAL: u32 = 500_000; // every 5us +const DEFAULT_SCALE_FACTOR: f64 = 1.0; /// Supported input flavors #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, ValueEnum)] @@ -60,6 +65,14 @@ struct Opts { #[clap(short = 'c', long, value_enum, default_value = "json")] codec: Codec, + + /// Path to graphql file containing Function schema; if omitted, defauls will be used to calculate limits. + #[clap(short = 's', long)] + schema_path: Option, + + /// Path to graphql file containing Function input query; if omitted, defauls will be used to calculate limits. + #[clap(short = 'q', long)] + query_path: Option, } impl Opts { @@ -89,6 +102,25 @@ impl Opts { path } + + pub fn read_schema_to_string(&self) -> Option> { + self.schema_path.as_ref().map(read_file_to_string) + } + + pub fn read_query_to_string(&self) -> Option> { + self.query_path.as_ref().map(read_file_to_string) + } +} + +fn read_file_to_string(file_path: &PathBuf) -> Result { + let mut file = File::open(file_path) + .map_err(|e| anyhow!("Couldn't open file {}: {}", file_path.to_string_lossy(), e))?; + + let mut contents = String::new(); + file.read_to_string(&mut contents) + .map_err(|e| anyhow!("Couldn't read file {}: {}", file_path.to_string_lossy(), e))?; + + Ok(contents) } fn main() -> Result<()> { @@ -109,27 +141,48 @@ fn main() -> Result<()> { let mut buffer = Vec::new(); input.read_to_end(&mut buffer)?; - let buffer = match opts.codec { + let schema_string = opts.read_schema_to_string().transpose()?; + + let query_string = opts.read_query_to_string().transpose()?; + + let (json_value, buffer) = match opts.codec { Codec::Json => { - let _ = serde_json::from_slice::(&buffer) + let json = serde_json::from_slice::(&buffer) .map_err(|e| anyhow!("Invalid input JSON: {}", e))?; - buffer + (Some(json), buffer) } - Codec::Raw => buffer, + Codec::Raw => (None, buffer), Codec::JsonToMessagepack => { let json: serde_json::Value = serde_json::from_slice(&buffer) .map_err(|e| anyhow!("Invalid input JSON: {}", e))?; - rmp_serde::to_vec(&json) - .map_err(|e| anyhow!("Couldn't convert JSON to MessagePack: {}", e))? + let bytes = rmp_serde::to_vec(&json) + .map_err(|e| anyhow!("Couldn't convert JSON to MessagePack: {}", e))?; + (Some(json), bytes) } }; + let scale_factor = if let (Some(schema_string), Some(query_string), Some(json_value)) = + (schema_string, query_string, json_value) + { + BluejaySchemaAnalyzer::analyze_schema_definition( + &schema_string, + opts.schema_path.as_ref().and_then(|p| p.to_str()), + &query_string, + opts.query_path.as_ref().and_then(|p| p.to_str()), + &json_value, + )? + } else { + DEFAULT_SCALE_FACTOR // Use default scale factor when schema or query is missing + }; + let profile_opts = opts.profile_opts(); + let function_run_result = run(FunctionRunParams { function_path: opts.function, input: buffer, export: opts.export.as_ref(), profile_opts: profile_opts.as_ref(), + scale_factor, })?; if opts.json { diff --git a/src/scale_limits_analyzer.rs b/src/scale_limits_analyzer.rs new file mode 100644 index 0000000..4a85cb1 --- /dev/null +++ b/src/scale_limits_analyzer.rs @@ -0,0 +1,169 @@ +use bluejay_core::{ + definition::{prelude::*, SchemaDefinition as CoreSchemaDefinition}, + AsIter, Directive, Value as CoreValue, ValueReference, +}; +use bluejay_parser::ast::{ + definition::FieldDefinition, + definition::{DefaultContext, SchemaDefinition}, + executable::ExecutableDocument, +}; +use serde_json::Value; +use std::collections::HashMap; + +pub type ScaleLimitsAnalyzer<'a> = bluejay_validator::executable::operation::Orchestrator< + 'a, + ExecutableDocument<'a>, + SchemaDefinition<'a>, + serde_json::Map, + ScaleLimits<'a>, +>; + +#[derive(Hash, PartialEq, Eq, Debug)] +struct PathWithIndex<'a> { + path: Vec<&'a str>, + index: usize, +} + +pub struct ScaleLimits<'a> { + value_stack: Vec>, + path_stack: Vec<&'a str>, + rates: HashMap, f64>, +} + +impl<'a> + bluejay_validator::executable::operation::Visitor< + 'a, + ExecutableDocument<'a>, + SchemaDefinition<'a>, + serde_json::Map, + > for ScaleLimits<'a> +{ + type ExtraInfo = &'a Value; + + fn new( + _operation_definition: &'a ::OperationDefinition, + _schema_definition: &'a SchemaDefinition<'a>, + _variable_values: &'a serde_json::Map, + _cache: &'a bluejay_validator::executable::Cache<'a, ExecutableDocument, SchemaDefinition>, + extra_info: &'a Value, + ) -> Self { + Self { + value_stack: vec![vec![extra_info]], + path_stack: Vec::new(), + rates: Default::default(), + } + } + + fn visit_field( + &mut self, + field: &'a as bluejay_core::executable::ExecutableDocument>::Field, + field_definition: &'_ ::FieldDefinition, + _scoped_type: bluejay_core::definition::TypeDefinitionReference< + '_, + as CoreSchemaDefinition>::TypeDefinition, + >, + _included: bool, + ) { + self.path_stack.push(field.response_key()); + let rate = Self::rate_for_field_definition(field_definition); + let values = self.value_stack.last().unwrap(); + let mut nested_values = Vec::new(); + + values.iter().enumerate().for_each(|(index, value)| { + let value_for_field = match value { + Value::Object(object) => object.get(field.response_key()), + Value::Null => None, + _ => None, + }; + if let Some(rate) = rate { + let length = match value_for_field { + Some(Value::String(s)) => s.len(), + Some(Value::Array(arr)) => arr.len(), + _ => 1, + }; + let increment = length as f64 * rate; + + let path_with_index = PathWithIndex { + path: self.path_stack.clone(), + index, + }; + + let entry = self.rates.entry(path_with_index).or_default(); + + *entry = entry.max(increment); + } + + match value_for_field { + Some(Value::Array(values)) => nested_values.extend(values), + Some(value) => nested_values.push(value), + None => {} + } + }); + + self.value_stack.push(nested_values); + } + + fn leave_field( + &mut self, + _field: &'a as bluejay_core::executable::ExecutableDocument>::Field, + _field_definition: &'a as CoreSchemaDefinition>::FieldDefinition, + _scoped_type: bluejay_core::definition::TypeDefinitionReference< + 'a, + as CoreSchemaDefinition>::TypeDefinition, + >, + _included: bool, + ) { + self.path_stack.pop().unwrap(); + self.value_stack.pop().unwrap(); + } +} + +impl<'a> + bluejay_validator::executable::operation::Analyzer< + 'a, + ExecutableDocument<'a>, + SchemaDefinition<'a>, + serde_json::Map, + > for ScaleLimits<'a> +{ + type Output = f64; + + fn into_output(self) -> Self::Output { + let normalized_rates = self.rates.into_iter().fold( + HashMap::new(), + |mut normalized_rates, (PathWithIndex { path, .. }, rate)| { + *normalized_rates.entry(path).or_default() += rate; + normalized_rates + }, + ); + + normalized_rates + .into_values() + .fold(Self::MIN_SCALE_FACTOR, f64::max) + .clamp(Self::MIN_SCALE_FACTOR, Self::MAX_SCALE_FACTOR) + } +} + +impl<'a> ScaleLimits<'a> { + const MIN_SCALE_FACTOR: f64 = 1.0; + const MAX_SCALE_FACTOR: f64 = 10.0; + + fn rate_for_field_definition( + field_definition: &FieldDefinition, + ) -> Option { + field_definition + .directives() + .iter() + .flat_map(|directives| directives.iter()) + .find(|directive| directive.name() == "scaleLimits") + .and_then(|directive| directive.arguments()) + .and_then(|arguments| arguments.iter().find(|argument| argument.name() == "rate")) + .and_then(|argument| { + if let ValueReference::Float(rate) = argument.value().as_ref() { + Some(rate) + } else { + None + } + }) + } +} diff --git a/tests/fixtures/query/query.graphql b/tests/fixtures/query/query.graphql new file mode 100644 index 0000000..ff0aefe --- /dev/null +++ b/tests/fixtures/query/query.graphql @@ -0,0 +1,7 @@ +query { + cart { + lines { + quantity + } + } +} diff --git a/tests/fixtures/schema/schema.graphql b/tests/fixtures/schema/schema.graphql new file mode 100644 index 0000000..75bdbc6 --- /dev/null +++ b/tests/fixtures/schema/schema.graphql @@ -0,0 +1,22 @@ +schema { + query: Query +} + +directive @scaleLimits(rate: Float!) on FIELD_DEFINITION + +type Attribute { + key: String! + value: String +} + +type Cart { + lines: [CartLine!]! @scaleLimits(rate: 0.005) +} + +type CartLine { + quantity: Int! +} + +type Query { + cart: Cart +} diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index dcf03a9..1d48846 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -192,6 +192,17 @@ mod tests { Ok(()) } + // output is expected + // pass in correct inputs + // add some with small input and use high scale rate (or just use a big one) + // test if there an easy way to see output and limits + + // errors test + // test 1: schema file does not exist, do we get right error + // test 2: query error + + // Nick: error or not in cases where one of query / schema are invalid or missing + #[test] fn exports() -> Result<(), Box> { let mut cmd = Command::cargo_bin("function-runner")?; @@ -246,4 +257,61 @@ mod tests { Ok(file) } + + #[test] + fn test_scale_limits_analyzer_with_missing_paths() -> Result<(), Box> { + let mut cmd = Command::cargo_bin("function-runner")?; + let input_file = temp_input(json!({"cart": { + "lines": [ + {"quantity": 2} + ] + }}))?; + + cmd.args(["--function", "tests/fixtures/build/exit_code.wasm"]) + .arg("--input") + .arg(input_file.as_os_str()); + cmd.assert().success(); + + cmd.assert() + .success() + .stdout(contains("Input Size: 62.50KB")) + .stdout(contains("Output Size: 19.53KB")) + .stdout(contains("Instructions: 11.00M")); + + Ok(()) + } + + #[test] + fn test_scale_limits_analyzer_with_scaled_limits() -> Result<(), Box> { + let mut cmd = Command::cargo_bin("function-runner")?; + + let input_data = vec![json!({"quantity": 2}); 400]; + let json_data = json!({ + "cart": { + "lines": input_data + } + }); + let input_file = temp_input(json_data)?; + + // Define paths to the schema, query, and input JSON files + let schema_path = "tests/fixtures/schema/yolo-schema_willnotmerge.graphql"; + let query_path = "tests/fixtures/query/query.graphql"; + + cmd.args(["--function", "tests/fixtures/build/exit_code.wasm"]) + .arg("--input") + .arg(input_file.as_os_str()) + .arg("--schema-path") + .arg(schema_path) + .arg("--query-path") + .arg(query_path); + cmd.assert().success(); + + cmd.assert() + .success() + .stdout(contains("Input Size: 125.00KB")) + .stdout(contains("Output Size: 39.06KB")) + .stdout(contains("Instructions: 22.00M")); + + Ok(()) + } }