Make port-forward endpoint to wait for the VM (#65)

* Make port-forward endpoint to wait for the VM

Fixes #62

* Fixes after rebase
This commit is contained in:
Fedor Korotkov
2023-03-27 23:52:21 +04:00
committed by GitHub
parent 357a042937
commit 5eaf6b24d4
7 changed files with 72 additions and 31 deletions
+8
View File
@@ -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
+6 -1
View File
@@ -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)
+4 -1
View File
@@ -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)
+4 -1
View File
@@ -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)
+34 -12
View File
@@ -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
+6 -14
View File
@@ -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{
+10 -2
View File
@@ -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 {