Small timout for worker notification (#242)

* Small timout for worker notification

It seems at the moment if a worker re-establishes notify stream (for example, if network flips or proxy breaks the connection) then we can see "no worker registered with this name" errors.

This change makes Notifier to wait for 30 seconds before failing, at the time of calling `Notifier#Notify` we know such worker exists.

PS not sure if we need to make the timeout configurable.

* Wait via context

* Make sure all `context`s for `Notify` is time bounded

* Lint issues

api/openapi.yaml

@@ -197,6 +197,15 @@ paths:
             minimum: 1
             maximum: 65535
           required: true
+        - in: query
+          name: wait
+          description: Duration in seconds for the worker to become available if it's not available already
+          schema:
+            type: integer
+            minimum: 0
+            maximum: 65535
+            default: 10
+          required: false
         - in: header
           name: Connection
           description: WebSocket protocol required header

internal/controller/api_vms_ip.go

@@ -26,7 +26,8 @@ func (controller *Controller) ip(ctx *gin.Context) responder.Responder {
 	if err != nil {
 		return responder.Code(http.StatusBadRequest)
 	}
-	waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
+	waitDuration := time.Duration(wait) * time.Second
+	waitContext, waitContextCancel := context.WithTimeout(ctx, waitDuration)
 	defer waitContextCancel()
 
 	// Look-up the VM
@@ -40,7 +41,7 @@ func (controller *Controller) ip(ctx *gin.Context) responder.Responder {
 	boomerangConnCh, cancel := controller.ipRendezvous.Request(ctx, session)
 	defer cancel()
 
-	err = controller.workerNotifier.Notify(ctx, vm.Worker, &rpc.WatchInstruction{
+	err = controller.workerNotifier.Notify(waitContext, vm.Worker, &rpc.WatchInstruction{
 		Action: &rpc.WatchInstruction_ResolveIpAction{
 			ResolveIpAction: &rpc.WatchInstruction_ResolveIP{
 				Session: session,

internal/controller/api_vms_portforward.go

@@ -37,12 +37,13 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
 		return responder.Code(http.StatusBadRequest)
 	}
 
-	waitRaw := ctx.DefaultQuery("wait", "0")
+	waitRaw := ctx.DefaultQuery("wait", "10")
 	wait, err := strconv.ParseUint(waitRaw, 10, 16)
 	if err != nil {
 		return responder.Code(http.StatusBadRequest)
 	}
-	waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
+	waitDuration := time.Duration(wait) * time.Second
+	waitContext, waitContextCancel := context.WithTimeout(ctx, waitDuration)
 	defer waitContextCancel()
 
 	// Look-up the VM
@@ -52,7 +53,7 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
 	}
 
 	// Commence port-forwarding
-	return controller.portForward(ctx, vm.Worker, vm.UID, uint32(port))
+	return controller.portForward(ctx, vm.Worker, vm.UID, uint32(port), waitDuration)
 }
 
 func (controller *Controller) portForward(
@@ -60,6 +61,7 @@ func (controller *Controller) portForward(
 	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)
@@ -71,7 +73,9 @@ func (controller *Controller) portForward(
 	defer cancel()
 
 	// send request to worker to initiate port-forwarding connection back to us
-	err := controller.workerNotifier.Notify(ctx, workerName, &rpc.WatchInstruction{
+	waitContext, waitContextCancel := context.WithTimeout(ctx, waitTimeout)
+	defer waitContextCancel()
+	err := controller.workerNotifier.Notify(waitContext, workerName, &rpc.WatchInstruction{
 		Action: &rpc.WatchInstruction_PortForwardAction{
 			PortForwardAction: &rpc.WatchInstruction_PortForward{
 				Session: session,

internal/controller/api_workers_portforward.go

@@ -7,6 +7,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"net/http"
 	"strconv"
+	"time"
 )
 
 func (controller *Controller) portForwardWorker(ctx *gin.Context) responder.Responder {
@@ -26,6 +27,13 @@ func (controller *Controller) portForwardWorker(ctx *gin.Context) responder.Resp
 		return responder.Code(http.StatusBadRequest)
 	}
 
+	waitRaw := ctx.DefaultQuery("wait", "10")
+	wait, err := strconv.ParseUint(waitRaw, 10, 16)
+	if err != nil {
+		return responder.Code(http.StatusBadRequest)
+	}
+	waitDuration := time.Duration(wait) * time.Second
+
 	var worker *v1.Worker
 
 	if responder := controller.storeView(func(txn storepkg.Transaction) responder.Responder {
@@ -40,5 +48,5 @@ func (controller *Controller) portForwardWorker(ctx *gin.Context) responder.Resp
 	}
 
 	// Commence port-forwarding
-	return controller.portForward(ctx, worker.Name, "", uint32(port))
+	return controller.portForward(ctx, worker.Name, "", uint32(port), waitDuration)
 }

internal/controller/notifier/notifier.go

@@ -3,10 +3,10 @@ package notifier
 import (
 	"context"
 	"errors"
-	"fmt"
 	"github.com/cirruslabs/orchard/internal/concurrentmap"
 	"github.com/cirruslabs/orchard/rpc"
 	"go.uber.org/zap"
+	"time"
 )
 
 var ErrNoWorker = errors.New("no worker registered with this name")
@@ -47,8 +47,16 @@ func (watcher *Notifier) Register(ctx context.Context, worker string) (chan *rpc
 
 func (watcher *Notifier) Notify(ctx context.Context, worker string, msg *rpc.WatchInstruction) error {
 	slot, ok := watcher.workers.Load(worker)
-	if !ok {
-		return fmt.Errorf("%w: %s", ErrNoWorker, worker)
+
+	for !ok {
+		select {
+		case <-ctx.Done():
+			watcher.logger.Warnf("failed to notify watcher of worker %s due to timeout: %v", worker, ctx.Err())
+			return ctx.Err()
+		case <-time.After(time.Second):
+			watcher.logger.Infof("waiting for worker %s to re-connect...", worker)
+		}
+		slot, ok = watcher.workers.Load(worker)
 	}
 
 	select {

internal/controller/scheduler/scheduler.go

@@ -350,11 +350,13 @@ NextVM:
 		// It's fine to not treat the error as fatal here,
 		// since the worker will sync the VMs on the next
 		// scheduling iteration
-		if err := scheduler.notifier.Notify(context.Background(), affectedWorker, &rpc.WatchInstruction{
+		notifyContext, notifyContextCancel := context.WithTimeout(context.Background(), time.Second)
+		if err := scheduler.notifier.Notify(notifyContext, affectedWorker, &rpc.WatchInstruction{
 			Action: &rpc.WatchInstruction_SyncVmsAction{},
 		}); err != nil {
 			scheduler.logger.Errorf("Failed to reactively sync VMs on worker %s: %v", affectedWorker, err)
 		}
+		notifyContextCancel()
 	}
 
 	return nil

internal/controller/sshserver/sshserver.go

@@ -17,6 +17,7 @@ import (
 	"golang.org/x/crypto/ssh"
 	"net"
 	"strings"
+	"time"
 )
 
 const (
@@ -235,7 +236,9 @@ func (server *SSHServer) handleDirectTCPIP(ctx context.Context, newChannel ssh.N
 	boomerangConnCh, cancel := server.connRendezvous.Request(ctx, session)
 	defer cancel()
 
-	err = server.workerNotifier.Notify(ctx, vm.Worker, &rpc.WatchInstruction{
+	notifyContext, notifyContextCancel := context.WithTimeout(ctx, time.Second)
+	defer notifyContextCancel()
+	err = server.workerNotifier.Notify(notifyContext, vm.Worker, &rpc.WatchInstruction{
 		Action: &rpc.WatchInstruction_PortForwardAction{
 			PortForwardAction: &rpc.WatchInstruction_PortForward{
 				Session: session,