78 lines
1.4 KiB
Go
78 lines
1.4 KiB
Go
package pause
|
|
|
|
import (
|
|
"context"
|
|
"github.com/cirruslabs/orchard/pkg/client"
|
|
v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
|
|
"github.com/spf13/cobra"
|
|
"golang.org/x/exp/slices"
|
|
"time"
|
|
)
|
|
|
|
var wait uint64
|
|
|
|
func newPauseWorkerCommand() *cobra.Command {
|
|
command := &cobra.Command{
|
|
Use: "worker NAME",
|
|
Short: "Pause a worker",
|
|
RunE: runPauseWorker,
|
|
Args: cobra.ExactArgs(1),
|
|
}
|
|
|
|
command.Flags().Uint64Var(&wait, "wait", 0,
|
|
"wait the specified amount of seconds for the worker to stop running any VMs")
|
|
|
|
return command
|
|
}
|
|
|
|
func runPauseWorker(cmd *cobra.Command, args []string) error {
|
|
name := args[0]
|
|
|
|
client, err := client.New()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
worker, err := client.Workers().Get(cmd.Context(), name)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if !worker.SchedulingPaused {
|
|
worker.SchedulingPaused = true
|
|
|
|
_, err = client.Workers().Update(cmd.Context(), *worker)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if wait != 0 {
|
|
subCtx, cancel := context.WithTimeout(cmd.Context(), time.Duration(wait)*time.Second)
|
|
defer cancel()
|
|
|
|
for {
|
|
vms, err := client.VMs().FindForWorker(cmd.Context(), worker.Name)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
hasActiveVMs := slices.ContainsFunc(vms, func(vm v1.VM) bool {
|
|
return !vm.TerminalState()
|
|
})
|
|
if !hasActiveVMs {
|
|
break
|
|
}
|
|
|
|
select {
|
|
case <-time.After(time.Second):
|
|
continue
|
|
case <-subCtx.Done():
|
|
return subCtx.Err()
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|