fix ws closing

This commit is contained in:
yyasha 2024-06-05 01:08:40 +03:00
parent 978ab02672
commit 784d4a6006
3 changed files with 35 additions and 20 deletions

View File

@ -5,28 +5,30 @@ import (
"crypto/hmac"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"time"
"github.com/google/uuid"
"github.com/gorilla/websocket"
"golang.org/x/sync/errgroup"
)
type MessageHandler func(message string) error
func (b *WebSocket) handleIncomingMessages() {
func (b *WebSocket) handleIncomingMessages() error {
for {
_, message, err := b.conn.ReadMessage()
if err != nil {
fmt.Println("Error reading:", err)
return
return err
}
if b.onMessage != nil {
err := b.onMessage(string(message))
if err != nil {
fmt.Println("Error handling message:", err)
return
return err
}
}
}
@ -112,35 +114,45 @@ func (b *WebSocket) Connect(args []string) error {
}
}
go b.handleIncomingMessages()
eg := errgroup.Group{}
b.ctx, b.cancel = context.WithCancel(context.Background())
Ping(b)
return b.sendSubscription(args)
eg.Go(b.Ping)
eg.Go(b.handleIncomingMessages)
if err := eg.Wait(); err != nil {
b.Disconnect()
return fmt.Errorf("failed to handle message: %w", err)
}
func Ping(b *WebSocket) {
return nil
}
var ErrPingFailed = errors.New("failed to send ping")
func (b *WebSocket) Ping() error {
ticker := time.NewTicker(time.Duration(b.pingInterval) * time.Second)
go func() {
defer ticker.Stop() // Ensure the ticker is stopped when this goroutine ends
for {
select {
case <-ticker.C: // Wait until the ticker sends a signal
if err := b.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
fmt.Println("Failed to send ping:", err)
return ErrPingFailed
}
case <-b.ctx.Done():
fmt.Println("Exit ping")
return
return nil
}
}
}()
}
func (b *WebSocket) Disconnect() error {
b.cancel()
return b.conn.Close()
err := b.conn.Close()
fmt.Println("WebSocket disconnected, err =", err)
return err
}
func (b *WebSocket) Send(message string) error {

1
go.mod
View File

@ -8,6 +8,7 @@ require (
github.com/gorilla/websocket v1.5.1
github.com/json-iterator/go v1.1.12
github.com/stretchr/testify v1.8.4
golang.org/x/sync v0.7.0
)
require (

2
go.sum
View File

@ -30,6 +30,8 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
golang.org/x/net v0.19.0 h1:zTwKpTd2XuCqf8huc7Fo2iSy+4RHPd10s4KzeTnVr1c=
golang.org/x/net v0.19.0/go.mod h1:CfAk/cbD4CthTvqiEl8NpboMuiuOYsAr/7NOjZJtv1U=
golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M=
golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=