diff --git a/crates/integration-tests/source.lua b/crates/integration-tests/source.lua index 2022dbf0..d0dcdd3b 100644 --- a/crates/integration-tests/source.lua +++ b/crates/integration-tests/source.lua @@ -246,10 +246,32 @@ kumo.on('get_queue_config', function(domain, _tenant, _campaign) protocol = protocol, retry_interval = os.getenv 'KUMOD_RETRY_INTERVAL', strategy = os.getenv 'KUMOD_QUEUE_STRATEGY', + egress_pool = os.getenv 'KUMOD_POOL_NAME', } end) -kumo.on('get_egress_path_config', function(domain, _source_name, _site_name) +kumo.on('get_egress_pool', function(pool_name) + if pool_name == 'warming' then + -- coupled with source_selection_rate_pool.rs + return kumo.make_egress_pool { + name = pool_name, + entries = { + { name = 'warming_a' }, + { name = 'warming_b' }, + }, + } + end + + error('integration-tests/source.lua: unhandled pool ' .. pool_name) +end) + +kumo.on('get_egress_source', function(source_name) + return kumo.make_egress_source { + name = source_name, + } +end) + +kumo.on('get_egress_path_config', function(domain, source_name, _site_name) -- Allow sending to a sink local params = { enable_tls = os.getenv 'KUMOD_ENABLE_TLS' or 'OpportunisticInsecure', @@ -260,8 +282,21 @@ kumo.on('get_egress_path_config', function(domain, _source_name, _site_name) opportunistic_tls_reconnect_on_failed_handshake = ( (os.getenv 'KUMOD_OPPORTUNISTIC_TLS_RECONNECT') and true ) or false, + source_selection_rate = os.getenv 'KUMOD_SOURCE_SELECTION_RATE', } + -- See if there is a source-specific rate exported to us via the environment. + -- We assign this using additional_source_selection_rates regardless of + -- whether we have a more generate rate specified above so that we can + -- excercise the additional_source_selection_rates collection logic in the core. + local source_rate_name = 'KUMOD_SOURCE_SELECTION_RATE_' + .. source_name:upper() + local source_rate = os.getenv(source_rate_name) + if source_rate then + params.additional_source_selection_rates = + { [source_rate_name] = source_rate } + end + local username = os.getenv 'KUMOD_SMTP_AUTH_USERNAME' local password = os.getenv 'KUMOD_SMTP_AUTH_PASSWORD' diff --git a/crates/integration-tests/src/test/mod.rs b/crates/integration-tests/src/test/mod.rs index c97aa11d..a676d210 100644 --- a/crates/integration-tests/src/test/mod.rs +++ b/crates/integration-tests/src/test/mod.rs @@ -16,6 +16,8 @@ mod rebind; mod rebind_event_defined; mod rebind_event_missing; mod retry_schedule; +mod source_selection_rate; +mod source_selection_rate_pool; mod spf_basic; mod suspend_delivery_ready_q; mod suspend_delivery_ready_q_and_deliver; diff --git a/crates/integration-tests/src/test/source_selection_rate.rs b/crates/integration-tests/src/test/source_selection_rate.rs new file mode 100644 index 00000000..ef11ea8a --- /dev/null +++ b/crates/integration-tests/src/test/source_selection_rate.rs @@ -0,0 +1,77 @@ +use crate::kumod::{generate_message_text, DaemonWithMaildir, MailGenParams}; +use anyhow::Context; +use k9::assert_equal; +use kumo_log_types::RecordType; +use std::time::Duration; + +/// Check that setting source_selection_rate to 2/day limits us to delivering +/// only 2 out of the 10 messages we plan to send here, and that we get the +/// expected TransientFailure log for the messages that can't be delivered. +/// We have just a single source in this example. +#[tokio::test] +async fn source_selection_rate() -> anyhow::Result<()> { + let mut daemon = + DaemonWithMaildir::start_with_env(vec![("KUMOD_SOURCE_SELECTION_RATE", "2/day")]) + .await + .context("DaemonWithMaildir::start")?; + + eprintln!("sending message"); + let mut client = daemon.smtp_client().await.context("make smtp_client")?; + + let body = generate_message_text(1024, 78); + + const NUM_MSGS: usize = 10; + + for _ in 0..NUM_MSGS { + let response = MailGenParams { + body: Some(&body), + ..Default::default() + } + .send(&mut client) + .await + .context("send message")?; + eprintln!("{response:?}"); + anyhow::ensure!(response.code == 250); + } + + daemon + .wait_for_maildir_count(2, Duration::from_secs(10)) + .await; + + daemon.stop_both().await.context("stop_both")?; + println!("Stopped!"); + + let records = daemon.source.collect_logs()?; + let mut receptions = 0; + let mut delivery = 0; + let mut trans_fail = 0; + + for record in &records { + match record.kind { + RecordType::Reception => { + receptions += 1; + } + RecordType::Delivery => { + delivery += 1; + } + RecordType::TransientFailure => { + trans_fail += 1; + eprintln!("Considering TransientFailure: {record:#?}"); + assert_equal!( + record.response.content, + "KumoMTA internal: no sources for example.com pool=`unspecified` \ + are eligible for selection at this time" + ); + } + _ => { + panic!("unexpected record: {record:#?}"); + } + } + } + + assert_equal!(receptions, 10); + assert_equal!(delivery, 2); + assert_equal!(trans_fail, 8); + + Ok(()) +} diff --git a/crates/integration-tests/src/test/source_selection_rate_pool.rs b/crates/integration-tests/src/test/source_selection_rate_pool.rs new file mode 100644 index 00000000..8d256232 --- /dev/null +++ b/crates/integration-tests/src/test/source_selection_rate_pool.rs @@ -0,0 +1,77 @@ +use crate::kumod::{generate_message_text, DaemonWithMaildir, MailGenParams}; +use anyhow::Context; +use k9::assert_equal; +use kumo_log_types::RecordType; +use std::time::Duration; + +/// Check that setting source_selection_rate to 2/day for just one of the +/// sources in the pool results in only 2 messages being sent via that +/// source. The remainder should go through the other source in that pool +/// We use `warming` as the pool name here; source.lua knows that it should +/// generate two sources named `warming_a` and `warming_b` for that source. +/// We limit `warming_a` to `2/day`. +#[tokio::test] +async fn source_selection_rate_pool() -> anyhow::Result<()> { + let mut daemon = DaemonWithMaildir::start_with_env(vec![ + ("KUMOD_SOURCE_SELECTION_RATE_WARMING_A", "2/day"), + ("KUMOD_POOL_NAME", "warming"), + ]) + .await + .context("DaemonWithMaildir::start")?; + + let mut client = daemon.smtp_client().await.context("make smtp_client")?; + + let body = generate_message_text(1024, 78); + + const NUM_MSGS: usize = 10; + + for _ in 0..NUM_MSGS { + let response = MailGenParams { + body: Some(&body), + ..Default::default() + } + .send(&mut client) + .await + .context("send message")?; + eprintln!("{response:?}"); + anyhow::ensure!(response.code == 250); + } + + daemon + .wait_for_maildir_count(NUM_MSGS, Duration::from_secs(10)) + .await; + + daemon.stop_both().await.context("stop_both")?; + println!("Stopped!"); + + let records = daemon.source.collect_logs()?; + let mut receptions = 0; + let mut delivery_a = 0; + let mut delivery_b = 0; + + for record in &records { + match record.kind { + RecordType::Reception => { + receptions += 1; + } + RecordType::Delivery => match record.egress_source.as_deref().unwrap() { + "warming_a" => { + delivery_a += 1; + } + "warming_b" => { + delivery_b += 1; + } + wat => panic!("unexpected source {wat} in {record:#?}"), + }, + _ => { + panic!("unexpected record: {record:#?}"); + } + } + } + + assert_equal!(receptions, 10); + assert_equal!(delivery_a, 2); + assert_equal!(delivery_b, 8); + + Ok(()) +} diff --git a/crates/kumod/src/egress_source.rs b/crates/kumod/src/egress_source.rs index 332e8227..7a3d7a0e 100644 --- a/crates/kumod/src/egress_source.rs +++ b/crates/kumod/src/egress_source.rs @@ -521,6 +521,7 @@ impl EgressPoolSourceSelector { let mut entries = vec![]; let mut min_delay = None; + let mut is_full = false; // filter to non-suspended pathways for entry in &self.entries { @@ -571,17 +572,31 @@ impl EgressPoolSourceSelector { Ok(site) => { match site.make_reservation() { Some(reservation) => { - if !is_source_selection_throttled(deadline, &site, &source_name) - .await? + match get_source_selection_throttle_delay( + deadline, + &site, + &source_name, + ) + .await? { - site.redeem_reservation(msg, reservation).await; - return Ok(SourceInsertResult::Inserted); - } + None => { + site.redeem_reservation(msg, reservation).await; + return Ok(SourceInsertResult::Inserted); + } + Some(delay) => { + // Throttled; revise min delay to match throttle + if let Ok(delay) = chrono::Duration::from_std(delay) { + min_delay + .replace(min_delay.unwrap_or(delay).min(delay)); + } - // Throttled; fall through. + // fall through + } + } } None => { - // Not usable; fall through. + // Not usable; it is too full fall through. + is_full = true; } } } @@ -596,30 +611,37 @@ impl EgressPoolSourceSelector { // by going around again once we've filtered this // particular source out of the set entries.retain(|(entry, _)| entry.name != source_name); - - if entries.is_empty() { - // If there are no more sources, then we just - // report that the queue (whichever of the ones - // that we tried) is full. - return Err(ReadyQueueFull.into()); - } } None => { - return Ok(match min_delay { - Some(duration) => SourceInsertResult::Delay(duration), - None => SourceInsertResult::NoSources, - }) + // There are no more sources left to consider. + // + // If we definitively hit a full queue as one + // of the candidates, let's return that we are + // full + return if is_full { + Err(ReadyQueueFull.into()) + } else { + // If we got a delay value, it means that at least one + // of the candidates was either suspended until that duration, + // or was subject to a source_selection_rate with a duration. + // Let our response reflect that delay. + Ok(match min_delay { + Some(duration) => SourceInsertResult::Delay(duration), + None => SourceInsertResult::NoSources, + }) + }; } } } } } -async fn is_source_selection_throttled( +/// If selection is throttled, return Some(delay) +async fn get_source_selection_throttle_delay( deadline: Option, site: &ReadyQueueHandle, source_name: &str, -) -> anyhow::Result { +) -> anyhow::Result> { let path_config = site.get_path_config().borrow(); let mut throttles = Vec::with_capacity( @@ -645,7 +667,7 @@ async fn is_source_selection_throttled( } if throttles.is_empty() { - return Ok(false); + return Ok(None); } Box::pin(async move { @@ -659,11 +681,11 @@ async fn is_source_selection_throttled( opt_timeout_at(deadline, async { for (key, throttle) in throttles { let result = throttle.throttle(&key).await?; - if result.retry_after.is_some() { - return Ok(true); + if let Some(delay) = result.retry_after { + return Ok(Some(delay)); } } - Ok(false) + Ok(None) }) .await }) diff --git a/crates/kumod/src/queue.rs b/crates/kumod/src/queue.rs index 35c824e3..3fcdcbd4 100644 --- a/crates/kumod/src/queue.rs +++ b/crates/kumod/src/queue.rs @@ -2054,12 +2054,12 @@ impl Queue { detail: 4, }), content: format!( - "all possible sources for {} are suspended", - self.name + "KumoMTA internal: no sources for {} pool=`{}` are eligible for selection at this time", + self.name, source_selector.name ), command: None, }, - egress_pool: None, + egress_pool: Some(&source_selector.name), egress_source: None, relay_disposition: None, delivery_protocol: None, @@ -2070,7 +2070,11 @@ impl Queue { }) .await; context.note(InsertReason::LoggedTransientFailure); - anyhow::bail!("all possible sources for {} are suspended", self.name); + anyhow::bail!( + "no sources for {} pool=`{}` are eligible for selection at this time", + self.name, + source_selector.name + ); } SourceInsertResult::NoSources => { log_disposition(LogDisposition { @@ -2086,12 +2090,12 @@ impl Queue { detail: 4, }), content: format!( - "no non-zero-weighted sources available for {}. {:?}", - self.name, self.source_selector, + "KumoMTA internal: no sources available for {} pool=`{}`", + self.name, source_selector.name, ), command: None, }, - egress_pool: None, + egress_pool: Some(&source_selector.name), egress_source: None, relay_disposition: None, delivery_protocol: None, @@ -2102,7 +2106,11 @@ impl Queue { }) .await; context.note(InsertReason::LoggedTransientFailure); - anyhow::bail!("no non-zero-weighted sources available for {}", self.name); + anyhow::bail!( + "no sources available for {} pool=`{}`", + self.name, + source_selector.name + ); } SourceInsertResult::FailedResolve(err) => { log_disposition(LogDisposition {