refactor: move execution context and timeout management into pool

common/lua/pool.go

@@ -2,7 +2,9 @@ package lua
 
 import (
 	"context"
+	"errors"
 	"sync"
+	"time"
 
 	glua "github.com/yuin/gopher-lua"
 )
@@ -13,8 +15,9 @@ const maxIdleStates = 16
 // keeps up to maxIdleStates idle states until Close. Acquire/Release callers
 // decide reusability; WithState uses its callback's error.
 type Pool struct {
-	ctx    context.Context
-	cancel context.CancelFunc
+	ctx     context.Context
+	cancel  context.CancelFunc
+	timeout time.Duration
 
 	factory LStateFactory
 	idle    []*glua.LState
@@ -26,7 +29,11 @@ type Pool struct {
 }
 
 // NewPool tests the factory by creating one state during initialization.
-func NewPool(ctx context.Context, factory LStateFactory) (*Pool, error) {
+func NewPool(ctx context.Context, timeout time.Duration, factory LStateFactory) (*Pool, error) {
+	if timeout <= 0 {
+		return nil, errors.New("Lua pool timeout must be positive")
+	}
+
 	poolCtx, cancel := context.WithCancel(ctx)
 
 	state, err := factory(poolCtx)
@@ -35,20 +42,26 @@ func NewPool(ctx context.Context, factory LStateFactory) (*Pool, error) {
 		return nil, err
 	}
 
-	return &Pool{ctx: poolCtx, cancel: cancel, factory: factory, idle: []*glua.LState{state}, top: state.GetTop()}, nil
-}
-
-// Context is cancelled by Close. Query contexts should derive from it.
-func (p *Pool) Context() context.Context {
-	return p.ctx
+	return &Pool{ctx: poolCtx, cancel: cancel, timeout: timeout, factory: factory, idle: []*glua.LState{state}, top: state.GetTop()}, nil
 }
 
 // Acquire returns an initialized exclusive state, growing the pool if necessary.
-func (p *Pool) Acquire() (*glua.LState, error) {
+// ctx is passed to the factory for state creation; nil uses the pool context.
+func (p *Pool) Acquire(ctx context.Context) (*glua.LState, error) {
 	p.mu.Lock()
-	if p.closed || p.ctx.Err() != nil {
+	if p.closed {
 		p.mu.Unlock()
-		return nil, p.ctx.Err()
+		return nil, errors.New("Lua pool is closed")
+	}
+	if err := p.ctx.Err(); err != nil {
+		p.mu.Unlock()
+		return nil, err
+	}
+	if ctx == nil {
+		ctx = p.ctx
+	} else if err := ctx.Err(); err != nil {
+		p.mu.Unlock()
+		return nil, err
 	}
 
 	p.active.Add(1)
@@ -65,7 +78,7 @@ func (p *Pool) Acquire() (*glua.LState, error) {
 	// TODO: Limit the total number of states. When the limit is reached, wait
 	// for a Release instead of creating another state; allow the wait to be
 	// cancelled by the caller or by Close.
-	state, err := p.factory(p.ctx)
+	state, err := p.factory(ctx)
 	if err != nil {
 		p.active.Done()
 		return nil, err
@@ -74,15 +87,24 @@ func (p *Pool) Acquire() (*glua.LState, error) {
 	return state, nil
 }
 
-// WithState runs work on an exclusive state and releases it afterward. A state
-// is reusable only when work succeeds; a panic closes it before propagating.
-func (p *Pool) WithState(work func(*glua.LState) error) error {
-	state, err := p.Acquire()
+// WithState runs work on an exclusive state and releases it afterward.
+// Nil ctx and zero timeout use pool defaults. The timeout starts after acquisition.
+func (p *Pool) WithState(ctx context.Context, timeout time.Duration, work func(*glua.LState) error) error {
+	state, err := p.Acquire(ctx)
 	if err != nil {
 		return err
 	}
+	if ctx == nil {
+		ctx = p.ctx
+	}
+	if timeout == 0 {
+		timeout = p.timeout
+	}
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	state.SetContext(ctx)
 	reusable := false
 	defer func() {
+		cancel()
 		p.Release(state, reusable)
 	}()
 	err = work(state)
@@ -90,9 +112,10 @@ func (p *Pool) WithState(work func(*glua.LState) error) error {
 	return err
 }
 
-// Release returns a healthy state to the pool and closes a failed or cancelled one.
+// Release resets a state for reuse or closes it.
 func (p *Pool) Release(state *glua.LState, reusable bool) {
 	if reusable {
+		state.RemoveContext()
 		state.SetTop(p.top)
 		p.mu.Lock()
 		if !p.closed && p.ctx.Err() == nil && len(p.idle) < maxIdleStates {
@@ -110,7 +133,7 @@ func (p *Pool) Release(state *glua.LState, reusable bool) {
 	p.active.Done()
 }
 
-// Close cancels active work, closes idle states, and waits for borrowed states.
+// Close cancels the pool context, closes idle states, and waits for borrowed states.
 func (p *Pool) Close() {
 	p.mu.Lock()
 	if !p.closed {

common/lua/pool_test.go

@@ -9,186 +9,458 @@ import (
 	glua "github.com/yuin/gopher-lua"
 )
 
+func newTestPool(t testing.TB, ctx context.Context, timeout time.Duration, factory LStateFactory) *Pool {
+	t.Helper()
+	pool, err := NewPool(ctx, timeout, factory)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(pool.Close)
+	return pool
+}
+
+func assertPoolCloseBlocked(t *testing.T, done <-chan struct{}) {
+	t.Helper()
+	select {
+	case <-done:
+		t.Fatal("Close returned while work was still active")
+	case <-time.After(20 * time.Millisecond):
+	}
+}
+
+func TestPoolTimeoutValidation(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		timeout time.Duration
+		wantErr bool
+	}{
+		{"zero", 0, true},
+		{"negative", -time.Nanosecond, true},
+		{"positive", time.Nanosecond, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			called := false
+			pool, err := NewPool(context.Background(), tc.timeout, func(context.Context) (*glua.LState, error) {
+				called = true
+				return glua.NewState(), nil
+			})
+			if pool != nil {
+				t.Cleanup(pool.Close)
+			}
+			if (err != nil) != tc.wantErr {
+				t.Fatalf("NewPool error = %v, want error %t", err, tc.wantErr)
+			}
+			if tc.wantErr && (pool != nil || called) {
+				t.Fatal("invalid timeout created a pool or called the factory")
+			}
+		})
+	}
+}
+
 func TestPoolFactoryFailure(t *testing.T) {
 	failure := errors.New("factory failed")
-	_, err := NewPool(context.Background(), func(context.Context) (*glua.LState, error) {
+	_, err := NewPool(context.Background(), time.Second, func(context.Context) (*glua.LState, error) {
 		return nil, failure
 	})
 	if !errors.Is(err, failure) {
-		t.Fatalf("NewPool error = %v, want %v", err, failure)
+		t.Fatalf("NewPool error = %v, want original factory error", err)
 	}
 
 	calls := 0
-	pool, err := NewPool(context.Background(), func(context.Context) (*glua.LState, error) {
+	pool := newTestPool(t, context.Background(), time.Second, func(context.Context) (*glua.LState, error) {
 		calls++
 		if calls == 1 {
 			return glua.NewState(), nil
 		}
 		return nil, failure
 	})
+	state, err := pool.Acquire(nil)
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer pool.Close()
-	borrowed, err := pool.Acquire()
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer pool.Release(borrowed, true)
-	_, err = pool.Acquire()
+	defer pool.Release(state, true)
+	err = pool.WithState(nil, 0, func(*glua.LState) error {
+		t.Error("work ran after factory failure")
+		return nil
+	})
 	if !errors.Is(err, failure) {
-		t.Fatalf("Acquire error = %v, want %v", err, failure)
+		t.Fatalf("WithState error = %v, want original factory error", err)
 	}
 }
 
 func TestPoolReusesStatesAndLimitsIdle(t *testing.T) {
 	created := 0
-	pool, err := NewPool(context.Background(), func(context.Context) (*glua.LState, error) {
+	pool := newTestPool(t, context.Background(), time.Second, func(context.Context) (*glua.LState, error) {
 		created++
 		return glua.NewState(), nil
 	})
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer pool.Close()
-
-	states := make([]*glua.LState, maxIdleStates+3)
-	for i := range states {
-		states[i], err = pool.Acquire()
+	var borrowed []*glua.LState
+	defer func() {
+		for _, state := range borrowed {
+			pool.Release(state, false)
+		}
+	}()
+	for range maxIdleStates + 3 {
+		state, err := pool.Acquire(nil)
 		if err != nil {
 			t.Fatal(err)
 		}
+		borrowed = append(borrowed, state)
+		state.SetContext(context.Background())
 	}
+	states := borrowed
 	for _, state := range states {
 		pool.Release(state, true)
 	}
-	for i, state := range states {
-		if got, want := state.IsClosed(), i >= maxIdleStates; got != want {
-			t.Fatalf("state %d closed = %t, want %t", i, got, want)
+	borrowed = nil
+	open := 0
+	for _, state := range states {
+		if !state.IsClosed() {
+			if state.Context() != nil {
+				t.Fatal("Release left a context on a reusable state")
+			}
+			open++
 		}
 	}
-	borrowed, err := pool.Acquire()
-	if err != nil {
+	if open != maxIdleStates {
+		t.Fatalf("retained %d states, want %d", open, maxIdleStates)
+	}
+	if err := pool.WithState(nil, 0, func(*glua.LState) error { return nil }); err != nil {
 		t.Fatal(err)
 	}
 	if created != len(states) {
-		t.Fatalf("Acquire created %d states, want %d", created, len(states))
+		t.Fatalf("created %d states, want %d", created, len(states))
+	}
+	pool.Close()
+	for _, state := range states {
+		if !state.IsClosed() {
+			t.Fatal("Close left an idle state open")
+		}
 	}
-	pool.Release(borrowed, true)
 }
 
-func TestPoolCloseCancelsAndWaitsForBorrowedState(t *testing.T) {
-	pool, err := NewPool(context.Background(), func(context.Context) (*glua.LState, error) {
+func TestPoolWithStateOptions(t *testing.T) {
+	key := struct{}{}
+	parent := context.WithValue(context.Background(), key, "pool")
+	caller := context.WithValue(context.Background(), key, "caller")
+	pool := newTestPool(t, parent, time.Second, func(context.Context) (*glua.LState, error) {
 		return glua.NewState(), nil
 	})
-	if err != nil {
-		t.Fatal(err)
+	for _, tc := range []struct {
+		name        string
+		ctx         context.Context
+		timeout     time.Duration
+		wantValue   string
+		wantTimeout time.Duration
+	}{
+		{"defaults", nil, 0, "pool", time.Second},
+		{"context", caller, 0, "caller", time.Second},
+		{"timeout", nil, 2 * time.Second, "pool", 2 * time.Second},
+		{"both", caller, 2 * time.Second, "caller", 2 * time.Second},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			started := time.Now()
+			err := pool.WithState(tc.ctx, tc.timeout, func(L *glua.LState) error {
+				ctx := L.Context()
+				if ctx.Value(key) != tc.wantValue {
+					t.Errorf("context value = %v, want %q", ctx.Value(key), tc.wantValue)
+				}
+				deadline, ok := ctx.Deadline()
+				if !ok || deadline.Before(started.Add(tc.wantTimeout)) || deadline.After(time.Now().Add(tc.wantTimeout)) {
+					t.Errorf("deadline = %v, want timeout %v", deadline, tc.wantTimeout)
+				}
+				return nil
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+		})
 	}
-	state, err := pool.Acquire()
-	if err != nil {
-		t.Fatal(err)
+}
+
+func TestPoolFactoryContext(t *testing.T) {
+	caller, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+	for _, tc := range []struct {
+		name string
+		ctx  context.Context
+	}{
+		{"default", nil},
+		{"caller", caller},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var contexts []context.Context
+			pool := newTestPool(t, context.Background(), time.Second, func(ctx context.Context) (*glua.LState, error) {
+				contexts = append(contexts, ctx)
+				return glua.NewState(), nil
+			})
+			state, err := pool.Acquire(nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer pool.Release(state, true)
+			if err := pool.WithState(tc.ctx, 2*time.Second, func(*glua.LState) error { return nil }); err != nil {
+				t.Fatal(err)
+			}
+			want := tc.ctx
+			if want == nil {
+				want = pool.ctx
+			}
+			if len(contexts) != 2 || contexts[0] != pool.ctx || contexts[1] != want {
+				t.Fatal("factory did not receive the initialization and acquisition contexts unchanged")
+			}
+		})
 	}
-	done := make(chan struct{})
+}
+
+func TestPoolWithStateLifecycle(t *testing.T) {
+	failure := errors.New("work failed")
+	for _, tc := range []struct {
+		name      string
+		work      func(*glua.LState, context.CancelFunc) error
+		reusable  bool
+		wantPanic bool
+		wantErr   error
+	}{
+		{"success", func(*glua.LState, context.CancelFunc) error { return nil }, true, false, nil},
+		{"canceled success", func(_ *glua.LState, cancel context.CancelFunc) error {
+			cancel()
+			return nil
+		}, true, false, nil},
+		{"error", func(*glua.LState, context.CancelFunc) error { return failure }, false, false, failure},
+		{"timeout", func(L *glua.LState, _ context.CancelFunc) error { return L.DoString("while true do end") }, false, false, nil},
+		{"panic", func(*glua.LState, context.CancelFunc) error { panic(failure) }, false, true, nil},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			pool := newTestPool(t, context.Background(), 10*time.Millisecond, func(context.Context) (*glua.LState, error) {
+				state := glua.NewState()
+				state.Push(glua.LTrue)
+				return state, nil
+			})
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			var state *glua.LState
+			var workCtx context.Context
+			var recovered any
+			err := func() (err error) {
+				defer func() { recovered = recover() }()
+				return pool.WithState(ctx, 0, func(L *glua.LState) error {
+					state, workCtx = L, L.Context()
+					L.Push(glua.LFalse)
+					return tc.work(L, cancel)
+				})
+			}()
+			if tc.wantPanic {
+				if recovered != failure {
+					t.Fatalf("panic = %v, want original panic", recovered)
+				}
+			} else {
+				if recovered != nil || (err == nil) != tc.reusable {
+					t.Fatalf("WithState error = %v, panic = %v", err, recovered)
+				}
+				if tc.wantErr != nil && !errors.Is(err, tc.wantErr) {
+					t.Fatalf("WithState error = %v, want %v", err, tc.wantErr)
+				}
+			}
+			if workCtx.Err() == nil {
+				t.Fatal("WithState did not cancel the execution context")
+			}
+			if closed := state.IsClosed(); closed == tc.reusable {
+				t.Fatalf("state closed = %t, want %t", closed, !tc.reusable)
+			}
+			if tc.reusable && (state.Context() != nil || state.GetTop() != 1 || state.Get(1) != glua.LTrue) {
+				t.Fatal("WithState did not reset the state for reuse")
+			}
+			if err := pool.WithState(nil, 0, func(L *glua.LState) error {
+				if (L == state) != tc.reusable {
+					t.Error("unexpected state reuse")
+				}
+				return nil
+			}); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}
+
+func TestPoolClose(t *testing.T) {
+	pool := newTestPool(t, context.Background(), time.Second, func(context.Context) (*glua.LState, error) {
+		return glua.NewState(), nil
+	})
+	finishCtx, finish := context.WithCancel(context.Background())
+	t.Cleanup(finish)
+	started, done := make(chan *glua.LState, 1), make(chan error, 1)
+	var workCtx context.Context
+	go func() {
+		done <- pool.WithState(nil, 0, func(L *glua.LState) error {
+			workCtx = L.Context()
+			started <- L
+			<-finishCtx.Done()
+			return nil
+		})
+	}()
+	var state *glua.LState
+	select {
+	case state = <-started:
+	case <-time.After(time.Second):
+		t.Fatal("WithState did not start")
+	}
+	closed := make(chan struct{})
 	go func() {
 		pool.Close()
-		close(done)
+		close(closed)
 	}()
 	select {
-	case <-pool.Context().Done():
+	case <-workCtx.Done():
 	case <-time.After(time.Second):
-		pool.Release(state, false)
-		t.Fatal("Close did not cancel the pool context")
+		t.Fatal("Close did not cancel work using the pool context")
+	}
+	if !errors.Is(workCtx.Err(), context.Canceled) {
+		t.Fatalf("work context error = %v, want context.Canceled", workCtx.Err())
+	}
+	assertPoolCloseBlocked(t, closed)
+	finish()
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("successful work returned an error: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("WithState did not finish")
 	}
 	select {
-	case <-done:
-		pool.Release(state, false)
-		t.Fatal("Close returned while a state was borrowed")
-	default:
-	}
-	pool.Release(state, true)
-	select {
-	case <-done:
+	case <-closed:
 	case <-time.After(time.Second):
-		t.Fatal("Close did not finish after Release")
+		t.Fatal("Close did not finish after WithState")
 	}
 	if !state.IsClosed() {
-		t.Fatal("borrowed state was not closed")
+		t.Fatal("Release returned a state to a closed pool")
 	}
-	if _, err := pool.Acquire(); !errors.Is(err, context.Canceled) {
-		t.Fatalf("Acquire after Close = %v, want context.Canceled", err)
+	if state, err := pool.Acquire(nil); state != nil || err == nil || errors.Is(err, context.Canceled) {
+		t.Fatalf("Acquire after Close = %v, %v; want closed pool error", state, err)
 	}
 	pool.Close()
 }
 
-func TestPoolCloseCancelsStateCreation(t *testing.T) {
-	started := make(chan struct{})
+func TestPoolCloseWaitsForFactory(t *testing.T) {
+	finishCtx, finish := context.WithCancel(context.Background())
+	started, canceled := make(chan struct{}), make(chan struct{})
 	first := true
-	pool, err := NewPool(context.Background(), func(ctx context.Context) (*glua.LState, error) {
+	pool := newTestPool(t, context.Background(), time.Second, func(ctx context.Context) (*glua.LState, error) {
 		if first {
 			first = false
 			return glua.NewState(), nil
 		}
 		close(started)
 		<-ctx.Done()
+		close(canceled)
+		<-finishCtx.Done()
 		return nil, ctx.Err()
 	})
+	t.Cleanup(finish)
+	state, err := pool.Acquire(nil)
 	if err != nil {
 		t.Fatal(err)
 	}
-	borrowed, err := pool.Acquire()
-	if err != nil {
-		t.Fatal(err)
-	}
+	pool.Release(state, false)
 	acquireDone := make(chan error, 1)
 	go func() {
-		_, err := pool.Acquire()
+		_, err := pool.Acquire(nil)
 		acquireDone <- err
 	}()
-	<-started
-	closeDone := make(chan struct{})
+	select {
+	case <-started:
+	case <-time.After(time.Second):
+		t.Fatal("state creation did not start")
+	}
+	closed := make(chan struct{})
 	go func() {
 		pool.Close()
-		close(closeDone)
+		close(closed)
 	}()
 	select {
+	case <-canceled:
+	case <-time.After(time.Second):
+		t.Fatal("Close did not cancel state creation")
+	}
+	assertPoolCloseBlocked(t, closed)
+	finish()
+	select {
 	case err := <-acquireDone:
 		if !errors.Is(err, context.Canceled) {
-			t.Fatalf("Acquire during Close = %v, want context.Canceled", err)
+			t.Fatalf("Acquire error = %v, want context.Canceled", err)
 		}
 	case <-time.After(time.Second):
-		t.Fatal("state creation did not stop after Close")
+		t.Fatal("state creation did not finish")
 	}
 	select {
-	case <-closeDone:
-		pool.Release(borrowed, false)
-		t.Fatal("Close returned while the initial state was borrowed")
-	default:
-	}
-	pool.Release(borrowed, true)
-	select {
-	case <-closeDone:
+	case <-closed:
 	case <-time.After(time.Second):
-		t.Fatal("Close did not finish after Release")
+		t.Fatal("Close did not finish after state creation")
+	}
+}
+
+func TestPoolCloseWaitsForCallerContext(t *testing.T) {
+	pool := newTestPool(t, context.Background(), time.Minute, func(context.Context) (*glua.LState, error) {
+		return glua.NewState(), nil
+	})
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	started, done := make(chan context.Context, 1), make(chan error, 1)
+	go func() {
+		done <- pool.WithState(ctx, 0, func(L *glua.LState) error {
+			started <- L.Context()
+			<-L.Context().Done()
+			return L.Context().Err()
+		})
+	}()
+	var workCtx context.Context
+	select {
+	case workCtx = <-started:
+	case <-time.After(time.Second):
+		t.Fatal("WithState did not start")
+	}
+	closed := make(chan struct{})
+	go func() {
+		pool.Close()
+		close(closed)
+	}()
+	select {
+	case <-pool.ctx.Done():
+	case <-time.After(time.Second):
+		t.Fatal("Close did not cancel the pool context")
+	}
+	assertPoolCloseBlocked(t, closed)
+	if workCtx.Err() != nil || ctx.Err() != nil {
+		t.Fatal("Close canceled the caller's execution context")
+	}
+	cancel()
+	select {
+	case err := <-done:
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("WithState error = %v, want context.Canceled", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("WithState did not stop after caller cancellation")
+	}
+	select {
+	case <-closed:
+	case <-time.After(time.Second):
+		t.Fatal("Close did not finish after WithState")
 	}
 }
 
 func BenchmarkPoolAcquireRelease(b *testing.B) {
-	pool, err := NewPool(context.Background(), func(context.Context) (*glua.LState, error) {
+	pool := newTestPool(b, context.Background(), time.Second, func(context.Context) (*glua.LState, error) {
 		return glua.NewState(), nil
 	})
-	if err != nil {
-		b.Fatal(err)
-	}
-	defer pool.Close()
-
 	b.ReportAllocs()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		state, err := pool.Acquire()
+		state, err := pool.Acquire(nil)
 		if err != nil {
 			b.Fatal(err)
 		}
 		pool.Release(state, true)
 	}
-	b.StopTimer()
 }