From 5eaf6b24d43afe6566215e40996f527c96fd3cc8 Mon Sep 17 00:00:00 2001 From: Fedor Korotkov Date: Mon, 27 Mar 2023 15:52:21 -0400 Subject: [PATCH] Make port-forward endpoint to wait for the VM (#65) * Make port-forward endpoint to wait for the VM Fixes #62 * Fixes after rebase --- api/openapi.yaml | 8 ++++ internal/command/portforward/vm.go | 7 +++- internal/command/ssh/vm.go | 5 ++- internal/command/vnc/vm.go | 5 ++- internal/controller/api_vms_portforward.go | 46 ++++++++++++++++------ internal/tests/integration_test.go | 20 +++------- pkg/client/vms.go | 12 +++++- 7 files changed, 72 insertions(+), 31 deletions(-) diff --git a/api/openapi.yaml b/api/openapi.yaml index 1d33865..a5ff648 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -234,6 +234,14 @@ paths: minimum: 1 maximum: 65535 required: true + - in: query + name: wait + description: Duration in seconds to wait for the VM to transition into "running" state if not already running. + schema: + type: integer + minimum: 0 + maximum: 65535 + required: false - in: header name: Connection description: WebSocket protocol required header diff --git a/internal/command/portforward/vm.go b/internal/command/portforward/vm.go index 3f103c2..b493a6f 100644 --- a/internal/command/portforward/vm.go +++ b/internal/command/portforward/vm.go @@ -8,6 +8,8 @@ import ( "net" ) +var wait uint16 + func newPortForwardVMCommand() *cobra.Command { command := &cobra.Command{ Use: "vm VM_NAME [LOCAL_PORT]:REMOTE_PORT", @@ -16,6 +18,9 @@ func newPortForwardVMCommand() *cobra.Command { RunE: runPortForwardVMCommand, } + command.PersistentFlags().Uint16VarP(&wait, "wait", "t", 60, + "Amount of seconds to wait for the VM to start running if it's not running already") + return command } @@ -59,7 +64,7 @@ func runPortForwardVMCommand(cmd *cobra.Command, args []string) (err error) { go func() { defer conn.Close() - wsConn, err := client.VMs().PortForward(cmd.Context(), name, portSpec.RemotePort) + wsConn, err := client.VMs().PortForward(cmd.Context(), name, portSpec.RemotePort, wait) if err != nil { fmt.Printf("failed to forward port: %v\n", err) diff --git a/internal/command/ssh/vm.go b/internal/command/ssh/vm.go index 7d15bcd..d2e3809 100644 --- a/internal/command/ssh/vm.go +++ b/internal/command/ssh/vm.go @@ -18,6 +18,7 @@ var ErrFailed = errors.New("ssh command failed") var username string var password string +var wait uint16 func newSSHVMCommand() *cobra.Command { command := &cobra.Command{ @@ -31,6 +32,8 @@ func newSSHVMCommand() *cobra.Command { "SSH username") command.PersistentFlags().StringVarP(&password, "password", "p", "", "SSH password") + command.PersistentFlags().Uint16VarP(&wait, "wait", "t", 60, + "Amount of seconds to wait for the VM to start running if it's not running already") return command } @@ -43,7 +46,7 @@ func runSSHVM(cmd *cobra.Command, args []string) error { return err } - wsConn, err := client.VMs().PortForward(cmd.Context(), name, 22) + wsConn, err := client.VMs().PortForward(cmd.Context(), name, 22, wait) if err != nil { fmt.Printf("failed to forward an SSH port to VM %s: %v\n", name, err) diff --git a/internal/command/vnc/vm.go b/internal/command/vnc/vm.go index 2373ace..019b784 100644 --- a/internal/command/vnc/vm.go +++ b/internal/command/vnc/vm.go @@ -14,6 +14,7 @@ const vncPort = 5900 var username string var password string +var wait uint16 func newVNCVMCommand() *cobra.Command { command := &cobra.Command{ @@ -27,6 +28,8 @@ func newVNCVMCommand() *cobra.Command { "VNC username") command.PersistentFlags().StringVarP(&password, "password", "p", "", "VNC password") + command.PersistentFlags().Uint16VarP(&wait, "wait", "w", 60, + "Amount of seconds to wait for the VM to start running if it's not running already") return command } @@ -63,7 +66,7 @@ func runVNCVM(cmd *cobra.Command, args []string) (err error) { } go func() { - wsConn, err := client.VMs().PortForward(cmd.Context(), name, vncPort) + wsConn, err := client.VMs().PortForward(cmd.Context(), name, vncPort, wait) if err != nil { fmt.Printf("failed to forward port: %v\n", err) diff --git a/internal/controller/api_vms_portforward.go b/internal/controller/api_vms_portforward.go index 94e9142..2f3fff9 100644 --- a/internal/controller/api_vms_portforward.go +++ b/internal/controller/api_vms_portforward.go @@ -1,6 +1,7 @@ package controller import ( + "context" storepkg "github.com/cirruslabs/orchard/internal/controller/store" "github.com/cirruslabs/orchard/internal/proxy" "github.com/cirruslabs/orchard/internal/responder" @@ -11,6 +12,7 @@ import ( "golang.org/x/net/websocket" "net/http" "strconv" + "time" ) func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responder { @@ -30,23 +32,43 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde return responder.Code(http.StatusBadRequest) } + waitRaw := ctx.Query("wait") + 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) + defer waitContextCancel() + // Look-up the VM var vm *v1.VM - if responder := controller.storeView(func(txn storepkg.Transaction) responder.Responder { - vm, err = txn.GetVM(name) - if err != nil { - return responder.Error(err) + for { + if lookupResponder := controller.storeView(func(txn storepkg.Transaction) responder.Responder { + vm, err = txn.GetVM(name) + if err != nil { + return responder.Error(err) + } + + return nil + }); lookupResponder != nil { + return lookupResponder } - return nil - }); responder != nil { - return responder - } - - // Sanity-check - if vm.Worker == "" { - return responder.Code(http.StatusServiceUnavailable) + if vm.TerminalState() { + return responder.JSON(http.StatusExpectationFailed, NewErrorResponse("VM is in a terminal state '%s'", vm.Status)) + } + if vm.Status == v1.VMStatusRunning { + // VM is running, proceed + break + } + select { + case <-waitContext.Done(): + return responder.JSON(http.StatusRequestTimeout, NewErrorResponse("VM is not running on '%s' worker", vm.Worker)) + case <-time.After(1 * time.Second): + // try again + continue + } } // Request and wait for a connection with a worker diff --git a/internal/tests/integration_test.go b/internal/tests/integration_test.go index 3faeb02..6e1e3a6 100644 --- a/internal/tests/integration_test.go +++ b/internal/tests/integration_test.go @@ -200,24 +200,16 @@ func TestPortForwarding(t *testing.T) { }) require.NoError(t, err) - // Wait for the VM to start - var vm *v1.VM - - require.True(t, Wait(2*time.Minute, func() bool { - vm, err = devClient.VMs().Get(ctx, "test-vm") - require.NoError(t, err) - - t.Logf("Waiting for the VM to start, current status: %s", vm.Status) - - return vm.Status == v1.VMStatusRunning || vm.Status == v1.VMStatusFailed - }), "failed to start a VM") + // Establish port forwarding to VMs SSH port + wsConn, err := devClient.VMs().PortForward(ctx, "test-vm", 22, 120) + require.NoError(t, err) + vm, err := devClient.VMs().Get(ctx, "test-vm") + require.NoError(t, err) require.Equal(t, v1.VMStatusRunning, vm.Status) require.Empty(t, vm.StatusMessage) - // Establish port forwarding to VMs SSH port - wsConn, err := devClient.VMs().PortForward(ctx, "test-vm", 22) - require.NoError(t, err) + t.Logf("Waiting for the VM to start, current status: %s", vm.Status) // Make sure we can connect to the VM over SSH via the forwarded port sshConfig := &ssh.ClientConfig{ diff --git a/pkg/client/vms.go b/pkg/client/vms.go index c4aac4e..f1985e3 100644 --- a/pkg/client/vms.go +++ b/pkg/client/vms.go @@ -101,9 +101,17 @@ func (service *VMsService) Delete(ctx context.Context, name string) error { return nil } -func (service *VMsService) PortForward(ctx context.Context, name string, port uint16) (*websocket.Conn, error) { +func (service *VMsService) PortForward( + ctx context.Context, + name string, + port uint16, + waitSeconds uint16, +) (*websocket.Conn, error) { return service.client.wsRequest(ctx, fmt.Sprintf("vms/%s/port-forward", name), - map[string]string{"port": strconv.FormatUint(uint64(port), 10)}) + map[string]string{ + "port": strconv.FormatUint(uint64(port), 10), + "wait": strconv.FormatUint(uint64(waitSeconds), 10), + }) } func (service *VMsService) StreamEvents(name string) *EventStreamer {