Merge branch 'master' of github.com:The-K-R-O-K/flow-go into illia-ma…

…lachyn/6635-tests-for-websocket-controller

Makefile

@@ -205,6 +205,7 @@ generate-mocks: install-mock-generators
 	mockery --name 'BlockTracker' --dir="./engine/access/subscription" --case=underscore --output="./engine/access/subscription/mock"  --outpkg="mock"
 	mockery --name 'DataProvider' --dir="./engine/access/rest/websockets/data_providers" --case=underscore --output="./engine/access/rest/websockets/data_providers/mock"  --outpkg="mock"
 	mockery --name 'DataProviderFactory' --dir="./engine/access/rest/websockets/data_providers" --case=underscore --output="./engine/access/rest/websockets/data_providers/mock"  --outpkg="mock"
+	mockery --name 'WebsocketConnection' --dir="./engine/access/rest/websockets" --case=underscore --output="./engine/access/rest/websockets/mock"  --outpkg="mock"
 	mockery --name 'ExecutionDataTracker' --dir="./engine/access/subscription" --case=underscore --output="./engine/access/subscription/mock"  --outpkg="mock"
 	mockery --name 'ConnectionFactory' --dir="./engine/access/rpc/connection" --case=underscore --output="./engine/access/rpc/connection/mock" --outpkg="mock"
 	mockery --name 'Communicator' --dir="./engine/access/rpc/backend" --case=underscore --output="./engine/access/rpc/backend/mock" --outpkg="mock"

engine/access/rest/websockets/config.go

@@ -4,6 +4,33 @@ import (
 	"time"
 )
 
+const (
+	// PingPeriod defines the interval at which ping messages are sent to the client.
+	// This value must be less than pongWait, cause it that case the server ensures it sends a ping well before the PongWait
+	// timeout elapses. Each new pong message resets the server's read deadline, keeping the connection alive as long as
+	// the client is responsive.
+	//
+	// Example:
+	// At t=9, the server sends a ping, initial read deadline is t=10 (for the first message)
+	// At t=10, the client responds with a pong. The server resets its read deadline to t=20.
+	// At t=18, the server sends another ping. If the client responds with a pong at t=19, the read deadline is extended to t=29.
+	//
+	// In case of failure:
+	// If the client stops responding, the server will send a ping at t=9 but won't receive a pong by t=10. The server then closes the connection.
+	PingPeriod = (PongWait * 9) / 10
+
+	// PongWait specifies the maximum time to wait for a pong response message from the peer
+	// after sending a ping
+	PongWait = 10 * time.Second
+
+	// WriteWait specifies a timeout for the write operation. If the write
+	// isn't completed within this duration, it fails with a timeout error.
+	// SetWriteDeadline ensures the write operation does not block indefinitely
+	// if the client is slow or unresponsive. This prevents resource exhaustion
+	// and allows the server to gracefully handle timeouts for delayed writes.
+	WriteWait = 10 * time.Second
+)
+
 type Config struct {
 	MaxSubscriptionsPerConnection uint64
 	MaxResponsesPerSecond         uint64

engine/access/rest/websockets/connection.go

@@ -1,42 +1,57 @@
 package websockets
 
 import (
+	"time"
+
 	"github.com/gorilla/websocket"
 )
 
-// We wrap gorilla's websocket connection with interface
-// to be able to mock it in order to test the types dependent on it
-
 type WebsocketConnection interface {
 	ReadJSON(v interface{}) error
 	WriteJSON(v interface{}) error
+	WriteControl(messageType int, deadline time.Time) error
 	Close() error
+	SetReadDeadline(deadline time.Time) error
+	SetWriteDeadline(deadline time.Time) error
+	SetPongHandler(h func(string) error)
 }
 
-type GorillaWebsocketConnection struct {
+type WebsocketConnectionImpl struct {
 	conn *websocket.Conn
 }
 
-func NewGorillaWebsocketConnection(conn *websocket.Conn) *GorillaWebsocketConnection {
-	return &GorillaWebsocketConnection{
+func NewWebsocketConnection(conn *websocket.Conn) *WebsocketConnectionImpl {
+	return &WebsocketConnectionImpl{
 		conn: conn,
 	}
 }
 
-var _ WebsocketConnection = (*GorillaWebsocketConnection)(nil)
+var _ WebsocketConnection = (*WebsocketConnectionImpl)(nil)
+
+func (c *WebsocketConnectionImpl) ReadJSON(v interface{}) error {
+	return c.conn.ReadJSON(v)
+}
+
+func (c *WebsocketConnectionImpl) WriteJSON(v interface{}) error {
+	return c.conn.WriteJSON(v)
+}
+
+func (c *WebsocketConnectionImpl) WriteControl(messageType int, deadline time.Time) error {
+	return c.conn.WriteControl(messageType, nil, deadline)
+}
 
-func (m *GorillaWebsocketConnection) ReadJSON(v interface{}) error {
-	return m.conn.ReadJSON(v)
+func (c *WebsocketConnectionImpl) Close() error {
+	return c.conn.Close()
 }
 
-func (m *GorillaWebsocketConnection) WriteJSON(v interface{}) error {
-	return m.conn.WriteJSON(v)
+func (c *WebsocketConnectionImpl) SetReadDeadline(deadline time.Time) error {
+	return c.conn.SetReadDeadline(deadline)
 }
 
-func (m *GorillaWebsocketConnection) SetCloseHandler(handler func(code int, text string) error) {
-	m.conn.SetCloseHandler(handler)
+func (c *WebsocketConnectionImpl) SetWriteDeadline(deadline time.Time) error {
+	return c.conn.SetWriteDeadline(deadline)
 }
 
-func (m *GorillaWebsocketConnection) Close() error {
-	return m.conn.Close()
+func (c *WebsocketConnectionImpl) SetPongHandler(h func(string) error) {
+	c.conn.SetPongHandler(h)
 }

engine/access/rest/websockets/controller.go

@@ -4,30 +4,34 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/gorilla/websocket"
 	"github.com/rs/zerolog"
+	"golang.org/x/sync/errgroup"
 
 	dp "github.com/onflow/flow-go/engine/access/rest/websockets/data_providers"
 	"github.com/onflow/flow-go/engine/access/rest/websockets/models"
 	"github.com/onflow/flow-go/utils/concurrentmap"
 )
 
 type Controller struct {
-	logger               zerolog.Logger
-	config               Config
-	conn                 WebsocketConnection
-	communicationChannel chan interface{}
-	dataProviders        *concurrentmap.Map[uuid.UUID, dp.DataProvider]
-	dataProviderFactory  dp.DataProviderFactory
+	logger zerolog.Logger
+	config Config
+	conn   WebsocketConnection
+
+	communicationChannel chan interface{} // Channel for sending messages to the client.
+
+	dataProviders       *concurrentmap.Map[uuid.UUID, dp.DataProvider]
+	dataProviderFactory dp.DataProviderFactory
 }
 
 func NewWebSocketController(
 	logger zerolog.Logger,
 	config Config,
-	dataProviderFactory dp.DataProviderFactory,
 	conn WebsocketConnection,
+	dataProviderFactory dp.DataProviderFactory,
 ) *Controller {
 	return &Controller{
 		logger:               logger.With().Str("component", "websocket-controller").Logger(),
@@ -39,40 +43,99 @@ func NewWebSocketController(
 	}
 }
 
-// HandleConnection manages the WebSocket connection, adding context and error handling.
+// HandleConnection manages the lifecycle of a WebSocket connection,
+// including setup, message processing, and graceful shutdown.
+//
+// Parameters:
+// - ctx: The context for controlling cancellation and timeouts.
 func (c *Controller) HandleConnection(ctx context.Context) {
-	//TODO: configure the connection with ping-pong and deadlines
-	//TODO: spin up a response limit tracker routine
 	defer c.shutdownConnection()
-	go c.readMessages(ctx)
-	c.writeMessages(ctx)
+
+	// configuring the connection with appropriate read/write deadlines and handlers.
+	err := c.configureKeepalive()
+	if err != nil {
+		// TODO: add error handling here
+		c.logger.Error().Err(err).Msg("error configuring keepalive connection")
+
+		return
+	}
+
+	// for track all goroutines and error handling
+	g, gCtx := errgroup.WithContext(ctx)
+
+	g.Go(func() error {
+		return c.readMessagesFromClient(gCtx)
+	})
+
+	g.Go(func() error {
+		return c.keepalive(gCtx)
+	})
+
+	g.Go(func() error {
+		return c.writeMessagesToClient(gCtx)
+	})
+
+	if err = g.Wait(); err != nil {
+		//TODO: add error handling here
+		c.logger.Error().Err(err).Msg("error detected in one of the goroutines")
+	}
+}
+
+// configureKeepalive sets up the WebSocket connection with a read deadline
+// and a handler for receiving pong messages from the client.
+//
+// The function does the following:
+//  1. Sets an initial read deadline to ensure the server doesn't wait indefinitely
+//     for a pong message from the client. If no message is received within the
+//     specified `pongWait` duration, the connection will be closed.
+//  2. Establishes a Pong handler that resets the read deadline every time a pong
+//     message is received from the client, allowing the server to continue waiting
+//     for further pong messages within the new deadline.
+//
+// No errors are expected during normal operation.
+func (c *Controller) configureKeepalive() error {
+	// Set the initial read deadline for the first pong message
+	// The Pong handler itself only resets the read deadline after receiving a Pong.
+	// It doesn't set an initial deadline. The initial read deadline is crucial to prevent the server from waiting
+	// forever if the client doesn't send Pongs.
+	if err := c.conn.SetReadDeadline(time.Now().Add(PongWait)); err != nil {
+		return fmt.Errorf("failed to set the initial read deadline: %w", err)
+	}
+	// Establish a Pong handler which sets the handler for pong messages received from the peer.
+	c.conn.SetPongHandler(func(string) error {
+		return c.conn.SetReadDeadline(time.Now().Add(PongWait))
+	})
+
+	return nil
 }
 
 // writeMessages reads a messages from communication channel and passes them on to a client WebSocket connection.
 // The communication channel is filled by data providers. Besides, the response limit tracker is involved in
 // write message regulation
-func (c *Controller) writeMessages(ctx context.Context) {
+//
+// No errors are expected during normal operation. All errors are considered benign.
+func (c *Controller) writeMessagesToClient(ctx context.Context) error {
 	for {
 		select {
 		case <-ctx.Done():
-			return
+			return nil
 		case msg, ok := <-c.communicationChannel:
 			if !ok {
-				return
+				return fmt.Errorf("communication channel closed, no error occurred")
 			}
-			c.logger.Debug().Msgf("read message from communication channel: %s", msg)
-
 			// TODO: handle 'response per second' limits
 
+			// Specifies a timeout for the write operation. If the write
+			// isn't completed within this duration, it fails with a timeout error.
+			// SetWriteDeadline ensures the write operation does not block indefinitely
+			// if the client is slow or unresponsive. This prevents resource exhaustion
+			// and allows the server to gracefully handle timeouts for delayed writes.
+			if err := c.conn.SetWriteDeadline(time.Now().Add(WriteWait)); err != nil {
+				return fmt.Errorf("failed to set the write deadline: %w", err)
+			}
 			err := c.conn.WriteJSON(msg)
 			if err != nil {
-				if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) ||
-					websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
-					return
-				}
-
-				c.logger.Error().Err(err).Msg("error writing to connection")
-				return
+				return fmt.Errorf("failed to write message to connection: %w", err)
 			}
 
 			c.logger.Debug().Msg("written message to client")
@@ -82,30 +145,30 @@ func (c *Controller) writeMessages(ctx context.Context) {
 
 // readMessages continuously reads messages from a client WebSocket connection,
 // processes each message, and handles actions based on the message type.
-func (c *Controller) readMessages(ctx context.Context) {
+//
+// No errors are expected during normal operation. All errors are considered benign.
+func (c *Controller) readMessagesFromClient(ctx context.Context) error {
 	for {
-		msg, err := c.readMessage()
-		if err != nil {
-			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseAbnormalClosure) ||
-				websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
-				return
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+			msg, err := c.readMessage()
+			if err != nil {
+				if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseAbnormalClosure) {
+					return nil
+				}
+				return fmt.Errorf("failed to read message from client: %w", err)
 			}
 
-			c.logger.Debug().Err(err).Msg("error reading message from client")
-			continue
-		}
-
-		baseMsg, validatedMsg, err := c.parseAndValidateMessage(msg)
-		if err != nil {
-			c.logger.Debug().Err(err).Msg("error parsing and validating client message")
-			//TODO: write error to error channel
-			continue
-		}
+			_, validatedMsg, err := c.parseAndValidateMessage(msg)
+			if err != nil {
+				return fmt.Errorf("failed to parse and validate client message: %w", err)
+			}
 
-		if err := c.handleAction(ctx, validatedMsg); err != nil {
-			c.logger.Debug().Err(err).Str("action", baseMsg.Action).Msg("error handling action")
-			//TODO: write error to error channel
-			continue
+			if err := c.handleAction(ctx, validatedMsg); err != nil {
+				return fmt.Errorf("failed to handle message action: %w", err)
+			}
 		}
 	}
 }
@@ -150,7 +213,6 @@ func (c *Controller) parseAndValidateMessage(message json.RawMessage) (models.Ba
 		validatedMsg = listMsg
 
 	default:
-		c.logger.Debug().Str("action", baseMsg.Action).Msg("unknown action type")
 		return baseMsg, nil, fmt.Errorf("unknown action type: %s", baseMsg.Action)
 	}
 
@@ -229,6 +291,7 @@ func (c *Controller) shutdownConnection() {
 		if err := c.conn.Close(); err != nil {
 			c.logger.Warn().Err(err).Msg("error closing connection")
 		}
+		// TODO: safe closing communicationChannel will be included as a part of PR #6642
 	}()
 
 	c.logger.Debug().Msg("shutting down connection")
@@ -243,3 +306,24 @@ func (c *Controller) shutdownConnection() {
 
 	c.dataProviders.Clear()
 }
+
+// keepalive sends a ping message periodically to keep the WebSocket connection alive
+// and avoid timeouts.
+//
+// No errors are expected during normal operation. All errors are considered benign.
+func (c *Controller) keepalive(ctx context.Context) error {
+	pingTicker := time.NewTicker(PingPeriod)
+	defer pingTicker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		case <-pingTicker.C:
+			err := c.conn.WriteControl(websocket.PingMessage, time.Now().Add(WriteWait))
+			if err != nil {
+				return fmt.Errorf("failed to write ping message: %w", err)
+			}
+		}
+	}
+}