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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
package cancellablewebsocket
import (
"context"
"net/http"
"time"
"github.com/gorilla/websocket"
)
type CancellableWebSocket struct {
conn *websocket.Conn
ctx context.Context
cancel context.CancelFunc
closeCode int // 0 means killed
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,
closeCode: 0,
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) Kill() error {
cws.cancel()
err, ok := <-cws.errors
if ok && err != nil {
return err
}
return nil
}
func (cws *CancellableWebSocket) Close(code int) error {
cws.closeCode = code
if err := cws.Kill(); err != nil {
cws.closeCode = 0 // Reset close code
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 cws.closeCode == 0 {
// Close without sending close.
if err = cws.conn.Close(); err != nil {
cws.errors <- err
}
} else {
// Send close with code.
// NOTE: The SetWriteDeadline above does not affect this. :)
err = cws.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(cws.closeCode, ""), time.Now().Add(time.Second))
if err != nil {
cws.errors <- err
}
}
close(cws.errors)
}
|
