Skip to content

Commit 81850c1

Browse files
超渡法師brettchien
authored andcommitted
fix(dispatch): idle eviction, config validation, avoid clone, timestamp precision
- Add 5-min idle timeout to consumer_loop to prevent per-thread handle/task leak (unbounded growth from one-shot thread keys like Slack non-thread msgs) - Validate max_buffered_messages > 0 at config load time (prevents panic from tokio::sync::mpsc::channel(0)) - Use into_iter() in dispatch_batch to avoid deep-copying extra_blocks (may contain base64 image data) - Add TODO comment for gateway multibot detection - Use real milliseconds in now_iso8601() via dur.subsec_millis() Co-authored-by: 超渡法師 <chaodu@openab.dev>
1 parent 66fcf3f commit 81850c1

4 files changed

Lines changed: 54 additions & 20 deletions

File tree

src/config.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -504,6 +504,18 @@ fn parse_config(raw: &str, source: &str) -> anyhow::Result<Config> {
504504
let expanded = expand_env_vars(raw);
505505
let config: Config = toml::from_str(&expanded)
506506
.map_err(|e| anyhow::anyhow!("failed to parse config from {source}: {e}"))?;
507+
508+
// Validate max_buffered_messages > 0 (tokio::sync::mpsc::channel panics on 0).
509+
if let Some(ref d) = config.discord {
510+
anyhow::ensure!(d.max_buffered_messages > 0, "discord.max_buffered_messages must be > 0");
511+
}
512+
if let Some(ref s) = config.slack {
513+
anyhow::ensure!(s.max_buffered_messages > 0, "slack.max_buffered_messages must be > 0");
514+
}
515+
if let Some(ref g) = config.gateway {
516+
anyhow::ensure!(g.max_buffered_messages > 0, "gateway.max_buffered_messages must be > 0");
517+
}
518+
507519
Ok(config)
508520
}
509521

src/dispatch.rs

Lines changed: 34 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,10 @@
1111
use std::collections::HashMap;
1212
use std::sync::atomic::{AtomicU64, Ordering};
1313
use std::sync::{Arc, Mutex};
14-
use std::time::Instant;
14+
use std::time::{Duration, Instant};
1515

1616
use anyhow::Result;
17-
use tracing::{error, info, info_span, warn};
17+
use tracing::{debug, error, info, info_span, warn};
1818

1919
use crate::adapter::{AdapterRouter, ChannelRef, ChatAdapter, MessageRef};
2020
use crate::acp::ContentBlock;
@@ -255,6 +255,12 @@ impl Dispatcher {
255255
// consumer_loop
256256
// ---------------------------------------------------------------------------
257257

258+
/// Idle timeout for per-thread consumer tasks. When no message arrives within
259+
/// this window the consumer exits, allowing `per_thread` map cleanup on the
260+
/// next `submit` (via `SendError` → `try_evict_locked`). Prevents unbounded
261+
/// task/memory growth from one-shot thread keys (e.g. Slack non-thread messages).
262+
const CONSUMER_IDLE_TIMEOUT: Duration = Duration::from_secs(300); // 5 min
263+
258264
#[allow(clippy::too_many_arguments)]
259265
async fn consumer_loop(
260266
thread_key: String,
@@ -271,13 +277,26 @@ async fn consumer_loop(
271277

272278
loop {
273279
// I1: block until at least one message arrives (zero latency for first message).
280+
// Idle timeout: if no message arrives within CONSUMER_IDLE_TIMEOUT the
281+
// consumer exits, freeing the task and mpsc. The next `submit` for this
282+
// thread_key will observe `SendError`, evict the stale entry, and lazily
283+
// spawn a fresh consumer (§2.5 generation check prevents mis-eviction).
274284
let first = match pending.take() {
275285
Some(msg) => msg,
276-
None => match rx.recv().await {
277-
Some(msg) => msg,
278-
// All senders dropped → either shutdown() cleared the map, or
279-
// cancel_buffered() removed our entry. Exit the loop.
280-
None => break,
286+
None => match tokio::time::timeout(CONSUMER_IDLE_TIMEOUT, rx.recv()).await {
287+
Ok(Some(msg)) => msg,
288+
Ok(None) => {
289+
// All senders dropped → shutdown() or cancel_buffered().
290+
break;
291+
}
292+
Err(_elapsed) => {
293+
debug!(
294+
thread_key = %thread_key,
295+
channel = %thread_channel.channel_id,
296+
"consumer idle timeout, exiting"
297+
);
298+
break;
299+
}
281300
},
282301
};
283302

@@ -340,19 +359,23 @@ async fn dispatch_batch(
340359
)
341360
.await;
342361

343-
// Collect per-event observability data.
362+
// Collect per-event observability data (before consuming the batch).
344363
let tokens_per_event: Vec<usize> = batch.iter().map(|m| m.estimated_tokens).collect();
345364
let wait_ms: Vec<u128> = batch
346365
.iter()
347366
.map(|m| m.arrived_at.elapsed().as_millis())
348367
.collect();
349-
let senders: Vec<&str> = batch.iter().map(|m| m.sender_name.as_str()).collect();
368+
let senders: Vec<String> = batch.iter().map(|m| m.sender_name.clone()).collect();
369+
370+
// Anchor reactions on the last message in the batch (before consuming).
371+
let trigger_msg = batch.last().unwrap().trigger_msg.clone();
350372

351373
// Pack all arrival events into one Vec<ContentBlock> (§3.3).
374+
// Uses into_iter() to avoid deep-copying extra_blocks (may contain base64 image data).
352375
let mut content_blocks: Vec<ContentBlock> = Vec::new();
353-
for msg in &batch {
376+
for msg in batch {
354377
let mut event_blocks =
355-
AdapterRouter::pack_arrival_event(&msg.sender_json, &msg.prompt, msg.extra_blocks.clone());
378+
AdapterRouter::pack_arrival_event(&msg.sender_json, &msg.prompt, msg.extra_blocks);
356379
content_blocks.append(&mut event_blocks);
357380
}
358381
let packed_block_count = content_blocks.len();
@@ -367,8 +390,6 @@ async fn dispatch_batch(
367390
return;
368391
}
369392

370-
// Anchor reactions on the last message in the batch.
371-
let trigger_msg = batch.last().unwrap().trigger_msg.clone();
372393
let reactions_config = router.reactions_config().clone();
373394
let reactions = Arc::new(StatusReactionController::new(
374395
reactions_config.enabled,

src/gateway.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -564,6 +564,7 @@ pub async fn run_gateway_adapter(
564564
trigger_msg,
565565
arrived_at: std::time::Instant::now(),
566566
estimated_tokens,
567+
// TODO: implement gateway multibot detection
567568
other_bot_present: false,
568569
};
569570
if let Err(e) = dispatcher

src/timestamp.rs

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,13 +42,12 @@ pub fn slack_ts_to_iso8601(ts: &str) -> String {
4242
unix_to_iso8601(secs, ms)
4343
}
4444

45-
/// Current wall-clock instant as ISO 8601 UTC (millisecond precision, always `.000Z`).
45+
/// Current wall-clock instant as ISO 8601 UTC with millisecond precision.
4646
pub fn now_iso8601() -> String {
47-
let secs = SystemTime::now()
47+
let dur = SystemTime::now()
4848
.duration_since(UNIX_EPOCH)
49-
.unwrap_or_default()
50-
.as_secs();
51-
unix_to_iso8601(secs, 0)
49+
.unwrap_or_default();
50+
unix_to_iso8601(dur.as_secs(), (dur.subsec_millis()) as u64)
5251
}
5352

5453
#[cfg(test)]
@@ -79,10 +78,11 @@ mod tests {
7978
#[test]
8079
fn now_iso8601_has_expected_shape() {
8180
let s = now_iso8601();
82-
// YYYY-MM-DDTHH:MM:SS.000Z = 24 chars
81+
// YYYY-MM-DDTHH:MM:SS.mmmZ = 24 chars
8382
assert_eq!(s.len(), 24);
84-
assert!(s.ends_with(".000Z"));
83+
assert!(s.ends_with('Z'));
8584
assert_eq!(&s[4..5], "-");
8685
assert_eq!(&s[10..11], "T");
86+
assert_eq!(&s[19..20], ".");
8787
}
8888
}

0 commit comments

Comments
 (0)