1use std::sync::LazyLock;
8use std::time::{Duration, Instant};
9
10use crate::cid_collections::CidHashMap;
11use crate::prelude::*;
12use crate::utils::misc::env::env_or_default_logged;
13use ahash::HashSet;
14use futures::StreamExt;
15use libp2p::PeerId;
16use nonzero_ext::nonzero;
17use parking_lot::RwLock;
18use tokio::sync::Semaphore;
19
20use crate::libp2p_bitswap::{event_handlers::*, *};
21
22const BITSWAP_BLOCK_REQUEST_INTERVAL: Duration = Duration::from_millis(500);
23
24const SERVE_RESPONSE_CHANNEL_CAP: usize = 128;
31
32pub(in crate::libp2p_bitswap) const MAX_WANTLIST_ENTRIES_SERVED: usize = 1024;
40
41pub(in crate::libp2p_bitswap) static MAX_CONCURRENT_INBOUND_WANTLIST_SERVES: LazyLock<usize> =
45 LazyLock::new(|| {
46 env_or_default_logged(
47 "FOREST_MAX_CONCURRENT_INBOUND_WANTLIST_SERVES",
48 nonzero!(8_usize),
49 )
50 .get()
51 });
52
53pub type ValidatePeerCallback = dyn Fn(PeerId) -> bool + Send + Sync;
54
55#[derive(Debug, Clone)]
56struct ResponseChannels {
57 block_have: flume::Sender<PeerId>,
58 block_received: flume::Sender<Option<Vec<u8>>>,
59}
60
61pub struct BitswapRequestManager {
64 outbound_have_request_tx: flume::Sender<(PeerId, Cid)>,
66 outbound_have_request_rx: flume::Receiver<(PeerId, Cid)>,
67 outbound_cancel_request_tx: flume::Sender<(PeerId, Cid)>,
69 outbound_cancel_request_rx: flume::Receiver<(PeerId, Cid)>,
70 outbound_block_request_tx: flume::Sender<(PeerId, Cid)>,
72 outbound_block_request_rx: flume::Receiver<(PeerId, Cid)>,
73 serve_response_tx: flume::Sender<(PeerId, Cid, BitswapResponse)>,
75 serve_response_rx: flume::Receiver<(PeerId, Cid, BitswapResponse)>,
76 inbound_serve_limiter: Arc<Semaphore>,
78 peers: RwLock<HashSet<PeerId>>,
79 response_channels: RwLock<CidHashMap<ResponseChannels>>,
80}
81
82impl BitswapRequestManager {
83 pub fn outbound_request_stream(
88 &self,
89 ) -> impl futures::stream::Stream<Item = (PeerId, BitswapRequest)> + '_ {
90 type MapperType = fn((libp2p::PeerId, Cid)) -> (libp2p::PeerId, BitswapRequest);
91
92 fn new_block((peer, cid): (PeerId, Cid)) -> (PeerId, BitswapRequest) {
93 (peer, BitswapRequest::new_block(cid).send_dont_have(false))
94 }
95
96 fn new_have((peer, cid): (PeerId, Cid)) -> (PeerId, BitswapRequest) {
97 (peer, BitswapRequest::new_have(cid).send_dont_have(false))
98 }
99
100 fn new_cancel((peer, cid): (PeerId, Cid)) -> (PeerId, BitswapRequest) {
101 (peer, BitswapRequest::new_cancel(cid).send_dont_have(false))
102 }
103
104 let streams = vec![
106 self.outbound_block_request_rx
107 .stream()
108 .map(new_block as MapperType),
109 self.outbound_have_request_rx
110 .stream()
111 .map(new_have as MapperType),
112 self.outbound_cancel_request_rx
113 .stream()
114 .map(new_cancel as MapperType),
115 ];
116 futures::stream::select_all(streams)
117 }
118
119 pub fn outbound_serve_response_stream(
123 &self,
124 ) -> impl futures::stream::Stream<Item = (PeerId, Cid, BitswapResponse)> + '_ {
125 self.serve_response_rx.stream()
126 }
127
128 pub(in crate::libp2p_bitswap) fn serve_inbound_requests<S>(
134 self: &Arc<Self>,
135 store: &S,
136 peer: PeerId,
137 requests: Vec<BitswapRequest>,
138 ) where
139 S: BitswapStoreRead + ShallowClone + Send + Sync + 'static,
140 {
141 if requests.is_empty() {
142 return;
143 }
144
145 let Ok(permit) = self
146 .inbound_serve_limiter
147 .shallow_clone()
148 .try_acquire_owned()
149 else {
150 debug!(%peer, "dropping inbound bitswap wantlist: too many serves in flight");
151 return;
152 };
153 if requests.len() > MAX_WANTLIST_ENTRIES_SERVED {
154 debug!(
155 %peer,
156 "truncating inbound bitswap wantlist from {} to {MAX_WANTLIST_ENTRIES_SERVED} entries",
157 requests.len(),
158 );
159 }
160 let store = store.shallow_clone();
161 let serve_response_tx = self.serve_response_tx.clone();
162 task::spawn_blocking(move || {
163 let _permit = permit;
164 for request in requests.into_iter().take(MAX_WANTLIST_ENTRIES_SERVED) {
165 if let Some(response) = handle_inbound_request(&store, &request)
166 && serve_response_tx
167 .send((peer, request.cid, response))
168 .is_err()
169 {
170 break; }
172 }
173 });
174 }
175}
176
177impl Default for BitswapRequestManager {
178 fn default() -> Self {
179 let (outbound_have_request_tx, outbound_have_request_rx) = flume::unbounded();
180 let (outbound_cancel_request_tx, outbound_cancel_request_rx) = flume::unbounded();
181 let (outbound_block_request_tx, outbound_block_request_rx) = flume::unbounded();
182 let (serve_response_tx, serve_response_rx) = flume::bounded(SERVE_RESPONSE_CHANNEL_CAP);
183 Self {
184 outbound_have_request_tx,
185 outbound_have_request_rx,
186 outbound_cancel_request_tx,
187 outbound_cancel_request_rx,
188 outbound_block_request_tx,
189 outbound_block_request_rx,
190 serve_response_tx,
191 serve_response_rx,
192 inbound_serve_limiter: Arc::new(Semaphore::new(
193 *MAX_CONCURRENT_INBOUND_WANTLIST_SERVES,
194 )),
195 peers: RwLock::new(HashSet::new()),
196 response_channels: RwLock::new(CidHashMap::new()),
197 }
198 }
199}
200
201impl BitswapRequestManager {
202 pub fn handle_event<S: BitswapStoreRead + ShallowClone + Send + Sync + 'static>(
204 self: &Arc<Self>,
205 bitswap: &mut BitswapBehaviour,
206 store: &S,
207 event: BitswapBehaviourEvent,
208 ) -> anyhow::Result<()> {
209 handle_event_impl(self, bitswap, store, event)
210 }
211
212 #[cfg(not(target_arch = "wasm32"))]
216 pub fn get_block(
217 self: Arc<Self>,
218 store: impl BitswapStoreReadWrite + ShallowClone,
219 cid: Cid,
220 timeout: Duration,
221 responder: Option<flume::Sender<bool>>,
222 validate_peer: Option<Arc<ValidatePeerCallback>>,
223 ) {
224 let start = Instant::now();
225 task::spawn(async move {
226 let mut success = store.contains(&cid).unwrap_or_default();
227 if !success {
228 let deadline = start.checked_add(timeout).expect("Infallible");
229 success = self
230 .get_block_inner(&store, cid, deadline, validate_peer)
231 .await;
232 while !success && Instant::now() < deadline {
235 task::sleep(BITSWAP_BLOCK_REQUEST_INTERVAL).await;
236 success = store.contains(&cid).unwrap_or_default();
237 }
238 }
239
240 if success {
241 metrics::message_counter_get_block_success().inc();
242 } else {
243 metrics::message_counter_get_block_failure().inc();
244 }
245
246 if let Some(responder) = responder
247 && let Err(e) = responder.send_async(success).await
248 {
249 debug!("{e}");
250 }
251
252 metrics::GET_BLOCK_TIME.observe((Instant::now() - start).as_secs_f64());
253 });
254 }
255
256 async fn get_block_inner(
257 &self,
258 store: &(impl BitswapStoreReadWrite + ShallowClone),
259 cid: Cid,
260 deadline: Instant,
261 validate_peer: Option<Arc<ValidatePeerCallback>>,
262 ) -> bool {
263 if self.response_channels.read().contains_key(&cid) {
265 return false;
266 }
267
268 let (block_have_tx, block_have_rx) = flume::unbounded();
269 let (block_saved_tx, block_saved_rx) = flume::unbounded();
270 let channels = ResponseChannels {
271 block_have: block_have_tx,
272 block_received: block_saved_tx,
273 };
274 {
275 self.response_channels.write().insert(cid, channels);
276 }
277
278 let peers: Vec<_> = self.peers.read().iter().cloned().collect();
279 let validated_peers: Vec<_> = peers
280 .iter()
281 .filter(|&&p| validate_peer.as_ref().map(|f| f(p)).unwrap_or(true))
282 .cloned()
283 .collect();
284
285 debug!("Found {} valid peers for {cid}", validated_peers.len());
286 let selected_peers = if validated_peers.is_empty() {
287 peers
289 } else {
290 validated_peers
291 };
292
293 for peer in selected_peers {
294 if let Err(e) = self.outbound_have_request_tx.send((peer, cid)) {
295 debug!("{e}");
296 }
297 }
298
299 let timeout = tokio::time::sleep_until(tokio::time::Instant::from_std(deadline));
305 tokio::pin!(timeout);
306 let mut have_open = true;
307 let response = loop {
308 tokio::select! {
309 biased;
310 saved = block_saved_rx.recv_async() => break saved.ok(),
311 () = &mut timeout => break None,
312 have = block_have_rx.recv_async(), if have_open => match have {
313 Ok(peer) => {
314 _ = self.outbound_block_request_tx.send((peer, cid));
315 }
316 Err(_) => have_open = false,
317 },
318 }
319 };
320
321 let success = match response {
322 Some(None) => true,
324 None => false,
326 Some(Some(data)) => {
329 let store = store.shallow_clone();
330 task::spawn_blocking(move || match Block::new(cid, data) {
331 Ok(block) => match store.insert(&block) {
332 Ok(()) => {
333 metrics::message_counter_inbound_response_block_update_db().inc();
334 true
335 }
336 Err(e) => {
337 metrics::message_counter_inbound_response_block_update_db_failure()
338 .inc();
339 warn!(
340 "Failed to update db, cid: {cid}, data: {:?}, error: {e:#}",
341 block.data()
342 );
343 false
344 }
345 },
346 Err(e) => {
347 warn!("Failed to construct block, cid: {cid}, error: {e:#}");
348 false
349 }
350 })
351 .await
352 .unwrap_or(false)
353 }
354 };
355
356 {
358 let mut response_channels = self.response_channels.write();
359 if response_channels.remove(&cid).is_some() {
360 response_channels.shrink_to_fit();
361 metrics::response_channel_container_capacity()
362 .set(response_channels.total_capacity() as _);
363 }
364 }
365
366 success
367 }
368
369 pub(in crate::libp2p_bitswap) fn on_inbound_response_event<S: BitswapStoreRead>(
370 &self,
371 store: &S,
372 response: BitswapInboundResponseEvent,
373 ) {
374 use BitswapInboundResponseEvent::*;
375
376 match response {
377 HaveBlock(peer, cid) => {
378 if let Some(chans) = self.response_channels.read().get(&cid) {
379 _ = chans.block_have.send(peer);
380 }
381 }
382 DataBlock(_peer, cid, data) => {
383 if let Some(chans) = self.response_channels.read().get(&cid) {
384 if let Ok(true) = store.contains(&cid) {
385 metrics::message_counter_inbound_response_block_already_exists_in_db()
387 .inc();
388 _ = chans.block_received.send(None);
389 } else {
390 _ = chans.block_received.send(Some(data));
391 }
392
393 for &peer in self.peers.read().iter() {
398 if let Err(e) = self.outbound_cancel_request_tx.send((peer, cid)) {
399 debug!("{e}");
400 }
401 }
402 } else {
403 metrics::message_counter_inbound_response_block_not_requested().inc();
404 }
405 }
406 }
407 }
408
409 pub(in crate::libp2p_bitswap) fn on_peer_connected(&self, peer: PeerId) -> bool {
410 let mut peers = self.peers.write();
411 let success = peers.insert(peer);
412 if success {
413 metrics::peer_container_capacity().set(peers.capacity() as _);
414 }
415 success
416 }
417
418 pub(in crate::libp2p_bitswap) fn on_peer_disconnected(&self, peer: &PeerId) -> bool {
419 let mut peers = self.peers.write();
420 let success = peers.remove(peer);
421 if success {
422 peers.shrink_to_fit();
423 metrics::peer_container_capacity().set(peers.capacity() as _);
424 }
425 success
426 }
427}