admission_gate_test.go RawHistoryFolder

1package process
2
3import (
4	"context"
5	"sync"
6	"testing"
7	"time"
8
9	"github.com/stretchr/testify/require"
10)
11
12func TestAdmissionGateAcquireRelease(t *testing.T) {
13	gate := NewAdmissionGate()
14	require.True(t, gate.Acquire())
15	require.True(t, gate.Acquire())
16	require.Equal(t, 2, gate.Active())
17	gate.Release()
18	require.Equal(t, 1, gate.Active())
19	gate.Release()
20	require.Equal(t, 0, gate.Active())
21}
22
23func TestAdmissionGateCloseStopsAdmission(t *testing.T) {
24	gate := NewAdmissionGate()
25	require.True(t, gate.Acquire())
26	gate.Close()
27	require.True(t, gate.Closed())
28	require.False(t, gate.Acquire(), "closed gate must refuse new leases")
29	require.Equal(t, 1, gate.Active(), "in-flight lease survives close")
30	gate.Release()
31	require.Equal(t, 0, gate.Active())
32}
33
34func TestAdmissionGateWaitDrained(t *testing.T) {
35	gate := NewAdmissionGate()
36	gate.Acquire()
37	done := make(chan bool, 1)
38	go func() {
39		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
40		defer cancel()
41		done <- gate.WaitDrained(ctx)
42	}()
43	select {
44	case <-done:
45		t.Fatal("drain must not complete while a lease is held")
46	case <-time.After(50 * time.Millisecond):
47	}
48	gate.Release()
49	select {
50	case drained := <-done:
51		require.True(t, drained)
52	case <-time.After(5 * time.Second):
53		t.Fatal("drain must complete after release")
54	}
55}
56
57func TestAdmissionGateWaitDrainedExpiry(t *testing.T) {
58	gate := NewAdmissionGate()
59	gate.Acquire()
60	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
61	defer cancel()
62	require.False(t, gate.WaitDrained(ctx), "expiry must not hang shutdown")
63	gate.Release()
64}
65
66func TestAdmissionGateOnDrained(t *testing.T) {
67	gate := NewAdmissionGate()
68	var mu sync.Mutex
69	calls := 0
70	gate.OnDrained(func() {
71		mu.Lock()
72		calls++
73		mu.Unlock()
74	})
75	gate.Acquire()
76	gate.Release()
77	mu.Lock()
78	defer mu.Unlock()
79	require.Equal(t, 1, calls)
80}
81
82func TestAdmissionGateConcurrent(t *testing.T) {
83	gate := NewAdmissionGate()
84	var wg sync.WaitGroup
85	for range 32 {
86		wg.Add(1)
87		go func() {
88			defer wg.Done()
89			if gate.Acquire() {
90				gate.Release()
91			}
92		}()
93	}
94	wg.Wait()
95	require.Equal(t, 0, gate.Active())
96}