From 8dd74db44623c8025f05cedaddd07d8117e886ae Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Fri, 7 Feb 2025 00:38:04 +0400
Subject: [PATCH] Worker notification improvements (#246)

* OpenAPI: document all default "wait" values

* Re-use waitContext instead of instantiating it anew
---
 api/openapi.yaml                               |  2 ++
 internal/controller/api_vms_ip.go              |  3 +--
 internal/controller/api_vms_portforward.go     | 11 ++++-------
 internal/controller/api_workers_portforward.go |  6 ++++--
 4 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/api/openapi.yaml b/api/openapi.yaml
index 0afe4e0..9b173c7 100644
--- a/api/openapi.yaml
+++ b/api/openapi.yaml
@@ -354,6 +354,7 @@ paths:
             type: integer
             minimum: 0
             maximum: 65535
+            default: 10
           required: false
         - in: header
           name: Connection
@@ -393,6 +394,7 @@ paths:
             type: integer
             minimum: 0
             maximum: 65535
+            default: 0
           required: false
       responses:
         '200':
diff --git a/internal/controller/api_vms_ip.go b/internal/controller/api_vms_ip.go
index a8595dc..477bb1a 100644
--- a/internal/controller/api_vms_ip.go
+++ b/internal/controller/api_vms_ip.go
@@ -26,8 +26,7 @@ func (controller *Controller) ip(ctx *gin.Context) responder.Responder {
 	if err != nil {
 		return responder.Code(http.StatusBadRequest)
 	}
-	waitDuration := time.Duration(wait) * time.Second
-	waitContext, waitContextCancel := context.WithTimeout(ctx, waitDuration)
+	waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
 	defer waitContextCancel()
 
 	// Look-up the VM
diff --git a/internal/controller/api_vms_portforward.go b/internal/controller/api_vms_portforward.go
index 7f91ce3..e0e9224 100644
--- a/internal/controller/api_vms_portforward.go
+++ b/internal/controller/api_vms_portforward.go
@@ -42,8 +42,7 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
 	if err != nil {
 		return responder.Code(http.StatusBadRequest)
 	}
-	waitDuration := time.Duration(wait) * time.Second
-	waitContext, waitContextCancel := context.WithTimeout(ctx, waitDuration)
+	waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
 	defer waitContextCancel()
 
 	// Look-up the VM
@@ -53,15 +52,15 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
 	}
 
 	// Commence port-forwarding
-	return controller.portForward(ctx, vm.Worker, vm.UID, uint32(port), waitDuration)
+	return controller.portForward(ctx, waitContext, vm.Worker, vm.UID, uint32(port))
 }
 
 func (controller *Controller) portForward(
 	ctx *gin.Context,
+	notifyContext context.Context,
 	workerName string,
 	vmUID string,
 	port uint32,
-	waitTimeout time.Duration,
 ) responder.Responder {
 	// Request and wait for a connection with a worker
 	rendezvousCtx, rendezvousCtxCancel := context.WithCancel(ctx)
@@ -73,9 +72,7 @@ func (controller *Controller) portForward(
 	defer cancel()
 
 	// send request to worker to initiate port-forwarding connection back to us
-	waitContext, waitContextCancel := context.WithTimeout(ctx, waitTimeout)
-	defer waitContextCancel()
-	err := controller.workerNotifier.Notify(waitContext, workerName, &rpc.WatchInstruction{
+	err := controller.workerNotifier.Notify(notifyContext, workerName, &rpc.WatchInstruction{
 		Action: &rpc.WatchInstruction_PortForwardAction{
 			PortForwardAction: &rpc.WatchInstruction_PortForward{
 				Session: session,
diff --git a/internal/controller/api_workers_portforward.go b/internal/controller/api_workers_portforward.go
index f5700af..1ef62b1 100644
--- a/internal/controller/api_workers_portforward.go
+++ b/internal/controller/api_workers_portforward.go
@@ -1,6 +1,7 @@
 package controller
 
 import (
+	"context"
 	storepkg "github.com/cirruslabs/orchard/internal/controller/store"
 	"github.com/cirruslabs/orchard/internal/responder"
 	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
@@ -32,7 +33,8 @@ func (controller *Controller) portForwardWorker(ctx *gin.Context) responder.Resp
 	if err != nil {
 		return responder.Code(http.StatusBadRequest)
 	}
-	waitDuration := time.Duration(wait) * time.Second
+	waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
+	defer waitContextCancel()
 
 	var worker *v1.Worker
 
@@ -48,5 +50,5 @@ func (controller *Controller) portForwardWorker(ctx *gin.Context) responder.Resp
 	}
 
 	// Commence port-forwarding
-	return controller.portForward(ctx, worker.Name, "", uint32(port), waitDuration)
+	return controller.portForward(ctx, waitContext, worker.Name, "", uint32(port))
 }