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.
This commit is contained in:
Max 2025-04-26 21:25:38 +08:00
parent 77efc03582
commit c76024fda1
2 changed files with 151 additions and 47 deletions

View file

@ -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)
}

145
neo/message/queue.go Normal file
View file

@ -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
}