From 2f0cd0f23f4deb15ef4a44504896ef8e0be11cac Mon Sep 17 00:00:00 2001
From: iBug <git@ibugone.com>
Date: Sat, 20 Jan 2024 02:12:51 +0800
Subject: [PATCH] Add killall subcommand

---
 cmd/killall/killall.go | 72 ++++++++++++++++++++++++++++++++++++++++++
 cmd/root.go            |  4 ++-
 go.sum                 |  2 --
 pkg/pve/pct.go         | 13 ++++++++
 4 files changed, 88 insertions(+), 3 deletions(-)
 create mode 100644 cmd/killall/killall.go
 create mode 100644 pkg/pve/pct.go

diff --git a/cmd/killall/killall.go b/cmd/killall/killall.go
new file mode 100644
index 0000000..104cd06
--- /dev/null
+++ b/cmd/killall/killall.go
@@ -0,0 +1,72 @@
+package killall
+
+import (
+	"strconv"
+	"sync"
+
+	"github.com/USTC-vlab/vct/pkg/cgroup"
+	"github.com/USTC-vlab/vct/pkg/pve"
+	"github.com/spf13/cobra"
+)
+
+func killWorker(ch <-chan string, errCh chan<- error, wg *sync.WaitGroup) {
+	defer wg.Done()
+	for id := range ch {
+		cmd := pve.Stop(id)
+		if err := cmd.Run(); err != nil {
+			errCh <- err
+		}
+	}
+}
+
+func killallMain(n int, minID int) error {
+	ch := make(chan string)
+	chErr := make(chan error)
+	wg := &sync.WaitGroup{}
+	wg.Add(n)
+	for i := 0; i < n; i++ {
+		go killWorker(ch, chErr, wg)
+	}
+	defer func() {
+		close(ch)
+		wg.Wait()
+		close(chErr)
+	}()
+	// discard errors for now
+	go func() {
+		for range chErr {
+		}
+	}()
+
+	ids, err := cgroup.ListLXC()
+	if err != nil {
+		return err
+	}
+	for _, id := range ids {
+		numID, err := strconv.Atoi(id)
+		if err != nil {
+			return err
+		}
+		if numID < minID {
+			continue
+		}
+		ch <- id
+	}
+	return nil
+}
+
+func MakeCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "killall",
+		Short: "Kill all running containers",
+		Args:  cobra.NoArgs,
+	}
+	flags := cmd.Flags()
+	pN := flags.IntP("n", "n", 5, "max number of parallel killing containers")
+	pS := flags.IntP("min", "m", 1000, "minimum ID of containers to kill")
+
+	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		return killallMain(*pN, *pS)
+	}
+	return cmd
+}
diff --git a/cmd/root.go b/cmd/root.go
index 075cf83..6d5a846 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -4,6 +4,7 @@ import (
 	"github.com/USTC-vlab/vct/cmd/df"
 	"github.com/USTC-vlab/vct/cmd/findpid"
 	"github.com/USTC-vlab/vct/cmd/iostat"
+	"github.com/USTC-vlab/vct/cmd/killall"
 	"github.com/USTC-vlab/vct/cmd/pressure"
 	"github.com/spf13/cobra"
 )
@@ -29,6 +30,7 @@ func MakeCmd() *cobra.Command {
 		df.MakeCmd(),
 		findpid.MakeCmd(),
 		iostat.MakeCmd(),
+		killall.MakeCmd(),
 		pressure.MakeCmd(),
 		versionCmd,
 	)
@@ -36,7 +38,7 @@ func MakeCmd() *cobra.Command {
 
 	cmd.Run = func(cmd *cobra.Command, args []string) {
 		if *pVersion {
-			versionCmd.Run(cmd, args)
+			versionCmd.Run(versionCmd, args)
 		} else {
 			cmd.Help()
 		}
diff --git a/go.sum b/go.sum
index 8f0e7f9..ceaf173 100644
--- a/go.sum
+++ b/go.sum
@@ -8,8 +8,6 @@ github.com/spf13/cobra v1.8.0 h1:7aJaZx1B85qltLMc546zn58BxxfZdR/W22ej9CFoEf0=
 github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
-golang.org/x/exp v0.0.0-20240103183307-be819d1f06fc h1:ao2WRsKSzW6KuUY9IWPwWahcHCgR0s52IfwutMfEbdM=
-golang.org/x/exp v0.0.0-20240103183307-be819d1f06fc/go.mod h1:iRJReGqOEeBhDZGkGbynYwcHlctCvnjTYIamk7uXpHI=
 golang.org/x/exp v0.0.0-20240112132812-db7319d0e0e3 h1:hNQpMuAJe5CtcUqCXaWga3FHu+kQvCqcsoVaQgSV60o=
 golang.org/x/exp v0.0.0-20240112132812-db7319d0e0e3/go.mod h1:idGWGoKP1toJGkd5/ig9ZLuPcZBC3ewk7SzmH0uou08=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
diff --git a/pkg/pve/pct.go b/pkg/pve/pct.go
new file mode 100644
index 0000000..654bada
--- /dev/null
+++ b/pkg/pve/pct.go
@@ -0,0 +1,13 @@
+package pve
+
+import "os/exec"
+
+const PctPath = "/usr/sbin/pct"
+
+func Pct(args ...string) *exec.Cmd {
+	return exec.Command(PctPath, args...)
+}
+
+func Stop(vmid string) *exec.Cmd {
+	return Pct("stop", vmid)
+}