118 lines
3.0 KiB
Go
118 lines
3.0 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// Coalescing is safe because each snapshot includes the persisted detection
|
|
// history. Subscribers register before reading their initial snapshot.
|
|
var groupStatusEvents = struct {
|
|
sync.Mutex
|
|
clients map[chan struct{}]struct{}
|
|
}{clients: make(map[chan struct{}]struct{})}
|
|
|
|
func subscribeGroupStatus() (chan struct{}, func(), error) {
|
|
groupStatusEvents.Lock()
|
|
defer groupStatusEvents.Unlock()
|
|
if len(groupStatusEvents.clients) >= 32 {
|
|
return nil, nil, errors.New("too many status subscribers")
|
|
}
|
|
ch := make(chan struct{}, 1)
|
|
groupStatusEvents.clients[ch] = struct{}{}
|
|
return ch, func() {
|
|
groupStatusEvents.Lock()
|
|
delete(groupStatusEvents.clients, ch)
|
|
groupStatusEvents.Unlock()
|
|
}, nil
|
|
}
|
|
|
|
func publishGroupStatusChange() {
|
|
groupStatusEvents.Lock()
|
|
defer groupStatusEvents.Unlock()
|
|
for ch := range groupStatusEvents.clients {
|
|
select {
|
|
case ch <- struct{}{}:
|
|
default:
|
|
}
|
|
}
|
|
}
|
|
|
|
func StreamGroupStatus(c *gin.Context) {
|
|
streamGroupStatus(c, func(ctx context.Context) (any, error) {
|
|
results, err := currentGroupMonitors(ctx, true)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for i := range results {
|
|
for j := range results[i].Monitors {
|
|
monitor := &results[i].Monitors[j]
|
|
monitor.Name, monitor.Alias, monitor.Hidden = monitor.Group, "", false
|
|
}
|
|
}
|
|
return gin.H{"success": true, "data": results}, nil
|
|
})
|
|
}
|
|
|
|
func streamGroupStatus(c *gin.Context, snapshot func(context.Context) (any, error)) {
|
|
changed, unsubscribe, err := subscribeGroupStatus()
|
|
if err != nil {
|
|
c.JSON(http.StatusTooManyRequests, gin.H{"success": false})
|
|
return
|
|
}
|
|
defer unsubscribe()
|
|
read := func() ([]byte, error) {
|
|
ctx, cancel := context.WithTimeout(c.Request.Context(), requestTimeout)
|
|
defer cancel()
|
|
payload, err := snapshot(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return common.Marshal(payload)
|
|
}
|
|
data, err := read()
|
|
if err != nil {
|
|
c.JSON(http.StatusServiceUnavailable, gin.H{"success": false})
|
|
return
|
|
}
|
|
c.Header("Content-Type", "text/event-stream")
|
|
c.Header("Cache-Control", "no-cache, no-transform")
|
|
c.Header("X-Accel-Buffering", "no")
|
|
write := func(frame string) error {
|
|
_ = http.NewResponseController(c.Writer).SetWriteDeadline(time.Now().Add(10 * time.Second))
|
|
if _, err := fmt.Fprint(c.Writer, frame); err != nil {
|
|
return err
|
|
}
|
|
c.Writer.Flush()
|
|
return nil
|
|
}
|
|
if write("retry: 5000\nevent: status\ndata: "+string(data)+"\n\n") != nil {
|
|
return
|
|
}
|
|
heartbeat := time.NewTicker(20 * time.Second)
|
|
defer heartbeat.Stop()
|
|
for {
|
|
select {
|
|
case <-c.Request.Context().Done():
|
|
return
|
|
case <-heartbeat.C:
|
|
// Keep proxies alive without querying the database or polling status.
|
|
if write(": keepalive\n\n") != nil {
|
|
return
|
|
}
|
|
case <-changed:
|
|
data, err := read()
|
|
if err != nil || write("event: status\ndata: "+string(data)+"\n\n") != nil {
|
|
return // Reconnect replays persisted history after transient failures.
|
|
}
|
|
}
|
|
}
|
|
}
|