Controller: periodically send PINGs on all WebSocket connections (#315)

internal/controller/api_rpc_portforward.go

@@ -1,12 +1,14 @@
 package controller
 
 import (
+	"context"
 	"github.com/cirruslabs/orchard/internal/controller/rendezvous"
 	"github.com/cirruslabs/orchard/internal/responder"
 	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/coder/websocket"
 	"github.com/gin-gonic/gin"
 	"net"
+	"time"
 )
 
 func (controller *Controller) rpcPortForward(ctx *gin.Context) responder.Responder {
@@ -36,13 +38,24 @@ func (controller *Controller) rpcPortForward(ctx *gin.Context) responder.Respond
 			"failure to respond with the established WebSocket connection", err)
 	}
 
-	select {
-	case <-proxyCtx.Done():
-		// Do not close the WebSocket connection as it should be already closed by our rendezvous party
-		return responder.Empty()
-	case <-ctx.Done():
-		// Connection should normally be closed by our rendezvous party, not by the worker
-		return controller.wsError(wsConn, websocket.StatusAbnormalClosure, "port forwarding RPC",
-			"unexpectedly disconnected worker", err)
+	for {
+		select {
+		case <-proxyCtx.Done():
+			// Do not close the WebSocket connection as it should be already closed by our rendezvous party
+			return responder.Empty()
+		case <-time.After(30 * time.Second):
+			pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second)
+
+			if err := wsConn.Ping(pingCtx); err != nil {
+				controller.logger.Warnf("port forwarding RPC: failed to ping the worker, "+
+					"connection might time out: %v", err)
+			}
+
+			pingCtxCancel()
+		case <-ctx.Done():
+			// Connection shouldn't be normally closed by the worker
+			return controller.wsErrorNoClose("watch RPC",
+				"worker unexpectedly disconnected", ctx.Err())
+		}
 	}
 }

internal/controller/api_vms_portforward.go

@@ -51,7 +51,7 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
 		return responderImpl
 	}
 
-	// Commence port-forwarding
+	// Commence port forwarding
 	return controller.portForward(ctx, waitContext, vm.Worker, vm.UID, uint32(port))
 }
 
@@ -71,7 +71,7 @@ func (controller *Controller) portForward(
 	boomerangConnCh, cancel := controller.connRendezvous.Request(rendezvousCtx, session)
 	defer cancel()
 
-	// send request to worker to initiate port-forwarding connection back to us
+	// Send request to worker to initiate port forwarding connection back to us
 	err := controller.workerNotifier.Notify(notifyContext, workerName, &rpc.WatchInstruction{
 		Action: &rpc.WatchInstruction_PortForwardAction{
 			PortForwardAction: &rpc.WatchInstruction_PortForward{
@@ -82,13 +82,13 @@ func (controller *Controller) portForward(
 		},
 	})
 	if err != nil {
-		controller.logger.Warnf("failed to request port-forwarding from the worker %s: %v",
+		controller.logger.Warnf("failed to request port forwarding from the worker %s: %v",
 			workerName, err)
 
 		return responder.Code(http.StatusServiceUnavailable)
 	}
 
-	// worker will asynchronously start port-forwarding so we wait
+	// Worker will asynchronously start port forwarding, so we wait
 	select {
 	case rendezvousResponse := <-boomerangConnCh:
 		if rendezvousResponse.ErrorMessage != "" {
@@ -111,29 +111,48 @@ func (controller *Controller) portForward(
 			expectedMsgType = websocket.MessageText
 		}
 
+		proxyConnectionsErrCh := make(chan error, 1)
 		wsConnAsNetConn := websocket.NetConn(ctx, wsConn, expectedMsgType)
 		fromWorkerConnectionWithCancel := netconncancel.New(rendezvousResponse.Result, rendezvousCtxCancel)
 
-		if err := proxy.Connections(wsConnAsNetConn, fromWorkerConnectionWithCancel); err != nil {
-			var websocketCloseError websocket.CloseError
+		go func() {
+			proxyConnectionsErrCh <- proxy.Connections(wsConnAsNetConn, fromWorkerConnectionWithCancel)
+		}()
+
+		for {
+			select {
+			case err := <-proxyConnectionsErrCh:
+				if err != nil {
+					var websocketCloseError websocket.CloseError
+
+					// Normal closure from the user
+					if errors.As(err, &websocketCloseError) && websocketCloseError.Code == websocket.StatusNormalClosure {
+						return responder.Empty()
+					}
+
+					if errors.Is(err, context.Canceled) {
+						return responder.Empty()
+					}
+
+					if status, ok := status.FromError(err); ok && status.Code() == codes.Canceled {
+						return responder.Empty()
+					}
+
+					controller.logger.Warnf("port forwarding: failed to proxy connections: %v", err)
+				}
 
-			// Normal closure from the user
-			if errors.As(err, &websocketCloseError) && websocketCloseError.Code == websocket.StatusNormalClosure {
 				return responder.Empty()
-			}
+			case <-time.After(30 * time.Second):
+				pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second)
 
-			if errors.Is(err, context.Canceled) {
-				return responder.Empty()
-			}
+				if err := wsConn.Ping(pingCtx); err != nil {
+					controller.logger.Warnf("port forwarding: failed to ping the client, "+
+						"connection might time out: %v", err)
+				}
 
-			if status, ok := status.FromError(err); ok && status.Code() == codes.Canceled {
-				return responder.Empty()
+				pingCtxCancel()
 			}
-
-			controller.logger.Warnf("failed to port-forward: %v", err)
 		}
-
-		return responder.Empty()
 	case <-ctx.Done():
 		return responder.Error(ctx.Err())
 	}