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

     1  package augs
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/Rookout/GoSDK/pkg/rookoutErrors"
     8  	"github.com/Rookout/GoSDK/pkg/types"
     9  )
    10  
    11  type AugTimeLimiter struct {
    12  	disableAug bool
    13  	maxAugTime time.Duration
    14  
    15  	executionStartTimesLock sync.Mutex
    16  	executionStartTimes     map[string]time.Time
    17  }
    18  
    19  func NewAugTimeLimiter(maxAugTime time.Duration) *AugTimeLimiter {
    20  	return &AugTimeLimiter{
    21  		disableAug:          false,
    22  		maxAugTime:          maxAugTime,
    23  		executionStartTimes: make(map[string]time.Time),
    24  	}
    25  }
    26  
    27  func (t *AugTimeLimiter) BeforeRun(executionId string) (types.AugStatus, rookoutErrors.RookoutError) {
    28  	t.executionStartTimesLock.Lock()
    29  	defer t.executionStartTimesLock.Unlock()
    30  
    31  	if t.disableAug {
    32  		return types.Error, rookoutErrors.NewRookRuleMaxExecutionTimeReached()
    33  	}
    34  
    35  	t.executionStartTimes[executionId] = time.Now()
    36  	return types.Active, nil
    37  }
    38  
    39  func (t *AugTimeLimiter) CancelRun(executionId string) {
    40  	t.executionStartTimesLock.Lock()
    41  	defer t.executionStartTimesLock.Unlock()
    42  
    43  	delete(t.executionStartTimes, executionId)
    44  }
    45  
    46  func (t *AugTimeLimiter) AfterRun(executionId string) (types.AugStatus, rookoutErrors.RookoutError) {
    47  	t.executionStartTimesLock.Lock()
    48  	defer t.executionStartTimesLock.Unlock()
    49  
    50  	augTime := time.Since(t.executionStartTimes[executionId])
    51  	delete(t.executionStartTimes, executionId)
    52  
    53  	if t.maxAugTime > 0 && augTime > t.maxAugTime {
    54  		t.disableAug = true
    55  		return types.Error, rookoutErrors.NewRookRuleMaxExecutionTimeReached()
    56  	}
    57  
    58  	return types.Active, nil
    59  }