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,