iop/apps/node/internal/node/run_manager.go

package node

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

// ErrConcurrencyLimitExceeded is returned when a run is rejected because
// admission could not be granted under the current concurrency policy.
var ErrConcurrencyLimitExceeded = errors.New("concurrency limit exceeded")

type runHandle struct {
	runID     string
	adapter   string
	target    string
	sessionID string
	cancel    context.CancelFunc
	done      chan struct{}
}

type runManager struct {
	mu   sync.RWMutex
	runs map[string]*runHandle
}

func newRunManager() *runManager {
	return &runManager{runs: make(map[string]*runHandle)}
}

func (m *runManager) register(h *runHandle) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.runs[h.runID] = h
}

func (m *runManager) deregister(runID string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.runs, runID)
}

func (m *runManager) cancelRun(runID string) bool {
	m.mu.RLock()
	h, ok := m.runs[runID]
	m.mu.RUnlock()
	if !ok {
		return false
	}
	h.cancel()
	return true
}

func (m *runManager) hasActiveRunsForAdapter(adapterKey string) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	for _, r := range m.runs {
		if r.adapter == adapterKey {
			return true
		}
	}
	return false
}

func (m *runManager) hasAnyActiveRuns() bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.runs) > 0
}

// snapshotActive returns a copy of all currently active run handles so callers
// can wait for pre-swap runs to drain without holding the manager lock.
func (m *runManager) snapshotActive() []*runHandle {
	m.mu.RLock()
	defer m.mu.RUnlock()
	out := make([]*runHandle, 0, len(m.runs))
	for _, h := range m.runs {
		out = append(out, h)
	}
	return out
}

// waitHandles blocks until every handle's run has finished or ctx is cancelled.
func waitHandles(ctx context.Context, handles []*runHandle) {
	for _, h := range handles {
		select {
		case <-h.done:
		case <-ctx.Done():
			return
		}
	}
}

// rejectReason classifies why a run could not be admitted.
type rejectReason string

const (
	// reasonConcurrencyUnavailable means admission failed because the concurrency limit was exceeded.
	reasonConcurrencyUnavailable rejectReason = "concurrency_unavailable"
)

// admissionError carries the concrete rejection reason while still satisfying
// errors.Is(err, ErrConcurrencyLimitExceeded) for callers relying on the legacy sentinel.
type admissionError struct {
	reason rejectReason
	msg    string
}

func (e *admissionError) Error() string { return e.msg }

func (e *admissionError) Is(target error) bool { return target == ErrConcurrencyLimitExceeded }

func newConcurrencyUnavailableError() *admissionError {
	return &admissionError{
		reason: reasonConcurrencyUnavailable,
		msg:    "concurrency unavailable: admission rejected because safety limit exceeded",
	}
}

// admissionRejectReason extracts the rejectReason from an admission error.
func admissionRejectReason(err error) rejectReason {
	var ae *admissionError
	if errors.As(err, &ae) {
		return ae.reason
	}
	return reasonConcurrencyUnavailable
}

// admissionMessage extracts the human-readable rejection message.
func admissionMessage(err error) string {
	var ae *admissionError
	if errors.As(err, &ae) {
		return ae.msg
	}
	if err != nil {
		return err.Error()
	}
	return string(reasonConcurrencyUnavailable)
}

// fifoGate is a simple in-flight concurrency semaphore (no queue).
type fifoGate struct {
	mu       sync.Mutex
	capacity int
	inFlight int
}

func newFifoGate(capacity int) *fifoGate {
	return &fifoGate{
		capacity: capacity,
	}
}

func (g *fifoGate) acquire() bool {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.capacity <= 0 || g.inFlight < g.capacity {
		g.inFlight++
		return true
	}
	return false
}

func (g *fifoGate) release() {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.inFlight > 0 {
		g.inFlight--
	}
}

func (g *fifoGate) activeCount() int {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.inFlight
}

func (g *fifoGate) updateCapacity(capacity int) {
	g.mu.Lock()
	defer g.mu.Unlock()
	g.capacity = capacity
}

// currentCapacity returns the current concurrency capacity (0 means unlimited).
func (g *fifoGate) currentCapacity() int {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.capacity
}

func (g *fifoGate) queuedCount() int {
	return 0
}

// admissionManager tracks the optional gates held by a run admission.
type admissionManager struct {
	global  *fifoGate
	adapter *fifoGate
}

type admissionTicket struct {
	mgr         *admissionManager
	adapterHeld bool
	globalHeld  bool
}

func (a *admissionManager) acquire() (*admissionTicket, error) {
	if a.adapter != nil {
		if !a.adapter.acquire() {
			return nil, newConcurrencyUnavailableError()
		}
	}
	if a.global != nil {
		if !a.global.acquire() {
			if a.adapter != nil {
				a.adapter.release()
			}
			return nil, newConcurrencyUnavailableError()
		}
	}
	return &admissionTicket{
		mgr:         a,
		adapterHeld: a.adapter != nil,
		globalHeld:  a.global != nil,
	}, nil
}

func (t *admissionTicket) release() {
	if t.globalHeld && t.mgr.global != nil {
		t.mgr.global.release()
		t.globalHeld = false
	}
	if t.adapterHeld && t.mgr.adapter != nil {
		t.mgr.adapter.release()
		t.adapterHeld = false
	}
}