gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/pkg/test/dockerutil/container.go (about)

     1  // Copyright 2020 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package dockerutil
    16  
    17  import (
    18  	"bytes"
    19  	"context"
    20  	"encoding/json"
    21  	"errors"
    22  	"fmt"
    23  	"io/ioutil"
    24  	"net"
    25  	"os"
    26  	"path"
    27  	"path/filepath"
    28  	"regexp"
    29  	"sort"
    30  	"strconv"
    31  	"strings"
    32  	"time"
    33  
    34  	"github.com/docker/docker/api/types"
    35  	"github.com/docker/docker/api/types/container"
    36  	"github.com/docker/docker/api/types/mount"
    37  	"github.com/docker/docker/api/types/network"
    38  	"github.com/docker/go-connections/nat"
    39  	"github.com/moby/moby/client"
    40  	"github.com/moby/moby/pkg/stdcopy"
    41  	"gvisor.dev/gvisor/pkg/sync"
    42  	"gvisor.dev/gvisor/pkg/test/testutil"
    43  )
    44  
    45  // Container represents a Docker Container allowing
    46  // user to configure and control as one would with the 'docker'
    47  // client. Container is backed by the official golang docker API.
    48  // See: https://pkg.go.dev/github.com/docker/docker.
    49  type Container struct {
    50  	Name     string
    51  	runtime  string
    52  	logger   testutil.Logger
    53  	client   *client.Client
    54  	id       string
    55  	mounts   []mount.Mount
    56  	links    []string
    57  	copyErr  error
    58  	cleanups []func()
    59  
    60  	// profile is the profiling hook associated with this container.
    61  	profile *profile
    62  }
    63  
    64  // RunOpts are options for running a container.
    65  type RunOpts struct {
    66  	// Image is the image relative to images/. This will be mangled
    67  	// appropriately, to ensure that only first-party images are used.
    68  	Image string
    69  
    70  	// Memory is the memory limit in bytes.
    71  	Memory int
    72  
    73  	// Cpus in which to allow execution. ("0", "1", "0-2").
    74  	CpusetCpus string
    75  
    76  	// Ports are the ports to be allocated.
    77  	Ports []int
    78  
    79  	// WorkDir sets the working directory.
    80  	WorkDir string
    81  
    82  	// ReadOnly sets the read-only flag.
    83  	ReadOnly bool
    84  
    85  	// Env are additional environment variables.
    86  	Env []string
    87  
    88  	// User is the user to use.
    89  	User string
    90  
    91  	// Optional argv to override the ENTRYPOINT specified in the image.
    92  	Entrypoint []string
    93  
    94  	// Privileged enables privileged mode.
    95  	Privileged bool
    96  
    97  	// Sets network mode for the container. See container.NetworkMode for types. Several options will
    98  	// not work w/ gVisor. For example, you can't set the "sandbox" network option for gVisor using
    99  	// this handle.
   100  	NetworkMode string
   101  
   102  	// CapAdd are the extra set of capabilities to add.
   103  	CapAdd []string
   104  
   105  	// CapDrop are the extra set of capabilities to drop.
   106  	CapDrop []string
   107  
   108  	// Mounts is the list of directories/files to be mounted inside the container.
   109  	Mounts []mount.Mount
   110  
   111  	// Links is the list of containers to be connected to the container.
   112  	Links []string
   113  
   114  	// DeviceRequests are device requests on the container itself.
   115  	DeviceRequests []container.DeviceRequest
   116  
   117  	Devices []container.DeviceMapping
   118  }
   119  
   120  func makeContainer(ctx context.Context, logger testutil.Logger, runtime string) *Container {
   121  	// Slashes are not allowed in container names.
   122  	name := testutil.RandomID(logger.Name())
   123  	name = strings.ReplaceAll(name, "/", "-")
   124  	client, err := client.NewClientWithOpts(client.FromEnv)
   125  	if err != nil {
   126  		return nil
   127  	}
   128  	client.NegotiateAPIVersion(ctx)
   129  	return &Container{
   130  		logger:  logger,
   131  		Name:    name,
   132  		runtime: runtime,
   133  		client:  client,
   134  	}
   135  }
   136  
   137  // MakeContainer constructs a suitable Container object.
   138  //
   139  // The runtime used is determined by the runtime flag.
   140  //
   141  // Containers will check flags for profiling requests.
   142  func MakeContainer(ctx context.Context, logger testutil.Logger) *Container {
   143  	return makeContainer(ctx, logger, *runtime)
   144  }
   145  
   146  // MakeContainerWithRuntime is like MakeContainer, but allows for a runtime
   147  // to be specified by suffix.
   148  func MakeContainerWithRuntime(ctx context.Context, logger testutil.Logger, suffix string) *Container {
   149  	return makeContainer(ctx, logger, *runtime+suffix)
   150  }
   151  
   152  // MakeNativeContainer constructs a suitable Container object.
   153  //
   154  // The runtime used will be the system default.
   155  //
   156  // Native containers aren't profiled.
   157  func MakeNativeContainer(ctx context.Context, logger testutil.Logger) *Container {
   158  	unsandboxedRuntime := "runc"
   159  	if override, found := os.LookupEnv("UNSANDBOXED_RUNTIME"); found {
   160  		unsandboxedRuntime = override
   161  	}
   162  	return makeContainer(ctx, logger, unsandboxedRuntime)
   163  }
   164  
   165  // Spawn is analogous to 'docker run -d'.
   166  func (c *Container) Spawn(ctx context.Context, r RunOpts, args ...string) error {
   167  	if err := c.create(ctx, r.Image, c.config(r, args), c.hostConfig(r), nil); err != nil {
   168  		return err
   169  	}
   170  	return c.Start(ctx)
   171  }
   172  
   173  // SpawnProcess is analogous to 'docker run -it'. It returns a process
   174  // which represents the root process.
   175  func (c *Container) SpawnProcess(ctx context.Context, r RunOpts, args ...string) (Process, error) {
   176  	config, hostconf, netconf := c.ConfigsFrom(r, args...)
   177  	config.Tty = true
   178  	config.OpenStdin = true
   179  
   180  	if err := c.CreateFrom(ctx, r.Image, config, hostconf, netconf); err != nil {
   181  		return Process{}, err
   182  	}
   183  
   184  	// Open a connection to the container for parsing logs and for TTY.
   185  	stream, err := c.client.ContainerAttach(ctx, c.id,
   186  		types.ContainerAttachOptions{
   187  			Stream: true,
   188  			Stdin:  true,
   189  			Stdout: true,
   190  			Stderr: true,
   191  		})
   192  	if err != nil {
   193  		return Process{}, fmt.Errorf("connect failed container id %s: %v", c.id, err)
   194  	}
   195  
   196  	c.cleanups = append(c.cleanups, func() { stream.Close() })
   197  
   198  	if err := c.Start(ctx); err != nil {
   199  		return Process{}, err
   200  	}
   201  
   202  	return Process{container: c, conn: stream}, nil
   203  }
   204  
   205  // Run is analogous to 'docker run'.
   206  func (c *Container) Run(ctx context.Context, r RunOpts, args ...string) (string, error) {
   207  	if err := c.create(ctx, r.Image, c.config(r, args), c.hostConfig(r), nil); err != nil {
   208  		return "", err
   209  	}
   210  
   211  	if err := c.Start(ctx); err != nil {
   212  		logs, _ := c.Logs(ctx)
   213  		return logs, err
   214  	}
   215  
   216  	if err := c.Wait(ctx); err != nil {
   217  		logs, _ := c.Logs(ctx)
   218  		return logs, err
   219  	}
   220  
   221  	return c.Logs(ctx)
   222  }
   223  
   224  // ConfigsFrom returns container configs from RunOpts and args. The caller should call 'CreateFrom'
   225  // and Start.
   226  func (c *Container) ConfigsFrom(r RunOpts, args ...string) (*container.Config, *container.HostConfig, *network.NetworkingConfig) {
   227  	return c.config(r, args), c.hostConfig(r), &network.NetworkingConfig{}
   228  }
   229  
   230  // MakeLink formats a link to add to a RunOpts.
   231  func (c *Container) MakeLink(target string) string {
   232  	return fmt.Sprintf("%s:%s", c.Name, target)
   233  }
   234  
   235  // CreateFrom creates a container from the given configs.
   236  func (c *Container) CreateFrom(ctx context.Context, profileImage string, conf *container.Config, hostconf *container.HostConfig, netconf *network.NetworkingConfig) error {
   237  	return c.create(ctx, profileImage, conf, hostconf, netconf)
   238  }
   239  
   240  // Create is analogous to 'docker create'.
   241  func (c *Container) Create(ctx context.Context, r RunOpts, args ...string) error {
   242  	return c.create(ctx, r.Image, c.config(r, args), c.hostConfig(r), nil)
   243  }
   244  
   245  func (c *Container) create(ctx context.Context, profileImage string, conf *container.Config, hostconf *container.HostConfig, netconf *network.NetworkingConfig) error {
   246  	if c.runtime != "" && c.runtime != "runc" {
   247  		// Use the image name as provided here; which normally represents the
   248  		// unmodified "basic/alpine" image name. This should be easy to grok.
   249  		c.profileInit(profileImage)
   250  	}
   251  	if hostconf == nil || hostconf.LogConfig.Type == "" {
   252  		// If logging config is not explicitly specified, set it to a fairly
   253  		// generous default. This makes large volumes of log more reliably
   254  		// captured, which is useful for profiling metrics.
   255  		if hostconf == nil {
   256  			hostconf = &container.HostConfig{}
   257  		}
   258  		hostconf.LogConfig.Type = "local"
   259  		hostconf.LogConfig.Config = map[string]string{
   260  			"mode":     "blocking",
   261  			"max-size": "1g",
   262  			"max-file": "10",
   263  			"compress": "false",
   264  		}
   265  	}
   266  	cont, err := c.client.ContainerCreate(ctx, conf, hostconf, nil, nil, c.Name)
   267  	if err != nil {
   268  		return err
   269  	}
   270  	c.id = cont.ID
   271  	return nil
   272  }
   273  
   274  func (c *Container) config(r RunOpts, args []string) *container.Config {
   275  	ports := nat.PortSet{}
   276  	for _, p := range r.Ports {
   277  		port := nat.Port(fmt.Sprintf("%d", p))
   278  		ports[port] = struct{}{}
   279  	}
   280  	env := append(r.Env, fmt.Sprintf("RUNSC_TEST_NAME=%s", c.Name))
   281  
   282  	return &container.Config{
   283  		Image:        testutil.ImageByName(r.Image),
   284  		Cmd:          args,
   285  		Entrypoint:   r.Entrypoint,
   286  		ExposedPorts: ports,
   287  		Env:          env,
   288  		WorkingDir:   r.WorkDir,
   289  		User:         r.User,
   290  	}
   291  }
   292  
   293  func (c *Container) hostConfig(r RunOpts) *container.HostConfig {
   294  	c.mounts = append(c.mounts, r.Mounts...)
   295  
   296  	return &container.HostConfig{
   297  		Runtime:         c.runtime,
   298  		Mounts:          c.mounts,
   299  		PublishAllPorts: true,
   300  		Links:           r.Links,
   301  		CapAdd:          r.CapAdd,
   302  		CapDrop:         r.CapDrop,
   303  		Privileged:      r.Privileged,
   304  		ReadonlyRootfs:  r.ReadOnly,
   305  		NetworkMode:     container.NetworkMode(r.NetworkMode),
   306  		Resources: container.Resources{
   307  			Memory:         int64(r.Memory), // In bytes.
   308  			CpusetCpus:     r.CpusetCpus,
   309  			DeviceRequests: r.DeviceRequests,
   310  			Devices:        r.Devices,
   311  		},
   312  	}
   313  }
   314  
   315  // Start is analogous to 'docker start'.
   316  func (c *Container) Start(ctx context.Context) error {
   317  	if err := c.client.ContainerStart(ctx, c.id, types.ContainerStartOptions{}); err != nil {
   318  		return fmt.Errorf("ContainerStart failed: %v", err)
   319  	}
   320  
   321  	if c.profile != nil {
   322  		if err := c.profile.Start(c); err != nil {
   323  			c.logger.Logf("profile.Start failed: %v", err)
   324  		}
   325  	}
   326  
   327  	return nil
   328  }
   329  
   330  // Stop is analogous to 'docker stop'.
   331  func (c *Container) Stop(ctx context.Context) error {
   332  	return c.client.ContainerStop(ctx, c.id, container.StopOptions{})
   333  }
   334  
   335  // Pause is analogous to 'docker pause'.
   336  func (c *Container) Pause(ctx context.Context) error {
   337  	return c.client.ContainerPause(ctx, c.id)
   338  }
   339  
   340  // Unpause is analogous to 'docker unpause'.
   341  func (c *Container) Unpause(ctx context.Context) error {
   342  	return c.client.ContainerUnpause(ctx, c.id)
   343  }
   344  
   345  // Checkpoint is analogous to 'docker checkpoint'.
   346  func (c *Container) Checkpoint(ctx context.Context, name string) error {
   347  	return c.client.CheckpointCreate(ctx, c.Name, types.CheckpointCreateOptions{CheckpointID: name, Exit: true})
   348  }
   349  
   350  // Restore is analogous to 'docker start --checkpoint [name]'.
   351  func (c *Container) Restore(ctx context.Context, name string) error {
   352  	return c.client.ContainerStart(ctx, c.id, types.ContainerStartOptions{CheckpointID: name})
   353  }
   354  
   355  // CheckpointResume is analogous to 'docker checkpoint'.
   356  func (c *Container) CheckpointResume(ctx context.Context, name string) error {
   357  	return c.client.CheckpointCreate(ctx, c.Name, types.CheckpointCreateOptions{CheckpointID: name, Exit: false})
   358  }
   359  
   360  // Logs is analogous 'docker logs'.
   361  func (c *Container) Logs(ctx context.Context) (string, error) {
   362  	var out bytes.Buffer
   363  	err := c.logs(ctx, &out, &out)
   364  	return out.String(), err
   365  }
   366  
   367  func (c *Container) logs(ctx context.Context, stdout, stderr *bytes.Buffer) error {
   368  	opts := types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true}
   369  	writer, err := c.client.ContainerLogs(ctx, c.id, opts)
   370  	if err != nil {
   371  		return err
   372  	}
   373  	defer writer.Close()
   374  	_, err = stdcopy.StdCopy(stdout, stderr, writer)
   375  
   376  	return err
   377  }
   378  
   379  // ID returns the container id.
   380  func (c *Container) ID() string {
   381  	return c.id
   382  }
   383  
   384  // RootDirectory returns an educated guess about the container's root directory.
   385  func (c *Container) RootDirectory() (string, error) {
   386  	// The root directory of this container's runtime.
   387  	rootDir := fmt.Sprintf("/var/run/docker/runtime-%s/moby", c.runtime)
   388  	_, err := os.Stat(rootDir)
   389  	if err == nil {
   390  		return rootDir, nil
   391  	}
   392  	// In docker v20+, due to https://github.com/moby/moby/issues/42345 the
   393  	// rootDir seems to always be the following.
   394  	const defaultDir = "/var/run/docker/runtime-runc/moby"
   395  	_, derr := os.Stat(defaultDir)
   396  	if derr == nil {
   397  		return defaultDir, nil
   398  	}
   399  
   400  	return "", fmt.Errorf("cannot stat %q: %v or %q: %v", rootDir, err, defaultDir, derr)
   401  }
   402  
   403  // SandboxPid returns the container's pid.
   404  func (c *Container) SandboxPid(ctx context.Context) (int, error) {
   405  	resp, err := c.client.ContainerInspect(ctx, c.id)
   406  	if err != nil {
   407  		return -1, err
   408  	}
   409  	return resp.ContainerJSONBase.State.Pid, nil
   410  }
   411  
   412  // ErrNoIP indicates that no IP address is available.
   413  var ErrNoIP = errors.New("no IP available")
   414  
   415  // FindIP returns the IP address of the container.
   416  func (c *Container) FindIP(ctx context.Context, ipv6 bool) (net.IP, error) {
   417  	resp, err := c.client.ContainerInspect(ctx, c.id)
   418  	if err != nil {
   419  		return nil, err
   420  	}
   421  
   422  	var ip net.IP
   423  	if ipv6 {
   424  		ip = net.ParseIP(resp.NetworkSettings.DefaultNetworkSettings.GlobalIPv6Address)
   425  	} else {
   426  		ip = net.ParseIP(resp.NetworkSettings.DefaultNetworkSettings.IPAddress)
   427  	}
   428  	if ip == nil {
   429  		return net.IP{}, ErrNoIP
   430  	}
   431  	return ip, nil
   432  }
   433  
   434  // FindPort returns the host port that is mapped to 'sandboxPort'.
   435  func (c *Container) FindPort(ctx context.Context, sandboxPort int) (int, error) {
   436  	desc, err := c.client.ContainerInspect(ctx, c.id)
   437  	if err != nil {
   438  		return -1, fmt.Errorf("error retrieving port: %v", err)
   439  	}
   440  
   441  	format := fmt.Sprintf("%d/tcp", sandboxPort)
   442  	ports, ok := desc.NetworkSettings.Ports[nat.Port(format)]
   443  	if !ok {
   444  		return -1, fmt.Errorf("error retrieving port: %v", err)
   445  
   446  	}
   447  
   448  	port, err := strconv.Atoi(ports[0].HostPort)
   449  	if err != nil {
   450  		return -1, fmt.Errorf("error parsing port %q: %v", port, err)
   451  	}
   452  	return port, nil
   453  }
   454  
   455  // CopyFiles copies in and mounts the given files. They are always ReadOnly.
   456  func (c *Container) CopyFiles(opts *RunOpts, target string, sources ...string) {
   457  	dir, err := ioutil.TempDir("", c.Name)
   458  	if err != nil {
   459  		c.copyErr = fmt.Errorf("ioutil.TempDir failed: %v", err)
   460  		return
   461  	}
   462  	c.cleanups = append(c.cleanups, func() { os.RemoveAll(dir) })
   463  	if err := os.Chmod(dir, 0755); err != nil {
   464  		c.copyErr = fmt.Errorf("os.Chmod(%q, 0755) failed: %v", dir, err)
   465  		return
   466  	}
   467  	for _, name := range sources {
   468  		src := name
   469  		if !filepath.IsAbs(src) {
   470  			src, err = testutil.FindFile(name)
   471  			if err != nil {
   472  				c.copyErr = fmt.Errorf("testutil.FindFile(%q) failed: %w", name, err)
   473  				return
   474  			}
   475  		}
   476  		dst := path.Join(dir, path.Base(name))
   477  		if err := testutil.Copy(src, dst); err != nil {
   478  			c.copyErr = fmt.Errorf("testutil.Copy(%q, %q) failed: %v", src, dst, err)
   479  			return
   480  		}
   481  		c.logger.Logf("copy: %s -> %s", src, dst)
   482  	}
   483  	opts.Mounts = append(opts.Mounts, mount.Mount{
   484  		Type:     mount.TypeBind,
   485  		Source:   dir,
   486  		Target:   target,
   487  		ReadOnly: false,
   488  	})
   489  }
   490  
   491  // Stats returns a snapshot of container stats similar to `docker stats`.
   492  func (c *Container) Stats(ctx context.Context) (*types.StatsJSON, error) {
   493  	responseBody, err := c.client.ContainerStats(ctx, c.id, false /*stream*/)
   494  	if err != nil {
   495  		return nil, fmt.Errorf("ContainerStats failed: %v", err)
   496  	}
   497  	defer responseBody.Body.Close()
   498  	var v types.StatsJSON
   499  	if err := json.NewDecoder(responseBody.Body).Decode(&v); err != nil {
   500  		return nil, fmt.Errorf("failed to decode container stats: %v", err)
   501  	}
   502  	return &v, nil
   503  }
   504  
   505  // Status inspects the container returns its status.
   506  func (c *Container) Status(ctx context.Context) (types.ContainerState, error) {
   507  	resp, err := c.client.ContainerInspect(ctx, c.id)
   508  	if err != nil {
   509  		return types.ContainerState{}, err
   510  	}
   511  	return *resp.State, err
   512  }
   513  
   514  // Wait waits for the container to exit.
   515  func (c *Container) Wait(ctx context.Context) error {
   516  	defer c.stopProfiling()
   517  	statusChan, errChan := c.client.ContainerWait(ctx, c.id, container.WaitConditionNotRunning)
   518  	select {
   519  	case err := <-errChan:
   520  		return err
   521  	case res := <-statusChan:
   522  		if res.StatusCode != 0 {
   523  			var msg string
   524  			if res.Error != nil {
   525  				msg = res.Error.Message
   526  			}
   527  			return fmt.Errorf("container returned non-zero status: %d, msg: %q", res.StatusCode, msg)
   528  		}
   529  		return nil
   530  	}
   531  }
   532  
   533  // WaitTimeout waits for the container to exit with a timeout.
   534  func (c *Container) WaitTimeout(ctx context.Context, timeout time.Duration) error {
   535  	ctx, cancel := context.WithTimeout(ctx, timeout)
   536  	defer cancel()
   537  	statusChan, errChan := c.client.ContainerWait(ctx, c.id, container.WaitConditionNotRunning)
   538  	select {
   539  	case <-ctx.Done():
   540  		if ctx.Err() == context.DeadlineExceeded {
   541  			return fmt.Errorf("container %s timed out after %v seconds", c.Name, timeout.Seconds())
   542  		}
   543  		return nil
   544  	case err := <-errChan:
   545  		return err
   546  	case <-statusChan:
   547  		return nil
   548  	}
   549  }
   550  
   551  // WaitForOutput searches container logs for pattern and returns or timesout.
   552  func (c *Container) WaitForOutput(ctx context.Context, pattern string, timeout time.Duration) (string, error) {
   553  	matches, err := c.WaitForOutputSubmatch(ctx, pattern, timeout)
   554  	if err != nil {
   555  		return "", err
   556  	}
   557  	if len(matches) == 0 {
   558  		return "", fmt.Errorf("didn't find pattern %s logs", pattern)
   559  	}
   560  	return matches[0], nil
   561  }
   562  
   563  // WaitForOutputSubmatch searches container logs for the given
   564  // pattern or times out. It returns any regexp submatches as well.
   565  func (c *Container) WaitForOutputSubmatch(ctx context.Context, pattern string, timeout time.Duration) ([]string, error) {
   566  	ctx, cancel := context.WithTimeout(ctx, timeout)
   567  	defer cancel()
   568  	re := regexp.MustCompile(pattern)
   569  	for {
   570  		logs, err := c.Logs(ctx)
   571  		if err != nil {
   572  			return nil, fmt.Errorf("failed to get logs: %v logs: %s", err, logs)
   573  		}
   574  		if matches := re.FindStringSubmatch(logs); matches != nil {
   575  			return matches, nil
   576  		}
   577  		time.Sleep(50 * time.Millisecond)
   578  	}
   579  }
   580  
   581  // stopProfiling stops profiling.
   582  func (c *Container) stopProfiling() {
   583  	if c.profile != nil {
   584  		if err := c.profile.Stop(c); err != nil {
   585  			// This most likely means that the runtime for the container
   586  			// was too short to connect and actually get a profile.
   587  			c.logger.Logf("warning: profile.Stop failed: %v", err)
   588  		}
   589  	}
   590  }
   591  
   592  // Kill kills the container.
   593  func (c *Container) Kill(ctx context.Context) error {
   594  	defer c.stopProfiling()
   595  	return c.client.ContainerKill(ctx, c.id, "")
   596  }
   597  
   598  // Remove is analogous to 'docker rm'.
   599  func (c *Container) Remove(ctx context.Context) error {
   600  	// Remove the image.
   601  	remove := types.ContainerRemoveOptions{
   602  		RemoveVolumes: c.mounts != nil,
   603  		RemoveLinks:   c.links != nil,
   604  		Force:         true,
   605  	}
   606  	return c.client.ContainerRemove(ctx, c.Name, remove)
   607  }
   608  
   609  // CleanUp kills and deletes the container (best effort).
   610  func (c *Container) CleanUp(ctx context.Context) {
   611  	// Execute all cleanups. We execute cleanups here to close any
   612  	// open connections to the container before closing. Open connections
   613  	// can cause Kill and Remove to hang.
   614  	for _, c := range c.cleanups {
   615  		c()
   616  	}
   617  	c.cleanups = nil
   618  
   619  	// Kill the container.
   620  	if err := c.Kill(ctx); err != nil && !strings.Contains(err.Error(), "is not running") {
   621  		// Just log; can't do anything here.
   622  		c.logger.Logf("error killing container %q: %v", c.Name, err)
   623  	}
   624  
   625  	// Remove the image.
   626  	if err := c.Remove(ctx); err != nil {
   627  		c.logger.Logf("error removing container %q: %v", c.Name, err)
   628  	}
   629  
   630  	// Forget all mounts.
   631  	c.mounts = nil
   632  }
   633  
   634  // ContainerPool represents a pool of reusable containers.
   635  // Callers may request a container from the pool, and must release it back
   636  // when they are done with it.
   637  //
   638  // This is useful for large tests which can `exec` individual test cases
   639  // inside the same set of reusable containers, to avoid the cost of creating
   640  // and destroying containers for each test.
   641  //
   642  // It also supports reserving the whole pool ("exclusive"), which locks out
   643  // all other callers from getting any other container from the pool.
   644  // This is useful for tests where running in parallel may induce unexpected
   645  // errors which running serially would not cause. This allows the test to
   646  // first try to run in parallel, and then re-run failing tests exclusively
   647  // to make sure their failure is not due to parallel execution.
   648  type ContainerPool struct {
   649  	// numContainers is the total number of containers in the pool, whether
   650  	// reserved or not.
   651  	numContainers int
   652  
   653  	// containersCh is the main container queue channel.
   654  	// It is buffered with as many items as there are total containers in the
   655  	// pool (i.e. `numContainers`).
   656  	containersCh chan *Container
   657  
   658  	// exclusiveLockCh is a lock-like channel for exclusive locking.
   659  	// It is buffered with a single element seeded in it.
   660  	// Whichever goroutine claims this element now holds the exclusive lock.
   661  	exclusiveLockCh chan struct{}
   662  
   663  	// containersExclusiveCh is an alternative container queue channel.
   664  	// It is preferentially written to over containerCh, but is unbuffered.
   665  	// A goroutine may only listen on this channel if they hold the exclusive
   666  	// lock.
   667  	// This allows exclusive locking to get released containers preferentially
   668  	// over non-exclusive waiters, to avoid needless starvation.
   669  	containersExclusiveCh chan *Container
   670  
   671  	// shutdownCh is always empty, and is closed when the pool is shutting down.
   672  	shutdownCh chan struct{}
   673  
   674  	// mu protects the fields below.
   675  	mu sync.Mutex
   676  
   677  	// statuses maps container name to their current status.
   678  	statuses map[*Container]containerPoolStatus
   679  
   680  	// firstReservation is the time of the first container reservation event.
   681  	firstReservation time.Time
   682  
   683  	// productive is the duration containers have cumulatively spent in reserved
   684  	// state.
   685  	productive time.Duration
   686  }
   687  
   688  // containerPoolStatus represents the status of a container in the pool.
   689  type containerPoolStatus struct {
   690  	when      time.Time
   691  	state     containerPoolState
   692  	userLabel string
   693  }
   694  
   695  // containerPoolState is the state of a container in the pool.
   696  type containerPoolState int
   697  
   698  // Set of containerPoolState used by ContainerPool.
   699  const (
   700  	stateNeverUsed containerPoolState = iota
   701  	stateIdle
   702  	stateReserved
   703  	stateReservedExclusive
   704  	stateShutdown
   705  	stateHeldForExclusive
   706  )
   707  
   708  // IsProductive returns true if the container is in a state where it is
   709  // actively being used.
   710  func (cps containerPoolState) IsProductive() bool {
   711  	return cps == stateReserved || cps == stateReservedExclusive
   712  }
   713  
   714  // String returns a human-friendly representation of the container state in
   715  // the pool.
   716  func (cps containerPoolState) String() string {
   717  	switch cps {
   718  	case stateNeverUsed:
   719  		return "never used"
   720  	case stateIdle:
   721  		return "idle"
   722  	case stateReserved:
   723  		return "reserved"
   724  	case stateReservedExclusive:
   725  		return "reserved in exclusive mode"
   726  	case stateHeldForExclusive:
   727  		return "held back to allow another container to be exclusively-reserved"
   728  	case stateShutdown:
   729  		return "shutdown"
   730  	default:
   731  		return fmt.Sprintf("unknownstate(%d)", cps)
   732  	}
   733  }
   734  
   735  // NewContainerPool returns a new ContainerPool holding the given set of
   736  // containers.
   737  func NewContainerPool(containers []*Container) *ContainerPool {
   738  	if len(containers) == 0 {
   739  		panic("cannot create an empty pool")
   740  	}
   741  	containersCh := make(chan *Container, len(containers))
   742  	for _, c := range containers {
   743  		containersCh <- c
   744  	}
   745  	exclusiveCh := make(chan struct{}, 1)
   746  	exclusiveCh <- struct{}{}
   747  	statuses := make(map[*Container]containerPoolStatus)
   748  	poolStarted := time.Now()
   749  	for _, c := range containers {
   750  		statuses[c] = containerPoolStatus{
   751  			when:  poolStarted,
   752  			state: stateNeverUsed,
   753  		}
   754  	}
   755  	return &ContainerPool{
   756  		containersCh:          containersCh,
   757  		exclusiveLockCh:       exclusiveCh,
   758  		containersExclusiveCh: make(chan *Container),
   759  		shutdownCh:            make(chan struct{}),
   760  		numContainers:         len(containers),
   761  		statuses:              statuses,
   762  	}
   763  }
   764  
   765  // releaseFn returns a function to release a container back to the pool.
   766  func (cp *ContainerPool) releaseFn(c *Container) func() {
   767  	return func() {
   768  		cp.setContainerState(c, stateIdle)
   769  		// Preferentially release to the exclusive channel.
   770  		select {
   771  		case cp.containersExclusiveCh <- c:
   772  			return
   773  		default:
   774  			// Otherwise, release to either channel.
   775  			select {
   776  			case cp.containersExclusiveCh <- c:
   777  			case cp.containersCh <- c:
   778  			}
   779  		}
   780  	}
   781  }
   782  
   783  // Get returns a free container and a function to release it back to the pool.
   784  func (cp *ContainerPool) Get(ctx context.Context) (*Container, func(), error) {
   785  	select {
   786  	case c := <-cp.containersCh:
   787  		cp.setContainerState(c, stateReserved)
   788  		return c, cp.releaseFn(c), nil
   789  	case <-cp.shutdownCh:
   790  		return nil, func() {}, errors.New("pool's closed")
   791  	case <-ctx.Done():
   792  		return nil, func() {}, ctx.Err()
   793  	}
   794  }
   795  
   796  // GetExclusive ensures all pooled containers are in the pool, reserves all
   797  // of them, and returns one of them, along with a function to release them all
   798  // back to the pool.
   799  func (cp *ContainerPool) GetExclusive(ctx context.Context) (*Container, func(), error) {
   800  	select {
   801  	case <-cp.exclusiveLockCh:
   802  		// Proceed.
   803  	case <-cp.shutdownCh:
   804  		return nil, func() {}, errors.New("pool's closed")
   805  	case <-ctx.Done():
   806  		return nil, func() {}, ctx.Err()
   807  	}
   808  	var reserved *Container
   809  	releaseFuncs := make([]func(), 0, cp.numContainers)
   810  	releaseAll := func() {
   811  		for i := len(releaseFuncs) - 1; i >= 0; i-- {
   812  			releaseFuncs[i]()
   813  		}
   814  		cp.exclusiveLockCh <- struct{}{}
   815  	}
   816  	for i := 0; i < cp.numContainers; i++ {
   817  		var got *Container
   818  		select {
   819  		case c := <-cp.containersExclusiveCh:
   820  			got = c
   821  		case c := <-cp.containersCh:
   822  			got = c
   823  		case <-cp.shutdownCh:
   824  			releaseAll()
   825  			return nil, func() {}, errors.New("pool's closed")
   826  		case <-ctx.Done():
   827  			releaseAll()
   828  			return nil, func() {}, ctx.Err()
   829  		}
   830  		cp.setContainerState(got, stateHeldForExclusive)
   831  		releaseFuncs = append(releaseFuncs, cp.releaseFn(got))
   832  		if reserved == nil {
   833  			reserved = got
   834  		}
   835  	}
   836  	cp.setContainerState(reserved, stateReservedExclusive)
   837  	return reserved, releaseAll, nil
   838  }
   839  
   840  // CleanUp waits for all containers to be back into the pool, and cleans up
   841  // each container as soon as it gets back in the pool.
   842  func (cp *ContainerPool) CleanUp(ctx context.Context) {
   843  	close(cp.shutdownCh)
   844  	for i := 0; i < cp.numContainers; i++ {
   845  		c := <-cp.containersCh
   846  		cp.setContainerState(c, stateShutdown)
   847  		c.CleanUp(ctx)
   848  	}
   849  }
   850  
   851  // setContainerState sets the state of the given container.
   852  func (cp *ContainerPool) setContainerState(c *Container, state containerPoolState) {
   853  	isProductive := state.IsProductive()
   854  	when := time.Now()
   855  	cp.mu.Lock()
   856  	defer cp.mu.Unlock()
   857  	if isProductive && cp.firstReservation.IsZero() {
   858  		cp.firstReservation = when
   859  	}
   860  	status := cp.statuses[c]
   861  	wasProductive := status.state.IsProductive()
   862  	if wasProductive && !isProductive {
   863  		cp.productive += when.Sub(status.when)
   864  	}
   865  	status.when = when
   866  	status.state = state
   867  	status.userLabel = "" // Clear any existing user label.
   868  	cp.statuses[c] = status
   869  }
   870  
   871  // Utilization returns the utilization of the pool.
   872  // This is the ratio of the cumulative duration containers have been in a
   873  // productive state (reserved) divided by the maximum potential productive
   874  // container-duration since the first container was reserved.
   875  func (cp *ContainerPool) Utilization() float64 {
   876  	cp.mu.Lock()
   877  	defer cp.mu.Unlock()
   878  	return cp.getUtilizationLocked(time.Now())
   879  }
   880  
   881  // getUtilizationLocked actually computes the utilization of the pool.
   882  // Preconditions: cp.mu is held.
   883  func (cp *ContainerPool) getUtilizationLocked(now time.Time) float64 {
   884  	if cp.firstReservation.IsZero() {
   885  		return 0
   886  	}
   887  	maxUtilization := now.Sub(cp.firstReservation) * time.Duration(len(cp.statuses))
   888  	actualUtilization := cp.productive
   889  	for container := range cp.statuses {
   890  		if status := cp.statuses[container]; status.state.IsProductive() {
   891  			actualUtilization += now.Sub(status.when)
   892  		}
   893  	}
   894  	return float64(actualUtilization.Nanoseconds()) / float64(maxUtilization.Nanoseconds())
   895  }
   896  
   897  // SetContainerLabel sets the label for a container.
   898  // This is printed in `cp.String`, and wiped every time the container changes
   899  // state. It is useful for the ContainerPool user to track what each container
   900  // is doing.
   901  func (cp *ContainerPool) SetContainerLabel(c *Container, label string) {
   902  	cp.mu.Lock()
   903  	defer cp.mu.Unlock()
   904  	status := cp.statuses[c]
   905  	status.userLabel = label
   906  	cp.statuses[c] = status
   907  }
   908  
   909  // String returns a string representation of the pool.
   910  // It includes the state of each container, and their user-set label.
   911  func (cp *ContainerPool) String() string {
   912  	cp.mu.Lock()
   913  	defer cp.mu.Unlock()
   914  	if cp.firstReservation.IsZero() {
   915  		return "ContainerPool[never used]"
   916  	}
   917  	now := time.Now()
   918  	containers := make([]*Container, 0, len(cp.statuses))
   919  	for c := range cp.statuses {
   920  		containers = append(containers, c)
   921  	}
   922  	sort.Slice(containers, func(i, j int) bool { return containers[i].Name < containers[j].Name })
   923  	containersInUse := 0
   924  	for _, container := range containers {
   925  		if status := cp.statuses[container]; status.state.IsProductive() {
   926  			containersInUse++
   927  		}
   928  	}
   929  	utilizationPct := 100.0 * cp.getUtilizationLocked(now)
   930  	var sb strings.Builder
   931  	sb.WriteString(fmt.Sprintf("ContainerPool[%d/%d containers in use, utilization=%.1f%%]: ", containersInUse, len(containers), utilizationPct))
   932  	for i, container := range containers {
   933  		if i > 0 {
   934  			sb.WriteString(", ")
   935  		}
   936  		status := cp.statuses[container]
   937  		sb.WriteString(container.Name)
   938  		sb.WriteString("[")
   939  		sb.WriteString(status.state.String())
   940  		sb.WriteString("]")
   941  		if status.userLabel != "" {
   942  			sb.WriteString(": ")
   943  			sb.WriteString(status.userLabel)
   944  		}
   945  	}
   946  	return sb.String()
   947  }