fix: force-remove containers on pool cleanup and add ref-counting tests (#654)

Pool.cleanup() previously called resource.Close() which respects reference
counting, leaving reused containers running when the pool shuts down. Now
cleanup bypasses ref counting and force-removes all tracked containers.

Adds 5 tests covering the cleanup API and ref-counting behavior:
- TestResourceCloseRefCounting (integration)
- TestPoolCloseT (integration)
- TestPoolCleanupForceRemovesReusedContainers (integration)
- TestCheckForExistingIncrementsRefCount (unit)
- TestGetWithScopeDoesNotIncrementRefs (unit)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: aeneasr

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

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")

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,28 +97,71 @@ 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 := &registryEntry{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) {
 	val, ok := globalRegistry.Load(registryKey{scope: scope, reuseID: reuseID})
 	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
 	})