Gateway: Add graceful closing with CloseGracefully (#185)

* gateway: add the possibility of graceful closure

* wsutil: rename ConnGracefulCloser to GracefulCloser

* Gateway: rename Gateway.CloseSession to .CloseGracefully

gateway/gateway.go

@@ -194,10 +194,9 @@ func (g *Gateway) HasIntents(intents Intents) bool {
 // Close closes the underlying Websocket connection.
 func (g *Gateway) Close() error {
 	wsutil.WSDebug("Trying to close. Pacemaker check skipped.")
-
 	wsutil.WSDebug("Closing the Websocket...")
-	err := g.WS.Close()
 
+	err := g.WS.Close()
 	if errors.Is(err, wsutil.ErrWebsocketClosed) {
 		wsutil.WSDebug("Websocket already closed.")
 		return nil
@@ -207,7 +206,6 @@ func (g *Gateway) Close() error {
 	// the Start function exited before it could bind the event channel into the
 	// loop.
 	g.PacerLoop.Stop()
-
 	wsutil.WSDebug("Websocket closed; error:", err)
 
 	wsutil.WSDebug("Waiting for the Pacemaker loop to exit.")
@@ -220,6 +218,25 @@ func (g *Gateway) Close() error {
 	return err
 }
 
+// CloseGracefully attempts to close the gateway connection gracefully, by
+// sending a closing frame before ending the connection. This will cause the
+// gateway's session id to be rendered invalid.
+//
+// Note that a graceful closure is only possible, if the wsutil.Connection of
+// the Gateway's Websocket implements wsutil.GracefulCloser.
+func (g *Gateway) CloseGracefully() error {
+	err := g.WS.CloseGracefully()
+	if errors.Is(err, wsutil.ErrWebsocketClosed) {
+		wsutil.WSDebug("Websocket already closed.")
+		return nil
+	}
+
+	// Stop the pacemaker loop; This shouldn't error, so return is ignored
+	g.WS.Close()
+
+	return err
+}
+
 // SessionID returns the session ID received after Ready. This function is
 // concurrently safe.
 func (g *Gateway) SessionID() string {

utils/wsutil/conn.go

@@ -42,6 +42,8 @@ type Connection interface {
 
 	// Send allows the caller to send bytes. It does not need to clean itself
 	// up on errors, as the Websocket wrapper will do that.
+	//
+	// If the data is nil, it should send a close frame
 	Send(context.Context, []byte) error
 
 	// Close should close the websocket connection. The underlying connection
@@ -50,6 +52,14 @@ type Connection interface {
 	Close() error
 }
 
+// GracefulCloser is an interface used by Connections that support graceful
+// closure of their websocket connection.
+type GracefulCloser interface {
+	// CloseGracefully sends a close frame and then closes the websocket
+	// connection.
+	CloseGracefully() error
+}
+
 // Conn is the default Websocket connection. It tries to compresses all payloads
 // using zlib.
 type Conn struct {
@@ -72,7 +82,8 @@ func NewConn() *Conn {
 	})
 }
 
-// NewConn creates a new default websocket connection with a custom dialer.
+// NewConnWithDialer creates a new default websocket connection with a custom
+// dialer.
 func NewConnWithDialer(dialer websocket.Dialer) *Conn {
 	return &Conn{
 		Dialer: dialer,
@@ -136,7 +147,7 @@ func (c *Conn) Close() error {
 	c.Conn.SetWriteDeadline(resetDeadline)
 
 	WSDebug("Conn: Websocket closed; error:", err)
-	WSDebug("Conn: Flusing events...")
+	WSDebug("Conn: Flushing events...")
 
 	// Flush all events before closing the channel. This will return as soon as
 	// c.events is closed, or after closed.
@@ -148,6 +159,22 @@ func (c *Conn) Close() error {
 	return err
 }
 
+func (c *Conn) CloseGracefully() error {
+	WSDebug("Conn: CloseGracefully is called; sending close frame.")
+
+	c.Conn.SetWriteDeadline(time.Now().Add(CloseDeadline))
+
+	err := c.Conn.WriteMessage(websocket.CloseMessage,
+		websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+	if err != nil {
+		WSError(err)
+	}
+
+	WSDebug("Conn: Close frame sent; error:", err)
+
+	return c.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.

utils/wsutil/ws.go

@@ -168,6 +168,30 @@ func (ws *Websocket) Close() error {
 	return ws.close()
 }
 
+func (ws *Websocket) CloseGracefully() error {
+	WSDebug("Conn: Acquiring mutex lock to close gracefully...")
+
+	ws.mutex.Lock()
+	defer ws.mutex.Unlock()
+
+	WSDebug("Conn: Write mutex acquired")
+
+	if gc, ok := ws.conn.(GracefulCloser); ok {
+		if ws.closed {
+			WSDebug("Conn: Websocket is already closed.")
+			return ErrWebsocketClosed
+		}
+
+		WSDebug("Conn: closing gracefully")
+
+		ws.closed = true
+		return gc.CloseGracefully()
+	} else {
+		WSDebug("Conn: The Websocket's Connection does not provide graceful closure. Closing normally instead.")
+		return ws.close()
+	}
+}
+
 // close closes the Websocket without acquiring the mutex. Refer to Close for
 // more information.
 func (ws *Websocket) close() error {