From 376113a0b6ae6c3f3316424d79b3e5cd55b145c0 Mon Sep 17 00:00:00 2001
From: LamGC <lam827@lamgc.net>
Date: Mon, 25 Mar 2024 11:32:30 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E9=98=9F=E5=88=97?=
 =?UTF-8?q?=E5=B7=B2=E6=BB=A1=E5=AF=BC=E8=87=B4=20Task=20=E6=B7=BB?=
 =?UTF-8?q?=E5=8A=A0=E5=A4=B1=E8=B4=A5=E5=90=8E=E4=BB=8D=E5=88=A4=E5=AE=9A?=
 =?UTF-8?q?=E4=B8=BA=E5=8A=A0=E5=85=A5=E9=98=9F=E5=88=97=E6=88=90=E5=8A=9F?=
 =?UTF-8?q?=E7=9A=84=E9=97=AE=E9=A2=98.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/main/kotlin/OneDriveTransferCenter.kt | 25 ++++++++++++--------
 src/main/kotlin/Services.kt               | 28 ++++++++++++-----------
 2 files changed, 30 insertions(+), 23 deletions(-)

diff --git a/src/main/kotlin/OneDriveTransferCenter.kt b/src/main/kotlin/OneDriveTransferCenter.kt
index f9e7d03..3bdb775 100644
--- a/src/main/kotlin/OneDriveTransferCenter.kt
+++ b/src/main/kotlin/OneDriveTransferCenter.kt
@@ -29,11 +29,9 @@ import java.util.concurrent.*
 object OneDriveTransferCenter {
 
     val executor =
-        OneDriveTransferTaskExecutor(1, DefaultOneDriveTransferCallback, ArrayBlockingQueue(100))
+        OneDriveTransferTaskExecutor(1, DefaultOneDriveTransferCallback, ArrayBlockingQueue(Integer.MAX_VALUE))
 
-    fun submitUploadTask(task: OneDriveTransferTask) {
-        executor.submitTransferTask(task)
-    }
+    fun submitUploadTask(task: OneDriveTransferTask): Boolean = executor.submitTransferTask(task)
 
     fun cancelUploadTask(task: OneDriveTransferTask): Boolean = executor.cancelTransferTask(task)
 
@@ -291,9 +289,12 @@ class OneDriveTransferTaskExecutor(
         }
     }
 
-    fun submitTransferTask(task: OneDriveTransferTask) {
-        callback.onTransferTaskCreated(task)
-        taskQueue.offer(task)
+    fun submitTransferTask(task: OneDriveTransferTask): Boolean {
+        if (taskQueue.remainingCapacity() > 0) {
+            callback.onTransferTaskCreated(task)
+            return taskQueue.offer(task)
+        }
+        return false
     }
 
     fun cancelTransferTask(task: OneDriveTransferTask): Boolean {
@@ -348,9 +349,13 @@ class OneDriveTransferTaskExecutor(
                 }
                 logger.warn(e) { "OneDrive 中转任务执行失败: ${e.message}" }
                 progress.status = OneDriveTransferStatus.FAILURE
-                callback.onTransferFailure(task, progress.apply {
-                    this.exception = e
-                })
+                try {
+                    callback.onTransferFailure(task, progress.apply {
+                        this.exception = e
+                    })
+                } catch (e: Exception) {
+                    logger.warn(e) { "OneDrive 中转任务失败回调失败: ${e.message}" }
+                }
             } finally {
                 logger.info { "线程 $id 任务执行完毕: ${task.document.fileName}" }
                 threadStatusMap.remove(id)
diff --git a/src/main/kotlin/Services.kt b/src/main/kotlin/Services.kt
index 3a52fbc..37ea8fa 100644
--- a/src/main/kotlin/Services.kt
+++ b/src/main/kotlin/Services.kt
@@ -109,19 +109,21 @@ class OneDriveTransferService(
     fun submitUploadDocumentTask(userId: Long, messageId: Int, document: Document) {
         val transferSetting =
             accountManager.getTransferSetting(userId) ?: throw OneDriveNotLoginException()
-        OneDriveTransferCenter.submitUploadTask(
-            OneDriveTransferTask(
-                userId,
-                bot,
-                this,
-                document,
-                transferSetting.driveId,
-                transferSetting.storagePath,
-            ).apply {
-                extra["chatId"] = userId
-                extra["messageId"] = messageId
-            }.also { tasks.add(it) }
-        )
+        val task = OneDriveTransferTask(
+            userId,
+            bot,
+            this,
+            document,
+            transferSetting.driveId,
+            transferSetting.storagePath,
+        ).apply {
+            extra["chatId"] = userId
+            extra["messageId"] = messageId
+        }
+        val submitted = OneDriveTransferCenter.submitUploadTask(task)
+        if (submitted) {
+            tasks.add(task)
+        }
     }
 
     fun cancelTask(userId: Long, replyMessageId: Int): Boolean {