github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/cmd/pebble/random.go (about)

     1  // Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
     2  // of this source code is governed by a BSD-style license that can be found in
     3  // the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"strconv"
     9  	"strings"
    10  	"time"
    11  
    12  	"github.com/cockroachdb/errors"
    13  	"github.com/cockroachdb/pebble/internal/randvar"
    14  	"github.com/cockroachdb/pebble/internal/rate"
    15  )
    16  
    17  type rateFlag struct {
    18  	randvar.Flag
    19  	fluctuateDuration time.Duration
    20  	spec              string
    21  }
    22  
    23  func newRateFlag(spec string) *rateFlag {
    24  	f := &rateFlag{}
    25  	if err := f.Set(spec); err != nil {
    26  		panic(err)
    27  	}
    28  	return f
    29  }
    30  
    31  func (f *rateFlag) String() string {
    32  	return f.spec
    33  }
    34  
    35  // Type implements the Flag.Value interface.
    36  func (f *rateFlag) Type() string {
    37  	return "ratevar"
    38  }
    39  
    40  // Set implements the Flag.Value interface.
    41  func (f *rateFlag) Set(spec string) error {
    42  	if spec == "" {
    43  		if err := f.Flag.Set("0"); err != nil {
    44  			return err
    45  		}
    46  		f.fluctuateDuration = time.Duration(0)
    47  		f.spec = spec
    48  		return nil
    49  	}
    50  
    51  	parts := strings.Split(spec, "/")
    52  	if len(parts) == 0 || len(parts) > 2 {
    53  		return errors.Errorf("invalid ratevar spec: %s", errors.Safe(spec))
    54  	}
    55  	if err := f.Flag.Set(parts[0]); err != nil {
    56  		return err
    57  	}
    58  	// Don't fluctuate by default.
    59  	f.fluctuateDuration = time.Duration(0)
    60  	if len(parts) == 2 {
    61  		fluctuateDurationFloat, err := strconv.ParseFloat(parts[1], 64)
    62  		if err != nil {
    63  			return err
    64  		}
    65  		f.fluctuateDuration = time.Duration(fluctuateDurationFloat) * time.Second
    66  	}
    67  	f.spec = spec
    68  	return nil
    69  }
    70  
    71  func (f *rateFlag) newRateLimiter() *rate.Limiter {
    72  	if f.spec == "" {
    73  		return nil
    74  	}
    75  	rng := randvar.NewRand()
    76  	limiter := rate.NewLimiter(float64(f.Uint64(rng)), 1)
    77  	if f.fluctuateDuration != 0 {
    78  		go func(limiter *rate.Limiter) {
    79  			ticker := time.NewTicker(f.fluctuateDuration)
    80  			for range ticker.C {
    81  				limiter.SetRate(float64(f.Uint64(rng)))
    82  			}
    83  		}(limiter)
    84  	}
    85  	return limiter
    86  }
    87  
    88  func wait(l *rate.Limiter) {
    89  	if l != nil {
    90  		l.Wait(1)
    91  	}
    92  }