maq/sn0int
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Share the ratelimit state between main modules and all notifications

Browse Source
This commit is contained in:
kpcyrd
2020-06-23 05:41:24 +02:00
parent 63f226f68f
commit 2277089bda
6 changed files with 48 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
modules/harness/notify-ratelimit.lua Normal file
View File

@@ -0,0 +1,9 @@
-- Description: TODO your description here
-- Version: 0.1.0
-- License: GPL-3.0
-- Source: notifications
function run(arg)
ratelimit_throttle('foo', 3, 10000)
info('notication!: ' .. json_encode(arg))
end

5
src/cmd/notify_cmd.rs
View File

@@ -7,6 +7,7 @@ use crate::notify::{self, Notification};
use crate::options::{self, Opt};
use crate::shell::Shell;
use crate::term;
use sn0int_std::ratelimits::Ratelimiter;
use structopt::StructOpt;
use structopt::clap::AppSettings;
@@ -68,7 +69,7 @@ fn print_summary(module: &Module, sent: usize, errors: usize) {
fn send(args: SendArgs, rl: &mut Shell) -> Result<()> {
rl.signal_register().catch_ctrl();
notify::run_router(rl, &mut term::Term, args.dry_run, &args.topic, &args.notification)?;
notify::run_router(rl, &mut term::Term, &mut Ratelimiter::new(), args.dry_run, &args.topic, &args.notification)?;
rl.signal_register().reset_ctrlc();
Ok(())
}
@@ -78,7 +79,7 @@ fn exec(args: ExecArgs, rl: &mut Shell) -> Result<()> {
let options = Opt::collect(&args.options);
rl.signal_register().catch_ctrl();
let errors = notify::exec(rl, &module, options, args.verbose, &args.notification)?;
let errors = notify::exec(rl, &module, &mut Ratelimiter::new(), options, args.verbose, &args.notification)?;
rl.signal_register().reset_ctrlc();
print_summary(&module, 1, errors);

3
src/cmd/run_cmd.rs
View File

@@ -16,6 +16,7 @@ use crate::worker;
use serde::Serialize;
use serde_json;
use sn0int_common::metadata::Source;
use sn0int_std::ratelimits::Ratelimiter;
use std::collections::HashMap;
use structopt::StructOpt;
use structopt::clap::AppSettings;
@@ -189,7 +190,7 @@ pub fn execute(rl: &mut Shell, params: Params, options: HashMap<String, String>)
let args = get_args(rl, &module)?;
rl.signal_register().catch_ctrl();
let errors = worker::spawn(rl, &module, args, &params, rl.config().network.proxy.clone(), options);
let errors = worker::spawn(rl, &module, &mut Ratelimiter::new(), args, &params, rl.config().network.proxy.clone(), options);
rl.signal_register().reset_ctrlc();
if errors > 0 {

4
src/db/ttl.rs
View File

@@ -8,6 +8,7 @@ use crate::term::{self, Term};
use chrono::{NaiveDateTime, Duration, Utc};
use diesel;
use diesel::prelude::*;
use sn0int_std::ratelimits::Ratelimiter;
#[derive(Identifiable, Queryable, AsChangeset, PartialEq, Debug)]
@@ -137,13 +138,14 @@ impl Ttl {
pub fn reap_expired(rl: &mut Shell) -> Result<()> {
debug!("Reaping expired entities");
let mut ratelimit = Ratelimiter::new();
for expired in Ttl::expired(rl.db())? {
debug!("Expired: {:?}", expired);
expired.delete(rl.db())?;
let subject = format!("Deleted {} {:?}", &expired.family, &expired.value);
let topic = &format!("db:{}:{}:delete", &expired.family, &expired.value);
if let Err(err) = notify::trigger_notify_event(rl, &mut Term, &topic, &Notification {
if let Err(err) = notify::trigger_notify_event(rl, &mut Term, &mut ratelimit, &topic, &Notification {
subject,
body: None,
}) {

13
src/notify/mod.rs
View File

@@ -10,6 +10,7 @@ use serde::de::{self, Deserialize, Deserializer};
use serde::ser::{Serialize, Serializer};
use sn0int_common::metadata::Source;
use sn0int_std::blobs::Blob;
use sn0int_std::ratelimits::Ratelimiter;
use std::collections::HashMap;
use std::result;
use std::str::FromStr;
@@ -111,8 +112,8 @@ impl NotificationConfig {
}
}
pub fn trigger_notify_event<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, topic: &str, notification: &Notification) -> Result<()> {
run_router(rl, spinner, false, topic, notification)
pub fn trigger_notify_event<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, topic: &str, notification: &Notification) -> Result<()> {
run_router(rl, spinner, ratelimit, false, topic, notification)
}
fn prepare_arg(notification: &Notification) -> Result<(serde_json::Value, Option<String>, Vec<Blob>)> {
@@ -120,7 +121,7 @@ fn prepare_arg(notification: &Notification) -> Result<(serde_json::Value, Option
Ok((arg, None, vec![]))
}
pub fn exec(rl: &mut Shell, module: &Module, options: HashMap<String, String>, verbose: u64, notification: &Notification) -> Result<usize> {
pub fn exec(rl: &mut Shell, module: &Module, ratelimit: &mut Ratelimiter, options: HashMap<String, String>, verbose: u64, notification: &Notification) -> Result<usize> {
let module_name = module.canonical();
debug!("Setting up notification execution with {:?}", module_name);
@@ -142,13 +143,13 @@ pub fn exec(rl: &mut Shell, module: &Module, options: HashMap<String, String>, v
let args = vec![prepare_arg(&notification)?];
debug!("Executing notification module {:?}", module_name);
let errors = worker::spawn(rl, &module, args, &params, rl.config().network.proxy.clone(), options);
let errors = worker::spawn(rl, &module, ratelimit, args, &params, rl.config().network.proxy.clone(), options);
debug!("Notification module {:?} exited with {:?} errors", module_name, errors);
Ok(errors)
}
pub fn run_router<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, dry_run: bool, topic: &str, notification: &Notification) -> Result<()> {
pub fn run_router<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, dry_run: bool, topic: &str, notification: &Notification) -> Result<()> {
let configs = rl.config().notifications.clone();
debug!("Running notification router");
@@ -164,7 +165,7 @@ pub fn run_router<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, dry_run: bool,
spinner.success(&format!("Executed {} for {:?} (dry-run)", module.canonical(), name));
} else {
let options = options::Opt::collect(&config.options);
match exec(rl, &module, options, 0, notification) {
match exec(rl, &module, ratelimit, options, 0, notification) {
Ok(0) => {
let msg = format!("Executed {} for {:?}", module.canonical(), name);
spinner.success(&msg);

50
src/worker.rs
View File

@@ -169,8 +169,8 @@ impl EventWithCallback for DatabaseEvent {
}
impl DatabaseEvent {
fn notify<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, topic: &str, subject: String) {
if let Err(err) = notify::trigger_notify_event(rl, spinner, topic, &Notification {
fn notify<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, topic: &str, subject: String) {
if let Err(err) = notify::trigger_notify_event(rl, spinner, ratelimit, topic, &Notification {
subject,
body: None,
}) {
@@ -178,28 +178,28 @@ impl DatabaseEvent {
}
}
fn on_insert<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, family: &str, value: &str) {
fn on_insert<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, family: &str, value: &str) {
// TODO: also include fields, see update
let log = format!("Adding {} {:?}", family, value);
spinner.log(&log);
let subject = format!("Added {} {:?}", family, value);
let topic = format!("db:{}:{}:insert", family, value);
Self::notify(rl, spinner, &topic, subject);
Self::notify(rl, spinner, ratelimit, &topic, subject);
}
fn on_update<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, family: &str, value: &str, update: &Update) {
fn on_update<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, family: &str, value: &str, update: &Update) {
spinner.log(&format!("Updating {} {:?} ({})", family, value, update.to_term_str()));
// TODO: in the future we could consider firing multiple events, one for each column
// TODO: this would be super noisy if a lot of fields change though
let subject = format!("Updated {} {:?} ({})", family, value, update.to_plain_str());
let topic = format!("db:{}:{}:update", family, value);
Self::notify(rl, spinner, &topic, subject);
Self::notify(rl, spinner, ratelimit, &topic, subject);
}
fn on_activity<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, object: &NewActivity, verbose: u64) {
Self::spinner_log_new_activity(&object, spinner, verbose);
fn on_activity<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, object: &NewActivity, verbose: u64) {
Self::spinner_log_new_activity(spinner, &object, verbose);
// TODO: we don't want to copy the match arms everywhere
let mut subject = format!("New activity: {:?}", object.topic);
@@ -207,10 +207,10 @@ impl DatabaseEvent {
subject += &format!(" ({:?})", uniq);
}
let topic = format!("activity:{}", object.topic);
Self::notify(rl, spinner, &topic, subject);
Self::notify(rl, spinner, ratelimit, &topic, subject);
}
fn spinner_log_new_activity<T: SpinLogger>(object: &NewActivity, spinner: &mut T, verbose: u64) {
fn spinner_log_new_activity<T: SpinLogger>(spinner: &mut T, object: &NewActivity, verbose: u64) {
let mut log = format!("{:?} ", object.topic);
if let Some(uniq) = &object.uniq {
log.push_str(&format!("({:?}) ", uniq));
@@ -232,7 +232,7 @@ impl DatabaseEvent {
spinner.log(&log);
}
fn insert<T: SpinLogger>(rl: &mut Shell, object: Insert, ttl: Option<i32>, tx: DbSender, spinner: &mut T, verbose: u64) {
fn insert<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, object: Insert, ttl: Option<i32>, tx: DbSender, verbose: u64) {
let db = rl.db();
if verbose >= 1 {
spinner.debug(&format!("Inserting: {:?}", object));
@@ -251,7 +251,7 @@ impl DatabaseEvent {
}
}
Self::on_insert(rl, spinner, object.family(), &value);
Self::on_insert(rl, spinner, ratelimit, object.family(), &value);
}
Err(err) => {
spinner.error(&format!("Failed to query necessary fields for {:?}: {:?}", object, err));
@@ -268,7 +268,7 @@ impl DatabaseEvent {
}
match object.value(rl.db()) {
Ok(value) => Self::on_update(rl, spinner, object.family(), &value, &update),
Ok(value) => Self::on_update(rl, spinner, ratelimit, object.family(), &value, &update),
Err(err) => {
// TODO: this should be unreachable
spinner.error(&format!("Failed to get label for {:?}: {:?}", object, err));
@@ -298,14 +298,14 @@ impl DatabaseEvent {
}
pub fn activity<T: SpinLogger>(rl: &mut Shell, object: NewActivity, tx: DbSender, spinner: &mut T, verbose: u64) {
pub fn activity<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, object: NewActivity, tx: DbSender, verbose: u64) {
let db = rl.db();
let result = db.insert_activity(object.clone());
debug!("{:?} => {:?}", object, result);
let result = match result {
Ok(true) => {
Self::on_activity(rl, spinner, &object, verbose);
Self::on_activity(rl, spinner, ratelimit, &object, verbose);
Ok(DatabaseResponse::Inserted(0))
},
Ok(false) => Ok(DatabaseResponse::NoChange(0)),
@@ -319,7 +319,7 @@ impl DatabaseEvent {
tx.send(result).expect("Failed to send db result to channel");
}
pub fn update<T: SpinLogger>(rl: &mut Shell, family: &str, value: &str, update: &Update, tx: DbSender, spinner: &mut T, verbose: u64) {
pub fn update<T: SpinLogger>(rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, family: &str, value: &str, update: &Update, tx: DbSender, verbose: u64) {
let db = rl.db();
if verbose >= 1 {
spinner.debug(&format!("Updating: {:?}", update));
@@ -330,7 +330,7 @@ impl DatabaseEvent {
let result = match result {
Ok(id) => {
Self::on_update(rl, spinner, family, &value, &update);
Self::on_update(rl, spinner, ratelimit, family, &value, &update);
Ok(DatabaseResponse::Updated(id))
},
Err(err) => {
@@ -343,11 +343,11 @@ impl DatabaseEvent {
tx.send(result).expect("Failed to send db result to channel");
}
pub fn apply<T: SpinLogger>(self, rl: &mut Shell, tx: DbSender, spinner: &mut T, verbose: u64) {
pub fn apply<T: SpinLogger>(self, rl: &mut Shell, spinner: &mut T, ratelimit: &mut Ratelimiter, tx: DbSender, verbose: u64) {
match self {
DatabaseEvent::Insert(object) => Self::insert(rl, object, None, tx, spinner, verbose),
DatabaseEvent::InsertTtl((object, ttl)) => Self::insert(rl, object, Some(ttl), tx, spinner, verbose),
DatabaseEvent::Activity(object) => Self::activity(rl, object, tx, spinner, verbose),
DatabaseEvent::Insert(object) => Self::insert(rl, spinner, ratelimit, object, None, tx, verbose),
DatabaseEvent::InsertTtl((object, ttl)) => Self::insert(rl, spinner, ratelimit, object, Some(ttl), tx, verbose),
DatabaseEvent::Activity(object) => Self::activity(rl, spinner, ratelimit, object, tx, verbose),
DatabaseEvent::Select((family, value)) => {
let db = rl.db();
let result = match db.get_opt(&family, &value) {
@@ -358,7 +358,7 @@ impl DatabaseEvent {
tx.send(result).expect("Failed to send db result to channel");
},
DatabaseEvent::Update((family, value, update)) => Self::update(rl, family.as_str(), &value, &update, tx, spinner, verbose),
DatabaseEvent::Update((family, value, update)) => Self::update(rl, spinner, ratelimit, family.as_str(), &value, &update, tx, verbose),
}
}
}
@@ -437,7 +437,7 @@ impl RatelimitEvent {
}
}
pub fn spawn(rl: &mut Shell, module: &Module, args: Vec<(serde_json::Value, Option<String>, Vec<Blob>)>, params: &Params, proxy: Option<SocketAddr>, options: HashMap<String, String>) -> usize {
pub fn spawn(rl: &mut Shell, module: &Module, ratelimit: &mut Ratelimiter, args: Vec<(serde_json::Value, Option<String>, Vec<Blob>)>, params: &Params, proxy: Option<SocketAddr>, options: HashMap<String, String>) -> usize {
// This function hangs if args is empty, so return early if that's the case
if args.is_empty() {
return 0;
@@ -485,8 +485,6 @@ pub fn spawn(rl: &mut Shell, module: &Module, args: Vec<(serde_json::Value, Opti
expected += 1;
}
let mut ratelimit = Ratelimiter::new();
let mut errors = 0;
let mut failed = Vec::new();
let timeout = Duration::from_millis(100);
@@ -502,7 +500,7 @@ pub fn spawn(rl: &mut Shell, module: &Module, args: Vec<(serde_json::Value, Opti
stack.add(name, label);
},
Event2::Log(log) => log.apply(&mut stack.prefixed(name)),
Event2::Database((db, tx)) => db.apply(rl, tx, &mut stack.prefixed(name), verbose),
Event2::Database((db, tx)) => db.apply(rl, &mut stack.prefixed(name), ratelimit, tx, verbose),
Event2::Ratelimit((req, tx)) => ratelimit.pass(tx, &req.key, req.passes, req.time),
Event2::Blob((blob, tx)) => rl.store_blob(tx, &blob),
Event2::Exit(event) => {