david/nostr-rs-relay
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: rate limit event creation

Browse Source 
A configuration option, `messages_per_sec`, imposes a global limit on
the rate for which new events can be stored.

Fixes https://todo.sr.ht/~gheartsfield/nostr-rs-relay/6
This commit is contained in:
Greg Heartsfield
2021-12-30 21:07:21 -06:00
parent c60519de23
commit 20ee5a054c
7 changed files with 267 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config.rs
+19 -10
View File
@@ -1,4 +1,5 @@
use lazy_static::lazy_static;
use log::*;
use serde::{Deserialize, Serialize};
use std::sync::RwLock;
@@ -25,16 +26,16 @@ pub struct Options {
#[allow(unused)]
pub struct Retention {
// TODO: implement
pub max_events: Option<usize>, // max events
pub max_bytes: Option<usize>, // max size
pub persist_days: Option<usize>, // oldest message
pub whitelist_addresses: Vec<String>, // whitelisted addresses (never delete)
pub max_events: Option<usize>, // max events
pub max_bytes: Option<usize>, // max size
pub persist_days: Option<usize>, // oldest message
pub whitelist_addresses: Option<Vec<String>>, // whitelisted addresses (never delete)
}
#[derive(Debug, Serialize, Deserialize)]
#[allow(unused)]
pub struct Limits {
pub messages_per_sec: Option<usize>, // Artificially slow down event writing to limit disk consumption
pub messages_per_sec: Option<u32>, // Artificially slow down event writing to limit disk consumption (averaged over 1 minute)
pub max_event_bytes: Option<usize>,
pub max_ws_message_bytes: Option<usize>,
pub max_ws_frame_bytes: Option<usize>,
@@ -55,7 +56,15 @@ impl Settings {
pub fn new() -> Self {
let d = Self::default();
// attempt to construct settings with file
Self::new_from_default(&d).unwrap_or(d)
// Self::new_from_default(&d).unwrap_or(d)
let from_file = Self::new_from_default(&d);
match from_file {
Ok(f) => f,
Err(e) => {
warn!("Error reading config file ({:?})", e);
d
}
}
}
fn new_from_default(default: &Settings) -> Result<Self, config::ConfigError> {
@@ -86,10 +95,10 @@ impl Default for Settings {
event_persist_buffer: 16,
},
retention: Retention {
max_events: None, // max events
max_bytes: None, // max size
persist_days: None, // oldest message
whitelist_addresses: vec![], // whitelisted addresses (never delete)
max_events: None, // max events
max_bytes: None, // max size
persist_days: None, // oldest message
whitelist_addresses: None, // whitelisted addresses (never delete)
},
options: Options {
reject_future_seconds: Some(30 * 60), // Reject events 30min in the future or greater
src/db.rs
+35 -7
View File
@@ -2,12 +2,17 @@
use crate::error::Result;
use crate::event::Event;
use crate::subscription::Subscription;
use governor::clock::Clock;
use governor::{Quota, RateLimiter};
use hex;
use log::*;
use rusqlite::params;
use rusqlite::Connection;
use rusqlite::OpenFlags;
//use std::num::NonZeroU32;
use crate::config::SETTINGS;
use std::path::Path;
use std::thread;
use tokio::task;
/// Database file
@@ -114,6 +119,7 @@ PRAGMA user_version = 2;
pub async fn db_writer(
mut event_rx: tokio::sync::mpsc::Receiver<Event>,
bcast_tx: tokio::sync::broadcast::Sender<Event>,
mut shutdown: tokio::sync::broadcast::Receiver<()>,
) -> tokio::task::JoinHandle<Result<()>> {
task::spawn_blocking(move || {
let mut conn = Connection::open_with_flags(
@@ -122,26 +128,38 @@ pub async fn db_writer(
)?;
info!("opened database for writing");
upgrade_db(&mut conn)?;
// if version is zero, then we need to initialize from scratch.
// if version is one, we need to upgrade.
// if version is two, we are at the latest!
// TODO: determine if we need to execute the init script.
// TODO: check database app id / version before proceeding.
// get rate limit settings
let config = SETTINGS.read().unwrap();
let rps_setting = config.limits.messages_per_sec;
let mut lim_opt = None;
let clock = governor::clock::QuantaClock::default();
if let Some(rps) = rps_setting {
if rps > 0 {
info!("Enabling rate limits for event creation ({}/sec)", rps);
let quota = core::num::NonZeroU32::new(rps * 60).unwrap();
lim_opt = Some(RateLimiter::direct(Quota::per_minute(quota)));
}
}
loop {
if let Ok(_) = shutdown.try_recv() {
info!("shutting down database writer");
break;
}
// call blocking read on channel
let next_event = event_rx.blocking_recv();
// if the channel has closed, we will never get work
if next_event.is_none() {
break;
}
let mut event_write = false;
let event = next_event.unwrap();
match write_event(&mut conn, &event) {
Ok(updated) => {
if updated == 0 {
info!("nothing inserted (dupe?)");
debug!("ignoring duplicate event");
} else {
info!("persisted event: {}", event.get_event_id_prefix());
event_write = true;
// send this out to all clients
bcast_tx.send(event.clone()).ok();
}
@@ -150,6 +168,16 @@ pub async fn db_writer(
warn!("event insert failed: {}", err);
}
}
// use rate limit, if defined, and if an event was actually written.
if event_write {
if let Some(ref lim) = lim_opt {
if let Err(n) = lim.check() {
info!("Rate limiting event creation");
thread::sleep(n.wait_time_from(clock.now()));
continue;
}
}
}
}
conn.close().ok();
info!("database connection closed");
src/main.rs
+6 -5
View File
@@ -28,10 +28,10 @@ fn main() -> Result<(), Error> {
let mut settings = config::SETTINGS.write().unwrap();
// replace default settings with those read from config.toml
let c = config::Settings::new();
debug!("using settings: {:?}", c);
*settings = c;
}
let config = config::SETTINGS.read().unwrap();
debug!("config: {:?}", config);
let addr = format!("{}:{}", config.network.address.trim(), config.network.port);
// configure tokio runtime
let rt = Builder::new_multi_thread()
@@ -52,10 +52,6 @@ fn main() -> Result<(), Error> {
// validated events that need to be persisted are sent to the
// database on via this channel.
let (event_tx, event_rx) = mpsc::channel::<Event>(settings.limits.event_persist_buffer);
// start the database writer thread. Give it a channel for
// writing events, and for publishing events that have been
// written (to all connected clients).
db::db_writer(event_rx, bcast_tx.clone()).await;
// establish a channel for letting all threads now about a
// requested server shutdown.
let (invoke_shutdown, _) = broadcast::channel::<()>(1);
@@ -66,6 +62,11 @@ fn main() -> Result<(), Error> {
info!("shutting down due to SIGINT");
ctrl_c_shutdown.send(()).ok();
});
// start the database writer thread. Give it a channel for
// writing events, and for publishing events that have been
// written (to all connected clients).
db::db_writer(event_rx, bcast_tx.clone(), invoke_shutdown.subscribe()).await;
// track unique client connection count
let mut client_accept_count: usize = 0;
let mut stop_listening = invoke_shutdown.subscribe();
src/protostream.rs
+2
View File
@@ -54,6 +54,8 @@ impl Stream for NostrStream {
/// Convert Message to NostrMessage
fn convert(msg: String) -> Result<NostrMessage> {
debug!("raw msg: {}", msg);
let event_size = msg.len();
debug!("event size is {} bytes", event_size);
let parsed_res: Result<NostrMessage> = serde_json::from_str(&msg).map_err(|e| e.into());
match parsed_res {
Ok(m) => Ok(m),