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:
parent
77efc03582
commit
c76024fda1
2 changed files with 151 additions and 47 deletions
|
|
@ -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
145
neo/message/queue.go
Normal 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
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue