diff --git a/crates/partly-proxy-lib/src/listener.rs b/crates/partly-proxy-lib/src/listener.rs index 466815a..b451cd2 100644 --- a/crates/partly-proxy-lib/src/listener.rs +++ b/crates/partly-proxy-lib/src/listener.rs @@ -433,13 +433,24 @@ async fn record_success_exchange( return; } let (recorded_req, recorded_resp) = build_recorded(runtime, original_request, final_response); - persist_exchange( - runtime, + let exchange = RecordedExchange::new( + Some(runtime.name.clone()), recorded_req, ExchangeOutcome::Response(recorded_resp), duration, - ) - .await; + ); + // Promote a genuinely forwarded exchange into the live replay index so a + // later identical request replays it instead of re-forwarding and + // re-recording — this is what makes the cache deduplicate within a single + // run, even one started from an empty snapshot file. + if source == Some(ResponseSource::Upstream) { + if let Some(replay) = &runtime.replay { + replay.insert(exchange.clone()); + } + } + if let Err(e) = runtime.recorder.record(exchange).await { + tracing::warn!(name = %runtime.name, "recorder rejected exchange: {e}"); + } } async fn record_error_exchange( diff --git a/crates/partly-proxy-lib/src/replay.rs b/crates/partly-proxy-lib/src/replay.rs index e79d011..4733ac6 100644 --- a/crates/partly-proxy-lib/src/replay.rs +++ b/crates/partly-proxy-lib/src/replay.rs @@ -16,7 +16,10 @@ //! a live `Authorization` header still matches a snapshot recorded with //! that header stripped. -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::HashMap, + sync::{Arc, RwLock}, +}; // `ProxyError` is only constructed in the storage-error test below. #[cfg(test)] @@ -42,6 +45,14 @@ pub(crate) struct ReplaySource { } struct ReplaySourceInner { + /// Guarded so `Mode::Record` can promote freshly forwarded exchanges into + /// the index mid-run (see [`ReplaySource::insert`]). Critical sections are + /// short and never `.await`, so a `std::sync::RwLock` is appropriate even + /// on the async hot path. + state: RwLock, +} + +struct ReplayState { exchanges: Vec, /// Maps the lookup key to an index into `exchanges`. The *first* /// exchange written for a given key wins on collision — that way @@ -51,9 +62,10 @@ struct ReplaySourceInner { impl std::fmt::Debug for ReplaySource { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let state = self.inner.state.read().expect("replay lock poisoned"); f.debug_struct("ReplaySource") - .field("exchanges", &self.inner.exchanges.len()) - .field("index", &self.inner.index.len()) + .field("exchanges", &state.exchanges.len()) + .field("index", &state.index.len()) .finish_non_exhaustive() } } @@ -63,7 +75,9 @@ impl ReplaySource { pub(crate) fn new(exchanges: Vec) -> Self { let index = build_index(&exchanges); Self { - inner: Arc::new(ReplaySourceInner { exchanges, index }), + inner: Arc::new(ReplaySourceInner { + state: RwLock::new(ReplayState { exchanges, index }), + }), } } @@ -86,7 +100,41 @@ impl ReplaySource { /// Number of exchanges in the source. Test-only introspection. #[cfg(test)] pub(crate) fn len(&self) -> usize { - self.inner.exchanges.len() + self.inner + .state + .read() + .expect("replay lock poisoned") + .exchanges + .len() + } + + /// Add a freshly recorded exchange to the live index so a subsequent + /// identical request replays it instead of re-forwarding and re-recording. + /// + /// This is what makes `Mode::Record`'s deduplicating cache + /// (`SPECIFICATION.md` §8.3/§20.1) work *within* a single run — including + /// one that started from an empty snapshot file: the first forward of a + /// request is recorded and promoted here, and every later identical request + /// becomes a replay hit. Only `Response` outcomes are indexed (errors are + /// never replayed), and the first entry for a key wins, matching + /// [`build_index`]. + pub(crate) fn insert(&self, exchange: RecordedExchange) { + if !matches!(exchange.outcome, ExchangeOutcome::Response(_)) { + return; + } + let key = ( + exchange.request.method.clone(), + path_and_query_of_str(&exchange.request.uri), + exchange.request.body_sha256.clone(), + ); + let mut state = self.inner.state.write().expect("replay lock poisoned"); + if state.index.contains_key(&key) { + // First write wins — an entry for this key already replays. + return; + } + let idx = state.exchanges.len(); + state.exchanges.push(exchange); + state.index.insert(key, idx); } /// Look up a response for `req`. Returns `None` on miss or on a hit with @@ -108,11 +156,11 @@ impl ReplaySource { path_and_query_of_uri(&redacted.uri), sha256_hex(&redacted.body), ); - let exchange = self - .inner + let state = self.inner.state.read().expect("replay lock poisoned"); + let exchange = state .index .get(&key) - .and_then(|&i| self.inner.exchanges.get(i))?; + .and_then(|&i| state.exchanges.get(i))?; match &exchange.outcome { ExchangeOutcome::Response(r) => Some(ProxyResponse { status: r.status(), diff --git a/crates/partly-proxy-lib/tests/record.rs b/crates/partly-proxy-lib/tests/record.rs index e6791b7..1925d4d 100644 --- a/crates/partly-proxy-lib/tests/record.rs +++ b/crates/partly-proxy-lib/tests/record.rs @@ -206,9 +206,12 @@ async fn custom_storage_via_per_upstream_snapshots() { .unwrap(); let proxy = cluster.addr("upstream").unwrap(); - for _ in 0..3 { + // Distinct paths so each request is genuinely new: in `Mode::Record` the + // snapshot is a deduplicating cache, so three *identical* requests would + // collapse to a single append (the later two replay the first). + for n in 0..3 { let _ = http_client() - .get(format!("http://{proxy}/x")) + .get(format!("http://{proxy}/x/{n}")) .send() .await .unwrap() @@ -235,7 +238,7 @@ async fn custom_storage_via_per_upstream_snapshots() { let saved = storage.appended.lock().await.clone(); assert_eq!(saved.len(), 3); for (i, ex) in saved.iter().enumerate() { - assert_eq!(ex.request.uri, "/x", "exchange {i}"); + assert_eq!(ex.request.uri, format!("/x/{i}"), "exchange {i}"); } } @@ -311,6 +314,102 @@ async fn record_mode_does_not_re_record_request_already_in_snapshot() { ); } +/// Starting from an existing-but-empty snapshot file, the same request is sent +/// twice. The first forwards to the upstream and is recorded; the second must +/// hit that just-recorded snapshot and be served without a second recording. +/// If the freshly recorded exchange is not promoted into the live replay +/// index, the second request is treated as new and recorded again — leaving +/// two copies in the file instead of one. +#[tokio::test] +async fn record_mode_from_empty_file_dedupes_repeated_request() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("trace.ndjson"); + // File exists but is empty. + tokio::fs::File::create(&path).await.unwrap(); + assert_eq!(ndjson_line_count(&path).await, 0); + + let (echo_addr, _echo_task) = spawn_echo().await; + let storage: SharedStorage = Arc::new(JsonlStorage::open(&path).await.unwrap()); + let cfg = ProxyConfig::http( + "127.0.0.1:0".parse().unwrap(), + UpstreamTarget::new(format!("http://{echo_addr}")) + .with_connect_timeout(Duration::from_secs(1)) + .with_request_timeout(Duration::from_secs(5)), + ); + let cluster = ProxyClusterBuilder::new() + .recording(RecordingConfig::in_memory(100)) + .add_upstream_with("upstream", cfg, Vec::new(), Some(storage)) + .run() + .await + .unwrap(); + let proxy = cluster.addr("upstream").unwrap(); + + for _ in 0..2 { + let _ = http_client() + .get(format!("http://{proxy}/dup")) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(150)).await; + } + cluster.shutdown().await.unwrap(); + + assert_eq!( + ndjson_line_count(&path).await, + 1, + "the second identical request must be served from the freshly recorded \ + snapshot and not recorded again (SPECIFICATION.md §8.3 deduplicating cache)" + ); +} + +/// Guard (passes today): an existing-but-empty snapshot file must not suppress +/// recording of a genuinely new request. This pins the *other* half of the +/// dedup contract so a fix for the two regressions above cannot over-correct +/// into dropping new records. +#[tokio::test] +async fn record_mode_from_empty_file_still_records_new_request() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("trace.ndjson"); + tokio::fs::File::create(&path).await.unwrap(); + + let (echo_addr, _echo_task) = spawn_echo().await; + let storage: SharedStorage = Arc::new(JsonlStorage::open(&path).await.unwrap()); + let cfg = ProxyConfig::http( + "127.0.0.1:0".parse().unwrap(), + UpstreamTarget::new(format!("http://{echo_addr}")) + .with_connect_timeout(Duration::from_secs(1)) + .with_request_timeout(Duration::from_secs(5)), + ); + let cluster = ProxyClusterBuilder::new() + .recording(RecordingConfig::in_memory(100)) + .add_upstream_with("upstream", cfg, Vec::new(), Some(storage)) + .run() + .await + .unwrap(); + let proxy = cluster.addr("upstream").unwrap(); + + let _ = http_client() + .post(format!("http://{proxy}/brand-new")) + .body("payload") + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + let _ = wait_for_exchanges(cluster.recorder(), 1).await; + cluster.shutdown().await.unwrap(); + + assert_eq!( + ndjson_line_count(&path).await, + 1, + "a new request against an empty snapshot file must be recorded" + ); +} + /// Poll the recorder until it reaches at least `target` exchanges or times /// out. The lifecycle records exchanges *after* the response has been sent, /// so the client-side `await` returning is not sufficient by itself.