github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/thinktime/think.go (about)

     1  package thinktime
     2  
     3  import (
     4  	"crypto/rand"
     5  	"fmt"
     6  	"math/big"
     7  	"regexp"
     8  	"strings"
     9  	"time"
    10  )
    11  
    12  type ThinkTime struct {
    13  	Min, Max time.Duration
    14  }
    15  
    16  func (t ThinkTime) Think(thinkNow bool) (thinkTime time.Duration) {
    17  	if t.Min == t.Max {
    18  		thinkTime = t.Min
    19  	} else {
    20  		thinkTime = time.Duration(RandInt(int64(t.Max-t.Min))) + t.Min
    21  	}
    22  
    23  	if thinkNow {
    24  		time.Sleep(thinkTime)
    25  	}
    26  
    27  	return thinkTime
    28  }
    29  
    30  func RandInt(n int64) int64 {
    31  	result, _ := rand.Int(rand.Reader, big.NewInt(n))
    32  	return result.Int64()
    33  }
    34  
    35  func ParseThinkTime(think string) (t *ThinkTime, err error) {
    36  	if think == "" {
    37  		return nil, nil
    38  	}
    39  
    40  	t = &ThinkTime{}
    41  
    42  	rangePos := strings.Index(think, "-")
    43  	if rangePos < 0 {
    44  		if t.Min, err = time.ParseDuration(think); err != nil {
    45  			return nil, err
    46  		}
    47  		t.Max = t.Min
    48  		return t, nil
    49  	}
    50  
    51  	minThink, maxThink := think[0:rangePos], think[rangePos+1:]
    52  	if t.Max, err = time.ParseDuration(maxThink); err != nil {
    53  		return nil, err
    54  	}
    55  
    56  	if minThink == "" {
    57  		return t, nil
    58  	}
    59  
    60  	if regexp.MustCompile(`^\d+$`).MatchString(minThink) {
    61  		minThink += FindUnit(maxThink)
    62  	}
    63  
    64  	if t.Min, err = time.ParseDuration(minThink); err != nil {
    65  		return nil, err
    66  	}
    67  
    68  	if t.Min > t.Max {
    69  		return nil, fmt.Errorf("min think time should be less than max")
    70  	}
    71  
    72  	return t, nil
    73  }
    74  
    75  func FindUnit(s string) string {
    76  	pos := strings.LastIndexFunc(s, func(r rune) bool {
    77  		return r >= '0' && r <= '9'
    78  	})
    79  
    80  	if pos < 0 {
    81  		return s
    82  	}
    83  
    84  	return s[pos+1:]
    85  }