parent
942b7eb20b
commit
8db68101fd
@ -0,0 +1,37 @@ |
||||
package websocket |
||||
|
||||
import ( |
||||
"github.com/gorilla/websocket" |
||||
"time" |
||||
) |
||||
|
||||
const ( |
||||
// ReadTimeLimit is the maximum time the server is willing to wait after receiving a message before receiving another one.
|
||||
ReadTimeLimit = 60 * time.Second |
||||
// WriteTimeLimit is the maximum time the server is willing to wait to send a message.
|
||||
WriteTimeLimit = 10 * time.Second |
||||
// ControlTimeLimit is the maximum time the server is willing to wait to send a control message like Ping or Close.
|
||||
ControlTimeLimit = (WriteTimeLimit * 5) / 10 |
||||
// PingDelay is the time between pings.
|
||||
// It must be less than ReadTimeLimit to account for latency and delays on either side.
|
||||
PingDelay = (ReadTimeLimit * 7) / 10 |
||||
) |
||||
|
||||
// A Connection corresponds to a pair of actors.
|
||||
type Connection struct { |
||||
conn *websocket.Conn |
||||
r reader |
||||
w writer |
||||
} |
||||
|
||||
// ReadChannel returns the channel that can be used to read client messages from the connection.
|
||||
// After receiving SocketClosed, the reader will close its channel.
|
||||
func (c *Connection) ReadChannel() <-chan ClientMessage { |
||||
return c.r.channel |
||||
} |
||||
|
||||
// WriteChannel returns the channel that can be used to send server messages on the connection.
|
||||
// After sending SocketClosed, the writer will close its channel; do not send any further messages on the channel.
|
||||
func (c *Connection) WriteChannel() chan<- ServerMessage { |
||||
return c.w.channel |
||||
} |
||||
@ -0,0 +1,14 @@ |
||||
package websocket |
||||
|
||||
import ( |
||||
"github.com/gorilla/websocket" |
||||
"time" |
||||
) |
||||
|
||||
// Todo: Listen for pongs and extend the read deadline every time you get one
|
||||
|
||||
type reader struct { |
||||
conn *websocket.Conn |
||||
channel chan ClientMessage |
||||
readNotifications chan<- time.Duration |
||||
} |
||||
@ -0,0 +1,144 @@ |
||||
package websocket |
||||
|
||||
import ( |
||||
"encoding/json" |
||||
"fmt" |
||||
"github.com/gorilla/websocket" |
||||
"go.uber.org/zap" |
||||
"io" |
||||
"time" |
||||
) |
||||
|
||||
type writer struct { |
||||
conn *websocket.Conn |
||||
channel chan ServerMessage |
||||
readNotifications <-chan time.Duration |
||||
timer *time.Ticker |
||||
logger *zap.Logger |
||||
} |
||||
|
||||
func (w *writer) act() { |
||||
defer w.gracefulShutdown() |
||||
w.logger.Debug("Starting up") |
||||
w.timer = time.NewTicker(PingDelay) |
||||
for { |
||||
select { |
||||
case _, open := <-w.readNotifications: |
||||
if open { |
||||
w.logger.Debug("Received reader read, extending ping") |
||||
w.timer.Reset(PingDelay) |
||||
} else { |
||||
w.logger.Debug("Received reader close, shutting down") |
||||
w.readNotifications = nil |
||||
// bye bye, we'll graceful shutdown because we deferred it
|
||||
return |
||||
} |
||||
case raw := <-w.channel: |
||||
switch msg := raw.(type) { |
||||
case SocketClosed: |
||||
w.logger.Debug("Received close message, forwarding and shutting down", zap.Object("msg", msg)) |
||||
w.sendClose(msg) |
||||
// bye bye, we'll graceful shutdown because we deferred it
|
||||
return |
||||
default: |
||||
w.logger.Debug("Received message, forwarding", zap.Object("msg", msg)) |
||||
w.send(msg) |
||||
} |
||||
case <-w.timer.C: |
||||
w.sendPing() |
||||
} |
||||
w.logger.Debug("Awakening handled, resuming listening") |
||||
} |
||||
} |
||||
|
||||
func (w *writer) send(msg ServerMessage) { |
||||
writer, err := w.conn.NextWriter(websocket.TextMessage) |
||||
if err != nil { |
||||
w.logger.Error("error while getting writer from connection", zap.Error(err)) |
||||
return |
||||
} |
||||
defer func(writer io.WriteCloser) { |
||||
err := writer.Close() |
||||
if err != nil { |
||||
w.logger.Error("error while closing writer to send message", zap.Error(err)) |
||||
} |
||||
}(writer) |
||||
payload, err := json.Marshal(msg) |
||||
if err != nil { |
||||
w.logger.Error("error while rendering message payload to JSON", zap.Error(err)) |
||||
return |
||||
} |
||||
if len(payload) == 2 { |
||||
// This is an empty JSON message. We can leave it out.
|
||||
_, err = fmt.Fprintf(writer, "%s!", msg.ServerType()) |
||||
if err != nil { |
||||
w.logger.Error("error while writing command-only message", zap.Error(err)) |
||||
} |
||||
} else { |
||||
// Because we need to send this, we put in a space instead of an exclamation mark.
|
||||
_, err = fmt.Fprintf(writer, "%s %s", msg.ServerType(), payload) |
||||
if err != nil { |
||||
w.logger.Error("error while writing command-only message", zap.Error(err)) |
||||
} |
||||
} |
||||
} |
||||
|
||||
func (w *writer) sendClose(msg SocketClosed) { |
||||
w.logger.Debug("Shutting down the writer channel") |
||||
close(w.channel) |
||||
w.channel = nil |
||||
w.logger.Debug("Writing close message") |
||||
err := w.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(int(msg.Code), msg.Text), time.Now().Add(ControlTimeLimit)) |
||||
if err != nil { |
||||
w.logger.Error("Error while sending close", zap.Error(err)) |
||||
} |
||||
} |
||||
|
||||
func (w *writer) sendPing() { |
||||
w.logger.Debug("Sending ping") |
||||
err := w.conn.WriteControl(websocket.PingMessage, []byte("are you still there?"), time.Now().Add(ControlTimeLimit)) |
||||
if err != nil { |
||||
w.logger.Error("Error while sending ping", zap.Error(err)) |
||||
} |
||||
} |
||||
|
||||
// gracefulShutdown causes the writer to wait for the close handshake to finish and then shut down.
|
||||
// It waits for the reader's readNotifications to close, indicating that it has also shut down, and for the channel to
|
||||
// receive a SocketClosed message indicating that the client has shut down.
|
||||
// During this time, the writer ignores all other messages from the channel and sends no pings.
|
||||
func (w *writer) gracefulShutdown() { |
||||
defer w.finalShutdown() |
||||
w.logger.Debug("Waiting for all channels to shut down") |
||||
w.timer.Stop() |
||||
for { |
||||
if w.channel == nil && w.readNotifications == nil { |
||||
w.logger.Debug("All channels closed, beginning final shutdown") |
||||
// all done, we outta here, let the defer pick up the final shutdown
|
||||
return |
||||
} |
||||
select { |
||||
case _, open := <-w.readNotifications: |
||||
if !open { |
||||
w.logger.Debug("Received reader close while shutting down") |
||||
w.readNotifications = nil |
||||
} |
||||
case raw := <-w.channel: |
||||
switch msg := raw.(type) { |
||||
case SocketClosed: |
||||
w.logger.Debug("Received close message from channel while shutting down, forwarding", zap.Object("msg", msg)) |
||||
w.sendClose(msg) |
||||
default: |
||||
w.logger.Debug("Ignoring non-close message while shutting down", zap.Object("msg", msg)) |
||||
} |
||||
} |
||||
} |
||||
} |
||||
|
||||
func (w *writer) finalShutdown() { |
||||
w.logger.Debug("Closing WebSocket connection") |
||||
err := w.conn.Close() |
||||
if err != nil { |
||||
w.logger.Error("Received an error while closing", zap.Error(err)) |
||||
} |
||||
w.logger.Debug("Shut down") |
||||
} |
||||
Loading…
Reference in new issue