1 files changed, 65 insertions, 51 deletions
diff --git a/pkg/domain/infra/abi/containers.go b/pkg/domain/infra/abi/containers.go
index 21618f555..8b0d53940 100644
--- a/pkg/domain/infra/abi/containers.go
+++ b/pkg/domain/infra/abi/containers.go
@@ -1142,12 +1142,12 @@ func (ic *ContainerEngine) Shutdown(_ context.Context) {
 	})
 }
 
-func (ic *ContainerEngine) ContainerStats(ctx context.Context, namesOrIds []string, options entities.ContainerStatsOptions) error {
-	defer close(options.StatChan)
+func (ic *ContainerEngine) ContainerStats(ctx context.Context, namesOrIds []string, options entities.ContainerStatsOptions) (statsChan chan entities.ContainerStatsReport, err error) {
+	statsChan = make(chan entities.ContainerStatsReport, 1)
+
 	containerFunc := ic.Libpod.GetRunningContainers
+	queryAll := false
 	switch {
-	case len(namesOrIds) > 0:
-		containerFunc = func() ([]*libpod.Container, error) { return ic.Libpod.GetContainersByList(namesOrIds) }
 	case options.Latest:
 		containerFunc = func() ([]*libpod.Container, error) {
 			lastCtr, err := ic.Libpod.GetLatestContainer()
@@ -1156,62 +1156,76 @@ func (ic *ContainerEngine) ContainerStats(ctx context.Context, namesOrIds []stri
 			}
 			return []*libpod.Container{lastCtr}, nil
 		}
-	case options.All:
+	case len(namesOrIds) > 0:
+		containerFunc = func() ([]*libpod.Container, error) { return ic.Libpod.GetContainersByList(namesOrIds) }
+	default:
+		// No containers, no latest -> query all!
+		queryAll = true
 		containerFunc = ic.Libpod.GetAllContainers
 	}
 
-	ctrs, err := containerFunc()
-	if err != nil {
-		return errors.Wrapf(err, "unable to get list of containers")
-	}
-	containerStats := map[string]*define.ContainerStats{}
-	for _, ctr := range ctrs {
-		initialStats, err := ctr.GetContainerStats(&define.ContainerStats{})
-		if err != nil {
-			// when doing "all", don't worry about containers that are not running
-			cause := errors.Cause(err)
-			if options.All && (cause == define.ErrCtrRemoved || cause == define.ErrNoSuchCtr || cause == define.ErrCtrStateInvalid) {
-				continue
-			}
-			if cause == cgroups.ErrCgroupV1Rootless {
-				err = cause
-			}
-			return err
+	go func() {
+		defer close(statsChan)
+		var (
+			err            error
+			containers     []*libpod.Container
+			containerStats map[string]*define.ContainerStats
+		)
+		containerStats = make(map[string]*define.ContainerStats)
+
+	stream: // label to flatten the scope
+		select {
+		case <-ctx.Done():
+			// client cancelled
+			logrus.Debugf("Container stats stopped: context cancelled")
+			return
+		default:
+			// just fall through and do work
 		}
-		containerStats[ctr.ID()] = initialStats
-	}
-	for {
-		reportStats := []*define.ContainerStats{}
-		for _, ctr := range ctrs {
-			id := ctr.ID()
-			if _, ok := containerStats[ctr.ID()]; !ok {
-				initialStats, err := ctr.GetContainerStats(&define.ContainerStats{})
-				if errors.Cause(err) == define.ErrCtrRemoved || errors.Cause(err) == define.ErrNoSuchCtr || errors.Cause(err) == define.ErrCtrStateInvalid {
-					// skip dealing with a container that is gone
-					continue
+
+		// Anonymous func to easily use the return values for streaming.
+		computeStats := func() ([]define.ContainerStats, error) {
+			containers, err = containerFunc()
+			if err != nil {
+				return nil, errors.Wrapf(err, "unable to get list of containers")
+			}
+
+			reportStats := []define.ContainerStats{}
+			for _, ctr := range containers {
+				prev, ok := containerStats[ctr.ID()]
+				if !ok {
+					prev = &define.ContainerStats{}
 				}
+
+				stats, err := ctr.GetContainerStats(prev)
 				if err != nil {
-					return err
+					cause := errors.Cause(err)
+					if queryAll && (cause == define.ErrCtrRemoved || cause == define.ErrNoSuchCtr || cause == define.ErrCtrStateInvalid) {
+						continue
+					}
+					if cause == cgroups.ErrCgroupV1Rootless {
+						err = cause
+					}
+					return nil, err
 				}
-				containerStats[id] = initialStats
-			}
-			stats, err := ctr.GetContainerStats(containerStats[id])
-			if err != nil && errors.Cause(err) != define.ErrNoSuchCtr {
-				return err
+
+				containerStats[ctr.ID()] = stats
+				reportStats = append(reportStats, *stats)
 			}
-			// replace the previous measurement with the current one
-			containerStats[id] = stats
-			reportStats = append(reportStats, stats)
+			return reportStats, nil
 		}
-		ctrs, err = containerFunc()
-		if err != nil {
-			return err
-		}
-		options.StatChan <- reportStats
-		if options.NoStream {
-			break
+
+		report := entities.ContainerStatsReport{}
+		report.Stats, report.Error = computeStats()
+		statsChan <- report
+
+		if report.Error != nil || !options.Stream {
+			return
 		}
+
 		time.Sleep(time.Second)
-	}
-	return nil
+		goto stream
+	}()
+
+	return statsChan, nil
 }