Method AwaitAndTransition

internal/pool/conn_state.go:216–281

AwaitAndTransition waits for the connection to reach one of the valid states, then atomically transitions to the target state. Returns the current state after the transition attempt and an error if the operation failed. The returned state is the CURRENT state (after the attempt), not the previous st

(
	ctx context.Context,
	validFromStates []ConnState,
	targetState ConnState,
)

// - If already in a valid state, this is very fast (no allocation, no waiting)
// - If waiting is required, allocates one waiter struct and one channel
func (sm *ConnStateMachine) AwaitAndTransition(
	ctx context.Context,
	validFromStates []ConnState,
	targetState ConnState,
) (ConnState, error) {
	// Fast path: try immediate transition with CAS to prevent race conditions
	// BUT: only if there are no waiters in the queue (to maintain FIFO ordering)
	if sm.waiterCount.Load() == 0 {
		for _, fromState := range validFromStates {
			// Check if we're already in target state
			if fromState == targetState && sm.GetState() == targetState {
				return targetState, nil
			}

			// Try to atomically swap from fromState to targetState
			if sm.state.CompareAndSwap(uint32(fromState), uint32(targetState)) {
				// Success! We transitioned atomically
				sm.notifyWaiters()
				return targetState, nil
			}
		}
	}

	// Fast path failed - check if we should wait or fail
	currentState := sm.GetState()

	// Check if closed
	if currentState == StateClosed {
		return currentState, ErrStateMachineClosed
	}

	// Slow path: need to wait for state change
	// Create waiter with valid states map for fast lookup
	validStatesMap := make(map[ConnState]struct{}, len(validFromStates))
	for _, s := range validFromStates {
		validStatesMap[s] = struct{}{}
	}

	w := &waiter{
		validStates: validStatesMap,
		targetState: targetState,
		done:        make(chan error, 1), // Buffered to avoid goroutine leak
	}

	// Add to FIFO queue
	sm.mu.Lock()
	elem := sm.waiters.PushBack(w)
	sm.waiterCount.Add(1)
	sm.mu.Unlock()

	// Wait for state change or timeout
	select {
	case <-ctx.Done():
		// Timeout or cancellation - remove from queue
		sm.mu.Lock()
		sm.waiters.Remove(elem)
		sm.waiterCount.Add(-1)
		sm.mu.Unlock()