github.com/sunriselayer/sunrise-da@v0.13.1-sr3/cmd/cel-shed/eds_store_stress.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	_ "expvar"
     7  	"fmt"
     8  	"math"
     9  	"net/http"
    10  	"os"
    11  	"time"
    12  
    13  	"github.com/grafana/pyroscope-go"
    14  	logging "github.com/ipfs/go-log/v2"
    15  	"github.com/mitchellh/go-homedir"
    16  	"github.com/spf13/cobra"
    17  
    18  	"github.com/sunriselayer/sunrise-da/libs/edssser"
    19  	"github.com/sunriselayer/sunrise-da/nodebuilder"
    20  	"github.com/sunriselayer/sunrise-da/nodebuilder/node"
    21  )
    22  
    23  const (
    24  	edsStorePathFlag   = "path"
    25  	edsWritesFlag      = "writes"
    26  	edsSizeFlag        = "size"
    27  	edsDisableLogFlag  = "disable-log"
    28  	edsLogStatFreqFlag = "log-stat-freq"
    29  	edsCleanupFlag     = "cleanup"
    30  	edsFreshStartFlag  = "fresh"
    31  
    32  	pyroscopeEndpointFlag = "pyroscope"
    33  	putTimeoutFlag        = "timeout"
    34  	badgerLogLevelFlag    = "badger-log-level"
    35  )
    36  
    37  func init() {
    38  	edsStoreCmd.AddCommand(edsStoreStress)
    39  
    40  	defaultPath := "~/.edssser"
    41  	path, err := homedir.Expand(defaultPath)
    42  	if err != nil {
    43  		panic(err)
    44  	}
    45  
    46  	pathFlagUsage := fmt.Sprintf("Directory path to use for stress test. Uses %s by default.", defaultPath)
    47  	edsStoreStress.Flags().String(edsStorePathFlag, path, pathFlagUsage)
    48  	edsStoreStress.Flags().String(pyroscopeEndpointFlag, "",
    49  		"Pyroscope address. If no address provided, pyroscope will be disabled")
    50  	edsStoreStress.Flags().Int(edsWritesFlag, math.MaxInt, "Total EDS writes to make. MaxInt by default.")
    51  	edsStoreStress.Flags().Int(edsSizeFlag, 128, "Chooses EDS size. 128 by default.")
    52  	edsStoreStress.Flags().Bool(edsDisableLogFlag, false, "Disables logging. Enabled by default.")
    53  	edsStoreStress.Flags().Int(edsLogStatFreqFlag, 10, "Write statistic logging frequency. 10 by default.")
    54  	edsStoreStress.Flags().Bool(edsCleanupFlag, false, "Cleans up the store on stop. Disabled by default.")
    55  	edsStoreStress.Flags().Bool(edsFreshStartFlag, false, "Cleanup previous state on start. Disabled by default.")
    56  	edsStoreStress.Flags().Int(putTimeoutFlag, 30, "Sets put timeout in seconds. 30 sec by default.")
    57  	edsStoreStress.Flags().String(badgerLogLevelFlag, "INFO", "Badger log level, Defaults to INFO")
    58  
    59  	// kill redundant print
    60  	nodebuilder.PrintKeyringInfo = false
    61  }
    62  
    63  var edsStoreCmd = &cobra.Command{
    64  	Use:   "eds-store [subcommand]",
    65  	Short: "Collection of eds-store related utilities",
    66  }
    67  
    68  var edsStoreStress = &cobra.Command{
    69  	Use:          "stress",
    70  	Short:        `Runs eds.Store stress test over default node.Store Datastore backend (e.g. Badger).`,
    71  	SilenceUsage: true,
    72  	RunE: func(cmd *cobra.Command, args []string) (err error) {
    73  		// expose expvar vars over http
    74  		go http.ListenAndServe(":9999", http.DefaultServeMux) //nolint:errcheck,gosec
    75  
    76  		endpoint, _ := cmd.Flags().GetString(pyroscopeEndpointFlag)
    77  		if endpoint != "" {
    78  			_, err = pyroscope.Start(pyroscope.Config{
    79  				ApplicationName: "cel-shred.stresser",
    80  				ServerAddress:   endpoint,
    81  				ProfileTypes: []pyroscope.ProfileType{
    82  					pyroscope.ProfileCPU,
    83  					pyroscope.ProfileAllocObjects,
    84  					pyroscope.ProfileAllocSpace,
    85  					pyroscope.ProfileInuseObjects,
    86  					pyroscope.ProfileInuseSpace,
    87  				},
    88  			})
    89  			if err != nil {
    90  				fmt.Printf("failed to launch pyroscope with addr: %s err: %s\n", endpoint, err.Error())
    91  			} else {
    92  				fmt.Println("connected pyroscope to:", endpoint)
    93  			}
    94  		}
    95  
    96  		path, _ := cmd.Flags().GetString(edsStorePathFlag)
    97  		fmt.Printf("using %s\n", path)
    98  
    99  		freshStart, _ := cmd.Flags().GetBool(edsFreshStartFlag)
   100  		if freshStart {
   101  			err = os.RemoveAll(path)
   102  			if err != nil {
   103  				return err
   104  			}
   105  		}
   106  
   107  		cleanup, _ := cmd.Flags().GetBool(edsCleanupFlag)
   108  		if cleanup {
   109  			defer func() {
   110  				err = errors.Join(err, os.RemoveAll(path))
   111  			}()
   112  		}
   113  
   114  		loglevel, _ := cmd.Flags().GetString(badgerLogLevelFlag)
   115  		if err = logging.SetLogLevel("badger", loglevel); err != nil {
   116  			return err
   117  		}
   118  
   119  		disableLog, _ := cmd.Flags().GetBool(edsDisableLogFlag)
   120  		logFreq, _ := cmd.Flags().GetInt(edsLogStatFreqFlag)
   121  		edsWrites, _ := cmd.Flags().GetInt(edsWritesFlag)
   122  		edsSize, _ := cmd.Flags().GetInt(edsSizeFlag)
   123  		putTimeout, _ := cmd.Flags().GetInt(putTimeoutFlag)
   124  
   125  		cfg := edssser.Config{
   126  			EDSSize:     edsSize,
   127  			EDSWrites:   edsWrites,
   128  			EnableLog:   !disableLog,
   129  			LogFilePath: path,
   130  			StatLogFreq: logFreq,
   131  			OpTimeout:   time.Duration(putTimeout) * time.Second,
   132  		}
   133  
   134  		err = nodebuilder.Init(*nodebuilder.DefaultConfig(node.Full), path, node.Full)
   135  		if err != nil {
   136  			return err
   137  		}
   138  
   139  		nodestore, err := nodebuilder.OpenStore(path, nil)
   140  		if err != nil {
   141  			return err
   142  		}
   143  		defer func() {
   144  			err = errors.Join(err, nodestore.Close())
   145  		}()
   146  
   147  		datastore, err := nodestore.Datastore()
   148  		if err != nil {
   149  			return err
   150  		}
   151  
   152  		stresser, err := edssser.NewEDSsser(path, datastore, cfg)
   153  		if err != nil {
   154  			return err
   155  		}
   156  
   157  		stats, err := stresser.Run(cmd.Context())
   158  		if !errors.Is(err, context.Canceled) {
   159  			return err
   160  		}
   161  
   162  		fmt.Printf("%s", stats.Finalize())
   163  		return nil
   164  	},
   165  }