pd-edge: perf: decouple proxy forward wait from vm path

Author: fffonion

pd-edge/examples/http_proxy_perf_framework.rs

@@ -38,7 +38,6 @@ use tokio::{
 use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
 #[cfg(all(feature = "http2", feature = "tls"))]
 use tokio_rustls::{TlsAcceptor, rustls::ServerConfig};
-#[cfg(feature = "http3")]
 use url::Url;
 use vm::{compile_source, encode_program, validate_program};
 

pd-edge/src/abi_impl/http/state.rs

@@ -1867,6 +1867,7 @@ pub(crate) struct DownstreamState {
     pub(crate) post_response_plan: Option<DownstreamPostResponsePlan>,
     pub(crate) native_default_upstream_http_forward: bool,
     native_default_upstream_forward_response: Option<NativeDefaultUpstreamForwardResponse>,
+    native_default_upstream_forward_task: Option<NativeDefaultUpstreamForwardTask>,
     inline_http_response_sender: Option<InlineDownstreamHttpResponseSender>,
 }
 
@@ -1880,6 +1881,7 @@ impl DownstreamState {
             post_response_plan: None,
             native_default_upstream_http_forward: false,
             native_default_upstream_forward_response: None,
+            native_default_upstream_forward_task: None,
             inline_http_response_sender: None,
         }
     }
@@ -1918,6 +1920,7 @@ impl DownstreamState {
             post_response_plan: None,
             native_default_upstream_http_forward: false,
             native_default_upstream_forward_response: None,
+            native_default_upstream_forward_task: None,
             inline_http_response_sender: None,
         }
     }
@@ -2483,16 +2486,19 @@ impl ProxyVmContext {
     }
 
     pub(crate) fn clear_native_default_upstream_http_forward(&self) {
-        self.lock_downstream().native_default_upstream_http_forward = false;
+        let mut downstream = self.lock_downstream();
+        downstream.native_default_upstream_http_forward = false;
+        downstream.native_default_upstream_forward_response = None;
+        if let Some(task) = downstream.native_default_upstream_forward_task.take() {
+            task.abort();
+        }
     }
 
-    fn store_native_default_upstream_forward_response(
-        &self,
-        response: NativeDefaultUpstreamForwardResponse,
-    ) {
+    fn store_native_default_upstream_forward_task(&self, task: NativeDefaultUpstreamForwardTask) {
         let mut downstream = self.lock_downstream();
         downstream.native_default_upstream_http_forward = true;
-        downstream.native_default_upstream_forward_response = Some(response);
+        downstream.native_default_upstream_forward_response = None;
+        downstream.native_default_upstream_forward_task = Some(task);
     }
 
     fn take_native_default_upstream_forward_response(
@@ -2509,6 +2515,18 @@ impl ProxyVmContext {
             .is_some()
     }
 
+    fn take_native_default_upstream_forward_task(&self) -> Option<NativeDefaultUpstreamForwardTask> {
+        self.lock_downstream()
+            .native_default_upstream_forward_task
+            .take()
+    }
+
+    fn native_default_upstream_forward_task_pending(&self) -> bool {
+        self.lock_downstream()
+            .native_default_upstream_forward_task
+            .is_some()
+    }
+
     fn native_default_upstream_forward_latency_ms(&self) -> Option<u64> {
         self.lock_downstream()
             .native_default_upstream_forward_response
@@ -3303,6 +3321,10 @@ enum NativeDefaultUpstreamForwardBody {
     },
 }
 
+type NativeDefaultUpstreamForwardTask = tokio::task::JoinHandle<
+    Result<NativeDefaultUpstreamForwardResponse, UpstreamResponseStartError>,
+>;
+
 #[derive(Debug)]
 struct NativeDefaultUpstreamForwardResponse {
     status: u16,
@@ -4285,10 +4307,8 @@ async fn start_default_upstream_plain_http1_fast_path(
 
 fn materialize_native_default_upstream_forward_response(
     context: &SharedProxyVmContext,
+    response: NativeDefaultUpstreamForwardResponse,
 ) -> Result<Option<HttpUpstreamResponseSnapshot>, UpstreamResponseStartError> {
-    let Some(response) = context.take_native_default_upstream_forward_response() else {
-        return Ok(None);
-    };
     let NativeDefaultUpstreamForwardResponse {
         status,
         headers,
@@ -4342,6 +4362,53 @@ fn materialize_native_default_upstream_forward_response(
     Ok(Some(snapshot))
 }
 
+async fn take_or_await_native_default_upstream_forward_response(
+    context: &SharedProxyVmContext,
+) -> Result<Option<NativeDefaultUpstreamForwardResponse>, UpstreamResponseStartError> {
+    if let Some(response) = context.take_native_default_upstream_forward_response() {
+        return Ok(Some(response));
+    }
+
+    let Some(task) = context.take_native_default_upstream_forward_task() else {
+        return Ok(None);
+    };
+    match task.await {
+        Ok(Ok(response)) => Ok(Some(response)),
+        Ok(Err(err)) => Err(err),
+        Err(err) => Err(UpstreamResponseStartError::Protocol(format!(
+            "native default upstream forward task failed: {err}"
+        ))),
+    }
+}
+
+async fn try_materialize_ready_or_pending_native_default_upstream_forward_response(
+    context: &SharedProxyVmContext,
+) -> Result<Option<HttpUpstreamResponseSnapshot>, UpstreamResponseStartError> {
+    let Some(response) = take_or_await_native_default_upstream_forward_response(context).await?
+    else {
+        return Ok(None);
+    };
+    materialize_native_default_upstream_forward_response(context, response)
+}
+
+async fn try_resolve_ready_or_pending_native_default_upstream_forward_response(
+    context: &SharedProxyVmContext,
+    response_headers: HeaderMap,
+    response_status: Option<u16>,
+) -> Result<Option<ResolvedHttpGraphResponse>, UpstreamResponseStartError> {
+    let Some(response) = take_or_await_native_default_upstream_forward_response(context).await?
+    else {
+        return Ok(None);
+    };
+    let upstream_latency_ms = response.upstream_latency_ms;
+    Ok(Some(ResolvedHttpGraphResponse {
+        response: response_from_started_upstream_response(response, response_headers, response_status)
+            .await,
+        upstream_latency_ms,
+        post_response_plan: None,
+    }))
+}
+
 async fn forward_native_default_upstream_http_via_sender_pool(
     context: &SharedProxyVmContext,
     request: &DefaultUpstreamRequestSnapshot,
@@ -4487,10 +4554,23 @@ async fn forward_native_default_upstream_http_via_sender_pool(
     })
 }
 
+fn schedule_native_default_upstream_http_forward_response(
+    context: &SharedProxyVmContext,
+    request: DefaultUpstreamRequestSnapshot,
+) {
+    let task_context = context.clone();
+    let task = tokio::spawn(async move {
+        forward_native_default_upstream_http_via_sender_pool(&task_context, &request).await
+    });
+    context.store_native_default_upstream_forward_task(task);
+}
+
 pub(crate) async fn start_native_default_upstream_http_forward_response(
     context: &SharedProxyVmContext,
 ) -> Result<bool, VmError> {
-    if context.native_default_upstream_forward_response_ready() {
+    if context.native_default_upstream_forward_response_ready()
+        || context.native_default_upstream_forward_task_pending()
+    {
         return Ok(true);
     }
 
@@ -4521,13 +4601,8 @@ pub(crate) async fn start_native_default_upstream_http_forward_response(
         return Ok(false);
     }
 
-    match forward_native_default_upstream_http_via_sender_pool(context, &request).await {
-        Ok(response) => {
-            context.store_native_default_upstream_forward_response(response);
-            Ok(true)
-        }
-        Err(_) => Ok(false),
-    }
+    schedule_native_default_upstream_http_forward_response(context, request);
+    Ok(true)
 }
 
 async fn try_resolve_native_default_upstream_http_forward_response(
@@ -5096,10 +5171,13 @@ async fn start_outbound_exchange_response(
         }
     }
 
-    if handle == DEFAULT_UPSTREAM_EXCHANGE_HANDLE
-        && let Some(snapshot) = materialize_native_default_upstream_forward_response(context)?
-    {
-        return Ok(snapshot);
+    if handle == DEFAULT_UPSTREAM_EXCHANGE_HANDLE {
+        match try_materialize_ready_or_pending_native_default_upstream_forward_response(context)
+            .await
+        {
+            Ok(Some(snapshot)) => return Ok(snapshot),
+            Ok(None) | Err(_) => {}
+        }
     }
 
     if handle == DEFAULT_UPSTREAM_EXCHANGE_HANDLE
@@ -5690,7 +5768,6 @@ pub(crate) async fn resolve_http_graph_response(
         default_upstream_websocket_mode,
         upstream_response,
         native_default_upstream_http_forward,
-        native_default_upstream_forward_response_ready,
     ) = {
         let downstream = context.lock_downstream();
         let exchanges = context.lock_exchanges();
@@ -5710,9 +5787,6 @@ pub(crate) async fn resolve_http_graph_response(
                 HttpUpstreamResponseNode::NotStarted => None,
             },
             downstream.native_default_upstream_http_forward,
-            downstream
-                .native_default_upstream_forward_response
-                .is_some(),
         )
     };
 
@@ -5763,24 +5837,20 @@ pub(crate) async fn resolve_http_graph_response(
         };
     }
 
-    if native_default_upstream_forward_response_ready {
-        if let Some(response) = context.take_native_default_upstream_forward_response() {
-            let upstream_latency_ms = response.upstream_latency_ms;
-            context.clear_native_default_upstream_http_forward();
-            return ResolvedHttpGraphResponse {
-                response: response_from_started_upstream_response(
-                    response,
-                    response_headers,
-                    response_status,
-                )
-                .await,
-                upstream_latency_ms,
-                post_response_plan: None,
-            };
-        }
-    }
-
     if native_default_upstream_http_forward && upstream_response.is_none() {
+        match try_resolve_ready_or_pending_native_default_upstream_forward_response(
+            context,
+            response_headers.clone(),
+            response_status,
+        )
+        .await
+        {
+            Ok(Some(resolved)) => {
+                context.clear_native_default_upstream_http_forward();
+                return resolved;
+            }
+            Ok(None) | Err(_) => {}
+        }
         match try_resolve_native_default_upstream_http_forward_response(
             context,
             response_headers.clone(),
@@ -6029,10 +6099,13 @@ mod tests {
         resolve_http_graph_response, response_from_upstream_snapshot,
     };
     use crate::abi_impl::RateLimiterStore;
+    use crate::abi_impl::http2::{Http2DownstreamStreamAttachment, Http2StreamRef};
+    #[cfg(feature = "http2")]
     use crate::abi_impl::http2::{
-        Http2DownstreamStreamAttachment, Http2SendRequest, Http2StreamRef, Http2UpstreamMode,
-        new_shared_http_upstream_sessions, send_request, total_active_streams,
+        Http2SendRequest, Http2UpstreamMode, new_shared_http_upstream_sessions, send_request,
+        total_active_streams,
     };
+    #[cfg(feature = "http2")]
     use crate::abi_impl::transport::TlsFlowState;
 
     fn test_context() -> SharedProxyVmContext {

pd-edge/src/abi_impl/http2/mod.rs

@@ -12,7 +12,7 @@ pub(crate) use self::model::{
     Http2StreamRef, Http2UpstreamMode, configure_reqwest_builder, response_version_label,
     select_upstream_mode, supports_response_version,
 };
-#[cfg(test)]
+#[cfg(all(test, feature = "http2"))]
 pub(crate) use self::upstream::total_active_streams;
 #[cfg(feature = "http2")]
 pub(crate) use self::upstream::{Http2RequestError, Http2SendRequest, send_request};

pd-edge/src/runtime/http_plane/proxy_path.rs

@@ -1175,18 +1175,15 @@ fn text_response(status: StatusCode, text: &str) -> Response<Body> {
 
 #[cfg(test)]
 mod tests {
+    use super::*;
     #[cfg(feature = "http2")]
     use axum::http::Version;
     #[cfg(feature = "http2")]
     use http_body_util::{BodyExt, Full};
-    #[cfg(feature = "http2")]
     use vm::{compile_source, encode_program};
 
-    #[cfg(feature = "http2")]
-    use super::*;
     #[cfg(feature = "http2")]
     use crate::abi_impl::Http2SessionFrontier;
-    #[cfg(feature = "http2")]
     use crate::runtime::{VmExecutionConfig, VmExecutionMode, apply_program_from_bytes};
 
     #[cfg(feature = "http2")]

pd-vm/src/vm/jit/runtime.rs

@@ -282,13 +282,12 @@ impl Vm {
         let Some(trace) = self.jit.trace_clone(trace_id) else {
             return Ok(ExecOutcome::Continue);
         };
-        let step_indices_by_ip = trace
-            .step_ips
-            .iter()
-            .copied()
-            .enumerate()
-            .map(|(step_index, step_ip)| (step_ip, step_index))
-            .collect::<HashMap<_, _>>();
+        let step_indices_by_ip = self.jit.trace_step_index_lookup(trace_id).ok_or_else(|| {
+            VmError::JitNative(format!(
+                "trace {} is missing its step-index lookup",
+                trace_id
+            ))
+        })?;
         let mut step_index = 0usize;
         while step_index < trace.steps.len() {
             self.ip = trace