github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/runtime/profiling.go (about)

     1  package runtime
     2  
     3  import (
     4  	"runtime"
     5  
     6  	"github.com/jzelinskie/cobrautil/v2"
     7  	"github.com/spf13/cobra"
     8  	"github.com/spf13/pflag"
     9  )
    10  
    11  // RegisterFlags adds flags for configuring profiling rates.
    12  //
    13  // The following flags are added:
    14  // - "pprof-mutex-profile-rate"
    15  // - "pprof-block-profile-rate"
    16  func RegisterFlags(flags *pflag.FlagSet) {
    17  	flags.Int("pprof-mutex-profile-rate", 0, "sets the mutex profile sampling rate")
    18  	flags.Int("pprof-block-profile-rate", 0, "sets the block profile sampling rate")
    19  }
    20  
    21  // RunE returns a Cobra RunFunc that configures mutex and block profiles.
    22  //
    23  // The required flags can be added to a command by using RegisterFlags().
    24  func RunE() cobrautil.CobraRunFunc {
    25  	return func(cmd *cobra.Command, args []string) error {
    26  		if cobrautil.IsBuiltinCommand(cmd) {
    27  			return nil // No-op for builtins
    28  		}
    29  
    30  		runtime.SetMutexProfileFraction(cobrautil.MustGetInt(cmd, "pprof-mutex-profile-rate"))
    31  		runtime.SetBlockProfileRate(cobrautil.MustGetInt(cmd, "pprof-block-profile-rate"))
    32  		return nil
    33  	}
    34  }