diff --git a/crates/nailrater/src/futures.rs b/crates/nailrater/src/futures.rs index 1781a01..3e9fd9c 100644 --- a/crates/nailrater/src/futures.rs +++ b/crates/nailrater/src/futures.rs @@ -36,6 +36,7 @@ pin_project! { #[project = NailedStateProj] enum NailedState { RatePeer { + peer: IdentifiedPeer, req: Option, spicy_payloads: Option>, mode: LimitModes, @@ -122,10 +123,11 @@ impl NailedResponseFuture { ) -> Self { Self { state: NailedState::RatePeer { + peer: peer.clone(), + entry: boxed_future_within(|| PEERS.entry_async(peer)), req: Some(req), mode, spicy_payloads, - entry: boxed_future_within(|| PEERS.entry_async(peer)), inner, }, } @@ -156,6 +158,7 @@ where loop { match state.as_mut().project() { NailedStateProj::RatePeer { + peer, req, spicy_payloads, mode, @@ -164,14 +167,21 @@ where } => { let entry = ready!(entry.poll(cx)); let req = req.take().unwrap(); // If this panics, it is because the future was polled twice in the wrong state. - let peer = entry + let peer_record = entry .and_modify(|p| { p.count += 1; p.last_seen = Instant::now(); - p.state = mode.limit(&p.count); + + let next_mode = mode.limit(&p.count); + + if p.state != next_mode && matches!(next_mode, PeerState::SpicyDrop) { + tracing::info!("remote.peer" = %peer, "Spicy mode activated!"); + } + + p.state = next_mode; }) - .or_insert_with_key(|proxied| { - tracing::info!("remote.peer" = %proxied, "New remote peer"); + .or_insert_with(|| { + tracing::info!("remote.peer" = %peer, "New remote peer"); PeerRecord { count: 1, state: mode.limit(&1), @@ -180,13 +190,13 @@ where } }); - let delay = match peer.state { + let delay = match peer_record.state { PeerState::Ready => None, PeerState::Delay(delay) => Some(boxed_future_within(|| sleep(delay))), PeerState::SpicyDrop => { let new_state = spicy_payloads .as_deref() - .zip(peer.supports_spicy) + .zip(peer_record.supports_spicy) .and_then(|(payloads, kind)| { payloads.get(&kind).map(|payload| (payload.clone(), kind)) })