From 86f0afb5a3c0f61416a88427f84fb00b7a22460b Mon Sep 17 00:00:00 2001 From: Fedor Korotkov Date: Thu, 6 Feb 2025 12:30:09 -0500 Subject: [PATCH] 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 | 9 +++++++++ internal/controller/api_vms_ip.go | 5 +++-- internal/controller/api_vms_portforward.go | 12 ++++++++---- internal/controller/api_workers_portforward.go | 10 +++++++++- internal/controller/notifier/notifier.go | 14 +++++++++++--- internal/controller/scheduler/scheduler.go | 4 +++- internal/controller/sshserver/sshserver.go | 5 ++++- 7 files changed, 47 insertions(+), 12 deletions(-) diff --git a/api/openapi.yaml b/api/openapi.yaml index 669d2f3..0afe4e0 100644 --- a/api/openapi.yaml +++ b/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 diff --git a/internal/controller/api_vms_ip.go b/internal/controller/api_vms_ip.go index ac1c6a4..a8595dc 100644 --- a/internal/controller/api_vms_ip.go +++ b/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, diff --git a/internal/controller/api_vms_portforward.go b/internal/controller/api_vms_portforward.go index 47b8d03..7f91ce3 100644 --- a/internal/controller/api_vms_portforward.go +++ b/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, diff --git a/internal/controller/api_workers_portforward.go b/internal/controller/api_workers_portforward.go index 0336841..f5700af 100644 --- a/internal/controller/api_workers_portforward.go +++ b/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) } diff --git a/internal/controller/notifier/notifier.go b/internal/controller/notifier/notifier.go index ab15675..df7e6f7 100644 --- a/internal/controller/notifier/notifier.go +++ b/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 { diff --git a/internal/controller/scheduler/scheduler.go b/internal/controller/scheduler/scheduler.go index 93ba316..85995d6 100644 --- a/internal/controller/scheduler/scheduler.go +++ b/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 diff --git a/internal/controller/sshserver/sshserver.go b/internal/controller/sshserver/sshserver.go index 33bd3b9..7f5c84f 100644 --- a/internal/controller/sshserver/sshserver.go +++ b/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,