diff --git a/internal/net/gphttp/websocket/manager.go b/internal/net/gphttp/websocket/manager.go index 71cbd67d..571048b4 100644 --- a/internal/net/gphttp/websocket/manager.go +++ b/internal/net/gphttp/websocket/manager.go @@ -16,6 +16,7 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" + "github.com/rs/zerolog/log" "github.com/yusing/go-proxy/internal/common" ) @@ -121,10 +122,21 @@ func NewManagerWithUpgrade(c *gin.Context) (*Manager, error) { return cm, nil } -// Periodic writes data to the connection periodically. +func (cm *Manager) Context() context.Context { + return cm.ctx +} + +// Periodic writes data to the connection periodically, with deduplication. // If the connection is closed, the error is returned. // If the write timeout is reached, ErrWriteTimeout is returned. -func (cm *Manager) PeriodicWrite(interval time.Duration, getData func() (any, error)) error { +func (cm *Manager) PeriodicWrite(interval time.Duration, getData func() (any, error), deduplicate ...DeduplicateFunc) error { + var lastData any + + var equals DeduplicateFunc + if len(deduplicate) > 0 { + equals = deduplicate[0] + } + write := func() { data, err := getData() if err != nil { @@ -133,6 +145,13 @@ func (cm *Manager) PeriodicWrite(interval time.Duration, getData func() (any, er return } + // skip if the data is the same as the last data + if equals != nil && equals(data, lastData) { + return + } + + lastData = data + if err := cm.WriteJSON(data, interval); err != nil { cm.err = err cm.Close() @@ -230,6 +249,12 @@ func (cm *Manager) close() { cm.conn.Close() cm.pingCheckTicker.Stop() + + if cm.err != nil { + log.Debug().Caller(4).Msg("Closing WebSocket connection: " + cm.err.Error()) + } else { + log.Debug().Caller(4).Msg("Closing WebSocket connection") + } } // Done returns a channel that is closed when the context is done or the connection is closed