Refactor ProxyContext to remove upstream_override and update traffic split plugin to use context-based upstream selection

Author: zhu327

src/core/mod.rs

@@ -613,8 +613,6 @@ pub struct ProxyContext {
     /// Parameters extracted from the route pattern.
     /// Stored as Vec for better performance with small number of params (typical case).
     pub route_params: Option<Vec<(String, String)>>,
-    // Override upstream selector for traffic splitting and similar use cases.
-    pub upstream_override: Option<Arc<dyn UpstreamSelector>>,
     // Selected HTTP peer for the upstream request.
     pub peer: Option<Box<HttpPeer>>,
     /// Number of retry attempts so far.
@@ -636,7 +634,6 @@ impl Default for ProxyContext {
         Self {
             route: None,
             route_params: None,
-            upstream_override: None,
             peer: None,
             tries: 0,
             plugin: ProxyPluginExecutor::default_shared(),

src/plugin/traffic_split.rs

@@ -14,6 +14,8 @@ use crate::utils::request::request_selector_key;
 pub const PLUGIN_NAME: &str = "traffic-split";
 const PRIORITY: i32 = 966;
 
+pub const CTX_KEY_UPSTREAM_OVERRIDE: &str = "pingsix_upstream_override";
+
 #[derive(Debug, Serialize, Deserialize, Validate)]
 struct WeightedUpstream {
     pub upstream_id: Option<String>,
@@ -56,7 +58,7 @@ impl ProxyPlugin for PluginTrafficSplit {
             if self.match_vars(session, &rule.vars) {
                 // 命中规则，执行加权选择
                 if let Some(selected_ups) = self.pick_upstream(rule_idx, &rule.weighted_upstreams) {
-                    ctx.upstream_override = Some(selected_ups);
+                    ctx.set(CTX_KEY_UPSTREAM_OVERRIDE, selected_ups);
                 }
                 break; // 匹配到第一个规则即停止
             }

src/service/http.rs

@@ -24,8 +24,11 @@ use pingora_proxy::{ProxyHttp, Session};
 
 use crate::{
     config,
-    core::{ProxyContext, ProxyError, ProxyPlugin, RouteContext},
-    plugin::cache::{CacheSettings, CTX_KEY_CACHE_SETTINGS},
+    core::{ProxyContext, ProxyError, ProxyPlugin, RouteContext, UpstreamSelector},
+    plugin::{
+        cache::{CacheSettings, CTX_KEY_CACHE_SETTINGS},
+        traffic_split::CTX_KEY_UPSTREAM_OVERRIDE,
+    },
     proxy::{global_rule::global_plugin_fetch, route::global_route_match_fetch},
 };
 
@@ -120,7 +123,9 @@ impl ProxyHttp for HttpService {
         session: &mut Session,
         ctx: &mut Self::CTX,
     ) -> Result<Box<HttpPeer>> {
-        let peer = if let Some(ups_override) = ctx.upstream_override.as_ref() {
+        let peer = if let Some(ups_override) = ctx
+            .get::<Arc<dyn UpstreamSelector + 'static>>(CTX_KEY_UPSTREAM_OVERRIDE)
+        {
             let mut backend = ups_override.select_backend(session).ok_or_else(|| {
                 ProxyError::UpstreamSelection("Traffic-split selected no backend".to_string())
             })?;