From 08ed01b285870c50a7a4bae89210aa176cf66a20 Mon Sep 17 00:00:00 2001 From: Etienne Perot Date: Tue, 26 Mar 2024 18:49:45 -0700 Subject: [PATCH] `dockerutil`: Implement `ContainerPool`, a pool of reusable test containers. Callers may request a container from the pool, and must release it back when they are done with it. This is useful for large tests which can `exec` individual test cases inside the same set of reusable containers, to avoid the cost of creating and destroying containers for each test. It also supports reserving the whole pool ("exclusive"), which locks out all other callers from getting any other container from the pool. This is useful for tests where running in parallel may induce unexpected errors which running serially would not cause. This allows the test to first try to run in parallel, and then re-run failing tests exclusively to make sure their failure is not due to parallel execution. I plan to use this for CUDA sample tests. PiperOrigin-RevId: 619377512 --- pkg/test/dockerutil/container.go | 143 +++++++++++++++++++++++++++++++ 1 file changed, 143 insertions(+) diff --git a/pkg/test/dockerutil/container.go b/pkg/test/dockerutil/container.go index c31a357af..077a73057 100644 --- a/pkg/test/dockerutil/container.go +++ b/pkg/test/dockerutil/container.go @@ -606,3 +606,146 @@ func (c *Container) CleanUp(ctx context.Context) { // Forget all mounts. c.mounts = nil } + +// ContainerPool represents a pool of reusable containers. +// Callers may request a container from the pool, and must release it back +// when they are done with it. +// +// This is useful for large tests which can `exec` individual test cases +// inside the same set of reusable containers, to avoid the cost of creating +// and destroying containers for each test. +// +// It also supports reserving the whole pool ("exclusive"), which locks out +// all other callers from getting any other container from the pool. +// This is useful for tests where running in parallel may induce unexpected +// errors which running serially would not cause. This allows the test to +// first try to run in parallel, and then re-run failing tests exclusively +// to make sure their failure is not due to parallel execution. +type ContainerPool struct { + // numContainers is the total number of containers in the pool, whether + // reserved or not. + numContainers int + + // containersCh is the main container queue channel. + // It is buffered with as many items as there are total containers in the + // pool (i.e. `numContainers`). + containersCh chan *Container + + // exclusiveLockCh is a lock-like channel for exclusive locking. + // It is buffered with a single element seeded in it. + // Whichever goroutine claims this element now holds the exclusive lock. + exclusiveLockCh chan struct{} + + // containersExclusiveCh is an alternative container queue channel. + // It is preferentially written to over containerCh, but is unbuffered. + // A goroutine may only listen on this channel if they hold the exclusive + // lock. + // This allows exclusive locking to get released containers preferentially + // over non-exclusive waiters, to avoid needless starvation. + containersExclusiveCh chan *Container + + // shutdownCh is always empty, and is closed when the pool is shutting down. + shutdownCh chan struct{} +} + +// NewContainerPool returns a new ContainerPool holding the given set of +// containers. +func NewContainerPool(containers []*Container) *ContainerPool { + if len(containers) == 0 { + panic("cannot create an empty pool") + } + containersCh := make(chan *Container, len(containers)) + for _, c := range containers { + containersCh <- c + } + exclusiveCh := make(chan struct{}, 1) + exclusiveCh <- struct{}{} + return &ContainerPool{ + containersCh: containersCh, + exclusiveLockCh: exclusiveCh, + containersExclusiveCh: make(chan *Container), + shutdownCh: make(chan struct{}), + numContainers: len(containers), + } +} + +// releaseFn returns a function to release a container back to the pool. +func (cp *ContainerPool) releaseFn(c *Container) func() { + return func() { + // Preferentially release to the exclusive channel. + select { + case cp.containersExclusiveCh <- c: + return + default: + // Otherwise, release to either channel. + select { + case cp.containersExclusiveCh <- c: + case cp.containersCh <- c: + } + } + } +} + +// Get returns a free container and a function to release it back to the pool. +func (cp *ContainerPool) Get(ctx context.Context) (*Container, func(), error) { + select { + case c := <-cp.containersCh: + return c, cp.releaseFn(c), nil + case <-cp.shutdownCh: + return nil, func() {}, errors.New("pool's closed") + case <-ctx.Done(): + return nil, func() {}, ctx.Err() + } +} + +// GetExclusive ensures all pooled containers are in the pool, reserves all +// of them, and returns one of them, along with a function to release them all +// back to the pool. +func (cp *ContainerPool) GetExclusive(ctx context.Context) (*Container, func(), error) { + select { + case <-cp.exclusiveLockCh: + // Proceed. + case <-cp.shutdownCh: + return nil, func() {}, errors.New("pool's closed") + case <-ctx.Done(): + return nil, func() {}, ctx.Err() + } + var reserved *Container + releaseFuncs := make([]func(), 0, cp.numContainers) + releaseAll := func() { + for i := len(releaseFuncs) - 1; i >= 0; i-- { + releaseFuncs[i]() + } + cp.exclusiveLockCh <- struct{}{} + } + for i := 0; i < cp.numContainers; i++ { + var got *Container + select { + case c := <-cp.containersExclusiveCh: + got = c + case c := <-cp.containersCh: + got = c + case <-cp.shutdownCh: + releaseAll() + return nil, func() {}, errors.New("pool's closed") + case <-ctx.Done(): + releaseAll() + return nil, func() {}, ctx.Err() + } + releaseFuncs = append(releaseFuncs, cp.releaseFn(got)) + if reserved == nil { + reserved = got + } + } + return reserved, releaseAll, nil +} + +// CleanUp waits for all containers to be back into the pool, and cleans up +// each container as soon as it gets back in the pool. +func (cp *ContainerPool) CleanUp(ctx context.Context) { + close(cp.shutdownCh) + for i := 0; i < cp.numContainers; i++ { + c := <-cp.containersCh + c.CleanUp(ctx) + } +}