From a219a1f67e17cf3cfaf61e43ab23480e994fcfff Mon Sep 17 00:00:00 2001
From: chengqian <wuchengqian5@163.com>
Date: Mon, 8 Jun 2026 17:39:21 +0800
Subject: [PATCH] =?UTF-8?q?fix:=E7=B3=BB=E7=BB=9F=E4=BC=98=E9=9B=85?=
 =?UTF-8?q?=E9=80=80=E5=87=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/cmd/cmd.go                         |  4 +++-
 internal/logic/task/async_polling.go        | 10 ++++++++++
 internal/logic/tenant/webhook_dispatcher.go | 17 +++++++++++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/internal/cmd/cmd.go b/internal/cmd/cmd.go
index 0cfbce8..7b5eac2 100644
--- a/internal/cmd/cmd.go
+++ b/internal/cmd/cmd.go
@@ -302,7 +302,7 @@ var (
 
 			// Initialize active task count and start polling
 			task.InitActiveCount(ctx)
-			go task.StartAsyncPolling(ctx)
+			task.StartAsyncPolling(ctx)
 
 			// Start webhook dispatcher (event-driven delivery)
 			tenant.InitWebhookDispatcher(ctx)
@@ -312,6 +312,8 @@ var (
 
 			// Flush usage log writer on shutdown (s.Run blocks until server stops)
 			defer plugin.Shutdown(ctx)
+			defer tenant.ShutdownWebhookDispatcher()
+			defer task.StopAsyncPolling()
 			defer common.CloseChannelErrorWriter()
 			defer common.CloseUsageLogWriter()
 			defer response.CloseErrorLogWriter()
diff --git a/internal/logic/task/async_polling.go b/internal/logic/task/async_polling.go
index 12967d0..09392fa 100644
--- a/internal/logic/task/async_polling.go
+++ b/internal/logic/task/async_polling.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"io"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/gogf/gf/v2/frame/g"
@@ -24,10 +25,14 @@ const (
 	timeoutBatchSize = 100
 )
 
+var pollingWg sync.WaitGroup
+
 // StartAsyncPolling 启动异步任务轮询 goroutine
 func StartAsyncPolling(ctx context.Context) {
 	g.Log().Info(ctx, "Starting async task polling...")
+	pollingWg.Add(1)
 	go func() {
+		defer pollingWg.Done()
 		ticker := time.NewTicker(pollInterval)
 		defer ticker.Stop()
 
@@ -43,6 +48,11 @@ func StartAsyncPolling(ctx context.Context) {
 	}()
 }
 
+// StopAsyncPolling 等待轮询 goroutine 完全退出，在服务停机时调用。
+func StopAsyncPolling() {
+	pollingWg.Wait()
+}
+
 // pollOnce 执行一次轮询
 func pollOnce(ctx context.Context) {
 	if !HasActiveTasks() {
diff --git a/internal/logic/tenant/webhook_dispatcher.go b/internal/logic/tenant/webhook_dispatcher.go
index 080ed67..e929495 100644
--- a/internal/logic/tenant/webhook_dispatcher.go
+++ b/internal/logic/tenant/webhook_dispatcher.go
@@ -2,6 +2,7 @@ package tenant
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/gogf/gf/v2/frame/g"
@@ -27,6 +28,7 @@ const (
 type WebhookDispatcher struct {
 	eventCh chan int64
 	ctx     context.Context
+	wg      sync.WaitGroup
 }
 
 var defaultDispatcher *WebhookDispatcher
@@ -40,6 +42,7 @@ func InitWebhookDispatcher(ctx context.Context) {
 	defaultDispatcher = d
 
 	for i := 0; i < dispatcherWorkerCount; i++ {
+		d.wg.Add(1)
 		go d.worker()
 	}
 	go d.sweepLoop()
@@ -65,6 +68,7 @@ func NotifyNewEvent(eventID int64) {
 
 // worker consumes event IDs from the channel and delivers them.
 func (d *WebhookDispatcher) worker() {
+	defer d.wg.Done()
 	for {
 		select {
 		case <-d.ctx.Done():
@@ -209,3 +213,16 @@ func (d *WebhookDispatcher) cleanup() {
 		g.Log().Infof(ctx, "webhook cleanup: deleted %d exhausted events older than %v", rows, failedCutoff.Format("Y-m-d"))
 	}
 }
+
+// Shutdown waits for all in-flight deliveries to complete.
+func (d *WebhookDispatcher) Shutdown() {
+	d.wg.Wait()
+}
+
+// ShutdownWebhookDispatcher waits for all in-flight webhook deliveries to complete.
+// Call this via defer after s.Run() returns.
+func ShutdownWebhookDispatcher() {
+	if defaultDispatcher != nil {
+		defaultDispatcher.Shutdown()
+	}
+}