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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
|
package cancellablewebsocket
import (
"context"
"net/http"
"time"
"github.com/gorilla/websocket"
)
type CancellableWebSocket struct {
conn *websocket.Conn
ctx context.Context
cancel context.CancelFunc
errors chan error
}
func Dial(dialer *websocket.Dialer, ctx context.Context, url string, requestHeader http.Header) (*CancellableWebSocket, error) {
conn, _, err := dialer.Dial(url, requestHeader)
if err != nil {
return nil, err
}
childCtx, cancel := context.WithCancel(ctx)
cws := &CancellableWebSocket{
conn: conn,
ctx: childCtx,
cancel: cancel,
errors: make(chan error),
}
go cws.listenForCancel()
return cws, nil
}
func (cws *CancellableWebSocket) ReadJSON(v interface{}) error {
err := cws.conn.ReadJSON(v)
if err != nil {
// Check if context was not cancelled,
// if so, return error
if cws.ctx.Err() == nil {
return err
}
}
return nil
}
func (cws *CancellableWebSocket) WriteJSON(v interface{}) error {
err := cws.conn.WriteJSON(v)
if err != nil {
if cws.ctx.Err() == nil {
return err
}
}
return nil
}
func (cws *CancellableWebSocket) Close() error {
cws.cancel()
err, ok := <-cws.errors
if ok && err != nil {
return err
}
return nil
}
func (cws *CancellableWebSocket) listenForCancel() {
<-cws.ctx.Done()
var err error
aLongTimeAgo := time.Unix(0, 0)
if err = cws.conn.SetReadDeadline(aLongTimeAgo); err != nil {
cws.errors <- err
}
if err = cws.conn.SetWriteDeadline(aLongTimeAgo); err != nil {
cws.errors <- err
}
if err = cws.conn.Close(); err != nil {
cws.errors <- err
}
close(cws.errors)
}
|
