Add support for system adapter execution modes and JSON-RPC dispatching (#8)

* Add support for system adapter execution modes and JSON-RPC dispatching

* Refactor code formatting and improve readability in mod.rs

Author: lukekim

README.md

@@ -172,6 +172,11 @@ To benchmark a new platform, implement the adapter interface:
 
 - **stdio transport**: use `--system-adapter-stdio-cmd` (spicebench starts the child process).
 - **HTTP transport**: use `--system-adapter-http-url` (spicebench connects to a remote adapter endpoint).
+- **execution mode**: use `--system-adapter-execution-mode`:
+    - `adapter-command` (default): dispatches `spicebench run ...` to adapter JSON-RPC `run.load`
+    - `direct-query`: spicebench runs load/query path directly, while still connecting to adapter
+
+When adapter transport is configured, `spicebench run ...` is dispatched to spidapter method `run.load` over JSON-RPC and the adapter's `stdout`/`stderr` are streamed back.
 
 #### Stdio example (child process started by spicebench)
 
@@ -202,7 +207,18 @@ Notes:
 - Set **exactly one** of `--system-adapter-stdio-cmd` or `--system-adapter-http-url`.
 - `--system-adapter-stdio-args` passes CLI args to the stdio adapter command.
 - `--system-adapter-env` is only valid for stdio transport.
-- On connection, spicebench issues JSON-RPC `rpc.methods` to verify the adapter is reachable.
+- Spicebench validates the adapter supports JSON-RPC method `run.load` before dispatch.
+
+#### Direct-query example (run in spicebench, keep adapter connected)
+
+```bash
+spicebench run \
+    --query-set tpch \
+    --spicepod-path ./spicepod.yaml \
+    --system-adapter-name spidapter \
+    --system-adapter-execution-mode direct-query \
+    --system-adapter-http-url http://127.0.0.1:8080/jsonrpc
+```
 
 ## License
 

src/args/mod.rs

@@ -16,7 +16,7 @@ limitations under the License.
 
 use std::path::PathBuf;
 
-use clap::{ArgAction, Parser, Subcommand};
+use clap::{ArgAction, Parser, Subcommand, ValueEnum};
 
 mod dataset;
 pub use dataset::{DatasetTestArgs, LoadTestArgs};
@@ -82,6 +82,12 @@ pub struct CommonArgs {
     #[arg(long, default_value = "system_adapter")]
     pub(crate) system_adapter_name: String,
 
+    /// How to execute when a system adapter transport is configured.
+    /// - adapter-command: dispatch spicebench run as a JSON-RPC command (e.g. run.load)
+    /// - direct-query: execute load/query path in spicebench directly
+    #[arg(long, value_enum, default_value = "adapter-command")]
+    pub(crate) system_adapter_execution_mode: SystemAdapterExecutionMode,
+
     /// Command to run for a stdio JSON-RPC system adapter.
     #[arg(long, conflicts_with = "system_adapter_http_url")]
     pub(crate) system_adapter_stdio_cmd: Option<String>,
@@ -124,3 +130,9 @@ fn parse_key_val(s: &str) -> Result<(String, String), String> {
         .ok_or_else(|| "expected KEY=VALUE formatted header".to_string())?;
     Ok((s[..pos].to_string(), s[pos + 1..].to_string()))
 }
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
+pub enum SystemAdapterExecutionMode {
+    AdapterCommand,
+    DirectQuery,
+}

src/commands/mod.rs

@@ -14,9 +14,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use std::{collections::BTreeMap, sync::Arc, time::Duration};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+    time::Duration,
+};
 
-use crate::args::{CommonArgs, DatasetTestArgs};
+use crate::args::{CommonArgs, DatasetTestArgs, SystemAdapterExecutionMode};
 use test_framework::{
     anyhow,
     anyhow::Context,
@@ -31,7 +35,7 @@ use test_framework::{
 };
 use tokio::{
     io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
-    process::{ChildStdin, ChildStdout, Command},
+    process::{Child, ChildStdin, ChildStdout, Command},
 };
 
 pub(crate) mod load;
@@ -102,16 +106,6 @@ pub(crate) async fn build_test_with_validation(
 pub(crate) async fn run_or_connect_spiced(
     args: &CommonArgs,
 ) -> anyhow::Result<(App, SpicedInstance)> {
-    if let Some(mut adapter) = SystemAdapterClient::connect(args).await? {
-        let methods = adapter.rpc_methods().await?;
-        println!(
-            "Connected to system adapter '{}' via {} ({} methods)",
-            args.system_adapter_name,
-            adapter.transport_name(),
-            methods.len()
-        );
-    }
-
     let (app, mut instance) = if args.is_external_instance() {
         println!(
             "Connecting to external spiced instance at: {}",
@@ -170,8 +164,154 @@ pub(crate) async fn get_app_and_start_request(
     Ok((app, start_request))
 }
 
+pub(crate) async fn maybe_dispatch_run_to_system_adapter(
+    raw_cli_args: &[String],
+    common_args: &CommonArgs,
+) -> anyhow::Result<bool> {
+    if !has_system_adapter_transport(common_args) {
+        return Ok(false);
+    }
+
+    let mut adapter = SystemAdapterClient::connect(common_args)
+        .await?
+        .context("System adapter transport was configured but could not be initialized")?;
+
+    let methods = adapter.rpc_methods().await?;
+
+    if common_args.system_adapter_execution_mode == SystemAdapterExecutionMode::DirectQuery {
+        println!(
+            "Connected to system adapter '{}' via {} in direct-query mode (spicebench executes query/load path directly)",
+            common_args.system_adapter_name,
+            adapter.transport_name(),
+        );
+        return Ok(false);
+    }
+
+    let Some(method) = resolve_system_adapter_method(raw_cli_args) else {
+        anyhow::bail!(
+            "No JSON-RPC adapter method mapping for current command invocation: {:?}",
+            raw_cli_args
+        );
+    };
+
+    if !methods.iter().any(|available| available == method) {
+        anyhow::bail!(
+            "System adapter '{}' via {} does not support required method '{method}'",
+            common_args.system_adapter_name,
+            adapter.transport_name(),
+        );
+    }
+
+    let adapter_args = adapter_cli_args_for_run(raw_cli_args);
+    let mut params = serde_json::Map::new();
+    params.insert("args".to_string(), serde_json::to_value(adapter_args)?);
+
+    for (key, value) in &common_args.system_adapter_param {
+        params.insert(key.clone(), serde_json::Value::String(value.clone()));
+    }
+
+    let request = serde_json::json!({
+        "jsonrpc": "2.0",
+        "id": 2,
+        "method": method,
+        "params": serde_json::Value::Object(params),
+    });
+
+    let response = adapter.call(request).await?;
+    handle_adapter_execution_response(&response)?;
+
+    Ok(true)
+}
+
+fn resolve_system_adapter_method(raw_cli_args: &[String]) -> Option<&'static str> {
+    match raw_cli_args.first().map(String::as_str) {
+        Some("run") => Some("run.load"),
+        _ => None,
+    }
+}
+
+fn has_system_adapter_transport(args: &CommonArgs) -> bool {
+    args.system_adapter_stdio_cmd.is_some() || args.system_adapter_http_url.is_some()
+}
+
+fn adapter_cli_args_for_run(raw_cli_args: &[String]) -> Vec<String> {
+    let mut filtered = Vec::new();
+    let mut skip_next = false;
+
+    for (index, arg) in raw_cli_args.iter().enumerate() {
+        if index == 0 && arg == "run" {
+            continue;
+        }
+
+        if skip_next {
+            skip_next = false;
+            continue;
+        }
+
+        let takes_value = [
+            "--system-adapter-name",
+            "--system-adapter-execution-mode",
+            "--system-adapter-stdio-cmd",
+            "--system-adapter-stdio-args",
+            "--system-adapter-http-url",
+            "--system-adapter-param",
+            "--system-adapter-env",
+        ];
+
+        if takes_value.contains(&arg.as_str()) {
+            skip_next = true;
+            continue;
+        }
+
+        if takes_value
+            .iter()
+            .any(|flag| arg.starts_with(&format!("{flag}=")))
+        {
+            continue;
+        }
+
+        filtered.push(arg.clone());
+    }
+
+    filtered
+}
+
+fn handle_adapter_execution_response(response: &serde_json::Value) -> anyhow::Result<()> {
+    let result = response
+        .get("result")
+        .context("System adapter response missing JSON-RPC result payload")?;
+
+    if let Some(stdout) = result.get("stdout").and_then(|v| v.as_str())
+        && !stdout.is_empty()
+    {
+        print!("{stdout}");
+    }
+
+    if let Some(stderr) = result.get("stderr").and_then(|v| v.as_str())
+        && !stderr.is_empty()
+    {
+        eprint!("{stderr}");
+    }
+
+    let success = result
+        .get("success")
+        .and_then(|v| v.as_bool())
+        .unwrap_or(false);
+    let exit_code = result
+        .get("exit_code")
+        .and_then(|v| v.as_i64())
+        .unwrap_or(1);
+
+    if !success || exit_code != 0 {
+        anyhow::bail!("System adapter command failed (success={success}, exit_code={exit_code})");
+    }
+
+    Ok(())
+}
+
 enum SystemAdapterClient {
     Stdio {
+        child: Child,
         stdin: ChildStdin,
         stdout: BufReader<ChildStdout>,
     },
@@ -238,9 +378,8 @@ impl SystemAdapterClient {
                 .take()
                 .context("System adapter stdio child missing stdout")?;
 
-            std::mem::forget(child);
-
             return Ok(Some(Self::Stdio {
+                child,
                 stdin,
                 stdout: BufReader::new(stdout),
             }));
@@ -283,7 +422,11 @@ impl SystemAdapterClient {
 
     async fn call(&mut self, request: serde_json::Value) -> anyhow::Result<serde_json::Value> {
         match self {
-            Self::Stdio { stdin, stdout } => {
+            Self::Stdio {
+                child: _,
+                stdin,
+                stdout,
+            } => {
                 let payload = serde_json::to_string(&request)?;
                 stdin.write_all(payload.as_bytes()).await?;
                 stdin.write_all(b"\n").await?;

src/main.rs

@@ -37,10 +37,19 @@ async fn main() -> anyhow::Result<()> {
     let _ = rustls::crypto::CryptoProvider::install_default(
         rustls::crypto::aws_lc_rs::default_provider(),
     );
+    let raw_cli_args: Vec<String> = std::env::args().skip(1).collect();
     let cli = Cli::parse();
 
     match cli.subcommand {
-        Commands::Run(args) => commands::load::run(&args).await?,
+        Commands::Run(args) => {
+            if commands::maybe_dispatch_run_to_system_adapter(&raw_cli_args, &args.test_args.common)
+                .await?
+            {
+                return Ok(());
+            }
+
+            commands::load::run(&args).await?
+        }
     }
 
     Ok(())