Skip to content

Commit b6bbba8

Browse files
People-11People-11
committed
Fix id register overwrite
1 parent cb31b34 commit b6bbba8

File tree

1 file changed

+63
-38
lines changed

1 file changed

+63
-38
lines changed

‎bot/src/handler.rs‎

Lines changed: 63 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ pub enum Command {
5555
description = "Sync a gallery(e-hentai/exhentai/nhentai are supported now). 同步一个画廊(目前支持 EH/EX/NH)"
5656
)]
5757
Sync(String),
58-
#[command(description = "Cancel current sync operation. 取消当前正在进行的本子同步")]
58+
#[command(description = "Cancel all ongoing sync operations. 取消所有正在进行的同步操作。")]
5959
Cancel,
6060
}
6161

@@ -75,8 +75,8 @@ pub struct Handler<C> {
7575

7676
single_flight: singleflight_async::SingleFlight<String>,
7777

78-
// Add this field to track active syncs
79-
active_syncs: Arc<Mutex<HashMap<i64, oneshot::Sender<()>>>>,
78+
// One user can have multiple active syncs
79+
active_syncs: Arc<Mutex<HashMap<i64, Vec<(String, oneshot::Sender<()>)>>>>,
8080
}
8181

8282
impl<C> Handler<C>
@@ -126,28 +126,54 @@ where
126126
self.whitelist.contains(&chat_id)
127127
}
128128

129-
// Add these new helper methods for sync cancellation
130-
fn register_sync(&self, user_id: i64) -> oneshot::Receiver<()> {
129+
// Support Multiple Sync Task
130+
fn register_sync(&self, user_id: i64, url: &str) -> oneshot::Receiver<()> {
131131
let (tx, rx) = oneshot::channel();
132-
self.active_syncs.lock().unwrap().insert(user_id, tx);
132+
133+
let mut active_syncs = self.active_syncs.lock().unwrap();
134+
135+
let user_syncs = active_syncs.entry(user_id).or_insert_with(Vec::new);
136+
137+
user_syncs.push((url.to_string(), tx));
138+
133139
rx
134140
}
135141

136-
fn unregister_sync(&self, user_id: i64) {
137-
self.active_syncs.lock().unwrap().remove(&user_id);
142+
fn unregister_sync(&self, user_id: i64, url: &str) {
143+
let mut active_syncs = self.active_syncs.lock().unwrap();
144+
145+
if let Some(user_syncs) = active_syncs.get_mut(&user_id) {
146+
user_syncs.retain(|(sync_url, _)| sync_url != url);
147+
148+
if user_syncs.is_empty() {
149+
active_syncs.remove(&user_id);
150+
}
151+
}
138152
}
139153

140-
fn cancel_sync(&self, user_id: i64) -> bool {
141-
if let Some(tx) = self.active_syncs.lock().unwrap().remove(&user_id) {
142-
// Send cancellation signal
143-
let _ = tx.send(());
154+
fn cancel_all_syncs(&self, user_id: i64) -> usize {
155+
let mut active_syncs = self.active_syncs.lock().unwrap();
156+
157+
if let Some(user_syncs) = active_syncs.remove(&user_id) {
158+
let count = user_syncs.len();
159+
160+
// Send all cancellation signal
161+
for (url, tx) in user_syncs {
162+
info!(
163+
"[cancel handler] cancelling sync for user {} and url {}",
164+
user_id, url
165+
);
166+
let _ = tx.send(());
167+
}
168+
144169
info!(
145-
"[cancel handler] user {} cancelled their sync operation",
146-
user_id
170+
"[cancel handler] user {} cancelled {} sync operations",
171+
user_id, count
147172
);
148-
true
173+
174+
count
149175
} else {
150-
false
176+
0
151177
}
152178
}
153179

@@ -218,28 +244,30 @@ where
218244
.await
219245
);
220246

221-
// Register this sync with the user's ID
222-
let cancel_rx = self.register_sync(msg.chat.id.0);
247+
let url_clone = url.clone();
248+
let cancel_rx = self.register_sync(msg.chat.id.0, &url);
223249

224250
tokio::spawn(async move {
225251
let result = self.sync_response(&url, cancel_rx).await;
226252

227-
// Unregister sync when done
228-
self.unregister_sync(msg.chat.id.0);
253+
self.unregister_sync(msg.chat.id.0, &url_clone);
229254

230255
let _ = bot.edit_message_text(msg.chat.id, msg.id, result).await;
231256
});
232257
}
233258
Command::Cancel => {
234-
let cancelled = self.cancel_sync(msg.chat.id.0);
235-
if cancelled {
259+
let cancelled_count = self.cancel_all_syncs(msg.chat.id.0);
260+
if cancelled_count > 0 {
236261
let _ = bot
237-
.send_message(msg.chat.id, escape("Sync operation cancelled."))
262+
.send_message(
263+
msg.chat.id,
264+
escape(&format!("Cancelled {} sync operations.", cancelled_count)),
265+
)
238266
.reply_to_message_id(msg.id)
239267
.await;
240268
} else {
241269
let _ = bot
242-
.send_message(msg.chat.id, escape("No active sync operation to cancel."))
270+
.send_message(msg.chat.id, escape("No active sync operations to cancel."))
243271
.reply_to_message_id(msg.id)
244272
.await;
245273
}
@@ -313,14 +341,13 @@ where
313341
.await
314342
);
315343

316-
// Register this sync with the user's ID
317-
let cancel_rx = self.register_sync(msg.chat.id.0);
344+
let url_clone = url.clone();
345+
let cancel_rx = self.register_sync(msg.chat.id.0, &url);
318346

319347
tokio::spawn(async move {
320348
let result = self.sync_response(&url, cancel_rx).await;
321349

322-
// Unregister sync when done
323-
self.unregister_sync(msg.chat.id.0);
350+
self.unregister_sync(msg.chat.id.0, &url_clone);
324351

325352
let _ = bot.edit_message_text(msg.chat.id, msg.id, result).await;
326353
});
@@ -383,14 +410,13 @@ where
383410
.await
384411
);
385412

386-
// Register this sync with the user's ID
387-
let cancel_rx = self.register_sync(msg.chat.id.0);
413+
let url_clone = url.clone();
414+
let cancel_rx = self.register_sync(msg.chat.id.0, &url);
388415

389416
tokio::spawn(async move {
390417
let result = self.sync_response(&url, cancel_rx).await;
391418

392-
// Unregister sync when done
393-
self.unregister_sync(msg.chat.id.0);
419+
self.unregister_sync(msg.chat.id.0, &url_clone);
394420

395421
let _ = bot.edit_message_text(msg.chat.id, msg.id, result).await;
396422
});
@@ -467,14 +493,13 @@ where
467493
.reply_to_message_id(msg.id)
468494
.await
469495
{
470-
// Register this sync with the user's ID
471-
let cancel_rx = self.register_sync(msg.chat.id.0);
496+
let url_clone = url.clone();
497+
let cancel_rx = self.register_sync(msg.chat.id.0, &url);
472498

473499
tokio::spawn(async move {
474500
let result = self.sync_response(&url, cancel_rx).await;
475501

476-
// Unregister sync when done
477-
self.unregister_sync(msg.chat.id.0);
502+
self.unregister_sync(msg.chat.id.0, &url_clone);
478503

479504
let _ = bot.edit_message_text(msg.chat.id, msg.id, result).await;
480505
});
@@ -505,8 +530,8 @@ where
505530
tokio::select! {
506531
result = self.single_flight.work(url, || async {
507532
match self.route_sync(url).await {
508-
Ok(url) => {
509-
format!("Sync to telegraph finished: {}", link(&url, &escape(&url)))
533+
Ok(sync_url) => {
534+
format!("Sync to telegraph finished: {}", link(&sync_url, &escape(&sync_url)))
510535
}
511536
Err(e) => {
512537
format!("Sync to telegraph failed: {}", escape(&e.to_string()))

0 commit comments

Comments
 (0)