cli/command/container: collect: handle context-cancellation

construct the decoder inside the go-routine, including closing the body,
and add handling for context-cancellation.

Signed-off-by: Sebastiaan van Stijn <[email protected]>

Author: thaJeztah

cli/command/container/stats_helpers.go

@@ -49,7 +49,7 @@ func (s *stats) isKnownContainer(cid string) (int, bool) {
 	return -1, false
 }
 
-func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, streamStats bool, waitFirst *sync.WaitGroup) {
+func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, streamStats bool, waitFirst *sync.WaitGroup) { //nolint:gocyclo
 	var (
 		getFirst       bool
 		previousCPU    uint64
@@ -70,11 +70,14 @@ func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, strea
 		s.SetError(err)
 		return
 	}
-	defer response.Body.Close()
 
-	dec := json.NewDecoder(response.Body)
 	go func() {
+		defer response.Body.Close()
+		dec := json.NewDecoder(response.Body)
 		for {
+			if ctx.Err() != nil {
+				return
+			}
 			var (
 				v                      *container.StatsResponse
 				memPercent, cpuPercent float64
@@ -141,8 +144,8 @@ func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, strea
 			}
 		case err := <-u:
 			s.SetError(err)
-			if err == io.EOF {
-				break
+			if errors.Is(err, io.EOF) {
+				return
 			}
 			if err != nil {
 				continue
@@ -152,6 +155,9 @@ func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, strea
 				getFirst = true
 				waitFirst.Done()
 			}
+		case <-ctx.Done():
+			s.SetError(ctx.Err())
+			return
 		}
 		if !streamStats {
 			return