forked from coder/coder
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathwebsocket.go
More file actions
68 lines (60 loc) · 1.6 KB
/
websocket.go
File metadata and controls
68 lines (60 loc) · 1.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
package httpapi
import (
"context"
"errors"
"time"
"golang.org/x/xerrors"
"cdr.dev/slog"
"github.com/coder/websocket"
)
// Heartbeat loops to ping a WebSocket to keep it alive.
// Default idle connection timeouts are typically 60 seconds.
// See: https://docs.aws.amazon.com/elasticloadbalancing/latest/application/application-load-balancers.html#connection-idle-timeout
func Heartbeat(ctx context.Context, conn *websocket.Conn) {
ticker := time.NewTicker(15 * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
}
err := conn.Ping(ctx)
if err != nil {
return
}
}
}
// Heartbeat loops to ping a WebSocket to keep it alive. It calls `exit` on ping
// failure.
func HeartbeatClose(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn) {
interval := 15 * time.Second
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
}
err := pingWithTimeout(ctx, conn, interval)
if err != nil {
// context.DeadlineExceeded is expected when the client disconnects without sending a close frame
if !errors.Is(err, context.DeadlineExceeded) {
logger.Error(ctx, "failed to heartbeat ping", slog.Error(err))
}
_ = conn.Close(websocket.StatusGoingAway, "Ping failed")
exit()
return
}
}
}
func pingWithTimeout(ctx context.Context, conn *websocket.Conn, timeout time.Duration) error {
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
err := conn.Ping(ctx)
if err != nil {
return xerrors.Errorf("failed to ping: %w", err)
}
return nil
}