RPC v2: fix Ping() hanging due to PONG not being processed (#247)

internal/controller/api_rpc_portforward.go

@@ -4,9 +4,9 @@ import (
 	"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"
-	"nhooyr.io/websocket"
 )
 
 func (controller *Controller) rpcPortForward(ctx *gin.Context) responder.Responder {

internal/controller/api_rpc_watch.go

@@ -1,13 +1,14 @@
 package controller
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"github.com/cirruslabs/orchard/internal/responder"
 	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/cirruslabs/orchard/rpc"
+	"github.com/coder/websocket"
 	"github.com/gin-gonic/gin"
-	"nhooyr.io/websocket"
 	"time"
 )
 
@@ -36,6 +37,12 @@ func (controller *Controller) rpcWatch(ctx *gin.Context) responder.Responder {
 		return responder.Error(err)
 	}
 
+	// Ensure that pongs will be processed by reading
+	// from the connection in the background
+	//
+	// Otherwise the wsConn.Ping() will wait forever.
+	wsConn.CloseRead(ctx)
+
 	for {
 		select {
 		case msg := <-workerCh:
@@ -71,10 +78,14 @@ func (controller *Controller) rpcWatch(ctx *gin.Context) responder.Responder {
 					"failure to write the watch instruction", err)
 			}
 		case <-time.After(30 * time.Second):
-			if err := wsConn.Ping(ctx); err != nil {
+			pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second)
+
+			if err := wsConn.Ping(pingCtx); err != nil {
 				controller.logger.Warnf("watch RPC: failed to ping the worker %s, "+
 					"connection might time out: %v", workerName, err)
 			}
+
+			pingCtxCancel()
 		case <-ctx.Done():
 			// Connection shouldn't be normally closed by the worker
 			return controller.wsError(wsConn, websocket.StatusAbnormalClosure, "watch RPC",

internal/controller/api_vms_portforward.go

@@ -9,13 +9,13 @@ import (
 	"github.com/cirruslabs/orchard/internal/responder"
 	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/cirruslabs/orchard/rpc"
+	"github.com/coder/websocket"
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"net/http"
-	"nhooyr.io/websocket"
 	"strconv"
 	"time"
 )

internal/controller/wserror.go

@@ -3,7 +3,7 @@ package controller
 import (
 	"fmt"
 	"github.com/cirruslabs/orchard/internal/responder"
-	"nhooyr.io/websocket"
+	"github.com/coder/websocket"
 )
 
 func (controller *Controller) wsError(