Skip to content

Commit 29e3f00

Browse files
ackintoshackintosh
committed
WIP request size limiter
1 parent c0ae632 commit 29e3f00

File tree

3 files changed

+277
-45
lines changed

3 files changed

+277
-45
lines changed

‎beacon_node/lighthouse_network/src/rpc/mod.rs‎

Lines changed: 91 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ pub use methods::{
3232
};
3333
pub(crate) use outbound::OutboundRequest;
3434
pub use protocol::{max_rpc_size, Protocol, RPCError};
35+
use crate::rpc::rate_limiter::RequestSizeLimiter;
3536

3637
use self::config::{InboundRateLimiterConfig, OutboundRateLimiterConfig};
3738
use self::protocol::RPCProtocol;
@@ -122,6 +123,8 @@ pub struct RPC<Id: ReqId, E: EthSpec> {
122123
limiter: Option<RateLimiter>,
123124
/// Rate limiter for our responses. This is shared with RPCHandlers.
124125
response_limiter: Option<Arc<Mutex<RateLimiter>>>,
126+
///
127+
request_size_limiter: Option<RequestSizeLimiter>,
125128
/// Rate limiter for our own requests.
126129
self_limiter: Option<SelfRateLimiter<Id, E>>,
127130
/// Queue of events to be processed.
@@ -152,21 +155,27 @@ impl<Id: ReqId, E: EthSpec> RPC<Id, E> {
152155
// });
153156
let inbound_limiter = None; // TODO
154157

155-
let response_limiter = inbound_rate_limiter_config.map(|config| {
158+
let response_limiter = inbound_rate_limiter_config.clone().map(|config| {
156159
debug!(log, "Using response rate limiting params"; "config" => ?config);
157160
Arc::new(Mutex::new(
158161
RateLimiter::new_with_config(config.0)
159162
.expect("Inbound limiter configuration parameters are valid"),
160163
))
161164
});
162165

166+
let request_size_limiter = inbound_rate_limiter_config.map(|config| {
167+
RequestSizeLimiter::new_with_config(config.0)
168+
.expect("Inbound limiter configuration parameters are valid")
169+
});
170+
163171
let self_limiter = outbound_rate_limiter_config.map(|config| {
164172
SelfRateLimiter::new(config, log.clone()).expect("Configuration parameters are valid")
165173
});
166174

167175
RPC {
168176
limiter: inbound_limiter,
169177
response_limiter,
178+
request_size_limiter,
170179
self_limiter,
171180
events: Vec::new(),
172181
fork_context,
@@ -315,58 +324,95 @@ where
315324
) {
316325
match event {
317326
HandlerEvent::Ok(RPCReceived::Request(ref id, ref req)) => {
318-
if let Some(limiter) = self.limiter.as_mut() {
319-
// check if the request is conformant to the quota
320-
match limiter.allows(&peer_id, req) {
321-
Ok(()) => {
322-
// send the event to the user
323-
self.events.push(ToSwarm::GenerateEvent(RPCMessage {
324-
peer_id,
325-
conn_id,
326-
event,
327-
}))
328-
}
329-
Err(RateLimitedErr::TooLarge) => {
330-
// we set the batch sizes, so this is a coding/config err for most protocols
331-
let protocol = req.versioned_protocol().protocol();
332-
if matches!(
327+
if let Some(limiter) = self.request_size_limiter.as_ref() {
328+
println!("req limiter");
329+
// if let Some(limiter) = self.limiter.as_mut() {
330+
// Check if the request is conformant to the quota
331+
if limiter.allows(req) {
332+
println!("req limiter allowed");
333+
// Send the event to the user
334+
self.events.push(ToSwarm::GenerateEvent(RPCMessage {
335+
peer_id,
336+
conn_id,
337+
event,
338+
}))
339+
} else {
340+
println!("req limiter NOT allowed");
341+
// We set the batch sizes, so this is a coding/config err for most protocols
342+
let protocol = req.versioned_protocol().protocol();
343+
if matches!(
333344
protocol,
334345
Protocol::BlocksByRange
335346
| Protocol::BlobsByRange
336347
| Protocol::BlocksByRoot
337348
| Protocol::BlobsByRoot
338349
) {
339-
debug!(self.log, "Request too large to process"; "request" => %req, "protocol" => %protocol);
340-
} else {
341-
// Other protocols shouldn't be sending large messages, we should flag the peer kind
342-
crit!(self.log, "Request size too large to ever be processed"; "protocol" => %protocol);
343-
}
344-
// send an error code to the peer.
345-
// the handler upon receiving the error code will send it back to the behaviour
346-
self.send_response(
347-
peer_id,
348-
(conn_id, *id),
349-
RPCCodedResponse::Error(
350-
RPCResponseErrorCode::RateLimited,
351-
"Rate limited. Request too large".into(),
352-
),
353-
);
354-
}
355-
Err(RateLimitedErr::TooSoon(wait_time)) => {
356-
debug!(self.log, "Request exceeds the rate limit";
357-
"request" => %req, "peer_id" => %peer_id, "wait_time_ms" => wait_time.as_millis());
358-
// send an error code to the peer.
359-
// the handler upon receiving the error code will send it back to the behaviour
360-
self.send_response(
361-
peer_id,
362-
(conn_id, *id),
363-
RPCCodedResponse::Error(
364-
RPCResponseErrorCode::RateLimited,
365-
format!("Wait {:?}", wait_time).into(),
366-
),
367-
);
350+
debug!(self.log, "Request too large to process"; "request" => %req, "protocol" => %protocol);
351+
} else {
352+
// Other protocols shouldn't be sending large messages, we should flag the peer kind
353+
crit!(self.log, "Request size too large to ever be processed"; "protocol" => %protocol);
368354
}
355+
// Send an error code to the peer.
356+
// The handler upon receiving the error code will send it back to the behaviour
357+
self.send_response(
358+
peer_id,
359+
(conn_id, *id),
360+
RPCCodedResponse::Error(
361+
RPCResponseErrorCode::RateLimited,
362+
"Rate limited. Request too large".into(),
363+
),
364+
);
369365
}
366+
// match limiter.allows(&peer_id, req) {
367+
// Ok(()) => {
368+
// // send the event to the user
369+
// self.events.push(ToSwarm::GenerateEvent(RPCMessage {
370+
// peer_id,
371+
// conn_id,
372+
// event,
373+
// }))
374+
// }
375+
// Err(RateLimitedErr::TooLarge) => {
376+
// // we set the batch sizes, so this is a coding/config err for most protocols
377+
// let protocol = req.versioned_protocol().protocol();
378+
// if matches!(
379+
// protocol,
380+
// Protocol::BlocksByRange
381+
// | Protocol::BlobsByRange
382+
// | Protocol::BlocksByRoot
383+
// | Protocol::BlobsByRoot
384+
// ) {
385+
// debug!(self.log, "Request too large to process"; "request" => %req, "protocol" => %protocol);
386+
// } else {
387+
// // Other protocols shouldn't be sending large messages, we should flag the peer kind
388+
// crit!(self.log, "Request size too large to ever be processed"; "protocol" => %protocol);
389+
// }
390+
// // send an error code to the peer.
391+
// // the handler upon receiving the error code will send it back to the behaviour
392+
// self.send_response(
393+
// peer_id,
394+
// (conn_id, *id),
395+
// RPCCodedResponse::Error(
396+
// RPCResponseErrorCode::RateLimited,
397+
// "Rate limited. Request too large".into(),
398+
// ),
399+
// );
400+
// }
401+
// Err(RateLimitedErr::TooSoon(wait_time)) => {
402+
// debug!(self.log, "Request exceeds the rate limit";
403+
// "request" => %req, "peer_id" => %peer_id, "wait_time_ms" => wait_time.as_millis());
404+
// // send an error code to the peer.
405+
// // the handler upon receiving the error code will send it back to the behaviour
406+
// self.send_response(
407+
// peer_id,
408+
// (conn_id, *id),
409+
// RPCCodedResponse::Error(
410+
// RPCResponseErrorCode::RateLimited,
411+
// format!("Wait {:?}", wait_time).into(),
412+
// ),
413+
// );
414+
// }
415+
// }
370416
} else {
371417
// No rate limiting, send the event to the user
372418
self.events.push(ToSwarm::GenerateEvent(RPCMessage {

‎beacon_node/lighthouse_network/src/rpc/rate_limiter.rs‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -434,6 +434,89 @@ impl<Key: Hash + Eq + Clone> Limiter<Key> {
434434
}
435435
}
436436

437+
pub(super) struct RequestSizeLimiter {
438+
ping: (Nanosecs, Nanosecs),
439+
status: (Nanosecs, Nanosecs),
440+
meta_data: (Nanosecs, Nanosecs),
441+
goodbye: (Nanosecs, Nanosecs),
442+
blocks_by_range: (Nanosecs, Nanosecs),
443+
blocks_by_root: (Nanosecs, Nanosecs),
444+
}
445+
446+
impl RequestSizeLimiter {
447+
pub fn new_with_config(config: RateLimiterConfig) -> Result<Self, &'static str> {
448+
// Destructure to make sure every configuration value is used.
449+
let RateLimiterConfig {
450+
ping_quota,
451+
meta_data_quota,
452+
status_quota,
453+
goodbye_quota,
454+
blocks_by_range_quota,
455+
blocks_by_root_quota,
456+
blobs_by_range_quota,
457+
blobs_by_root_quota,
458+
light_client_bootstrap_quota,
459+
light_client_optimistic_update_quota,
460+
light_client_finality_update_quota,
461+
} = config;
462+
463+
let tau_and_t = |quota: &Quota| {
464+
let tau = quota.replenish_all_every.as_nanos();
465+
if tau == 0 {
466+
return Err("Replenish time must be positive");
467+
}
468+
let t = (tau / quota.max_tokens as u128).try_into().map_err(|_| "total replenish time is too long")?;
469+
let tau = tau
470+
.try_into()
471+
.map_err(|_| "total replenish time is too long")?;
472+
Ok((tau, t))
473+
};
474+
475+
Ok(Self {
476+
ping: tau_and_t(&ping_quota)?,
477+
meta_data: tau_and_t(&meta_data_quota)?,
478+
status: tau_and_t(&status_quota)?,
479+
goodbye: tau_and_t(&goodbye_quota)?,
480+
blocks_by_range: tau_and_t(&blocks_by_range_quota)?,
481+
blocks_by_root: tau_and_t(&blocks_by_root_quota)?,
482+
})
483+
}
484+
485+
pub fn allows<Item: RateLimiterItem>(&self, request: &Item) -> bool {
486+
let tokens = request.max_responses().max(1);
487+
let (tau, t) = match request.protocol() {
488+
Protocol::Ping => self.ping,
489+
Protocol::Status => self.status,
490+
Protocol::MetaData => self.meta_data,
491+
Protocol::Goodbye => self.goodbye,
492+
Protocol::BlocksByRange => self.blocks_by_range,
493+
Protocol::BlocksByRoot => self.blocks_by_root,
494+
// Protocol::BlobsByRange => &mut self.blbrange_rl,
495+
// Protocol::BlobsByRoot => &mut self.blbroot_rl,
496+
// Protocol::LightClientBootstrap => &mut self.lc_bootstrap_rl,
497+
// Protocol::LightClientOptimisticUpdate => &mut self.lc_optimistic_update_rl,
498+
// Protocol::LightClientFinalityUpdate => &mut self.lc_finality_update_rl,
499+
_ => todo!(),
500+
};
501+
502+
// how long does it take to replenish these tokens
503+
let additional_time = t * tokens;
504+
505+
// DEBUG
506+
println!("{}", request.protocol());
507+
println!("t:{}, tokens:{}, tau:{}", t, tokens, tau);
508+
println!("{}", additional_time);
509+
println!("{}", tau);
510+
511+
if additional_time > tau {
512+
// the time required to process this amount of tokens is longer than the time that
513+
// makes the bucket full. So, this batch can _never_ be processed
514+
return false;
515+
}
516+
true
517+
}
518+
}
519+
437520
#[cfg(test)]
438521
mod tests {
439522
use crate::rpc::rate_limiter::{Limiter, Quota};

‎beacon_node/lighthouse_network/tests/rpc_tests.rs‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1297,3 +1297,106 @@ fn test_delayed_rpc_response() {
12971297
}
12981298
})
12991299
}
1300+
1301+
#[test]
1302+
fn test_request_too_large() {
1303+
let rt = Arc::new(Runtime::new().unwrap());
1304+
let log = logging::test_logger();
1305+
let spec = E::default_spec();
1306+
1307+
rt.block_on(async {
1308+
// get sender/receiver
1309+
let (mut sender, mut receiver) = common::build_node_pair(
1310+
Arc::downgrade(&rt),
1311+
&log,
1312+
ForkName::Base,
1313+
&spec,
1314+
Protocol::Tcp,
1315+
// Configure quotas for requests.
1316+
Some("beacon_blocks_by_range:1/1024;beacon_blocks_by_root:1/128".parse().unwrap()),
1317+
)
1318+
.await;
1319+
1320+
let mut rpc_requests = vec![
1321+
Request::BlocksByRange(BlocksByRangeRequest::new(0, 2)),
1322+
Request::BlocksByRoot(BlocksByRootRequest::new(
1323+
vec![
1324+
Hash256::from_low_u64_be(0),
1325+
Hash256::from_low_u64_be(0),
1326+
],
1327+
&spec,
1328+
)),
1329+
];
1330+
// // BlocksByRange Request
1331+
// let rpc_request = Request::BlocksByRange(BlocksByRangeRequest::new(0, 2));
1332+
// // BlocksByRoot Request
1333+
// let rpc_request = Request::BlocksByRoot(BlocksByRootRequest::new(
1334+
// vec![
1335+
// Hash256::from_low_u64_be(0),
1336+
// Hash256::from_low_u64_be(0),
1337+
// ],
1338+
// &spec,
1339+
// ));
1340+
1341+
1342+
// build the sender future
1343+
let sender_future = async {
1344+
let mut request_id = 1;
1345+
loop {
1346+
match sender.next_event().await {
1347+
NetworkEvent::PeerConnectedOutgoing(peer_id) => {
1348+
let request = rpc_requests.pop().unwrap();
1349+
debug!(log, "Sending RPC request"; "request_id" => request_id, "request" => ?request);
1350+
sender.send_request(peer_id, request_id, request);
1351+
}
1352+
NetworkEvent::ResponseReceived { id, response, .. } => {
1353+
// TODO: 終端を示すレスポンスが返ってくるのでハンドリングする
1354+
debug!(log, "ResponseReceived"; "request_id" => id, "response" => ?response);
1355+
match response {
1356+
Response::BlobsByRoot(None) => {},
1357+
_ => unreachable!(),
1358+
}
1359+
}
1360+
NetworkEvent::RPCFailed { id, peer_id, error } => {
1361+
debug!(log, "RPC Failed"; "error" => ?error, "request_id" => id);
1362+
assert_eq!(id, request_id);
1363+
let v = matches!(error, RPCError::ErrorResponse(RPCResponseErrorCode::RateLimited, .. ));
1364+
assert!(v);
1365+
if let Some(request) = rpc_requests.pop() {
1366+
request_id += 1;
1367+
debug!(log, "Sending RPC request"; "request_id" => request_id, "request" => ?request);
1368+
sender.send_request(peer_id, request_id, request);
1369+
} else {
1370+
return
1371+
}
1372+
}
1373+
_ => {}
1374+
}
1375+
}
1376+
};
1377+
1378+
// build the receiver future
1379+
let receiver_future = async {
1380+
loop {
1381+
match receiver.next_event().await {
1382+
NetworkEvent::RequestReceived {
1383+
peer_id,
1384+
id,
1385+
request,
1386+
} => {
1387+
unreachable!();
1388+
}
1389+
_ => {} // Ignore other events
1390+
}
1391+
}
1392+
};
1393+
1394+
tokio::select! {
1395+
_ = sender_future => {}
1396+
_ = receiver_future => {}
1397+
_ = sleep(Duration::from_secs(30)) => {
1398+
panic!("Future timed out");
1399+
}
1400+
}
1401+
});
1402+
}

0 commit comments

Comments
 (0)