iop/packages/go/streamgate/commit_boundary.go

package streamgate

import (
	"context"
	"errors"
	"fmt"
	"sync"
)

// ErrBoundaryUncommitted is returned when an operation requires the
// boundary to be in transport_uncommitted state but it is not.
var ErrBoundaryUncommitted = errors.New("streamgate: commit boundary not in uncommitted state")

// ErrBoundaryAlreadyTerminal is returned when a terminal has already been
// committed.
var ErrBoundaryAlreadyTerminal = errors.New("streamgate: commit boundary already terminal")

// ErrBoundaryNoStagedStart is returned when a release is attempted without
// a staged response start.
var ErrBoundaryNoStagedStart = errors.New("streamgate: no staged response start")

// ErrBoundaryDuplicateAttempt is returned when BeginAttempt is called for an
// attempt already in progress.
var ErrBoundaryDuplicateAttempt = errors.New("streamgate: duplicate attempt id")

// ErrBoundaryInvalidAttempt is returned when the given attempt ID does not
// match the current active attempt.
var ErrBoundaryInvalidAttempt = errors.New("streamgate: attempt id mismatch")

// ErrBoundaryNoCurrentAttempt is returned when no attempt has been begun.
var ErrBoundaryNoCurrentAttempt = errors.New("streamgate: no current attempt")

// ErrBoundaryReplaced is returned when an operation is attempted with an
// old attempt ID after the boundary has been replaced.
var ErrBoundaryReplaced = errors.New("streamgate: boundary replaced by newer attempt")

// ErrBoundaryRecoveryNotPermitted is returned when the current state does
// not allow the requested recovery strategy.
var ErrBoundaryRecoveryNotPermitted = errors.New("streamgate: recovery strategy not permitted for current state")

// ErrBoundarySerializerConflict is returned when another operation on the boundary
// is currently in progress.
var ErrBoundarySerializerConflict = errors.New("streamgate: boundary operation in progress")

// ErrInvalidStateTransition is returned when a sink returns an invalid state transition.
var ErrInvalidStateTransition = errors.New("streamgate: invalid commit state transition")

// ErrBoundaryInvalidCallbackOp is returned when the callback operation does
// not match the reservation's expected operation type.
var ErrBoundaryInvalidCallbackOp = errors.New("streamgate: callback operation does not match reservation")

// ReleaseProgress records the number of events successfully released through
// the boundary during a ReleaseSafe call. It is immutable to callers.
type ReleaseProgress struct {
	releasedEvents int
}

// ReleasedEvents returns the count of events that the downstream sink
// accepted during this release.
func (p ReleaseProgress) ReleasedEvents() int { return p.releasedEvents }

type boundaryOpType int

const (
	boundaryOpNone boundaryOpType = iota
	boundaryOpRelease
	boundaryOpTerminal
)

type reservation struct {
	op         boundaryOpType
	attemptID  string
	generation uint64
}

// CommitBoundary owns the request-scope state machine for committing release
// payloads to the downstream sink. It enforces staged response-start,
// monotonic state transitions, partial progress tracking, and exactly-once
// terminal semantics.
type CommitBoundary struct {
	mu                sync.Mutex
	sink              ReleaseSink
	attemptID         string
	stagedStart       *ResponseStart
	state             CommitState
	confirmed         []ReleaseEvent
	generation        uint64
	activeReservation *reservation
}

// NewCommitBoundary creates a new CommitBoundary bound to the given sink.
// The boundary starts in transport_uncommitted state.
func NewCommitBoundary(sink ReleaseSink) (*CommitBoundary, error) {
	if sink == nil {
		return nil, errors.New("streamgate: commit boundary sink is required")
	}
	return &CommitBoundary{
		sink:  sink,
		state: CommitStateTransportUncommitted,
	}, nil
}

func (b *CommitBoundary) reserveLocked(op boundaryOpType, attemptID string) (*reservation, error) {
	if b.activeReservation != nil {
		return nil, ErrBoundarySerializerConflict
	}
	b.generation++
	res := &reservation{
		op:         op,
		attemptID:  attemptID,
		generation: b.generation,
	}
	b.activeReservation = res
	return res, nil
}

func (b *CommitBoundary) releaseReservationLocked(res *reservation) {
	if b.activeReservation == res {
		b.activeReservation = nil
	}
}

func validateStateTransition(current, returned CommitState) error {
	if returned != "" {
		if err := returned.Validate(); err != nil {
			return fmt.Errorf("%w: %v", ErrInvalidStateTransition, err)
		}
	}
	if current == returned || returned == "" {
		return nil
	}

	switch current {
	case CommitStateTransportUncommitted:
		return nil
	case CommitStateStreamOpen:
		if returned == CommitStateTerminalCommitted {
			return nil
		}
		return fmt.Errorf("%w: cannot transition from stream_open to %v", ErrInvalidStateTransition, returned)
	case CommitStateTerminalCommitted:
		return fmt.Errorf("%w: cannot transition from terminal_committed to %v", ErrInvalidStateTransition, returned)
	default:
		return fmt.Errorf("%w: invalid current state %v", ErrInvalidStateTransition, current)
	}
}

func isStateMoreAdvanced(a, b CommitState) bool {
	if a == b {
		return false
	}
	if a == CommitStateTerminalCommitted {
		return true
	}
	if a == CommitStateStreamOpen && b == CommitStateTransportUncommitted {
		return true
	}
	return false
}

// applySinkOutcomeLockedWithOpLocked validates the sink outcome against the
// reservation operation type. It preserves monotonic state transitions
// independent of sink error or role mismatch, requires role-specific success
// states (stream_open for response-start and release, terminal_committed for
// terminal), and short-circuits on callback error for response-start and
// release operations.
func (b *CommitBoundary) applySinkOutcomeLockedWithOpLocked(
	res *reservation,
	op boundaryOpType,
	returnedState CommitState,
	sinkErr error,
) error {
	if b.activeReservation != res || res.generation != b.generation {
		return ErrBoundarySerializerConflict
	}

	// 1. Validate attempt identity and callback operation match the reservation.
	//    This is defense-in-depth: public entry points (ReleaseSafe, CommitTerminal)
	//    check attemptID before reserving, but the callback path must also reject
	//    a mismatch in case b.attemptID is changed between reserve and callback.
	identityErr := errors.Join(
		validateAttemptIdentity(res.attemptID, b.attemptID),
		validateReservationCallbackOperation(res.op, op),
	)

	// 2. Validate state transition (independent of sink error, role, and identity).
	transitionErr := validateStateTransition(b.state, returnedState)

	// 3. Apply valid forward state regardless of sink error, role error, or identity error.
	//    This ensures that a callback returning a more advanced state (e.g.
	//    terminal_committed) is trusted even when the role does not match,
	//    a sink error is present, or the callback operation mismatches the reservation.
	if returnedState != "" && transitionErr == nil && isStateMoreAdvanced(returnedState, b.state) {
		b.state = returnedState
	}

	// 4. Clear staged start when state advances out of transport_uncommitted.
	if b.state == CommitStateStreamOpen || b.state == CommitStateTerminalCommitted {
		b.stagedStart = nil
	}

	// 5. Validate role-specific expected state.
	var roleErr error
	if returnedState != "" {
		switch op {
		case boundaryOpTerminal:
			if returnedState != CommitStateTerminalCommitted {
				roleErr = fmt.Errorf("%w: terminal callback returned %q, want %q",
					ErrInvalidStateTransition, returnedState, CommitStateTerminalCommitted)
			}
		case boundaryOpRelease, boundaryOpNone:
			if returnedState != CommitStateStreamOpen {
				roleErr = fmt.Errorf("%w: %s callback returned %q, want %q",
					ErrInvalidStateTransition, opLabel(op), returnedState, CommitStateStreamOpen)
			}
		}
	}

	// 6. Empty state on success is a transition error.
	if returnedState == "" && sinkErr == nil {
		transitionErr = fmt.Errorf("%w: empty commit state returned on success", ErrInvalidStateTransition)
	}

	// 7. Join all errors: sink, transition, identity, role.
	return errors.Join(sinkErr, transitionErr, identityErr, roleErr)
}

// opLabel returns a human-readable label for the operation type.
func opLabel(op boundaryOpType) string {
	switch op {
	case boundaryOpRelease:
		return "release"
	case boundaryOpTerminal:
		return "terminal"
	default:
		return "response-start"
	}
}

// validateAttemptIdentity returns ErrBoundaryInvalidAttempt when the
// reservation's attemptID does not match the boundary's current attemptID.
// This is defense-in-depth for the callback path.
func validateAttemptIdentity(resAttemptID, boundAttemptID string) error {
	if resAttemptID != boundAttemptID {
		return ErrBoundaryInvalidAttempt
	}
	return nil
}

// validateReservationCallbackOperation returns ErrBoundaryInvalidCallbackOp
// when the callback operation is not permitted by the reservation's expected
// operation type. A release reservation accepts response-start and release;
// a terminal reservation accepts response-start (optional) and terminal.
func validateReservationCallbackOperation(reservedOp, callbackOp boundaryOpType) error {
	switch reservedOp {
	case boundaryOpRelease:
		if callbackOp != boundaryOpNone && callbackOp != boundaryOpRelease {
			return fmt.Errorf("%w: reservation expects release, got %s",
				ErrBoundaryInvalidCallbackOp, opLabel(callbackOp))
		}
	case boundaryOpTerminal:
		if callbackOp != boundaryOpNone && callbackOp != boundaryOpTerminal {
			return fmt.Errorf("%w: reservation expects terminal, got %s",
				ErrBoundaryInvalidCallbackOp, opLabel(callbackOp))
		}
	}
	return nil
}

// BeginAttempt opens a new attempt identified by attemptID.
func (b *CommitBoundary) BeginAttempt(attemptID string) error {
	if attemptID == "" {
		return errors.New("streamgate: begin attempt id is required")
	}
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.activeReservation != nil {
		return ErrBoundarySerializerConflict
	}
	if b.state != CommitStateTransportUncommitted {
		return ErrBoundaryUncommitted
	}
	if b.attemptID != "" {
		return ErrBoundaryDuplicateAttempt
	}
	b.attemptID = attemptID
	b.stagedStart = nil
	b.confirmed = nil
	b.generation++
	return nil
}

// StageResponseStart buffers a response start for the given attempt.
func (b *CommitBoundary) StageResponseStart(attemptID string, start ResponseStart) error {
	if attemptID == "" {
		return errors.New("streamgate: stage response start attempt id is required")
	}
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.activeReservation != nil {
		return ErrBoundarySerializerConflict
	}
	if b.state != CommitStateTransportUncommitted {
		return ErrBoundaryUncommitted
	}
	if b.attemptID == "" {
		return ErrBoundaryNoCurrentAttempt
	}
	if b.attemptID != attemptID {
		return ErrBoundaryInvalidAttempt
	}
	if b.stagedStart != nil {
		return errors.New("streamgate: response start already staged")
	}
	b.stagedStart = &start
	return nil
}

// ReleaseSafe commits the staged response start (if any) exactly once, then
// releases each event in order.
func (b *CommitBoundary) ReleaseSafe(ctx context.Context, attemptID string, events []ReleaseEvent) (ReleaseProgress, error) {
	if attemptID == "" {
		return ReleaseProgress{}, errors.New("streamgate: release safe attempt id is required")
	}
	if len(events) == 0 {
		return ReleaseProgress{}, errors.New("streamgate: release safe requires at least one event")
	}

	b.mu.Lock()
	if b.state == CommitStateTerminalCommitted {
		b.mu.Unlock()
		return ReleaseProgress{}, ErrBoundaryAlreadyTerminal
	}
	if b.attemptID == "" {
		b.mu.Unlock()
		return ReleaseProgress{}, ErrBoundaryNoCurrentAttempt
	}
	if b.attemptID != attemptID {
		b.mu.Unlock()
		return ReleaseProgress{}, ErrBoundaryInvalidAttempt
	}

	res, err := b.reserveLocked(boundaryOpRelease, attemptID)
	if err != nil {
		b.mu.Unlock()
		return ReleaseProgress{}, err
	}
	defer func() {
		b.mu.Lock()
		b.releaseReservationLocked(res)
		b.mu.Unlock()
	}()

	var rsToCommit *ResponseStart
	if b.stagedStart != nil {
		rs := *b.stagedStart
		rsToCommit = &rs
	}

	eventsCopy := make([]ReleaseEvent, len(events))
	copy(eventsCopy, events)
	b.mu.Unlock()

	var firstErr error

	if rsToCommit != nil {
		st, sErr := b.sink.CommitResponseStart(ctx, *rsToCommit)
		b.mu.Lock()
		applyErr := b.applySinkOutcomeLockedWithOpLocked(res, boundaryOpNone, st, sErr)
		b.mu.Unlock()
		if applyErr != nil {
			firstErr = applyErr
			// Short-circuit: response-start callback error prevents proceeding
			// to release events. State is already applied by applySinkOutcomeLocked.
			return ReleaseProgress{}, firstErr
		}
	}

	if firstErr != nil && b.State() == CommitStateTransportUncommitted {
		return ReleaseProgress{}, firstErr
	}

	var releaseDone int
	for _, ev := range eventsCopy {
		st, sErr := b.sink.Release(ctx, ev)
		b.mu.Lock()
		applyErr := b.applySinkOutcomeLockedWithOpLocked(res, boundaryOpRelease, st, sErr)
		if applyErr == nil {
			b.confirmed = append(b.confirmed, ev)
			releaseDone++
		} else if firstErr == nil {
			firstErr = applyErr
		}
		b.mu.Unlock()
		if sErr != nil || applyErr != nil {
			break
		}
	}

	return ReleaseProgress{releasedEvents: releaseDone}, firstErr
}

// CommitTerminal commits the terminal result.
func (b *CommitBoundary) CommitTerminal(ctx context.Context, attemptID string, result TerminalResult) error {
	if attemptID == "" {
		return errors.New("streamgate: commit terminal attempt id is required")
	}

	b.mu.Lock()
	if b.state == CommitStateTerminalCommitted {
		b.mu.Unlock()
		return ErrBoundaryAlreadyTerminal
	}
	if b.attemptID == "" {
		b.mu.Unlock()
		return ErrBoundaryNoCurrentAttempt
	}
	if b.attemptID != attemptID {
		b.mu.Unlock()
		return ErrBoundaryInvalidAttempt
	}

	res, err := b.reserveLocked(boundaryOpTerminal, attemptID)
	if err != nil {
		b.mu.Unlock()
		return err
	}
	defer func() {
		b.mu.Lock()
		b.releaseReservationLocked(res)
		b.mu.Unlock()
	}()

	isError := result.Error()
	var rsCopy *ResponseStart
	if b.stagedStart != nil && !isError {
		cp := *b.stagedStart
		rsCopy = &cp
	}
	b.mu.Unlock()

	var firstErr error
	if rsCopy != nil {
		st, sErr := b.sink.CommitResponseStart(ctx, *rsCopy)
		b.mu.Lock()
		applyErr := b.applySinkOutcomeLockedWithOpLocked(res, boundaryOpNone, st, sErr)
		b.mu.Unlock()
		if applyErr != nil {
			firstErr = applyErr
			// Short-circuit: response-start callback error prevents proceeding
			// to terminal commit. State is already applied.
			return firstErr
		}
	}

	st, sErr := b.sink.CommitTerminal(ctx, result)
	b.mu.Lock()
	termErr := b.applySinkOutcomeLockedWithOpLocked(res, boundaryOpTerminal, st, sErr)
	b.mu.Unlock()

	if termErr != nil && firstErr == nil {
		firstErr = termErr
	} else if termErr != nil {
		firstErr = errors.Join(firstErr, termErr)
	}

	return firstErr
}

// ReplaceUncommittedAttempt discards the current attempt and its staged start.
func (b *CommitBoundary) ReplaceUncommittedAttempt(oldAttemptID, newAttemptID string) error {
	if newAttemptID == "" {
		return errors.New("streamgate: replace attempt new id is required")
	}
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.activeReservation != nil {
		return ErrBoundarySerializerConflict
	}
	if b.state == CommitStateTerminalCommitted {
		return ErrBoundaryAlreadyTerminal
	}
	if b.state != CommitStateTransportUncommitted {
		return ErrBoundaryUncommitted
	}
	if b.attemptID == "" {
		return ErrBoundaryNoCurrentAttempt
	}
	if oldAttemptID == "" || b.attemptID != oldAttemptID {
		return ErrBoundaryInvalidAttempt
	}

	b.attemptID = newAttemptID
	b.stagedStart = nil
	b.confirmed = nil
	b.generation++
	return nil
}

// CommitAllowsStrategy returns true when the current state allows the given recovery strategy.
func (b *CommitBoundary) CommitAllowsStrategy(strategy RecoveryStrategy) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	switch strategy {
	case RecoveryStrategyExactReplay, RecoveryStrategySchemaRepair:
		return b.state == CommitStateTransportUncommitted
	case RecoveryStrategyContinuationRepair:
		return b.state == CommitStateStreamOpen
	default:
		return false
	}
}

// State returns the current commit state.
func (b *CommitBoundary) State() CommitState {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.state
}

// CurrentAttempt returns the current attempt ID, or "" if none.
func (b *CommitBoundary) CurrentAttempt() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.attemptID
}