feat(websocket): enhance PeriodicWrite with deduplication support and add Context method

- Updated PeriodicWrite to accept a deduplication function for optimized data writing.
- Introduced Context method to retrieve the manager's context.
- Added logging for WebSocket connection closure with error details.
This commit is contained in:
yusing
2025-09-04 06:37:06 +08:00
parent 54ae580645
commit 90fb9f0dcc

View File

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