stats: cancel old streams when containers restart

cli/command/container/stats.go

@@ -107,6 +107,9 @@ var acceptedStatsFilters = map[string]bool{
 //
 //nolint:gocyclo
 func RunStats(ctx context.Context, dockerCLI command.Cli, options *StatsOptions) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	apiClient := dockerCLI.Client()
 
 	// Get the daemonOSType to handle platform-specific stats fields.
@@ -121,6 +124,22 @@ func RunStats(ctx context.Context, dockerCLI command.Cli, options *StatsOptions)
 	// after the reader has exited, preventing deadlocks.
 	closeChan := make(chan error, 4)
 	cStats := stats{}
+	startCollecting := func(ctx context.Context, id string, waitFirst *sync.WaitGroup) {
+		ctx, cancel := context.WithCancel(ctx)
+		s := NewStats(id)
+		if !cStats.add(s, cancel) {
+			cancel()
+			return
+		}
+		if waitFirst != nil {
+			waitFirst.Add(1)
+		}
+		log.G(ctx).WithField("container", id).Debug("collecting stats for container")
+		go func() {
+			defer cancel()
+			collect(ctx, s, apiClient, !options.NoStream, waitFirst)
+		}()
+	}
 
 	showAll := len(options.Containers) == 0
 	if showAll {
@@ -147,11 +166,9 @@ func RunStats(ctx context.Context, dockerCLI command.Cli, options *StatsOptions)
 			addEvents = append(addEvents, events.ActionCreate)
 		}
 		eh.setHandler(addEvents, func(ctx context.Context, e events.Message) {
-			if s := NewStats(e.Actor.ID); cStats.add(s) {
-				waitFirst.Add(1)
-				log.G(ctx).Debug("collecting stats for container")
-				go collect(ctx, s, apiClient, !options.NoStream, waitFirst)
-			}
+			// Only the initial collectors participate in waitFirst. Reusing it
+			// after the initial Wait can race with a new Add on a zero counter.
+			startCollecting(ctx, e.Actor.ID, nil)
 		})
 
 		// Remove containers when they are removed ("destroyed"); containers
@@ -224,13 +241,7 @@ func RunStats(ctx context.Context, dockerCLI command.Cli, options *StatsOptions)
 			return err
 		}
 		for _, ctr := range cs.Items {
-			if s := NewStats(ctr.ID); cStats.add(s) {
-				waitFirst.Add(1)
-				log.G(ctx).WithFields(log.Fields{
-					"container": ctr.ID,
-				}).Debug("collecting stats for container")
-				go collect(ctx, s, apiClient, !options.NoStream, waitFirst)
-			}
+			startCollecting(ctx, ctr.ID, waitFirst)
 		}
 
 		// make sure each container get at least one valid stat data
@@ -247,13 +258,7 @@ func RunStats(ctx context.Context, dockerCLI command.Cli, options *StatsOptions)
 		// Create the list of containers, and start collecting stats for all
 		// containers passed.
 		for _, ctr := range options.Containers {
-			if s := NewStats(ctr); cStats.add(s) {
-				waitFirst.Add(1)
-				log.G(ctx).WithFields(log.Fields{
-					"container": ctr,
-				}).Debug("collecting stats for container")
-				go collect(ctx, s, apiClient, !options.NoStream, waitFirst)
-			}
+			startCollecting(ctx, ctr, waitFirst)
 		}
 
 		// We don't expect any asynchronous errors: closeChan can be closed and disabled.

cli/command/container/stats_helpers.go

@@ -13,19 +13,24 @@ import (
 )
 
 type stats struct {
-	mu sync.RWMutex
-	cs []*Stats
+	mu      sync.RWMutex
+	cs      []*Stats
+	cancels map[string]context.CancelFunc
 }
 
 // daemonOSType is set once we have at least one stat for a container
 // from the daemon. It is used to ensure we print the right header based
 // on the daemon platform.
 var daemonOSType string
 
-func (s *stats) add(cs *Stats) bool {
+func (s *stats) add(cs *Stats, cancel context.CancelFunc) bool {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if _, exists := s.isKnownContainer(cs.Container); !exists {
+		if s.cancels == nil {
+			s.cancels = make(map[string]context.CancelFunc)
+		}
+		s.cancels[cs.Container] = cancel
 		s.cs = append(s.cs, cs)
 		return true
 	}
@@ -37,12 +42,19 @@ func (s *stats) remove(id string) {
 	if i, exists := s.isKnownContainer(id); exists {
 		s.cs = append(s.cs[:i], s.cs[i+1:]...)
 	}
+	cancel := s.cancels[id]
+	delete(s.cancels, id)
 	s.mu.Unlock()
+	if cancel != nil {
+		// Removing the display entry must also stop its streaming API request.
+		// Otherwise a subsequent start event adds another collector for this ID.
+		cancel()
+	}
 }
 
 func (s *stats) isKnownContainer(cid string) (int, bool) {
 	for i, c := range s.cs {
-		if c.Container == cid {
+		if c.GetStatistics().Container == cid {
 			return i, true
 		}
 	}
@@ -72,7 +84,9 @@ func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, strea
 		// if error happens, and we get nothing of stats, release wait group whatever
 		if !getFirst {
 			getFirst = true
-			waitFirst.Done()
+			if waitFirst != nil {
+				waitFirst.Done()
+			}
 		}
 	}()
 
@@ -96,7 +110,11 @@ func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, strea
 			var v container.StatsResponse
 			if err := dec.Decode(&v); err != nil {
 				dec = json.NewDecoder(io.MultiReader(dec.Buffered(), response.Body))
-				u <- err
+				select {
+				case u <- err:
+				case <-ctx.Done():
+					return
+				}
 				if err == io.EOF {
 					break
 				}
@@ -140,7 +158,11 @@ func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, strea
 					PidsCurrent:      v.PidsStats.Current,
 				})
 			}
-			u <- nil
+			select {
+			case u <- nil:
+			case <-ctx.Done():
+				return
+			}
 			if !streamStats {
 				return
 			}
@@ -155,7 +177,9 @@ func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, strea
 			// if this is the first stat you get, release WaitGroup
 			if !getFirst {
 				getFirst = true
-				waitFirst.Done()
+				if waitFirst != nil {
+					waitFirst.Done()
+				}
 			}
 		case err := <-u:
 			s.SetError(err)
@@ -168,7 +192,9 @@ func collect(ctx context.Context, s *Stats, cli client.ContainerAPIClient, strea
 			// if this is the first stat you get, release WaitGroup
 			if !getFirst {
 				getFirst = true
-				waitFirst.Done()
+				if waitFirst != nil {
+					waitFirst.Done()
+				}
 			}
 		case <-ctx.Done():
 			s.SetError(ctx.Err())

cli/command/container/stats_test.go

@@ -0,0 +1,97 @@
+package container
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/docker/cli/internal/test"
+	"github.com/moby/moby/api/types/container"
+	"github.com/moby/moby/api/types/events"
+	"github.com/moby/moby/client"
+	"gotest.tools/v3/assert"
+)
+
+// TestRunStatsRestart cancels each old request even when the daemon keeps its
+// stats stream open across container exits, as it does for stopped containers.
+func TestRunStatsRestart(t *testing.T) {
+	const id = "test-container"
+	actions := make(chan events.Action)
+	opened := make(chan struct{}, 10)
+	closed := make(chan struct{}, 10)
+	ctx, cancel := context.WithCancel(context.Background())
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		enc := json.NewEncoder(w)
+		switch strings.TrimPrefix(r.URL.Path, "/v1.54") {
+		case "/containers/json":
+			_ = enc.Encode([]container.Summary{{ID: id}})
+		case "/events":
+			w.WriteHeader(http.StatusOK)
+			w.(http.Flusher).Flush()
+			for {
+				select {
+				case action := <-actions:
+					if err := enc.Encode(events.Message{Type: events.ContainerEventType, Action: action, Actor: events.Actor{ID: id}}); err != nil {
+						return
+					}
+					w.(http.Flusher).Flush()
+				case <-r.Context().Done():
+					return
+				}
+			}
+		case "/containers/" + id + "/stats":
+			_ = enc.Encode(container.StatsResponse{ID: id, OSType: "linux"})
+			w.(http.Flusher).Flush()
+			opened <- struct{}{}
+			<-r.Context().Done()
+			closed <- struct{}{}
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	t.Cleanup(server.Close)
+	// Cancel before closing the server, including when an assertion fails.
+	t.Cleanup(cancel)
+	apiClient, err := client.New(client.WithHost(server.URL), client.WithAPIVersion("1.54"))
+	assert.NilError(t, err)
+	t.Cleanup(func() { assert.NilError(t, apiClient.Close()) })
+	done := make(chan error, 1)
+	go func() { done <- RunStats(ctx, test.NewFakeCli(apiClient), &StatsOptions{}) }()
+
+	wait := func(ch <-chan struct{}, description string) {
+		t.Helper()
+		select {
+		case <-ch:
+		case <-time.After(5 * time.Second):
+			t.Fatal(description)
+		}
+	}
+	send := func(action events.Action) {
+		t.Helper()
+		select {
+		case actions <- action:
+		case <-time.After(5 * time.Second):
+			t.Fatal("timed out sending container event")
+		}
+	}
+	wait(opened, "initial stats request was not opened")
+	for range 3 {
+		send(events.ActionDie)
+		wait(closed, "stats request survived the die event")
+		send(events.ActionStart)
+		wait(opened, "stats request was not reopened after restart")
+	}
+	send(events.ActionDestroy)
+	wait(closed, "stats request survived the destroy event")
+	cancel()
+	select {
+	case err := <-done:
+		assert.ErrorIs(t, err, context.Canceled)
+	case <-time.After(5 * time.Second):
+		t.Fatal("stats command did not stop after cancellation")
+	}
+}