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 }