hexmap/server/websocket/writer.go

package websocket

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
	"io"
	"time"
)

type writer struct {
	// conn is the websocket connection that this writer is responsible for writing on.
	conn *websocket.Conn
	// channel is the channel used to receive server messages to be sent to the client.
	// When it receives a SocketClosed, the process on the sending end promises not to send any further messages, as the writer will close it right after.
	channel chan ServerMessage
	// readNotifications is the channel used to receive pings when the reader receives a message, so that a ping will be sent out before the reader is ready to time out.
	// When it is closed, the reader has shut down.
	readNotifications <-chan time.Duration
	// timer is the timer used to send pings when the reader is close to timing out, to make sure the other end of the connection is still listening.
	timer *time.Timer
	// logger is the logger used to record the state of the writer, primarily in Debug level.
	logger *zap.Logger
}

// act is the function responsible for actually doing the writing.
func (w *writer) act() {
	defer w.gracefulShutdown()
	w.logger.Debug("Starting up")
	w.timer = time.NewTimer(PingDelay)
	for {
		select {
		case _, open := <-w.readNotifications:
			if open {
				w.logger.Debug("Received reader read, extending ping")
				if !w.timer.Stop() {
					// The timer went off while we were doing this, so drain the channel before resetting
					<-w.timer.C
				}
				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.timer.Reset(PingDelay)
		}
		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))
		}
	}
}

// sendClose sends a close message on the websocket connection, but does not actually close the connection.
// It does, however, close the incoming message channel.
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))
	}
}

// sendPing sends a ping message on the websocket connection. The content is arbitrary.
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 main process 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.timer = nil
	w.logger.Debug("Waiting for all channels to shut down")
	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))
			}
		}
	}
}

// finalShutdown closes the socket and finishes cleanup.
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")
}