wsutil: Refactored and decoupled structures for better thread safety
This commit is contained in:
parent
6c332ac145
commit
16a408bf30
|
|
@ -9,17 +9,17 @@ import (
|
|||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/diamondburned/arikawa/utils/json"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
// CopyBufferSize is used for the initial size of the internal WS' buffer.
|
||||
const CopyBufferSize = 2048
|
||||
// CopyBufferSize is used for the initial size of the internal WS' buffer. Its
|
||||
// size is 4KB.
|
||||
var CopyBufferSize = 4096
|
||||
|
||||
// MaxCapUntilReset determines the maximum capacity before the bytes buffer is
|
||||
// re-allocated. This constant is 4MB.
|
||||
const MaxCapUntilReset = 4 * (1 << 20)
|
||||
// re-allocated. It is roughly 16KB, quadruple CopyBufferSize.
|
||||
var MaxCapUntilReset = CopyBufferSize * 4
|
||||
|
||||
// CloseDeadline controls the deadline to wait for sending the Close frame.
|
||||
var CloseDeadline = time.Second
|
||||
|
|
@ -50,78 +50,54 @@ type Connection interface {
|
|||
// Conn is the default Websocket connection. It compresses all payloads using
|
||||
// zlib.
|
||||
type Conn struct {
|
||||
mutex sync.Mutex
|
||||
|
||||
Conn *websocket.Conn
|
||||
json.Driver
|
||||
|
||||
dialer *websocket.Dialer
|
||||
events chan Event
|
||||
|
||||
// write channels
|
||||
writeMu *sync.Mutex
|
||||
|
||||
buf bytes.Buffer
|
||||
zlib io.ReadCloser // (compress/zlib).reader
|
||||
|
||||
// nil until Dial().
|
||||
closeOnce *sync.Once
|
||||
|
||||
// zlib *zlib.Inflator // zlib.NewReader
|
||||
// buf []byte // io.Copy buffer
|
||||
}
|
||||
|
||||
var _ Connection = (*Conn)(nil)
|
||||
|
||||
// NewConn creates a new default websocket connection with a default dialer.
|
||||
func NewConn() *Conn {
|
||||
return NewConnWithDriver(json.Default)
|
||||
return NewConnWithDialer(&websocket.Dialer{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
HandshakeTimeout: WSTimeout,
|
||||
ReadBufferSize: CopyBufferSize,
|
||||
WriteBufferSize: CopyBufferSize,
|
||||
EnableCompression: true,
|
||||
})
|
||||
}
|
||||
|
||||
func NewConnWithDriver(driver json.Driver) *Conn {
|
||||
writeMu := sync.Mutex{}
|
||||
writeMu.Lock()
|
||||
|
||||
writeBuf := bytes.Buffer{}
|
||||
writeBuf.Grow(CopyBufferSize)
|
||||
|
||||
return &Conn{
|
||||
Driver: driver,
|
||||
dialer: &websocket.Dialer{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
HandshakeTimeout: WSTimeout,
|
||||
ReadBufferSize: CopyBufferSize,
|
||||
WriteBufferSize: CopyBufferSize,
|
||||
EnableCompression: true,
|
||||
},
|
||||
writeMu: &writeMu,
|
||||
buf: writeBuf,
|
||||
}
|
||||
// NewConn creates a new default websocket connection with a custom dialer.
|
||||
func NewConnWithDialer(dialer *websocket.Dialer) *Conn {
|
||||
return &Conn{dialer: dialer}
|
||||
}
|
||||
|
||||
func (c *Conn) Dial(ctx context.Context, addr string) error {
|
||||
var err error
|
||||
|
||||
// Enable compression:
|
||||
headers := http.Header{}
|
||||
headers.Set("Accept-Encoding", "zlib")
|
||||
headers := http.Header{
|
||||
"Accept-Encoding": {"zlib"},
|
||||
}
|
||||
|
||||
// BUG: https://github.com/golang/go/issues/31514
|
||||
// // Enable stream compression:
|
||||
// addr = InjectValues(addr, url.Values{
|
||||
// "compress": {"zlib-stream"},
|
||||
// })
|
||||
// BUG which prevents stream compression.
|
||||
// See https://github.com/golang/go/issues/31514.
|
||||
|
||||
c.Conn, _, err = c.dialer.DialContext(ctx, addr, headers)
|
||||
conn, _, err := c.dialer.DialContext(ctx, addr, headers)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "failed to dial WS")
|
||||
}
|
||||
|
||||
// Set up the closer.
|
||||
c.closeOnce = &sync.Once{}
|
||||
events := make(chan Event, WSBuffer)
|
||||
go startReadLoop(conn, events)
|
||||
|
||||
c.events = make(chan Event, WSBuffer)
|
||||
go c.readLoop()
|
||||
c.mutex.Lock()
|
||||
defer c.mutex.Unlock()
|
||||
|
||||
// Unlock the mutex that would otherwise be acquired in NewConn and Close.
|
||||
c.writeMu.Unlock()
|
||||
c.Conn = conn
|
||||
c.events = events
|
||||
|
||||
return err
|
||||
}
|
||||
|
|
@ -130,96 +106,12 @@ func (c *Conn) Listen() <-chan Event {
|
|||
return c.events
|
||||
}
|
||||
|
||||
func (c *Conn) readLoop() {
|
||||
// Clean up the events channel in the end.
|
||||
defer close(c.events)
|
||||
|
||||
for {
|
||||
b, err := c.handle()
|
||||
if err != nil {
|
||||
// Is the error an EOF?
|
||||
if errors.Is(err, io.EOF) {
|
||||
// Yes it is, exit.
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the error is a normal one:
|
||||
if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
|
||||
return
|
||||
}
|
||||
|
||||
// Unusual error; log and exit:
|
||||
c.events <- Event{nil, errors.Wrap(err, "WS error")}
|
||||
return
|
||||
}
|
||||
|
||||
// If the payload length is 0, skip it.
|
||||
if len(b) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
c.events <- Event{b, nil}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) handle() ([]byte, error) {
|
||||
// skip message type
|
||||
t, r, err := c.Conn.NextReader()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if t == websocket.BinaryMessage {
|
||||
// Probably a zlib payload
|
||||
|
||||
if c.zlib == nil {
|
||||
z, err := zlib.NewReader(r)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to create a zlib reader")
|
||||
}
|
||||
c.zlib = z
|
||||
} else {
|
||||
if err := c.zlib.(zlib.Resetter).Reset(r, nil); err != nil {
|
||||
return nil, errors.Wrap(err, "failed to reset zlib reader")
|
||||
}
|
||||
}
|
||||
|
||||
defer c.zlib.Close()
|
||||
r = c.zlib
|
||||
}
|
||||
|
||||
return readAll(&c.buf, r)
|
||||
|
||||
// if t is a text message, then handle it normally.
|
||||
// if t == websocket.TextMessage {
|
||||
// return readAll(&c.buf, r)
|
||||
// }
|
||||
|
||||
// // Write to the zlib writer.
|
||||
// c.zlib.Write(r)
|
||||
// // if _, err := io.CopyBuffer(c.zlib, r, c.buf); err != nil {
|
||||
// // return nil, errors.Wrap(err, "Failed to write to zlib")
|
||||
// // }
|
||||
|
||||
// if !c.zlib.CanFlush() {
|
||||
// return nil, nil
|
||||
// }
|
||||
|
||||
// // Flush and get the uncompressed payload.
|
||||
// b, err := c.zlib.Flush()
|
||||
// if err != nil {
|
||||
// return nil, errors.Wrap(err, "Failed to flush zlib")
|
||||
// }
|
||||
|
||||
// return nil, errors.New("Unexpected binary message.")
|
||||
}
|
||||
|
||||
// resetDeadline is used to reset the write deadline after using the context's.
|
||||
var resetDeadline = time.Time{}
|
||||
|
||||
func (c *Conn) Send(ctx context.Context, b []byte) error {
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
c.mutex.Lock()
|
||||
defer c.mutex.Unlock()
|
||||
|
||||
d, ok := ctx.Deadline()
|
||||
if ok {
|
||||
|
|
@ -230,38 +122,37 @@ func (c *Conn) Send(ctx context.Context, b []byte) error {
|
|||
return c.Conn.WriteMessage(websocket.TextMessage, b)
|
||||
}
|
||||
|
||||
func (c *Conn) Close() (err error) {
|
||||
func (c *Conn) Close() error {
|
||||
// Use a sync.Once to guarantee that other Close() calls block until the
|
||||
// main call is done. It also prevents future calls.
|
||||
c.closeOnce.Do(func() {
|
||||
WSDebug("Conn: Acquiring write lock...")
|
||||
WSDebug("Conn: Acquiring write lock...")
|
||||
|
||||
// Acquire the write lock forever.
|
||||
c.writeMu.Lock()
|
||||
// Acquire the write lock forever.
|
||||
c.mutex.Lock()
|
||||
defer c.mutex.Unlock()
|
||||
|
||||
WSDebug("Conn: Write lock acquired; closing.")
|
||||
WSDebug("Conn: Write lock acquired; closing.")
|
||||
|
||||
// Close the WS.
|
||||
err = c.closeWS()
|
||||
// Close the WS.
|
||||
err := c.closeWS()
|
||||
|
||||
WSDebug("Conn: Websocket closed; error:", err)
|
||||
WSDebug("Conn: Flusing events...")
|
||||
WSDebug("Conn: Websocket closed; error:", err)
|
||||
WSDebug("Conn: Flusing events...")
|
||||
|
||||
// Flush all events before closing the channel. This will return as soon as
|
||||
// c.events is closed, or after closed.
|
||||
for range c.events {
|
||||
}
|
||||
// Flush all events before closing the channel. This will return as soon as
|
||||
// c.events is closed, or after closed.
|
||||
for range c.events {
|
||||
}
|
||||
|
||||
WSDebug("Flushed events.")
|
||||
WSDebug("Flushed events.")
|
||||
|
||||
// Mark c.Conn as empty.
|
||||
c.Conn = nil
|
||||
})
|
||||
// Mark c.Conn as empty.
|
||||
c.Conn = nil
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *Conn) closeWS() (err error) {
|
||||
func (c *Conn) closeWS() error {
|
||||
// We can't close with a write control here, since it will invalidate the
|
||||
// old session, breaking resumes.
|
||||
|
||||
|
|
@ -275,30 +166,100 @@ func (c *Conn) closeWS() (err error) {
|
|||
// // checking this because it's not important.
|
||||
// err = c.Conn.WriteControl(websocket.CloseMessage, msg, deadline)
|
||||
|
||||
if err := c.Conn.Close(); err != nil {
|
||||
return err
|
||||
return c.Conn.Close()
|
||||
}
|
||||
|
||||
// loopState is a thread-unsafe disposable state container for the read loop.
|
||||
// It's made to completely separate the read loop of any synchronization that
|
||||
// doesn't involve the websocket connection itself.
|
||||
type loopState struct {
|
||||
conn *websocket.Conn
|
||||
zlib io.ReadCloser
|
||||
buf bytes.Buffer
|
||||
}
|
||||
|
||||
func startReadLoop(conn *websocket.Conn, eventCh chan<- Event) {
|
||||
// Clean up the events channel in the end.
|
||||
defer close(eventCh)
|
||||
|
||||
// Allocate the read loop its own private resources.
|
||||
state := loopState{conn: conn}
|
||||
state.buf.Grow(CopyBufferSize)
|
||||
|
||||
for {
|
||||
b, err := state.handle()
|
||||
if err != nil {
|
||||
// Is the error an EOF?
|
||||
if errors.Is(err, io.EOF) {
|
||||
// Yes it is, exit.
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the error is a normal one:
|
||||
if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
|
||||
return
|
||||
}
|
||||
|
||||
// Unusual error; log and exit:
|
||||
eventCh <- Event{nil, errors.Wrap(err, "WS error")}
|
||||
return
|
||||
}
|
||||
|
||||
// If the payload length is 0, skip it.
|
||||
if len(b) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
eventCh <- Event{b, nil}
|
||||
}
|
||||
}
|
||||
|
||||
func (state *loopState) handle() ([]byte, error) {
|
||||
// skip message type
|
||||
t, r, err := state.conn.NextReader()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return
|
||||
if t == websocket.BinaryMessage {
|
||||
// Probably a zlib payload
|
||||
|
||||
if state.zlib == nil {
|
||||
z, err := zlib.NewReader(r)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to create a zlib reader")
|
||||
}
|
||||
state.zlib = z
|
||||
} else {
|
||||
if err := state.zlib.(zlib.Resetter).Reset(r, nil); err != nil {
|
||||
return nil, errors.Wrap(err, "failed to reset zlib reader")
|
||||
}
|
||||
}
|
||||
|
||||
defer state.zlib.Close()
|
||||
r = state.zlib
|
||||
}
|
||||
|
||||
return state.readAll(r)
|
||||
}
|
||||
|
||||
// readAll reads bytes into an existing buffer, copy it over, then wipe the old
|
||||
// buffer.
|
||||
func readAll(buf *bytes.Buffer, r io.Reader) ([]byte, error) {
|
||||
defer buf.Reset()
|
||||
func (state *loopState) readAll(r io.Reader) ([]byte, error) {
|
||||
defer state.buf.Reset()
|
||||
|
||||
if _, err := buf.ReadFrom(r); err != nil {
|
||||
if _, err := state.buf.ReadFrom(r); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Copy the bytes so we could empty the buffer for reuse.
|
||||
cpy := make([]byte, buf.Len())
|
||||
copy(cpy, buf.Bytes())
|
||||
cpy := make([]byte, state.buf.Len())
|
||||
copy(cpy, state.buf.Bytes())
|
||||
|
||||
// If the buffer's capacity is over the limit, then re-allocate a new one.
|
||||
if buf.Cap() > MaxCapUntilReset {
|
||||
*buf = bytes.Buffer{}
|
||||
buf.Grow(CopyBufferSize)
|
||||
if state.buf.Cap() > MaxCapUntilReset {
|
||||
state.buf = bytes.Buffer{}
|
||||
state.buf.Grow(CopyBufferSize)
|
||||
}
|
||||
|
||||
return cpy, nil
|
||||
|
|
|
|||
Loading…
Reference in New Issue