go.mondoo.com/cnquery@v0.0.0-20231005093811-59568235f6ea/cli/prof/prof.go (about)

     1  // Copyright (c) Mondoo, Inc.
     2  // SPDX-License-Identifier: BUSL-1.1
     3  
     4  // Package prof is responsible for setting up the go profiler for commands
     5  package prof
     6  
     7  import (
     8  	"net/http"
     9  	"net/http/pprof"
    10  	"os"
    11  	"runtime"
    12  	"strconv"
    13  	"strings"
    14  
    15  	"github.com/cockroachdb/errors"
    16  	"github.com/rs/zerolog/log"
    17  )
    18  
    19  // InitProfiler sets up the go profiler based on the MONDOO_PROF environment
    20  // variable.
    21  // MONDO_PROF is a list of comma separated key/key=value.
    22  // Allowed keys:
    23  //   - `enable`:          Enables the profiler if no value is provided, or the value of
    24  //     `true` is provided
    25  //
    26  // - `enabled`:        Alias for `enable`
    27  //
    28  //   - `listen`:         Sets the listen address for the profiler http server. See
    29  //     https://golang.org/pkg/net/http/pprof for more info about the
    30  //     endpoints provided
    31  //
    32  // - `memprofilerate`: Sets runtime.MemProfileRate to the provided value
    33  //
    34  // Example:
    35  // MONDOO_PROF='enable,listen=localhost:7474,memprofilerate=1'
    36  func InitProfiler() {
    37  	if profVal := os.Getenv("MONDOO_PROF"); profVal != "" {
    38  		opts, err := parseProf(profVal)
    39  		if err != nil {
    40  			log.Warn().Err(err).Msg("failed to parse MONDOO_PROF")
    41  			return
    42  		}
    43  		setupProfiler(opts)
    44  	}
    45  }
    46  
    47  type profilerOpts struct {
    48  	Enabled        bool
    49  	Listen         string
    50  	MemProfileRate *int
    51  }
    52  
    53  var defaultOpts = profilerOpts{
    54  	Enabled:        false,
    55  	Listen:         "localhost:6060",
    56  	MemProfileRate: nil,
    57  }
    58  
    59  func parseProf(profVal string) (profilerOpts, error) {
    60  	opts := defaultOpts
    61  
    62  	sOpts := strings.Split(profVal, ",")
    63  	for _, sOpt := range sOpts {
    64  		keyval := strings.SplitN(sOpt, "=", 2)
    65  		key := ""
    66  		val := ""
    67  
    68  		if len(keyval) == 0 {
    69  			continue
    70  		}
    71  
    72  		key = strings.TrimSpace(keyval[0])
    73  		if len(keyval) == 2 {
    74  			val = strings.TrimSpace(keyval[1])
    75  		}
    76  
    77  		switch key {
    78  		case "enable", "enabled":
    79  			opts.Enabled = val == "" || val == "true"
    80  		case "listen":
    81  			if val != "" {
    82  				opts.Listen = val
    83  			}
    84  		case "memprofilerate":
    85  			if val != "" {
    86  				i, err := strconv.Atoi(val)
    87  				if err != nil {
    88  					return opts, errors.Wrapf(err, "invalid value %q for memprofilerate", val)
    89  				}
    90  				opts.MemProfileRate = &i
    91  			}
    92  		}
    93  	}
    94  	return opts, nil
    95  }
    96  
    97  func setupProfiler(opts profilerOpts) {
    98  	if !opts.Enabled {
    99  		return
   100  	}
   101  
   102  	log.Info().Interface("opts", opts).Msg("Enabling profiler")
   103  
   104  	if opts.MemProfileRate != nil {
   105  		runtime.MemProfileRate = *opts.MemProfileRate
   106  	}
   107  
   108  	mux := http.NewServeMux()
   109  	mux.HandleFunc("/debug/pprof/", pprof.Index)
   110  	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
   111  	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
   112  	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
   113  	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
   114  
   115  	go func() {
   116  		http.ListenAndServe(opts.Listen, mux)
   117  	}()
   118  }