From c76024fda1416046fc579fe27eb7435a1ea00ba4 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 26 Apr 2025 21:25:38 +0800 Subject: [PATCH] refactor: Update message writing logic to use message queue - Modify the Write method in the Message struct to utilize WriteMessageAsync for improved asynchronous message handling. - Update ignored message types to include "plan" alongside "loading", "error", "action", and "progress" for better message filtering. --- neo/message/message.go | 53 ++------------- neo/message/queue.go | 145 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 47 deletions(-) create mode 100644 neo/message/queue.go diff --git a/neo/message/message.go b/neo/message/message.go index 4afb92d3..f1a5c87b 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -432,7 +432,7 @@ func (m *Message) AppendTo(contents *Contents) *Message { } return m - case "loading", "error", "action", "progress": // Ignore progress, loading, action and error messages + case "loading", "error", "action", "progress", "plan": // Ignore progress, loading, plan and error messages return m default: @@ -621,52 +621,6 @@ func (m *Message) Callback(fn interface{}) *Message { return m } -// Write writes the message to response writer -func (m *Message) Write(w gin.ResponseWriter) bool { - - // Sync write to response writer - locker.Lock() - defer locker.Unlock() - - defer func() { - if r := recover(); r != nil { - - // Ignore if done is true - if m.IsDone { - return - } - - message := "Write Response Exception: (if client close the connection, it's normal) \n %s\n\n" - color.Red(message, r) - - // Print the message - raw, _ := jsoniter.MarshalToString(m) - color.White("Message:\n %s", raw) - } - }() - - // Ignore silent messages - if m.Silent { - return true - } - - data, err := jsoniter.Marshal(m) - if err != nil { - log.Error("%s", err.Error()) - return false - } - - data = append([]byte("data: "), data...) - data = append(data, []byte("\n\n")...) - - if _, err := w.Write(data); err != nil { - color.Red("Write JSON Message Error: %s", err.Error()) - return false - } - w.Flush() - return true -} - // WriteError writes an error message to response writer func (m *Message) WriteError(w gin.ResponseWriter, message string) { errMsg := strings.Trim(exception.New(message, 500).Message, "\"") @@ -730,3 +684,8 @@ func (a *Action) UnmarshalJSON(data []byte) error { } return nil } + +// Write writes the message to response writer using the message queue +func (m *Message) Write(w gin.ResponseWriter) bool { + return WriteMessageAsync(m, w) +} diff --git a/neo/message/queue.go b/neo/message/queue.go new file mode 100644 index 00000000..cff9de5b --- /dev/null +++ b/neo/message/queue.go @@ -0,0 +1,145 @@ +package message + +import ( + "sync" + "time" + + "github.com/fatih/color" + "github.com/gin-gonic/gin" + jsoniter "github.com/json-iterator/go" + "github.com/yaoapp/kun/log" +) + +// AsyncMessageQueue represents a queue for handling message writes +type AsyncMessageQueue struct { + queue chan *AsyncTask + workers int + wg sync.WaitGroup + shutdown chan struct{} +} + +// AsyncTask represents a task to write a message +type AsyncTask struct { + message *Message + writer gin.ResponseWriter + done chan bool +} + +var ( + defaultQueue *AsyncMessageQueue + queueOnce sync.Once +) + +// GetQueue returns the default message queue instance +func GetQueue() *AsyncMessageQueue { + queueOnce.Do(func() { + defaultQueue = NewAsyncQueue(10) // Initialize with 10 workers + defaultQueue.Start() + }) + return defaultQueue +} + +// NewAsyncQueue creates a new message queue with the specified number of workers +func NewAsyncQueue(workers int) *AsyncMessageQueue { + return &AsyncMessageQueue{ + queue: make(chan *AsyncTask, 1000), // Buffer size of 1000 + workers: workers, + shutdown: make(chan struct{}), + } +} + +// Start starts the message queue workers +func (mq *AsyncMessageQueue) Start() { + for i := 0; i < mq.workers; i++ { + mq.wg.Add(1) + go mq.worker() + } +} + +// Stop stops the message queue workers +func (mq *AsyncMessageQueue) Stop() { + close(mq.shutdown) + mq.wg.Wait() +} + +// worker processes messages from the queue +func (mq *AsyncMessageQueue) worker() { + defer mq.wg.Done() + + for { + select { + case task := <-mq.queue: + if task == nil { + continue + } + success := writeMessageToResponse(task.message, task.writer) + if task.done != nil { + task.done <- success + } + case <-mq.shutdown: + return + } + } +} + +// WriteMessageAsync writes the message to response writer using the message queue +func WriteMessageAsync(m *Message, w gin.ResponseWriter) bool { + task := &AsyncTask{ + message: m, + writer: w, + done: nil, // No need for done channel anymore + } + + // Try to send the task to the queue with a short timeout + select { + case GetQueue().queue <- task: + return true + case <-time.After(100 * time.Millisecond): // Reduced timeout since we don't wait for result + log.Error("Queue is full, message dropped") + return false + } +} + +// writeMessageToResponse writes the message directly to the response writer +func writeMessageToResponse(m *Message, w gin.ResponseWriter) bool { + // Sync write to response writer + locker.Lock() + defer locker.Unlock() + + defer func() { + if r := recover(); r != nil { + // Ignore if done is true + if m.IsDone { + return + } + + message := "Write Response Exception: (if client close the connection, it's normal) \n %s\n\n" + color.Red(message, r) + + // Print the message + raw, _ := jsoniter.MarshalToString(m) + color.White("Message:\n %s", raw) + } + }() + + // Ignore silent messages + if m.Silent { + return true + } + + data, err := jsoniter.Marshal(m) + if err != nil { + log.Error("%s", err.Error()) + return false + } + + data = append([]byte("data: "), data...) + data = append(data, []byte("\n\n")...) + + if _, err := w.Write(data); err != nil { + color.Red("Write JSON Message Error: %s", err.Error()) + return false + } + w.Flush() + return true +}