From 4d59319d1775bcecc752dbe1038026f18848dc4c Mon Sep 17 00:00:00 2001 From: luo jiyin Date: Fri, 16 Oct 2026 15:08:26 +0800 Subject: [PATCH] perf(agent): share rpc response wait across result types --- .github/workflows/ci.yml | 40 ++++++++++++++++++++ crates/dbx-driver-agent/src/agent_driver.rs | 41 +++++++++++++++++++-- 2 files changed, 77 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 805df286a5..1eb0678acc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -653,6 +653,46 @@ jobs: if-no-files-found: error retention-days: 7 + - name: Measure Agent RPC monomorphization + if: github.head_ref == 'perf/win7-agent-rpc-monomorphization' + shell: pwsh + run: | + $outputDir = Join-Path $env:RUNNER_TEMP "dbx-agent-rpc-mono" + New-Item -ItemType Directory -Force -Path $outputDir | Out-Null + $env:TAURI_CONFIG = Get-Content src-tauri/tauri.webview2-win7-fixed.conf.json -Raw + $env:RUSTC_WRAPPER = "" + $timer = [System.Diagnostics.Stopwatch]::StartNew() + cargo -Z host-config -Z target-applies-to-host -Z build-std=std,panic_abort ` + --config .github/fixtures/win7-host-repro-config.toml ` + rustc --locked --package dbx --lib --release --features custom-protocol ` + --target x86_64-win7-windows-msvc -- ` + "-Zdump-mono-stats=$outputDir" "-Zdump-mono-stats-format=json" + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } + $timer.Stop() + $monoFile = Get-ChildItem $outputDir -File -Filter "*.mono_items.json" | Select-Object -First 1 + if (!$monoFile) { + throw "dbx_lib monomorphization data is missing." + } + $items = Get-Content $monoFile.FullName -Raw | ConvertFrom-Json + "dbx_lib profile: $([math]::Round($timer.Elapsed.TotalSeconds, 2)) seconds" >> $env:GITHUB_STEP_SUMMARY + foreach ($symbol in @("AgentDriverClient::call_typed_with_timeout_and_cancel::{closure#0}", "AgentDriverClient::call_value_with_timeout_and_cancel::{closure#0}", "AgentRuntimeClient::call_typed::{closure#0}")) { + $item = $items | Where-Object { $_.name.EndsWith($symbol) } | Select-Object -First 1 + if ($item) { + "$symbol`: $($item.instantiation_count) instances, estimate $($item.total_estimate)" >> $env:GITHUB_STEP_SUMMARY + } + } + + - name: Upload Agent RPC monomorphization data + if: github.head_ref == 'perf/win7-agent-rpc-monomorphization' + uses: actions/upload-artifact@v6 + with: + name: DBX-win7-agent-rpc-mono + path: ${{ runner.temp }}/dbx-agent-rpc-mono/*.mono_items.json + if-no-files-found: error + retention-days: 7 + - name: Measure dbx library command registry expansion if: vars.WIN7_REGISTRY_AB == 'true' shell: pwsh diff --git a/crates/dbx-driver-agent/src/agent_driver.rs b/crates/dbx-driver-agent/src/agent_driver.rs index b2a56f6096..4bc023ea47 100644 --- a/crates/dbx-driver-agent/src/agent_driver.rs +++ b/crates/dbx-driver-agent/src/agent_driver.rs @@ -1,3 +1,4 @@ +use std::any::Any; use std::collections::{HashMap, VecDeque}; use std::fmt; use std::io::{BufRead, BufReader, BufWriter, Write}; @@ -17,6 +18,14 @@ use crate::agent_recovery::{RecoveryDecision, RecoveryPolicy, RecoveryScope}; use crate::models::connection::DatabaseConnectionInfo; type PendingAgentResponse = tokio::sync::oneshot::Sender>; +type AgentDecodedResult = Box; +type AgentResultDecoder = fn(Value) -> Result; + +fn deserialize_agent_result(value: Value) -> Result { + serde_json::from_value::(value) + .map(|result| Box::new(result) as AgentDecodedResult) + .map_err(|error| format!("Failed to deserialize agent result: {error}")) +} #[derive(Clone)] struct CachedAgentQuery { @@ -1809,6 +1818,26 @@ impl AgentDriverClient { timeout_duration: Option, cancel_token: Option, ) -> Result { + let result = self + .call_decoded_with_timeout_and_cancel( + method, + params, + timeout_duration, + cancel_token, + deserialize_agent_result::, + ) + .await?; + Ok(*result.downcast::().expect("Agent result decoder returned a different type")) + } + + async fn call_decoded_with_timeout_and_cancel( + &mut self, + method: &str, + params: Value, + timeout_duration: Option, + cancel_token: Option, + decode: AgentResultDecoder, + ) -> Result { if let Some(runtime) = &self.shared_runtime { let agent_session_id = self.agent_session_id.clone(); let mut params = params; @@ -1822,10 +1851,15 @@ impl AgentDriverClient { .ok_or_else(|| "Agent RPC parameters must be an object".to_string())? .insert("agentSessionId".to_string(), Value::String(session_id.clone())); } - return runtime + let result: Value = runtime .call_typed(method, params, timeout_duration, cancel_token) .await - .map_err(|error| error.with_legacy_session_id(agent_session_id.as_deref())); + .map_err(|error| error.with_legacy_session_id(agent_session_id.as_deref()))?; + return decode(result).map_err(|message| AgentCallError::ContractViolation { + rpc_code: None, + message, + reason: ContractViolationReason::InvalidResult, + }); } self.next_id += 1; let id = self.next_id; @@ -1924,8 +1958,7 @@ impl AgentDriverClient { let result = if let Some(err) = resp.get("error") { Err(format_agent_rpc_error(err)) } else if let Some(result_val) = resp.get("result") { - serde_json::from_value::(result_val.clone()) - .map_err(|e| format!("Failed to deserialize agent result: {e}")) + decode(result_val.clone()) } else { Err(format!("Agent response missing both 'result' and 'error': {line}")) };