go/internal/pool/want_conn_test.go

package pool

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"
)

func TestWantConn_getCtxForDial(t *testing.T) {
	ctx := context.Background()
	w := &wantConn{
		ctx:    ctx,
		result: make(chan wantConnResult, 1),
	}

	// Test getting context when not done
	gotCtx := w.getCtxForDial()
	if gotCtx != ctx {
		t.Errorf("getCtxForDial() = %v, want %v", gotCtx, ctx)
	}

	// Test getting context when done
	w.done = true
	w.ctx = nil
	gotCtx = w.getCtxForDial()
	if gotCtx != nil {
		t.Errorf("getCtxForDial() after done = %v, want nil", gotCtx)
	}
}

func TestWantConn_tryDeliver_Success(t *testing.T) {
	w := &wantConn{
		ctx:    context.Background(),
		result: make(chan wantConnResult, 1),
	}

	// Create a mock connection
	conn := &Conn{}

	// Test successful delivery
	delivered := w.tryDeliver(conn, nil)
	if !delivered {
		t.Error("tryDeliver() = false, want true")
	}

	// Check that wantConn is marked as done
	if !w.done {
		t.Error("wantConn.done = false, want true after delivery")
	}

	// Check that context is cleared
	if w.ctx != nil {
		t.Error("wantConn.ctx should be nil after delivery")
	}

	// Check that result is sent
	select {
	case result := <-w.result:
		if result.cn != conn {
			t.Errorf("result.cn = %v, want %v", result.cn, conn)
		}
		if result.err != nil {
			t.Errorf("result.err = %v, want nil", result.err)
		}
	case <-time.After(time.Millisecond):
		t.Error("Expected result to be sent to channel")
	}
}

func TestWantConn_tryDeliver_WithError(t *testing.T) {
	w := &wantConn{
		ctx:    context.Background(),
		result: make(chan wantConnResult, 1),
	}

	testErr := errors.New("test error")

	// Test delivery with error
	delivered := w.tryDeliver(nil, testErr)
	if !delivered {
		t.Error("tryDeliver() = false, want true")
	}

	// Check result
	select {
	case result := <-w.result:
		if result.cn != nil {
			t.Errorf("result.cn = %v, want nil", result.cn)
		}
		if result.err != testErr {
			t.Errorf("result.err = %v, want %v", result.err, testErr)
		}
	case <-time.After(time.Millisecond):
		t.Error("Expected result to be sent to channel")
	}
}

func TestWantConn_tryDeliver_AlreadyDone(t *testing.T) {
	w := &wantConn{
		ctx:    context.Background(),
		done:   true, // Already done
		result: make(chan wantConnResult, 1),
	}

	// Test delivery when already done
	delivered := w.tryDeliver(&Conn{}, nil)
	if delivered {
		t.Error("tryDeliver() = true, want false when already done")
	}

	// Check that no result is sent
	select {
	case <-w.result:
		t.Error("No result should be sent when already done")
	case <-time.After(time.Millisecond):
		// Expected
	}
}

func TestWantConn_cancel_NotDone(t *testing.T) {
	w := &wantConn{
		ctx:    context.Background(),
		result: make(chan wantConnResult, 1),
	}

	// Test cancel when not done
	cn := w.cancel()

	// Should return nil since no connection was not delivered
	if cn != nil {
		t.Errorf("cancel()= %v, want nil when no connection delivered", cn)
	}

	// Check that wantConn is marked as done
	if !w.done {
		t.Error("wantConn.done = false, want true after cancel")
	}

	// Check that context is cleared
	if w.ctx != nil {
		t.Error("wantConn.ctx should be nil after cancel")
	}

	// Check that channel is closed
	select {
	case _, ok := <-w.result:
		if ok {
			t.Error("result channel should be closed after cancel")
		}
	case <-time.After(time.Millisecond):
		t.Error("Expected channel to be closed")
	}
}

func TestWantConn_cancel_AlreadyDone(t *testing.T) {
	w := &wantConn{
		ctx:    context.Background(),
		done:   true,
		result: make(chan wantConnResult, 1),
	}

	// Put a result in the channel without connection (to avoid nil pointer issues)
	testErr := errors.New("test error")
	w.result <- wantConnResult{cn: nil, err: testErr}

	// Test cancel when already done
	cn := w.cancel()

	// Should return nil since the result had no connection
	if cn != nil {
		t.Errorf("cancel()= %v, want nil when result had no connection", cn)
	}

	// Check that wantConn remains done
	if !w.done {
		t.Error("wantConn.done = false, want true")
	}

	// Check that context is cleared
	if w.ctx != nil {
		t.Error("wantConn.ctx should be nil after cancel")
	}
}

func TestWantConnQueue_newWantConnQueue(t *testing.T) {
	q := newWantConnQueue()
	if q == nil {
		t.Fatal("newWantConnQueue() returned nil")
	}
	if q.items == nil {
		t.Error("queue items should be initialized")
	}
	if len(q.items) != 0 {
		t.Errorf("new queue length = %d, want 0", len(q.items))
	}
}

func TestWantConnQueue_enqueue_dequeue(t *testing.T) {
	q := newWantConnQueue()

	// Test dequeue from empty queue
	item, ok := q.dequeue()
	if ok {
		t.Error("dequeue() from empty queue should return false")
	}
	if item != nil {
		t.Error("dequeue() from empty queue should return nil")
	}

	// Create test wantConn items
	w1 := &wantConn{ctx: context.Background(), result: make(chan wantConnResult, 1)}
	w2 := &wantConn{ctx: context.Background(), result: make(chan wantConnResult, 1)}
	w3 := &wantConn{ctx: context.Background(), result: make(chan wantConnResult, 1)}

	// Test enqueue
	q.enqueue(w1)
	q.enqueue(w2)
	q.enqueue(w3)

	// Test FIFO behavior
	item, ok = q.dequeue()
	if !ok {
		t.Error("dequeue() should return true when queue has items")
	}
	if item != w1 {
		t.Errorf("dequeue() = %v, want %v (FIFO order)", item, w1)
	}

	item, ok = q.dequeue()
	if !ok {
		t.Error("dequeue() should return true when queue has items")
	}
	if item != w2 {
		t.Errorf("dequeue() = %v, want %v (FIFO order)", item, w2)
	}

	item, ok = q.dequeue()
	if !ok {
		t.Error("dequeue() should return true when queue has items")
	}
	if item != w3 {
		t.Errorf("dequeue() = %v, want %v (FIFO order)", item, w3)
	}

	// Test dequeue from empty queue again
	item, ok = q.dequeue()
	if ok {
		t.Error("dequeue() from empty queue should return false")
	}
	if item != nil {
		t.Error("dequeue() from empty queue should return nil")
	}
}

func TestWantConnQueue_ConcurrentAccess(t *testing.T) {
	q := newWantConnQueue()
	const numWorkers = 10
	const itemsPerWorker = 100

	var wg sync.WaitGroup

	// Start enqueuers
	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < itemsPerWorker; j++ {
				w := &wantConn{
					ctx:    context.Background(),
					result: make(chan wantConnResult, 1),
				}
				q.enqueue(w)
			}
		}()
	}

	// Start dequeuers
	dequeued := make(chan *wantConn, numWorkers*itemsPerWorker)
	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < itemsPerWorker; j++ {
				for {
					if item, ok := q.dequeue(); ok {
						dequeued <- item
						break
					}
					// Small delay to avoid busy waiting
					time.Sleep(time.Microsecond)
				}
			}
		}()
	}

	wg.Wait()
	close(dequeued)

	// Count dequeued items
	count := 0
	for range dequeued {
		count++
	}

	expectedCount := numWorkers * itemsPerWorker
	if count != expectedCount {
		t.Errorf("dequeued %d items, want %d", count, expectedCount)
	}

	// Queue should be empty
	if item, ok := q.dequeue(); ok {
		t.Errorf("queue should be empty but got item: %v", item)
	}
}

func TestWantConnQueue_ThreadSafety(t *testing.T) {
	q := newWantConnQueue()
	const numOperations = 1000

	var wg sync.WaitGroup
	errors := make(chan error, numOperations*2)

	// Concurrent enqueue operations
	wg.Add(1)
	go func() {
		defer wg.Done()
		for i := 0; i < numOperations; i++ {
			w := &wantConn{
				ctx:    context.Background(),
				result: make(chan wantConnResult, 1),
			}
			q.enqueue(w)
		}
	}()

	// Concurrent dequeue operations
	wg.Add(1)
	go func() {
		defer wg.Done()
		dequeued := 0
		for dequeued < numOperations {
			if _, ok := q.dequeue(); ok {
				dequeued++
			} else {
				// Small delay when queue is empty
				time.Sleep(time.Microsecond)
			}
		}
	}()

	// Wait for completion
	wg.Wait()
	close(errors)

	// Check for any errors
	for err := range errors {
		t.Error(err)
	}

	// Final queue should be empty
	if item, ok := q.dequeue(); ok {
		t.Errorf("queue should be empty but got item: %v", item)
	}
}

// Benchmark tests
func BenchmarkWantConnQueue_Enqueue(b *testing.B) {
	q := newWantConnQueue()

	// Pre-allocate a pool of wantConn to reuse
	const poolSize = 1000
	wantConnPool := make([]*wantConn, poolSize)
	for i := 0; i < poolSize; i++ {
		wantConnPool[i] = &wantConn{
			ctx:    context.Background(),
			result: make(chan wantConnResult, 1),
		}
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		w := wantConnPool[i%poolSize]
		q.enqueue(w)
	}
}

func BenchmarkWantConnQueue_Dequeue(b *testing.B) {
	q := newWantConnQueue()

	// Use a reasonable fixed size for pre-population to avoid memory issues
	const queueSize = 10000

	// Pre-populate queue with a fixed reasonable size
	for i := 0; i < queueSize; i++ {
		w := &wantConn{
			ctx:    context.Background(),
			result: make(chan wantConnResult, 1),
		}
		q.enqueue(w)
	}

	b.ResetTimer()

	// Benchmark dequeue operations, refilling as needed
	for i := 0; i < b.N; i++ {
		if _, ok := q.dequeue(); !ok {
			// Queue is empty, refill a batch
			for j := 0; j < 1000; j++ {
				w := &wantConn{
					ctx:    context.Background(),
					result: make(chan wantConnResult, 1),
				}
				q.enqueue(w)
			}
			// Dequeue again
			q.dequeue()
		}
	}
}

func BenchmarkWantConnQueue_EnqueueDequeue(b *testing.B) {
	q := newWantConnQueue()

	// Pre-allocate a pool of wantConn to reuse
	const poolSize = 1000
	wantConnPool := make([]*wantConn, poolSize)
	for i := 0; i < poolSize; i++ {
		wantConnPool[i] = &wantConn{
			ctx:    context.Background(),
			result: make(chan wantConnResult, 1),
		}
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		w := wantConnPool[i%poolSize]
		q.enqueue(w)
		q.dequeue()
	}
}

// TestWantConn_RaceConditionNilContext tests the race condition where
// getCtxForDial can return nil after the context is cancelled.
// This test verifies that the fix in newConn handles nil context gracefully.
func TestWantConn_RaceConditionNilContext(t *testing.T) {
	// This test simulates the race condition described in the issue:
	// 1. Main goroutine creates a wantConn with a context
	// 2. Background goroutine starts but hasn't called getCtxForDial yet
	// 3. Main goroutine times out and calls cancel(), setting w.ctx to nil
	// 4. Background goroutine calls getCtxForDial() and gets nil
	// 5. Background goroutine calls newConn(nil, true) which should not panic

	dialCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	w := &wantConn{
		ctx:       dialCtx,
		cancelCtx: cancel,
		result:    make(chan wantConnResult, 1),
	}

	// Simulate the race condition by canceling the context
	// and then trying to get it
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		// Small delay to ensure cancel happens first
		time.Sleep(10 * time.Millisecond)

		// This should return nil after cancel
		ctx := w.getCtxForDial()

		// Verify that we got nil context
		if ctx != nil {
			t.Errorf("Expected nil context after cancel, got %v", ctx)
		}
	}()

	// Cancel the context immediately
	w.cancel()

	wg.Wait()

	// Verify the wantConn state
	if !w.done {
		t.Error("wantConn should be marked as done after cancel")
	}
	if w.ctx != nil {
		t.Error("wantConn.ctx should be nil after cancel")
	}
}