iolimit: Print current limits if no settings

cmd/iolimit/iolimit.go

@@ -3,19 +3,22 @@ package iolimit
 import (
 	"errors"
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/USTC-vlab/vct/pkg/cgroup"
 	"github.com/USTC-vlab/vct/pkg/pve"
+	"github.com/ryanuber/columnize"
 	"github.com/spf13/cobra"
 )
 
 func MakeCmd() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "iolimit [--riops] [--wiops] [--rbps] [--wbps] CTID...",
+		Use:   "iolimit [--riops COUNT] [--wiops COUNT] [--rbps BYTES] [--wbps BYTES] CTID...",
 		Short: "Set I/O limits for an LXC container",
-		Long:  "Set I/O limits for an LXC container. Use -1 to remove existing limit.",
-		Args:  cobra.MinimumNArgs(1),
+		Long: `Set I/O limits for an LXC container.
+If no setting is given, print the current limits.
+Note that zero means "don't change". Use -1 to remove an existing limit (set to "max").`,
 	}
 	flags := cmd.Flags()
 	var iops cgroup.IOPS
@@ -24,9 +27,13 @@ func MakeCmd() *cobra.Command {
 	flags.Int64VarP(&iops.Riops, "riops", "r", 0, "set read IOPS limit")
 	flags.Int64VarP(&iops.Wiops, "wiops", "w", 0, "set write IOPS limit")
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		if len(args) == 0 {
+			return cmd.Help()
+		}
+
 		errs := make([]error, 0, len(args))
 		for _, arg := range args {
-			err := iolimitMain(arg, iops)
+			err := iolimitMain(cmd, arg, iops)
 			if err != nil {
 				errs = append(errs, err)
 			}
@@ -36,7 +43,31 @@ func MakeCmd() *cobra.Command {
 	return cmd
 }
 
-func iolimitMain(ctid string, iops cgroup.IOPS) error {
+func showIOLimit(w io.Writer, ctid string) error {
+	iopss, err := cgroup.GetIOPSForLXC(ctid)
+	if err != nil {
+		return err
+	}
+	if len(iopss) == 0 {
+		fmt.Fprintf(w, "No I/O limits for %s\n", ctid)
+		return nil
+	}
+
+	fmt.Fprintf(w, "I/O limits for %s:\n", ctid)
+	lines := []string{"Device | Rbps | Wbps | Riops | Wiops"}
+	for _, iops := range iopss {
+		line := fmt.Sprintf("%d:%d | %s | %s | %s | %s",
+			iops.Major, iops.Minor,
+			cgroup.ItoaZeroMax(iops.Rbps), cgroup.ItoaZeroMax(iops.Wbps),
+			cgroup.ItoaZeroMax(iops.Riops), cgroup.ItoaZeroMax(iops.Wiops))
+		lines = append(lines, line)
+	}
+	fmt.Fprintln(w, columnize.SimpleFormat(lines))
+	return nil
+}
+
+func setIOLimit(ctid string, iops cgroup.IOPS) error {
+
 	pveStorage, err := pve.GetStorage()
 	if err != nil {
 		return err
@@ -64,3 +95,10 @@ func iolimitMain(ctid string, iops cgroup.IOPS) error {
 	}
 	return cgroup.SetIOPSForLXC(ctid, iopsline)
 }
+
+func iolimitMain(cmd *cobra.Command, ctid string, iops cgroup.IOPS) error {
+	if iops.IsZero() {
+		return showIOLimit(cmd.OutOrStdout(), ctid)
+	}
+	return setIOLimit(ctid, iops)
+}

pkg/cgroup/io.go

@@ -1,6 +1,9 @@
 package cgroup
 
 import (
+	"bufio"
+	"fmt"
+	"io"
 	"os"
 	"strconv"
 	"strings"
@@ -13,7 +16,7 @@ type IOPS struct {
 	Rbps, Wbps, Riops, Wiops int64
 }
 
-func itoaZeroMax(i int64) string {
+func ItoaZeroMax(i int64) string {
 	if i <= 0 {
 		return "max"
 	}
@@ -23,30 +26,86 @@ func itoaZeroMax(i int64) string {
 func (i IOPS) String() string {
 	parts := make([]string, 0, 4)
 	if i.Rbps != 0 {
-		parts = append(parts, "rbps="+itoaZeroMax(i.Rbps))
+		parts = append(parts, "rbps="+ItoaZeroMax(i.Rbps))
 	}
 	if i.Wbps != 0 {
-		parts = append(parts, "wbps="+itoaZeroMax(i.Wbps))
+		parts = append(parts, "wbps="+ItoaZeroMax(i.Wbps))
 	}
 	if i.Riops != 0 {
-		parts = append(parts, "riops="+itoaZeroMax(i.Riops))
+		parts = append(parts, "riops="+ItoaZeroMax(i.Riops))
 	}
 	if i.Wiops != 0 {
-		parts = append(parts, "wiops="+itoaZeroMax(i.Wiops))
+		parts = append(parts, "wiops="+ItoaZeroMax(i.Wiops))
 	}
 	return strings.Join(parts, " ")
 }
 
+func (i IOPS) IsZero() bool {
+	return i.Rbps == 0 && i.Wbps == 0 && i.Riops == 0 && i.Wiops == 0
+}
+
+func ParseIOPS(line string) (i IOPS) {
+	parts := strings.Fields(line)
+	for _, part := range parts {
+		kv := strings.SplitN(part, "=", 2)
+		if len(kv) != 2 {
+			continue
+		}
+		switch kv[0] {
+		case "rbps":
+			i.Rbps, _ = strconv.ParseInt(kv[1], 10, 64)
+		case "wbps":
+			i.Wbps, _ = strconv.ParseInt(kv[1], 10, 64)
+		case "riops":
+			i.Riops, _ = strconv.ParseInt(kv[1], 10, 64)
+		case "wiops":
+			i.Wiops, _ = strconv.ParseInt(kv[1], 10, 64)
+		}
+	}
+	return
+}
+
 // IOPSLine represents a line in the io.max file.
 type IOPSLine struct {
 	Major, Minor uint64
-	IOPS         IOPS
+	IOPS
 }
 
 func (l IOPSLine) String() string {
 	return strconv.FormatUint(l.Major, 10) + ":" + strconv.FormatUint(l.Minor, 10) + " " + l.IOPS.String()
 }
 
+func parseIOPSLines(r io.Reader) ([]IOPSLine, error) {
+	scanner := bufio.NewScanner(r)
+	iopss := make([]IOPSLine, 0)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		fields := strings.Fields(line)
+		var major, minor uint64
+		_, err := fmt.Sscanf(fields[0], "%d:%d", &major, &minor)
+		if err != nil {
+			return iopss, err
+		}
+		iops := ParseIOPS(strings.Join(fields[1:], " "))
+		iopss = append(iopss, IOPSLine{
+			Major: major,
+			Minor: minor,
+			IOPS:  iops,
+		})
+	}
+	return iopss, nil
+}
+
+func GetIOPSForLXC(id string) ([]IOPSLine, error) {
+	filename := GetFilenameLXC(id, "io.max")
+	f, err := os.Open(filename)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	return parseIOPSLines(f)
+}
+
 func SetIOPSForLXC(id string, iops IOPSLine) error {
 	filename := GetFilenameLXC(id, "io.max")
 	return os.WriteFile(filename, []byte(iops.String()+"\n"), 0)