diff --git a/pool.go b/pool.go index a8a789b7..5901ccef 100644 --- a/pool.go +++ b/pool.go @@ -160,7 +160,7 @@ func NewPoolT(t TestingTB, endpoint string, opts ...PoolOption) *Pool { } t.Cleanup(func() { - if err := pool.Close(t.Context()); err != nil { + if err := pool.Close(context.WithoutCancel(t.Context())); err != nil { t.Logf("pool.Close() error: %v", err) } }) @@ -186,18 +186,31 @@ func (p *Pool) Close(ctx context.Context) error { return cleanupErr } +// CloseT cleans up all tracked containers and networks, then closes the Pool's +// Docker client. It calls t.Fatalf on error. +func (p *Pool) CloseT(t TestingTB) { + t.Helper() + if err := p.Close(context.WithoutCancel(t.Context())); err != nil { + t.Fatalf("Pool.CloseT failed: %v", err) + } +} + // cleanup removes all containers and networks tracked by this pool. // Containers are removed first, then networks. Errors during cleanup // do not stop the cleanup process. The first error encountered is returned. +// +// Unlike Resource.Close, cleanup force-removes containers regardless of +// reference count. This ensures that Pool.Close fully cleans up even when +// reused containers still have outstanding references. func (p *Pool) cleanup(ctx context.Context) error { cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 60*time.Second) defer cancel() var firstErr error - // Remove containers first (they may be connected to tracked networks) + // Force-remove containers (bypass ref counting — the pool is shutting down). for _, resource := range p.trackedResources() { - if err := resource.Close(cleanupCtx); err != nil && firstErr == nil { + if err := p.forceRemoveContainer(cleanupCtx, resource.Container.ID); err != nil && firstErr == nil { firstErr = err } p.untrackResource(resource.Container.ID) @@ -218,13 +231,26 @@ func (p *Pool) cleanup(ctx context.Context) error { return firstErr } +// forceRemoveContainer stops and removes a container, ignoring not-found errors. +func (p *Pool) forceRemoveContainer(ctx context.Context, containerID string) error { + _, _ = p.client.ContainerStop(ctx, containerID, mobyclient.ContainerStopOptions{}) + _, err := p.client.ContainerRemove(ctx, containerID, mobyclient.ContainerRemoveOptions{ + RemoveVolumes: true, + Force: true, + }) + if err != nil && !errdefs.IsNotFound(err) { + return err + } + return nil +} + // Run starts a container with the given repository and options. // // By default, containers are reused based on repository:tag to speed up tests. -// This means that two calls with the same repository and tag will return the -// same container, even if other options (env, cmd, etc.) differ. To ensure -// a fresh container, use WithoutReuse(). To control reuse with a custom key -// that accounts for your configuration, use WithReuseID(). +// Reused containers are reference-counted: the Docker container is only removed +// when the last caller closes its reference. To ensure a fresh container, use +// WithoutReuse(). To control reuse with a custom key that accounts for your +// configuration, use WithReuseID(). // // Example: // @@ -302,15 +328,16 @@ func computeReuseID(repository string, cfg *runConfig) string { return fmt.Sprintf("%s:%s", repository, cfg.tag) } -// checkForExisting looks up an existing container in the registry. -// The returned Resource is a shallow clone of the canonical registry entry -// with the pool pointer updated to the caller's pool. This is safe because -// Container (a value type) is copied, and reuseID is an immutable string. +// checkForExisting looks up an existing container in the registry and +// increments its reference count. The returned Resource is a shallow clone +// of the canonical registry entry with the pool pointer updated to the +// caller's pool. This is safe because Container (a value type) is copied, +// and reuseID is an immutable string. func checkForExisting(p *Pool, reuseID string) *Resource { if reuseID == "" { return nil } - if existing, ok := getWithScope(p.reuseScope, reuseID); ok { + if existing, ok := acquireWithScope(p.reuseScope, reuseID); ok { cloned := *existing cloned.pool = p return &cloned diff --git a/pool_test.go b/pool_test.go index e9520ca7..3fa797e2 100644 --- a/pool_test.go +++ b/pool_test.go @@ -243,6 +243,112 @@ func TestCustomClientPoolHasIsolatedScope(t *testing.T) { } } +func TestPoolCloseT(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test in short mode") + } + + ResetRegistry() + t.Cleanup(func() { + ResetRegistry() + }) + + pool := NewPoolT(t, "") + resource := pool.RunT(t, "alpine", + WithTag("latest"), + WithCmd([]string{"sleep", "300"}), + WithoutReuse(), + ) + containerID := resource.ID() + + pool.CloseT(t) + + // After CloseT, the container should be removed + newPool, err := NewPool(t.Context(), "") + if err != nil { + t.Fatalf("NewPool() error = %v", err) + } + t.Cleanup(func() { newPool.Close(t.Context()) }) + + _, inspectErr := newPool.client.ContainerInspect(t.Context(), containerID, mobyclient.ContainerInspectOptions{}) + if !errdefs.IsNotFound(inspectErr) { + t.Fatalf("ContainerInspect() error = %v, want not found", inspectErr) + } +} + +func TestCheckForExistingIncrementsRefCount(t *testing.T) { + ResetRegistry() + + resource := &Resource{Container: container.InspectResponse{ID: "ref-count-check"}} + // Register: refs=1 + registerWithScope("scope", "id", resource) + + pool := &Pool{reuseScope: "scope"} + // checkForExisting calls acquireWithScope: refs=2 + got := checkForExisting(pool, "id") + if got == nil || got.ID() != resource.ID() { + t.Fatalf("checkForExisting returned %v, want resource %q", got, resource.ID()) + } + + // First release: refs=1, not last + if releaseWithScope("scope", "id") { + t.Fatal("first releaseWithScope was last, want false") + } + + // Second release: refs=0, last + if !releaseWithScope("scope", "id") { + t.Fatal("second releaseWithScope was not last, want true") + } +} + +func TestPoolCleanupForceRemovesReusedContainers(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test in short mode") + } + + ResetRegistry() + t.Cleanup(func() { + ResetRegistry() + }) + + pool := NewPoolT(t, "") + + // Run first reused container: refs=1 + r1 := pool.RunT(t, "alpine", + WithTag("latest"), + WithCmd([]string{"sleep", "300"}), + ) + + // Run second reused container (same repo:tag): refs=2 + r2 := pool.RunT(t, "alpine", + WithTag("latest"), + WithCmd([]string{"sleep", "300"}), + ) + + if r1.ID() != r2.ID() { + t.Fatalf("expected same container ID, got %s and %s", r1.ID(), r2.ID()) + } + + containerID := r1.ID() + reuseScope := pool.reuseScope + + // cleanup() should force-remove everything regardless of ref count + if err := pool.cleanup(t.Context()); err != nil { + t.Fatalf("cleanup() error = %v", err) + } + + // Container should be gone from Docker + _, err := pool.client.ContainerInspect(t.Context(), containerID, mobyclient.ContainerInspectOptions{}) + if !errdefs.IsNotFound(err) { + t.Fatalf("ContainerInspect() error = %v, want not found", err) + } + + // Registry entry should be cleared + if _, ok := getWithScope(reuseScope, "alpine:latest"); ok { + t.Fatal("registry entry still present after cleanup, want it removed") + } +} + func TestPoolCleanupRemovesWithoutReuse(t *testing.T) { if testing.Short() { t.Skip("Skipping integration test in short mode") diff --git a/registry.go b/registry.go index 3b768199..774db09d 100644 --- a/registry.go +++ b/registry.go @@ -6,6 +6,7 @@ package dockertest import ( "fmt" "sync" + "sync/atomic" "github.com/moby/moby/api/types/container" ) @@ -29,6 +30,15 @@ type registryKey struct { reuseID string } +// registryEntry wraps a Resource with an atomic reference count. +// When multiple callers share a reused container, the ref count tracks +// how many are still using it. The container is only removed from Docker +// when the last reference is released. +type registryEntry struct { + resource *Resource + refs atomic.Int32 +} + const defaultRegistryScope = "default" // globalRegistry is the package-level container registry using sync.Map for thread-safety. @@ -87,16 +97,59 @@ func GetAll() []*Resource { func registerWithScope(scope, reuseID string, r *Resource) (*Resource, bool) { key := registryKey{scope: scope, reuseID: reuseID} - actual, loaded := globalRegistry.LoadOrStore(key, r) - resource, ok := actual.(*Resource) + entry := ®istryEntry{resource: r} + entry.refs.Store(1) + actual, loaded := globalRegistry.LoadOrStore(key, entry) + existing, ok := actual.(*registryEntry) if !ok { return r, loaded } - return resource, loaded + if loaded { + existing.refs.Add(1) + return existing.resource, true + } + return existing.resource, false } -func unregisterWithScope(scope, reuseID string) { - globalRegistry.Delete(registryKey{scope: scope, reuseID: reuseID}) +// acquireWithScope looks up an existing entry and increments its ref count. +// Returns the resource and true if found, nil and false otherwise. +func acquireWithScope(scope, reuseID string) (*Resource, bool) { + key := registryKey{scope: scope, reuseID: reuseID} + val, ok := globalRegistry.Load(key) + if !ok { + return nil, false + } + entry, ok := val.(*registryEntry) + if !ok { + return nil, false + } + entry.refs.Add(1) + // Verify the entry is still in the map (not deleted and replaced between Load and Add). + if current, ok := globalRegistry.Load(key); !ok || current != val { + entry.refs.Add(-1) + return nil, false + } + return entry.resource, true +} + +// releaseWithScope decrements the reference count for the given reuseID. +// Returns true if this was the last reference (caller should remove the container). +func releaseWithScope(scope, reuseID string) bool { + key := registryKey{scope: scope, reuseID: reuseID} + val, ok := globalRegistry.Load(key) + if !ok { + return true // Not found, treat as last reference + } + entry, ok := val.(*registryEntry) + if !ok { + globalRegistry.Delete(key) + return true + } + if entry.refs.Add(-1) <= 0 { + globalRegistry.Delete(key) + return true + } + return false } func getWithScope(scope, reuseID string) (*Resource, bool) { @@ -104,11 +157,11 @@ func getWithScope(scope, reuseID string) (*Resource, bool) { if !ok { return nil, false } - resource, ok := val.(*Resource) + entry, ok := val.(*registryEntry) if !ok { return nil, false } - return resource, true + return entry.resource, true } func getAllWithScope(scope string) []*Resource { @@ -119,8 +172,8 @@ func getAllWithScope(scope string) []*Resource { if !ok || regKey.scope != scope { return true } - if resource, ok := value.(*Resource); ok { - resources = append(resources, resource) + if entry, ok := value.(*registryEntry); ok { + resources = append(resources, entry.resource) } return true }) diff --git a/registry_test.go b/registry_test.go index 27b36ce2..a36569ad 100644 --- a/registry_test.go +++ b/registry_test.go @@ -219,3 +219,116 @@ func TestRegisterWithScopeLoadOrStore(t *testing.T) { t.Fatalf("second stored resource = %q, want %q", stored.ID(), first.ID()) } } + +func TestRegistryRefCounting(t *testing.T) { + ResetRegistry() + + r := &Resource{Container: container.InspectResponse{ID: "refcount-container"}} + + // Register: refs=1 + _, loaded := registerWithScope("scope", "rc-id", r) + if loaded { + t.Fatal("first register loaded = true, want false") + } + + // Acquire: refs=2 + got, ok := acquireWithScope("scope", "rc-id") + if !ok { + t.Fatal("acquireWithScope returned false, want true") + } + if got.ID() != r.ID() { + t.Fatalf("acquireWithScope returned %q, want %q", got.ID(), r.ID()) + } + + // Release once: refs=1, should NOT be last + if releaseWithScope("scope", "rc-id") { + t.Fatal("first releaseWithScope returned true (last ref), want false") + } + // Entry should still exist + if _, ok := getWithScope("scope", "rc-id"); !ok { + t.Fatal("entry removed after first release, want it to remain") + } + + // Release again: refs=0, should be last + if !releaseWithScope("scope", "rc-id") { + t.Fatal("second releaseWithScope returned false, want true (last ref)") + } + // Entry should be gone + if _, ok := getWithScope("scope", "rc-id"); ok { + t.Fatal("entry still present after last release, want it removed") + } +} + +func TestRegistryRefCountingRegisterIncrementsOnDuplicate(t *testing.T) { + ResetRegistry() + + r1 := &Resource{Container: container.InspectResponse{ID: "dup-1"}} + r2 := &Resource{Container: container.InspectResponse{ID: "dup-2"}} + + // Register r1: refs=1 + registerWithScope("scope", "dup-id", r1) + // Register r2 with same key: refs=2 (r1 is canonical) + stored, loaded := registerWithScope("scope", "dup-id", r2) + if !loaded { + t.Fatal("second register loaded = false, want true") + } + if stored.ID() != r1.ID() { + t.Fatalf("canonical resource = %q, want %q", stored.ID(), r1.ID()) + } + + // Need 2 releases to remove + if releaseWithScope("scope", "dup-id") { + t.Fatal("first release was last, want false") + } + if !releaseWithScope("scope", "dup-id") { + t.Fatal("second release was not last, want true") + } +} + +func TestAcquireWithScopeNonExistent(t *testing.T) { + ResetRegistry() + + _, ok := acquireWithScope("scope", "nonexistent") + if ok { + t.Fatal("acquireWithScope returned true for nonexistent entry, want false") + } +} + +func TestReleaseWithScopeNonExistent(t *testing.T) { + ResetRegistry() + + // Releasing a nonexistent entry should return true (treat as last reference) + if !releaseWithScope("scope", "nonexistent") { + t.Fatal("releaseWithScope returned false for nonexistent entry, want true") + } +} + +func TestGetWithScopeDoesNotIncrementRefs(t *testing.T) { + ResetRegistry() + + r := &Resource{Container: container.InspectResponse{ID: "get-no-inc"}} + + // Register: refs=1 + registerWithScope("s", "id", r) + + // getWithScope five times — should NOT increment refs + for range 5 { + got, ok := getWithScope("s", "id") + if !ok { + t.Fatal("getWithScope returned false, want true") + } + if got.ID() != r.ID() { + t.Fatalf("getWithScope returned %q, want %q", got.ID(), r.ID()) + } + } + + // Single release should be the last ref (still 1) + if !releaseWithScope("s", "id") { + t.Fatal("releaseWithScope returned false, want true (last ref)") + } + + // Entry should be gone + if _, ok := getWithScope("s", "id"); ok { + t.Fatal("entry still present after last release, want it removed") + } +} diff --git a/resource.go b/resource.go index e8408865..348fa1d7 100644 --- a/resource.go +++ b/resource.go @@ -73,11 +73,23 @@ func (r *Resource) GetHostPort(portID string) string { // Close stops and removes the container. // Anonymous volumes created by the container are also removed. +// +// For reused containers (those with a reuseID), Close only removes the Docker +// container when the last reference is released. If other callers still hold +// references, Close simply untracks the resource from this pool. func (r *Resource) Close(ctx context.Context) error { if r.pool == nil || r.pool.client == nil { return ErrClientClosed } + if r.reuseID != "" { + if !releaseWithScope(r.pool.reuseScope, r.reuseID) { + // Other callers still hold references; just untrack from this pool. + r.pool.untrackResource(r.Container.ID) + return nil + } + } + // Stop container (ignore errors if already stopped) _, _ = r.pool.client.ContainerStop(ctx, r.Container.ID, mobyclient.ContainerStopOptions{}) //nolint:errcheck // Best effort stop @@ -90,9 +102,6 @@ func (r *Resource) Close(ctx context.Context) error { return err } - if r.reuseID != "" { - unregisterWithScope(r.pool.reuseScope, r.reuseID) - } r.pool.untrackResource(r.Container.ID) return nil diff --git a/run_test.go b/run_test.go index 6140b0ea..472106fd 100644 --- a/run_test.go +++ b/run_test.go @@ -6,8 +6,10 @@ package dockertest_test import ( "testing" + "github.com/containerd/errdefs" "github.com/moby/moby/api/types/container" "github.com/moby/moby/api/types/network" + mobyclient "github.com/moby/moby/client" dockertest "github.com/ory/dockertest/v4" ) @@ -505,6 +507,65 @@ func TestRunWithPortBindings(t *testing.T) { } } +func TestResourceCloseRefCounting(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test in short mode") + } + + dockertest.ResetRegistry() + t.Cleanup(func() { + dockertest.ResetRegistry() + }) + + pool := dockertest.NewPoolT(t, "") + + // First run creates the container: refs=1 + r1 := pool.RunT(t, "alpine", + dockertest.WithTag("latest"), + dockertest.WithCmd([]string{"sleep", "300"}), + ) + + // Second run reuses the same container: refs=2 + r2 := pool.RunT(t, "alpine", + dockertest.WithTag("latest"), + dockertest.WithCmd([]string{"sleep", "300"}), + ) + + if r1.ID() != r2.ID() { + t.Fatalf("expected same container ID for reused container, got %s and %s", r1.ID(), r2.ID()) + } + + containerID := r1.ID() + + // Create a standalone Docker client for container inspection + // (pool.client is unexported from this external test package) + dc, err := mobyclient.New(mobyclient.FromEnv) + if err != nil { + t.Fatalf("mobyclient.New() error = %v", err) + } + t.Cleanup(func() { dc.Close() }) + + // Close first reference: refs=1, container should still exist + r1.CloseT(t) + + // Container must still be alive (one reference remains) + resp, err := dc.ContainerInspect(t.Context(), containerID, mobyclient.ContainerInspectOptions{}) + if err != nil { + t.Fatalf("ContainerInspect() after first close: error = %v, want container still alive", err) + } + if resp.Container.ID != containerID { + t.Fatalf("ContainerInspect() returned ID %s, want %s", resp.Container.ID, containerID) + } + + // Close second reference: refs=0, container should be removed + r2.CloseT(t) + + _, err = dc.ContainerInspect(t.Context(), containerID, mobyclient.ContainerInspectOptions{}) + if !errdefs.IsNotFound(err) { + t.Fatalf("ContainerInspect() after last close: error = %v, want not found", err) + } +} + func TestRunWithMounts(t *testing.T) { if testing.Short() { t.Skip("Skipping integration test in short mode")