proxy-pool/internal/domain/proxy/capacity_test.go

package proxy

import (
	"errors"
	"reflect"
	"sync"
	"sync/atomic"
	"testing"
)

func TestReservationCancelReleasesReservedCapacity(t *testing.T) {
	capacity := NewCapacity(1)
	reservation, ok := capacity.Reserve()
	if !ok {
		t.Fatal("Reserve() = false, want reservation")
	}
	if err := reservation.Cancel(); err != nil {
		t.Fatalf("Cancel() error = %v", err)
	}
	assertCapacityCounters(t, capacity, 0, 0)

	if err := reservation.Cancel(); !errors.Is(err, ErrReservationFinished) {
		t.Fatalf("second Cancel() error = %v, want ErrReservationFinished", err)
	}
	if err := reservation.Commit(); !errors.Is(err, ErrReservationFinished) {
		t.Fatalf("Commit() after Cancel error = %v, want ErrReservationFinished", err)
	}
	if err := reservation.Release(); !errors.Is(err, ErrReservationFinished) {
		t.Fatalf("Release() after Cancel error = %v, want ErrReservationFinished", err)
	}

	reused, ok := capacity.Reserve()
	if !ok {
		t.Fatal("Reserve() after Cancel = false, want released slot")
	}
	if err := reused.Cancel(); err != nil {
		t.Fatalf("reused Cancel() error = %v", err)
	}
	assertCapacityCounters(t, capacity, 0, 0)
}

func TestCapacityCountersReadsOnePackedSnapshot(t *testing.T) {
	capacity := NewCapacity(3)
	first, ok := capacity.Reserve()
	if !ok {
		t.Fatal("first Reserve() = false")
	}
	second, ok := capacity.Reserve()
	if !ok {
		t.Fatal("second Reserve() = false")
	}
	if err := first.Commit(); err != nil {
		t.Fatalf("Commit(): %v", err)
	}

	active, reserved, maximum := capacity.Counters()
	if active != 1 || reserved != 1 || maximum != 3 {
		t.Fatalf("Counters() = (%d, %d, %d), want (1, 1, 3)", active, reserved, maximum)
	}
	if err := first.Release(); err != nil {
		t.Fatalf("Release(): %v", err)
	}
	if err := second.Cancel(); err != nil {
		t.Fatalf("Cancel(): %v", err)
	}
}

func TestCapacityActivityObserverTracksOnlyNonzeroTransitions(t *testing.T) {
	var transitions []bool
	capacity := NewCapacityWithActivityObserver(2, func(nonzero bool) {
		transitions = append(transitions, nonzero)
	})
	first, ok := capacity.Reserve()
	if !ok {
		t.Fatal("Reserve(first) = false")
	}
	second, ok := capacity.Reserve()
	if !ok {
		t.Fatal("Reserve(second) = false")
	}
	if err := first.Cancel(); err != nil {
		t.Fatalf("Cancel(first): %v", err)
	}
	if err := second.Cancel(); err != nil {
		t.Fatalf("Cancel(second): %v", err)
	}
	if !reflect.DeepEqual(transitions, []bool{true, false}) {
		t.Fatalf("transitions = %v, want [true false]", transitions)
	}
}

func TestCapacityRetirementStopsNewReservationsAndRetainsActiveWork(t *testing.T) {
	capacity := NewCapacity(1)
	reservation, ok := capacity.Reserve()
	if !ok {
		t.Fatal("Reserve() = false")
	}
	if err := reservation.Commit(); err != nil {
		t.Fatalf("Commit(): %v", err)
	}
	capacity.SetReservationEnabled(false)
	if _, ok := capacity.Reserve(); ok {
		t.Fatal("Reserve() after retirement = true, want false")
	}
	if capacity.Reclaimable() {
		t.Fatal("Reclaimable() with active work = true, want false")
	}
	if err := reservation.Release(); err != nil {
		t.Fatalf("Release(): %v", err)
	}
	if !capacity.Reclaimable() {
		t.Fatal("Reclaimable() after release = false, want true")
	}
}

func TestReservationCommitAndReleaseAreSingleUse(t *testing.T) {
	capacity := NewCapacity(1)
	reservation, ok := capacity.Reserve()
	if !ok {
		t.Fatal("Reserve() = false, want reservation")
	}
	if err := reservation.Commit(); err != nil {
		t.Fatalf("Commit() error = %v", err)
	}
	assertCapacityCounters(t, capacity, 1, 0)

	if err := reservation.Commit(); !errors.Is(err, ErrReservationCommitted) {
		t.Fatalf("second Commit() error = %v, want ErrReservationCommitted", err)
	}
	if err := reservation.Cancel(); !errors.Is(err, ErrReservationFinished) {
		t.Fatalf("Cancel() after Commit error = %v, want ErrReservationFinished", err)
	}
	assertCapacityCounters(t, capacity, 1, 0)

	if err := reservation.Release(); err != nil {
		t.Fatalf("Release() error = %v", err)
	}
	if err := reservation.Release(); !errors.Is(err, ErrReservationFinished) {
		t.Fatalf("second Release() error = %v, want ErrReservationFinished", err)
	}
	if err := reservation.Commit(); !errors.Is(err, ErrReservationFinished) {
		t.Fatalf("Commit() after Release error = %v, want ErrReservationFinished", err)
	}
	assertCapacityCounters(t, capacity, 0, 0)
}

func TestReleaseBeforeCommitDoesNotConsumeReservation(t *testing.T) {
	capacity := NewCapacity(1)
	reservation, ok := capacity.Reserve()
	if !ok {
		t.Fatal("Reserve() = false, want reservation")
	}
	if err := reservation.Release(); !errors.Is(err, ErrReservationFinished) {
		t.Fatalf("Release() before Commit error = %v, want ErrReservationFinished", err)
	}
	assertCapacityCounters(t, capacity, 0, 1)

	if err := reservation.Cancel(); err != nil {
		t.Fatalf("Cancel() after rejected Release error = %v", err)
	}
	assertCapacityCounters(t, capacity, 0, 0)
}

func TestConcurrentReservationTerminationPreservesCounters(t *testing.T) {
	for iteration := range 1_000 {
		capacity := NewCapacity(1)
		reservation, ok := capacity.Reserve()
		if !ok {
			t.Fatalf("iteration %d Reserve() = false", iteration)
		}

		start := make(chan struct{})
		results := make(chan terminationResult, 2)
		var wait sync.WaitGroup
		wait.Add(2)
		go func() {
			defer wait.Done()
			<-start
			results <- terminationResult{operation: "commit", err: reservation.Commit()}
		}()
		go func() {
			defer wait.Done()
			<-start
			results <- terminationResult{operation: "cancel", err: reservation.Cancel()}
		}()
		close(start)
		wait.Wait()
		close(results)

		var succeeded string
		for result := range results {
			if result.err == nil {
				if succeeded != "" {
					t.Fatalf("iteration %d operations %s and %s both succeeded", iteration, succeeded, result.operation)
				}
				succeeded = result.operation
				continue
			}
			if !errors.Is(result.err, ErrReservationFinished) {
				t.Fatalf("iteration %d %s error = %v, want ErrReservationFinished", iteration, result.operation, result.err)
			}
		}

		switch succeeded {
		case "commit":
			assertCapacityCounters(t, capacity, 1, 0)
			if err := reservation.Release(); err != nil {
				t.Fatalf("iteration %d Release() error = %v", iteration, err)
			}
		case "cancel":
			assertCapacityCounters(t, capacity, 0, 0)
		default:
			t.Fatalf("iteration %d has no successful termination", iteration)
		}
		assertCapacityCounters(t, capacity, 0, 0)
	}
}

func TestConcurrentReleaseSucceedsOnce(t *testing.T) {
	capacity := NewCapacity(1)
	reservation, ok := capacity.Reserve()
	if !ok {
		t.Fatal("Reserve() = false, want reservation")
	}
	if err := reservation.Commit(); err != nil {
		t.Fatalf("Commit() error = %v", err)
	}

	var succeeded atomic.Int64
	var unexpected atomic.Int64
	var wait sync.WaitGroup
	for range 100 {
		wait.Add(1)
		go func() {
			defer wait.Done()
			err := reservation.Release()
			switch {
			case err == nil:
				succeeded.Add(1)
			case !errors.Is(err, ErrReservationFinished):
				unexpected.Add(1)
			}
		}()
	}
	wait.Wait()

	if succeeded.Load() != 1 || unexpected.Load() != 0 {
		t.Fatalf("Release() results = success:%d unexpected:%d, want 1 and 0", succeeded.Load(), unexpected.Load())
	}
	assertCapacityCounters(t, capacity, 0, 0)
}

type terminationResult struct {
	operation string
	err       error
}

func assertCapacityCounters(t *testing.T, capacity *Capacity, active, reserved int64) {
	t.Helper()
	if got := capacity.Active(); got != active {
		t.Fatalf("Active() = %d, want %d", got, active)
	}
	if got := capacity.Reserved(); got != reserved {
		t.Fatalf("Reserved() = %d, want %d", got, reserved)
	}
}