github.com/Rookout/GoSDK@v0.1.48/pkg/augs/aug_limits_manager.go (about)

     1  package augs
     2  
     3  import (
     4  	"github.com/Rookout/GoSDK/pkg/com_ws"
     5  	"github.com/Rookout/GoSDK/pkg/rookoutErrors"
     6  	"github.com/Rookout/GoSDK/pkg/types"
     7  )
     8  
     9  type LimitsManager interface {
    10  	AddLimiter(newLimiter Limiter)
    11  	GetAllLimiters() []Limiter
    12  	BeforeRun(executionId string, skipLimiters bool) bool
    13  	AfterRun(executionId string)
    14  }
    15  
    16  type Limiter interface {
    17  	BeforeRun(executionId string) (types.AugStatus, rookoutErrors.RookoutError)
    18  	CancelRun(executionId string)
    19  	AfterRun(executionId string) (types.AugStatus, rookoutErrors.RookoutError)
    20  }
    21  
    22  type limitsManager struct {
    23  	limiters  []Limiter
    24  	augID     types.AugID
    25  	augStatus types.AugStatus
    26  	output    com_ws.Output
    27  }
    28  
    29  func NewLimitsManager(augID types.AugID, output com_ws.Output) LimitsManager {
    30  	return &limitsManager{augID: augID, output: output, augStatus: types.Active}
    31  }
    32  
    33  func (l *limitsManager) AddLimiter(newLimiter Limiter) {
    34  	l.limiters = append(l.limiters, newLimiter)
    35  }
    36  
    37  func (l *limitsManager) GetAllLimiters() []Limiter {
    38  	return l.limiters
    39  }
    40  
    41  func (l *limitsManager) cancelAllLimiters(executionId string) {
    42  	for _, limiter := range l.limiters {
    43  		limiter.CancelRun(executionId)
    44  	}
    45  }
    46  
    47  func (l *limitsManager) setupAllLimiters(executionId string, skipLimiters bool) bool {
    48  	for _, limiter := range l.limiters {
    49  		status, err := limiter.BeforeRun(executionId)
    50  		if status == types.Active {
    51  			continue
    52  		}
    53  
    54  		if l.augStatus == types.Active {
    55  			l.augStatus = status
    56  			_ = l.output.SendRuleStatus(l.augID, l.augStatus, err)
    57  		}
    58  
    59  		if !skipLimiters {
    60  			return false
    61  		}
    62  	}
    63  
    64  	if l.augStatus != types.Active && !skipLimiters {
    65  		l.augStatus = types.Active
    66  		_ = l.output.SendRuleStatus(l.augID, l.augStatus, nil)
    67  	}
    68  
    69  	return true
    70  }
    71  
    72  func (l *limitsManager) BeforeRun(executionId string, skipLimiters bool) bool {
    73  	if ok := l.setupAllLimiters(executionId, skipLimiters); !ok {
    74  		l.cancelAllLimiters(executionId)
    75  		return false
    76  	}
    77  
    78  	return true
    79  }
    80  
    81  func (l *limitsManager) AfterRun(executionId string) {
    82  	for _, limiter := range l.limiters {
    83  		status, err := limiter.AfterRun(executionId)
    84  		if status == types.Active {
    85  			continue
    86  		}
    87  
    88  		if l.augStatus == types.Active {
    89  			l.augStatus = status
    90  			_ = l.output.SendRuleStatus(l.augID, l.augStatus, err)
    91  		}
    92  	}
    93  }