From 3bff930818d6d579fd78cae7229a935b472ca3a9 Mon Sep 17 00:00:00 2001
From: Jian Qiu <swordqiu@gmail.com>
Date: Sun, 13 Apr 2025 15:43:46 +0800
Subject: [PATCH] fix: handler support worker manager callback (#22404)

Co-authored-by: Qiu Jian <qiujian@yunionyun.com>
---
 pkg/appsrv/appsrv.go      |  4 +--
 pkg/appsrv/handlerinfo.go | 76 ++++++++++++++++++++++++++-------------
 2 files changed, 53 insertions(+), 27 deletions(-)

diff --git a/pkg/appsrv/appsrv.go b/pkg/appsrv/appsrv.go
index 9dd97067fd..f07e823d4e 100644
--- a/pkg/appsrv/appsrv.go
+++ b/pkg/appsrv/appsrv.go
@@ -393,7 +393,7 @@ func (app *Application) defaultHandle(w http.ResponseWriter, r *http.Request, ri
 				fw:     newResponseWriterChannel(w),
 				r:      r,
 				segs:   segs,
-				to:     hand.FetchProcessTimeout(r),
+				to:     hand.fetchProcessTimeout(r),
 				cancel: nil,
 			}
 
@@ -408,7 +408,7 @@ func (app *Application) defaultHandle(w http.ResponseWriter, r *http.Request, ri
 				defer task.cancel()
 			}
 			task.ctx = appctx.WithRequestLang(task.ctx, r)
-			session := hand.workerMan
+			session := hand.fetchWorkerManager(r)
 			if session == nil {
 				if r.Method == "GET" || r.Method == "HEAD" {
 					session = app.readSession
diff --git a/pkg/appsrv/handlerinfo.go b/pkg/appsrv/handlerinfo.go
index e457c9dbf3..f4097fd434 100644
--- a/pkg/appsrv/handlerinfo.go
+++ b/pkg/appsrv/handlerinfo.go
@@ -28,6 +28,8 @@ type handlerRequestCounter struct {
 
 type TProcessTimeoutCallback func(*SHandlerInfo, *http.Request) time.Duration
 
+type TWorkerManagerCallback func(*SHandlerInfo, *http.Request) *SWorkerManager
+
 type SHandlerInfo struct {
 	method     string
 	path       []string
@@ -38,27 +40,66 @@ type SHandlerInfo struct {
 	counter2XX handlerRequestCounter
 	counter4XX handlerRequestCounter
 	counter5XX handlerRequestCounter
-	workerMan  *SWorkerManager
-	skipLog    bool
+
+	skipLog bool
 
 	processTimeout         time.Duration
 	processTimeoutCallback TProcessTimeoutCallback
+
+	workerManager         *SWorkerManager
+	workerManagerCallback TWorkerManagerCallback
 }
 
-func (this *SHandlerInfo) FetchProcessTimeout(r *http.Request) time.Duration {
-	if this.processTimeoutCallback != nil {
-		tm := this.processTimeoutCallback(this, r)
-		if tm < this.processTimeout {
-			tm = this.processTimeout
+func (hi *SHandlerInfo) fetchProcessTimeout(r *http.Request) time.Duration {
+	if hi.processTimeoutCallback != nil {
+		tm := hi.processTimeoutCallback(hi, r)
+		if tm < hi.processTimeout {
+			tm = hi.processTimeout
 		}
 		return tm
 	} else {
-		return this.processTimeout
+		return hi.processTimeout
 	}
 }
 
-func (this *SHandlerInfo) SetProcessTimeoutCallback(callback TProcessTimeoutCallback) {
-	this.processTimeoutCallback = callback
+func (hi *SHandlerInfo) SetProcessTimeoutCallback(callback TProcessTimeoutCallback) *SHandlerInfo {
+	hi.processTimeoutCallback = callback
+	return hi
+}
+
+func (hi *SHandlerInfo) SetProcessTimeout(to time.Duration) *SHandlerInfo {
+	hi.processTimeout = to
+	return hi
+}
+
+func (hi *SHandlerInfo) SetProcessNoTimeout() *SHandlerInfo {
+	hi.processTimeout = -1
+	return hi
+}
+
+func (hi *SHandlerInfo) fetchWorkerManager(r *http.Request) *SWorkerManager {
+	if hi.workerManagerCallback != nil {
+		wm := hi.workerManagerCallback(hi, r)
+		if wm != nil {
+			return wm
+		}
+	}
+	return hi.workerManager
+}
+
+func (hi *SHandlerInfo) SetWorkerManagerCallback(callback TWorkerManagerCallback) *SHandlerInfo {
+	hi.workerManagerCallback = callback
+	return hi
+}
+
+func (hi *SHandlerInfo) SetWorkerManager(workerMan *SWorkerManager) *SHandlerInfo {
+	hi.workerManager = workerMan
+	return hi
+}
+
+func (hi *SHandlerInfo) ClearWorkerManager() *SHandlerInfo {
+	hi.workerManager = nil
+	return hi
 }
 
 func (this *SHandlerInfo) GetName(params map[string]string) string {
@@ -126,21 +167,6 @@ func (hi *SHandlerInfo) SetTags(tags map[string]string) *SHandlerInfo {
 	return hi
 }
 
-func (hi *SHandlerInfo) SetProcessTimeout(to time.Duration) *SHandlerInfo {
-	hi.processTimeout = to
-	return hi
-}
-
-func (hi *SHandlerInfo) SetProcessNoTimeout() *SHandlerInfo {
-	hi.processTimeout = -1
-	return hi
-}
-
-func (hi *SHandlerInfo) SetWorkerManager(workerMan *SWorkerManager) *SHandlerInfo {
-	hi.workerMan = workerMan
-	return hi
-}
-
 func (hi *SHandlerInfo) SetSkipLog(skip bool) *SHandlerInfo {
 	hi.skipLog = skip
 	return hi