mirror of
https://github.com/diamondburned/arikawa.git
synced 2024-09-28 21:29:25 +00:00
Gateway: Fixed the double Close and Gateway ReconnectOP bugs
This commit is contained in:
parent
e0a002cc6d
commit
3d47bada07
|
@ -6,6 +6,7 @@ import (
|
||||||
"context"
|
"context"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/diamondburned/arikawa/utils/json"
|
"github.com/diamondburned/arikawa/utils/json"
|
||||||
|
@ -18,6 +19,9 @@ const CopyBufferSize = 2048
|
||||||
// CloseDeadline controls the deadline to wait for sending the Close frame.
|
// CloseDeadline controls the deadline to wait for sending the Close frame.
|
||||||
var CloseDeadline = time.Second
|
var CloseDeadline = time.Second
|
||||||
|
|
||||||
|
// ErrWebsocketClosed is returned if the websocket is already closed.
|
||||||
|
var ErrWebsocketClosed = errors.New("Websocket is closed.")
|
||||||
|
|
||||||
// Connection is an interface that abstracts around a generic Websocket driver.
|
// Connection is an interface that abstracts around a generic Websocket driver.
|
||||||
// This connection expects the driver to handle compression by itself, including
|
// This connection expects the driver to handle compression by itself, including
|
||||||
// modifying the connection URL.
|
// modifying the connection URL.
|
||||||
|
@ -45,14 +49,17 @@ type Conn struct {
|
||||||
json.Driver
|
json.Driver
|
||||||
|
|
||||||
dialer *websocket.Dialer
|
dialer *websocket.Dialer
|
||||||
// mut sync.RWMutex
|
|
||||||
events chan Event
|
events chan Event
|
||||||
|
|
||||||
// write channels
|
// write channels
|
||||||
writes chan []byte
|
writes chan []byte
|
||||||
errors chan error
|
errors chan error
|
||||||
|
|
||||||
buf bytes.Buffer
|
buf bytes.Buffer
|
||||||
|
zlib io.ReadCloser // (compress/zlib).reader
|
||||||
|
|
||||||
|
// nil until Dial().
|
||||||
|
closeOnce *sync.Once
|
||||||
|
|
||||||
// zlib *zlib.Inflator // zlib.NewReader
|
// zlib *zlib.Inflator // zlib.NewReader
|
||||||
// buf []byte // io.Copy buffer
|
// buf []byte // io.Copy buffer
|
||||||
|
@ -91,6 +98,9 @@ func (c *Conn) Dial(ctx context.Context, addr string) error {
|
||||||
return errors.Wrap(err, "Failed to dial WS")
|
return errors.Wrap(err, "Failed to dial WS")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Set up the closer.
|
||||||
|
c.closeOnce = &sync.Once{}
|
||||||
|
|
||||||
c.events = make(chan Event)
|
c.events = make(chan Event)
|
||||||
go c.readLoop()
|
go c.readLoop()
|
||||||
|
|
||||||
|
@ -173,13 +183,21 @@ func (c *Conn) handle() ([]byte, error) {
|
||||||
|
|
||||||
if t == websocket.BinaryMessage {
|
if t == websocket.BinaryMessage {
|
||||||
// Probably a zlib payload
|
// Probably a zlib payload
|
||||||
z, err := zlib.NewReader(r)
|
|
||||||
if err != nil {
|
if c.zlib == nil {
|
||||||
return nil, errors.Wrap(err, "Failed to create a zlib reader")
|
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 z.Close()
|
defer c.zlib.Close()
|
||||||
r = z
|
r = c.zlib
|
||||||
}
|
}
|
||||||
|
|
||||||
return readAll(&c.buf, r)
|
return readAll(&c.buf, r)
|
||||||
|
@ -209,39 +227,40 @@ func (c *Conn) handle() ([]byte, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *Conn) Send(b []byte) error {
|
func (c *Conn) Send(b []byte) error {
|
||||||
// Don't send a nil byte slice. That would confuse the write loop.
|
|
||||||
if b == nil {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// If websocket is already closed.
|
// If websocket is already closed.
|
||||||
if c.writes == nil {
|
if c.writes == nil {
|
||||||
return errors.New("Websocket is closed.")
|
return ErrWebsocketClosed
|
||||||
}
|
}
|
||||||
|
|
||||||
c.writes <- b
|
c.writes <- b
|
||||||
return <-c.errors
|
return <-c.errors
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *Conn) Close() error {
|
func (c *Conn) Close() (err error) {
|
||||||
// Close c.writes. This should trigger the websocket to close itself.
|
// Use a sync.Once to guarantee that other Close() calls block until the
|
||||||
close(c.writes)
|
// main call is done. It also prevents future calls.
|
||||||
|
c.closeOnce.Do(func() {
|
||||||
|
// Close c.writes. This should trigger the websocket to close itself.
|
||||||
|
close(c.writes)
|
||||||
|
// Mark c.writes as empty.
|
||||||
|
c.writes = nil
|
||||||
|
|
||||||
// Wait for the write loop to exit by flusing the errors channel.
|
// Wait for the write loop to exit by flusing the errors channel.
|
||||||
var err = <-c.errors
|
err = <-c.errors // get close error
|
||||||
for range c.errors {
|
for range c.errors { // then flush
|
||||||
}
|
}
|
||||||
|
|
||||||
// Flush all events before closing the channel. This will return as soon as
|
// Flush all events before closing the channel. This will return as soon as
|
||||||
// c.events is closed, or after closed.
|
// c.events is closed, or after closed.
|
||||||
for range c.events {
|
for range c.events {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Mark c.events as empty.
|
// Mark c.events as empty.
|
||||||
c.events = nil
|
c.events = nil
|
||||||
|
|
||||||
// Mark c.Conn as empty.
|
// Mark c.Conn as empty.
|
||||||
c.Conn = nil
|
c.Conn = nil
|
||||||
|
})
|
||||||
|
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in a new issue