// Package cmdlimit provides concurrent command limiting to prevent resource exhaustion. package cmdlimit import ( "context" "errors" "sync" "time" ) // ErrLimitExceeded is returned when the concurrent command limit is reached. var ErrLimitExceeded = errors.New("concurrent command limit exceeded") // Config defines the limiter configuration. type Config struct { // MaxConcurrentPerProject is the maximum concurrent commands per project. // Defaults to 5. MaxConcurrentPerProject int // MaxConcurrentTotal is the maximum concurrent commands across all projects. // Defaults to 20. MaxConcurrentTotal int // CommandTimeout is the maximum duration a command can hold a slot. // After this duration, the slot is automatically released. // Defaults to 30 minutes. CommandTimeout time.Duration } // DefaultConfig returns sensible defaults. func DefaultConfig() Config { return Config{ MaxConcurrentPerProject: 5, MaxConcurrentTotal: 20, CommandTimeout: 30 * time.Minute, } } // Limiter tracks and enforces concurrent command limits. type Limiter struct { cfg Config mu sync.Mutex projectCounts map[string]int totalCount int activeCommands map[string]*activeCommand } type activeCommand struct { projectID string startedAt time.Time cancel context.CancelFunc } // New creates a new concurrent command limiter. func New(cfg Config) *Limiter { if cfg.MaxConcurrentPerProject <= 0 { cfg.MaxConcurrentPerProject = 5 } if cfg.MaxConcurrentTotal <= 0 { cfg.MaxConcurrentTotal = 20 } if cfg.CommandTimeout <= 0 { cfg.CommandTimeout = 30 * time.Minute } return &Limiter{ cfg: cfg, projectCounts: make(map[string]int), activeCommands: make(map[string]*activeCommand), } } // Acquire attempts to acquire a command slot for the given project. // Returns a release function that MUST be called when the command completes. // Returns ErrLimitExceeded if the limit is reached. func (l *Limiter) Acquire(ctx context.Context, projectID, commandID string) (release func(), err error) { l.mu.Lock() defer l.mu.Unlock() // Check total limit if l.totalCount >= l.cfg.MaxConcurrentTotal { return nil, ErrLimitExceeded } // Check per-project limit if l.projectCounts[projectID] >= l.cfg.MaxConcurrentPerProject { return nil, ErrLimitExceeded } // Acquire the slot l.totalCount++ l.projectCounts[projectID]++ // Create a context with timeout for automatic release cmdCtx, cancel := context.WithTimeout(ctx, l.cfg.CommandTimeout) l.activeCommands[commandID] = &activeCommand{ projectID: projectID, startedAt: time.Now(), cancel: cancel, } // Start a goroutine to auto-release on timeout go func() { <-cmdCtx.Done() l.release(commandID) }() // Return release function return func() { cancel() l.release(commandID) }, nil } // release decrements the counters for a command. func (l *Limiter) release(commandID string) { l.mu.Lock() defer l.mu.Unlock() cmd, exists := l.activeCommands[commandID] if !exists { return // Already released } delete(l.activeCommands, commandID) l.totalCount-- l.projectCounts[cmd.projectID]-- if l.projectCounts[cmd.projectID] <= 0 { delete(l.projectCounts, cmd.projectID) } } // Stats returns current usage statistics. func (l *Limiter) Stats() Stats { l.mu.Lock() defer l.mu.Unlock() projectStats := make(map[string]int) for k, v := range l.projectCounts { projectStats[k] = v } return Stats{ TotalActive: l.totalCount, MaxTotal: l.cfg.MaxConcurrentTotal, ProjectCounts: projectStats, MaxPerProject: l.cfg.MaxConcurrentPerProject, ActiveCommandIDs: l.getActiveCommandIDs(), } } func (l *Limiter) getActiveCommandIDs() []string { ids := make([]string, 0, len(l.activeCommands)) for id := range l.activeCommands { ids = append(ids, id) } return ids } // Stats contains current limiter statistics. type Stats struct { TotalActive int MaxTotal int ProjectCounts map[string]int MaxPerProject int ActiveCommandIDs []string } // IsProjectAtLimit checks if a project has reached its limit. func (l *Limiter) IsProjectAtLimit(projectID string) bool { l.mu.Lock() defer l.mu.Unlock() return l.projectCounts[projectID] >= l.cfg.MaxConcurrentPerProject } // IsTotalAtLimit checks if the total limit has been reached. func (l *Limiter) IsTotalAtLimit() bool { l.mu.Lock() defer l.mu.Unlock() return l.totalCount >= l.cfg.MaxConcurrentTotal } // ActiveCount returns the number of active commands for a project. func (l *Limiter) ActiveCount(projectID string) int { l.mu.Lock() defer l.mu.Unlock() return l.projectCounts[projectID] } // TotalActiveCount returns the total number of active commands. func (l *Limiter) TotalActiveCount() int { l.mu.Lock() defer l.mu.Unlock() return l.totalCount }