nik/chorus
1
0
Fork 0
mirror of https://github.com/mikedilger/chorus.git synced 2026-08-01 07:21:39 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Throttling: also consume burst_tokens with output bytes sent

Browse Source
This commit is contained in:
Mike Dilger 2024-07-14 12:09:39 +12:00
parent a041896336
commit 7a1b8e7ec5
2 changed files with 31 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
src/lib.rs
View File

@ -323,6 +323,18 @@ struct WebSocketService {
}
impl WebSocketService {
async fn send(&mut self, m: Message) -> Result<(), Error> {
// Throttling: we consume burst tokens, but we do not throttle on output
if m.len() > self.burst_tokens {
self.burst_tokens = 0;
} else {
self.burst_tokens = self.burst_tokens - m.len();
}
self.replied = true;
Ok(self.websocket.send(m).await?)
}
async fn handle_websocket_stream(&mut self) -> Result<(), Error> {
// Subscribe to the shutting down channel
let mut shutting_down = GLOBALS.shutting_down.subscribe();
@ -332,7 +344,7 @@ impl WebSocketService {
// Offer AUTH to clients right off the bat
let reply = NostrReply::Auth(self.challenge.clone());
self.websocket.send(Message::text(reply.as_json())).await?;
self.send(Message::text(reply.as_json())).await?;
let mut last_message_at = Instant::now();
@ -349,7 +361,7 @@ impl WebSocketService {
if self.subscriptions.is_empty() {
// And they are idle for timeout_seconds with no subscriptions
if last_message_at + Duration::from_secs(timeout_seconds) < instant {
self.websocket.send(Message::Close(None)).await?;
self.send(Message::Close(None)).await?;
return Err(ChorusError::TimedOut.into());
}
}
@ -375,7 +387,7 @@ impl WebSocketService {
},
_r = shutting_down.changed() => {
// Shutdown the websocket gracefully
self.websocket.send(Message::Close(None)).await?;
self.send(Message::Close(None)).await?;
break;
},
}
@ -405,6 +417,7 @@ impl WebSocketService {
&& nostr::screen_outgoing_event(event, &event_flags, authorized_user)
{
let message = NostrReply::Event(subid, event);
// note, this is not currently counted in throttling
self.websocket
.send(Message::text(message.as_json()))
.await?;
@ -429,9 +442,7 @@ impl WebSocketService {
self.last_message = Instant::now();
// Grant new tokens
let new_tokens = throttling_bytes_per_second
* elapsed.as_millis() as usize
/ 1_000;
let new_tokens = throttling_bytes_per_second * elapsed.as_millis() as usize / 1_000;
self.burst_tokens = self.burst_tokens + new_tokens;
// Cap tokens to a maximum
@ -443,7 +454,7 @@ impl WebSocketService {
if message.len() > self.burst_tokens {
log::info!(target: "Client", "{}: Rate limited exceeded", self.peer);
let reply = NostrReply::Notice("Rate limit exceeded.".into());
self.websocket.send(Message::text(reply.as_json())).await?;
self.send(Message::text(reply.as_json())).await?;
let error = ChorusError::RateLimitExceeded;
self.error_punishment += error.punishment();
return Err(error.into());
@ -469,11 +480,11 @@ impl WebSocketService {
}
if !self.replied {
let reply = NostrReply::Notice(format!("error: {}", e.inner));
self.websocket.send(Message::text(reply.as_json())).await?;
self.send(Message::text(reply.as_json())).await?;
}
if self.error_punishment >= 1.0 {
let reply = NostrReply::Notice("Closing due to error(s)".into());
self.websocket.send(Message::text(reply.as_json())).await?;
self.send(Message::text(reply.as_json())).await?;
return Err(ChorusError::ErrorClose.into());
}
}
@ -482,7 +493,7 @@ impl WebSocketService {
let reply = NostrReply::Notice(
"binary messages are not processed by this relay".to_owned(),
);
self.websocket.send(Message::text(reply.as_json())).await?;
self.send(Message::text(reply.as_json())).await?;
log::info!(target: "Client",
"{}: Received unhandled binary message: {:02X?}",
self.peer,

31
src/nostr.rs
View File

@ -2,7 +2,6 @@ use crate::error::{ChorusError, Error};
use crate::globals::GLOBALS;
use crate::reply::{NostrReply, NostrReplyPrefix};
use crate::WebSocketService;
use futures::SinkExt;
use hyper_tungstenite::tungstenite::Message;
use pocket_types::json::{eat_whitespace, json_unescape, verify_char};
use pocket_types::{Event, Filter, Kind, OwnedFilter, Pubkey, Time};
@ -34,8 +33,7 @@ impl WebSocketService {
} else {
log::warn!(target: "Client", "{}: Received unhandled text message: {}", self.peer, msg);
let reply = NostrReply::Notice("Command unrecognized".to_owned());
self.websocket.send(Message::text(reply.as_json())).await?;
self.replied = true;
self.send(Message::text(reply.as_json())).await?;
}
Ok(())
@ -94,8 +92,7 @@ impl WebSocketService {
}
_ => NostrReply::Closed(&subid, NostrReplyPrefix::Error, format!("{}", e.inner)),
};
self.websocket.send(Message::text(reply.as_json())).await?;
self.replied = true;
self.send(Message::text(reply.as_json())).await?;
Err(e)
} else {
Ok(())
@ -125,8 +122,7 @@ impl WebSocketService {
NostrReplyPrefix::AuthRequired,
"DM kinds were included in the REQ".to_owned(),
);
self.websocket.send(Message::text(reply.as_json())).await?;
self.replied = true;
self.send(Message::text(reply.as_json())).await?;
return Ok(());
}
}
@ -168,19 +164,17 @@ impl WebSocketService {
for event in events.drain(..) {
let reply = NostrReply::Event(subid, event);
self.websocket.send(Message::text(reply.as_json())).await?;
self.send(Message::text(reply.as_json())).await?;
}
// eose
let reply = NostrReply::Eose(subid);
self.websocket.send(Message::text(reply.as_json())).await?;
self.send(Message::text(reply.as_json())).await?;
}
// Store subscription
self.subscriptions.insert(subid.to_owned(), filters);
self.replied = true;
log::debug!(target: "Client",
"{}: new subscription \"{subid}\", {} total",
self.peer,
@ -250,13 +244,11 @@ impl WebSocketService {
},
_ => NostrReply::Ok(id, false, NostrReplyPrefix::Error, format!("{}", e.inner)),
};
self.websocket.send(Message::text(reply.as_json())).await?;
self.replied = true;
self.send(Message::text(reply.as_json())).await?;
Err(e)
} else {
let reply = NostrReply::Ok(id, true, NostrReplyPrefix::None, "".to_string());
self.websocket.send(Message::text(reply.as_json())).await?;
self.replied = true;
self.send(Message::text(reply.as_json())).await?;
Ok(())
}
}
@ -314,8 +306,7 @@ impl WebSocketService {
// Remove it, and let them know
self.subscriptions.remove(subid);
let reply = NostrReply::Closed(subid, NostrReplyPrefix::None, "".to_owned());
self.websocket.send(Message::text(reply.as_json())).await?;
self.replied = true;
self.send(Message::text(reply.as_json())).await?;
Ok(())
} else {
Err(ChorusError::NoSuchSubscription.into())
@ -341,13 +332,11 @@ impl WebSocketService {
}
_ => NostrReply::Ok(id, false, NostrReplyPrefix::Error, format!("{}", e.inner)),
};
self.websocket.send(Message::text(reply.as_json())).await?;
self.replied = true;
self.send(Message::text(reply.as_json())).await?;
Err(e)
} else {
let reply = NostrReply::Ok(id, true, NostrReplyPrefix::None, "".to_string());
self.websocket.send(Message::text(reply.as_json())).await?;
self.replied = true;
self.send(Message::text(reply.as_json())).await?;
Ok(())
}
}