From d52aa919275dda4b9a52c19890e0a4916dd7325d Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Thu, 15 May 2025 16:43:52 +0200 Subject: [PATCH] Controller: periodically send PINGs on all WebSocket connections (#315) --- internal/controller/api_rpc_portforward.go | 29 ++++++++---- internal/controller/api_vms_portforward.go | 55 +++++++++++++++------- 2 files changed, 58 insertions(+), 26 deletions(-) diff --git a/internal/controller/api_rpc_portforward.go b/internal/controller/api_rpc_portforward.go index f13efbd..a91faca 100644 --- a/internal/controller/api_rpc_portforward.go +++ b/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()) + } } } diff --git a/internal/controller/api_vms_portforward.go b/internal/controller/api_vms_portforward.go index f07f57b..0395879 100644 --- a/internal/controller/api_vms_portforward.go +++ b/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()) }