perf:(menu): 优化GetCTX方法(处理data race)

yqchilde · Dec 31, 2022 · 58a93fa · 58a93fa
1 parent 21fa847
commit 58a93fa
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 5 deletions.
diff --git a/engine/control/register.go b/engine/control/register.go
@@ -20,7 +20,7 @@ func Register(service string, o *Options[*robot.Ctx]) *Engine {
 		log.Fatalf("[%s]插件优先级 %d 已被 %s 占用", service, priority, s)
 	}
 	priorityMap[priority] = service
-	log.Debugf("[%s]插件已注册, 优先级: %d", service, priority)
+	log.Printf("[%s]插件已注册, 优先级: %d", service, priority)
 	return newEngine(service, o)
 }
 

diff --git a/engine/robot/robot.go b/engine/robot/robot.go
@@ -4,13 +4,15 @@ import (
 	"fmt"
 	"math/rand"
 	"runtime/debug"
+	"sync/atomic"
 	"time"
 
 	"github.com/yqchilde/wxbot/engine/pkg/log"
 )
 
 var (
 	BotConfig   *Config      // 机器人配置
+	Framework   atomic.Value // 当前机器人框架
 	eventBuffer *EventBuffer // 事件缓冲区
 )
 
@@ -238,11 +240,13 @@ func GetCTX() *Ctx {
 		case <-t.C:
 			log.Fatal("[robot] 获取CTX超时")
 		default:
-			if BotConfig != nil {
-				t.Stop()
-				return &Ctx{framework: BotConfig.Framework}
+			framework := Framework.Load()
+			if framework == nil {
+				time.Sleep(time.Second)
+				continue
 			}
-			time.Sleep(time.Second)
+			t.Stop()
+			return &Ctx{framework: framework.(IFramework)}
 		}
 	}
 }
diff --git a/framework/qianxun/callback.go b/framework/qianxun/callback.go
@@ -41,6 +41,7 @@ func New(botWxId, apiUrl, apiToken string, servePort uint) *Framework {
 }
 
 func (f *Framework) Callback(handler func(*robot.Event, robot.IFramework)) {
+	robot.Framework.Store(f)
 	http.HandleFunc("/wxbot/callback", func(w http.ResponseWriter, r *http.Request) {
 		recv, err := io.ReadAll(r.Body)
 		if err != nil {

diff --git a/framework/vlw/callback.go b/framework/vlw/callback.go
@@ -43,6 +43,7 @@ func New(botWxId, apiUrl, apiToken string, servePort uint) *Framework {
 }
 
 func (f *Framework) Callback(handler func(*robot.Event, robot.IFramework)) {
+	robot.Framework.Store(f)
 	http.HandleFunc("/wxbot/callback", func(w http.ResponseWriter, r *http.Request) {
 		recv, err := io.ReadAll(r.Body)
 		if err != nil {