1 files changed, 245 insertions, 0 deletions
diff --git a/cmd/kpod/stats.go b/cmd/kpod/stats.go
new file mode 100644
index 000000000..ac81212a1
--- /dev/null
+++ b/cmd/kpod/stats.go
@@ -0,0 +1,245 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"strings"
+	"text/template"
+	"time"
+
+	"github.com/docker/go-units"
+
+	tm "github.com/buger/goterm"
+	"github.com/kubernetes-incubator/cri-o/libkpod"
+	"github.com/kubernetes-incubator/cri-o/oci"
+	"github.com/pkg/errors"
+	"github.com/urfave/cli"
+)
+
+var printf func(format string, a ...interface{}) (n int, err error)
+var println func(a ...interface{}) (n int, err error)
+
+type statsOutputParams struct {
+	Container string
+	ID        string
+	CPUPerc   string
+	MemUsage  string
+	MemPerc   string
+	NetIO     string
+	BlockIO   string
+	PIDs      uint64
+}
+
+var (
+	statsFlags = []cli.Flag{
+		cli.BoolFlag{
+			Name:  "all, a",
+			Usage: "show all containers. Only running containers are shown by default. The default is false",
+		},
+		cli.BoolFlag{
+			Name:  "no-stream",
+			Usage: "disable streaming stats and only pull the first result, default setting is false",
+		},
+		cli.StringFlag{
+			Name:  "format",
+			Usage: "pretty-print container statistics using a Go template",
+		},
+		cli.BoolFlag{
+			Name:  "json",
+			Usage: "output container statistics in json format",
+		},
+	}
+
+	statsDescription = "display a live stream of one or more containers' resource usage statistics"
+	statsCommand     = cli.Command{
+		Name:        "stats",
+		Usage:       "Display percentage of CPU, memory, network I/O, block I/O and PIDs for one or more containers",
+		Description: statsDescription,
+		Flags:       statsFlags,
+		Action:      statsCmd,
+		ArgsUsage:   "",
+	}
+)
+
+func statsCmd(c *cli.Context) error {
+	if err := validateFlags(c, statsFlags); err != nil {
+		return err
+	}
+	config, err := getConfig(c)
+	if err != nil {
+		return errors.Wrapf(err, "could not read config")
+	}
+	containerServer, err := libkpod.New(config)
+	if err != nil {
+		return errors.Wrapf(err, "could not create container server")
+	}
+	defer containerServer.Shutdown()
+	err = containerServer.Update()
+	if err != nil {
+		return errors.Wrapf(err, "could not update list of containers")
+	}
+	times := -1
+	if c.Bool("no-stream") {
+		times = 1
+	}
+	statsChan := make(chan []*libkpod.ContainerStats)
+	// iterate over the channel until it is closed
+	go func() {
+		// print using goterm
+		printf = tm.Printf
+		println = tm.Println
+		for stats := range statsChan {
+			// Continually refresh statistics
+			tm.Clear()
+			tm.MoveCursor(1, 1)
+			outputStats(stats, c.String("format"), c.Bool("json"))
+			tm.Flush()
+			time.Sleep(time.Second)
+		}
+	}()
+	return getStats(containerServer, c.Args(), c.Bool("all"), statsChan, times)
+}
+
+func getStats(server *libkpod.ContainerServer, args []string, all bool, statsChan chan []*libkpod.ContainerStats, times int) error {
+	ctrs, err := server.ListContainers(isRunning, ctrInList(args))
+	if err != nil {
+		return err
+	}
+	containerStats := map[string]*libkpod.ContainerStats{}
+	for _, ctr := range ctrs {
+		initialStats, err := server.GetContainerStats(ctr, &libkpod.ContainerStats{})
+		if err != nil {
+			return err
+		}
+		containerStats[ctr.ID()] = initialStats
+	}
+	step := 1
+	if times == -1 {
+		times = 1
+		step = 0
+	}
+	for i := 0; i < times; i += step {
+		reportStats := []*libkpod.ContainerStats{}
+		for _, ctr := range ctrs {
+			id := ctr.ID()
+			if _, ok := containerStats[ctr.ID()]; !ok {
+				initialStats, err := server.GetContainerStats(ctr, &libkpod.ContainerStats{})
+				if err != nil {
+					return err
+				}
+				containerStats[id] = initialStats
+			}
+			stats, err := server.GetContainerStats(ctr, containerStats[id])
+			if err != nil {
+				return err
+			}
+			// replace the previous measurement with the current one
+			containerStats[id] = stats
+			reportStats = append(reportStats, stats)
+		}
+		statsChan <- reportStats
+
+		err := server.Update()
+		if err != nil {
+			return err
+		}
+		ctrs, err = server.ListContainers(isRunning, ctrInList(args))
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func outputStats(stats []*libkpod.ContainerStats, format string, json bool) error {
+	if format == "" {
+		outputStatsHeader()
+	}
+	if json {
+		return outputStatsAsJSON(stats)
+	}
+	var err error
+	for _, s := range stats {
+		if format == "" {
+			outputStatsUsingFormatString(s)
+		} else {
+			params := getStatsOutputParams(s)
+			err2 := outputStatsUsingTemplate(format, params)
+			if err2 != nil {
+				err = errors.Wrapf(err, err2.Error())
+			}
+		}
+	}
+	return err
+}
+
+func outputStatsHeader() {
+	printf("%-64s %-16s %-32s %-16s %-24s %-24s %s\n", "CONTAINER", "CPU %", "MEM USAGE / MEM LIMIT", "MEM %", "NET I/O", "BLOCK I/O", "PIDS")
+}
+
+func outputStatsUsingFormatString(stats *libkpod.ContainerStats) {
+	printf("%-64s %-16s %-32s %-16s %-24s %-24s %d\n", stats.Container, floatToPercentString(stats.CPU), combineHumanValues(stats.MemUsage, stats.MemLimit), floatToPercentString(stats.MemPerc), combineHumanValues(stats.NetInput, stats.NetOutput), combineHumanValues(stats.BlockInput, stats.BlockOutput), stats.PIDs)
+}
+
+func combineHumanValues(a, b uint64) string {
+	return fmt.Sprintf("%s / %s", units.HumanSize(float64(a)), units.HumanSize(float64(b)))
+}
+
+func floatToPercentString(f float64) string {
+	return fmt.Sprintf("%.2f %s", f, "%")
+}
+
+func getStatsOutputParams(stats *libkpod.ContainerStats) statsOutputParams {
+	return statsOutputParams{
+		Container: stats.Container,
+		ID:        stats.Container,
+		CPUPerc:   floatToPercentString(stats.CPU),
+		MemUsage:  combineHumanValues(stats.MemUsage, stats.MemLimit),
+		MemPerc:   floatToPercentString(stats.MemPerc),
+		NetIO:     combineHumanValues(stats.NetInput, stats.NetOutput),
+		BlockIO:   combineHumanValues(stats.BlockInput, stats.BlockOutput),
+		PIDs:      stats.PIDs,
+	}
+}
+
+func outputStatsUsingTemplate(format string, params statsOutputParams) error {
+	tmpl, err := template.New("stats").Parse(format)
+	if err != nil {
+		return errors.Wrapf(err, "template parsing error")
+	}
+
+	err = tmpl.Execute(os.Stdout, params)
+	if err != nil {
+		return err
+	}
+	println()
+	return nil
+}
+
+func outputStatsAsJSON(stats []*libkpod.ContainerStats) error {
+	s, err := json.Marshal(stats)
+	if err != nil {
+		return err
+	}
+	println(s)
+	return nil
+}
+
+func isRunning(ctr *oci.Container) bool {
+	return ctr.State().Status == "running"
+}
+
+func ctrInList(idsOrNames []string) func(ctr *oci.Container) bool {
+	if len(idsOrNames) == 0 {
+		return func(*oci.Container) bool { return true }
+	}
+	return func(ctr *oci.Container) bool {
+		for _, idOrName := range idsOrNames {
+			if strings.HasPrefix(ctr.ID(), idOrName) || strings.HasSuffix(ctr.Name(), idOrName) {
+				return true
+			}
+		}
+		return false
+	}
+}