From 818f4288c20c5edca7512a34215914f029957cc1 Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Thu, 20 Feb 2025 02:00:57 +0400
Subject: [PATCH] Controller API: correctly detect WebSocket closure in Watch
 RPC (#259)

---
 internal/controller/api_rpc_watch.go | 11 ++++++++---
 internal/controller/wserror.go       | 16 ++++++++++++----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/internal/controller/api_rpc_watch.go b/internal/controller/api_rpc_watch.go
index 37740c5..474a73e 100644
--- a/internal/controller/api_rpc_watch.go
+++ b/internal/controller/api_rpc_watch.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"github.com/cirruslabs/orchard/internal/responder"
 	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/cirruslabs/orchard/rpc"
@@ -41,7 +42,7 @@ func (controller *Controller) rpcWatch(ctx *gin.Context) responder.Responder {
 	// from the connection in the background
 	//
 	// Otherwise the wsConn.Ping() will wait forever.
-	wsConn.CloseRead(ctx)
+	closeReadCtx := wsConn.CloseRead(ctx)
 
 	for {
 		select {
@@ -86,10 +87,14 @@ func (controller *Controller) rpcWatch(ctx *gin.Context) responder.Responder {
 			}
 
 			pingCtxCancel()
+		case <-closeReadCtx.Done():
+			// Connection shouldn't be normally closed by the worker
+			return controller.wsErrorNoClose("watch RPC",
+				fmt.Sprintf("worker %s unexpectedly disconnected", workerName), closeReadCtx.Err())
 		case <-ctx.Done():
 			// Connection shouldn't be normally closed by the worker
-			return controller.wsError(wsConn, websocket.StatusAbnormalClosure, "watch RPC",
-				"unexpectedly disconnected worker", err)
+			return controller.wsErrorNoClose("watch RPC",
+				fmt.Sprintf("worker %s unexpectedly disconnected", workerName), ctx.Err())
 		}
 	}
 }
diff --git a/internal/controller/wserror.go b/internal/controller/wserror.go
index 86ce631..e9aedf1 100644
--- a/internal/controller/wserror.go
+++ b/internal/controller/wserror.go
@@ -13,14 +13,22 @@ func (controller *Controller) wsError(
 	reason string,
 	err error,
 ) responder.Responder {
-	message := fmt.Sprintf("%s: %v", reason, err)
+	responder := controller.wsErrorNoClose(component, reason, err)
 
-	controller.logger.Warn(message)
-
-	if err := wsConn.Close(code, message); err != nil {
+	if err := wsConn.Close(code, fmt.Sprintf("%s: %v", reason, err)); err != nil {
 		controller.logger.Warnf("%s: failed to close the WebSocket connection that entered error state"+
 			" due to %s: %v", component, reason, err)
 	}
 
+	return responder
+}
+
+func (controller *Controller) wsErrorNoClose(
+	component string,
+	reason string,
+	err error,
+) responder.Responder {
+	controller.logger.Warnf("%s: %s: %v", component, reason, err)
+
 	return responder.Empty()
 }