Collect cpu and memory metrics

cmd/logfile/logfile.go

@@ -19,6 +19,7 @@ const (
 // Params contains log file parameters.
 type Params struct {
 	File              string
+	Metrics           bool
 	SendDiagsOnErrors bool
 	ToStdout          bool
 	Verbose           bool
@@ -27,6 +28,11 @@ type Params struct {
 // LoadParams loads needed data from the configuration file.
 func LoadParams(v *viper.Viper) (Params, error) {
 	params := Params{
+		Metrics: vipertools.FirstNonEmptyBool(
+			v,
+			"metrics",
+			"settings.metrics",
+		),
 		SendDiagsOnErrors: vipertools.FirstNonEmptyBool(
 			v,
 			"send-diagnostics-on-errors",
@@ -40,6 +46,11 @@ func LoadParams(v *viper.Viper) (Params, error) {
 		),
 	}
 
+	// if debug is disabled, disable metrics as well.
+	if !params.Verbose {
+		params.Metrics = false
+	}
+
 	logFile, ok := vipertools.FirstNonEmptyString(v, "log-file", "logfile", "settings.log_file")
 	if ok {
 		p, err := homedir.Expand(logFile)

cmd/logfile/logfile_test.go

@@ -29,23 +29,25 @@ func TestLoadParams(t *testing.T) {
 	require.NoError(t, err)
 
 	tests := map[string]struct {
+		EnvVar             string
+		ViperDebug         bool
+		ViperDebugConfig   bool
 		ViperLogFile       string
 		ViperLogFileConfig string
 		ViperLogFileOld    string
+		ViperMetrics       bool
+		ViperMetricsConfig bool
 		ViperToStdout      bool
-		EnvVar             string
-		ViperDebug         bool
-		ViperDebugConfig   bool
 		Expected           logfile.Params
 	}{
-		"log file and verbose set": {
+		"verbose set": {
 			ViperDebug: true,
 			Expected: logfile.Params{
 				File:    filepath.Join(home, ".wakatime", "wakatime.log"),
 				Verbose: true,
 			},
 		},
-		"log file and verbose from config": {
+		"verbose from config": {
 			ViperDebugConfig: true,
 			Expected: logfile.Params{
 				File:    filepath.Join(home, ".wakatime", "wakatime.log"),
@@ -84,6 +86,22 @@ func TestLoadParams(t *testing.T) {
 				File: filepath.Join(home, ".wakatime", "wakatime.log"),
 			},
 		},
+		"metrics set and verbose false": {
+			ViperMetrics: true,
+			Expected: logfile.Params{
+				File:    filepath.Join(home, ".wakatime", "wakatime.log"),
+				Metrics: false,
+			},
+		},
+		"metrics set and verbose true": {
+			ViperDebug:   true,
+			ViperMetrics: true,
+			Expected: logfile.Params{
+				File:    filepath.Join(home, ".wakatime", "wakatime.log"),
+				Metrics: true,
+				Verbose: true,
+			},
+		},
 		"log to stdout": {
 			ViperToStdout: true,
 			Expected: logfile.Params{
@@ -99,6 +117,8 @@ func TestLoadParams(t *testing.T) {
 			v.Set("log-file", test.ViperLogFile)
 			v.Set("logfile", test.ViperLogFileOld)
 			v.Set("log-to-stdout", test.ViperToStdout)
+			v.Set("metrics", test.ViperMetrics)
+			v.Set("settings.metrics", test.ViperMetricsConfig)
 			v.Set("settings.log_file", test.ViperLogFileConfig)
 			v.Set("settings.debug", test.ViperDebug)
 			v.Set("verbose", test.ViperDebugConfig)

cmd/root.go

@@ -159,6 +159,12 @@ func setFlags(cmd *cobra.Command, v *viper.Viper) {
 	flags.String("log-file", "", "Optional log file. Defaults to '~/.wakatime/wakatime.log'.")
 	flags.String("logfile", "", "(deprecated) Optional log file. Defaults to '~/.wakatime/wakatime.log'.")
 	flags.Bool("log-to-stdout", false, "If enabled, logs will go to stdout. Will overwrite logfile configs.")
+	flags.Bool(
+		"metrics",
+		false,
+		"When set and --verbose or debug enabled, collects metrics usage in '~/.wakatime/metrics' folder."+
+			" Defaults to false.",
+	)
 	flags.Bool(
 		"no-ssl-verify",
 		false,

cmd/run.go

@@ -27,6 +27,7 @@ import (
 	"github.com/wakatime/wakatime-cli/pkg/heartbeat"
 	"github.com/wakatime/wakatime-cli/pkg/ini"
 	"github.com/wakatime/wakatime-cli/pkg/log"
+	"github.com/wakatime/wakatime-cli/pkg/metrics"
 	"github.com/wakatime/wakatime-cli/pkg/offline"
 	"github.com/wakatime/wakatime-cli/pkg/vipertools"
 	"github.com/wakatime/wakatime-cli/pkg/wakaerror"
@@ -68,72 +69,84 @@ func Run(cmd *cobra.Command, v *viper.Viper) {
 		log.Fatalf("failed to setup logging: %s", err)
 	}
 
+	shutdown := func() {}
+
+	// start profiling if enabled
+	if logFileParams.Metrics {
+		shutdown, err = metrics.StartProfiling()
+		if err != nil {
+			log.Errorf("failed to start profiling: %s", err)
+		}
+	}
+
 	if v.GetBool("user-agent") {
 		log.Debugln("command: user-agent")
 
 		fmt.Println(heartbeat.UserAgent(vipertools.GetString(v, "plugin")))
 
+		shutdown()
+
 		os.Exit(exitcode.Success)
 	}
 
 	if v.GetBool("version") {
 		log.Debugln("command: version")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, runVersion)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, runVersion, shutdown)
 	}
 
 	if v.IsSet("config-read") {
 		log.Debugln("command: config-read")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, configread.Run)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, configread.Run, shutdown)
 	}
 
 	if v.IsSet("config-write") {
 		log.Debugln("command: config-write")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, configwrite.Run)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, configwrite.Run, shutdown)
 	}
 
 	if v.GetBool("today") {
 		log.Debugln("command: today")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, today.Run)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, today.Run, shutdown)
 	}
 
 	if v.IsSet("today-goal") {
 		log.Debugln("command: today-goal")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, todaygoal.Run)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, todaygoal.Run, shutdown)
 	}
 
 	if v.GetBool("file-experts") {
 		log.Debugln("command: file-experts")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, fileexperts.Run)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, fileexperts.Run, shutdown)
 	}
 
 	if v.IsSet("entity") {
 		log.Debugln("command: heartbeat")
 
-		RunCmdWithOfflineSync(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, cmdheartbeat.Run)
+		RunCmdWithOfflineSync(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, cmdheartbeat.Run, shutdown)
 	}
 
 	if v.IsSet("sync-offline-activity") {
 		log.Debugln("command: sync-offline-activity")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlinesync.Run)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlinesync.Run, shutdown)
 	}
 
 	if v.GetBool("offline-count") {
 		log.Debugln("command: offline-count")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlinecount.Run)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlinecount.Run, shutdown)
 	}
 
 	if v.IsSet("print-offline-heartbeats") {
 		log.Debugln("command: print-offline-heartbeats")
 
-		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlineprint.Run)
+		RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlineprint.Run, shutdown)
 	}
 
 	log.Warnf("one of the following parameters has to be provided: %s", strings.Join([]string{
@@ -238,36 +251,46 @@ func SetupLogging(v *viper.Viper) (*logfile.Params, error) {
 	return &logfileParams, nil
 }
 
-// cmdFn represents a command function.
-type cmdFn func(v *viper.Viper) (int, error)
+type (
+	// cmdFn represents a command function.
+	cmdFn func(v *viper.Viper) (int, error)
+	// shutdownFn represents a shutdown function. It will be called before exiting.
+	shutdownFn func()
+)
 
 // RunCmd runs a command function and exits with the exit code returned by
 // the command function. Will send diagnostic on any errors or panics.
-func RunCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) {
+func RunCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn, shutdown shutdownFn) {
 	exitCode := runCmd(v, verbose, sendDiagsOnErrors, cmd)
 
+	shutdown()
+
 	os.Exit(exitCode)
 }
 
 // RunCmdWithOfflineSync runs a command function and exits with the exit code
 // returned by the command function. If command run was successful, it will execute
 // offline sync command afterwards. Will send diagnostic on any errors or panics.
-func RunCmdWithOfflineSync(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) {
+func RunCmdWithOfflineSync(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn, shutdown shutdownFn) {
 	exitCode := runCmd(v, verbose, sendDiagsOnErrors, cmd)
 	if exitCode != exitcode.Success {
+		shutdown()
+
 		os.Exit(exitCode)
 	}
 
 	exitCode = runCmd(v, verbose, sendDiagsOnErrors, offlinesync.Run)
 
+	shutdown()
+
 	os.Exit(exitCode)
 }
 
 // runCmd contains the main logic of RunCmd.
 // It will send diagnostic on any errors or panics.
 // On panic, it will send diagnostic and exit with ErrGeneric exit code.
 // On error, it will only send diagnostic if sendDiagsOnErrors and verbose is true.
-func runCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) int {
+func runCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) (exitCode int) {
 	logs := bytes.NewBuffer(nil)
 	resetLogs := captureLogs(logs)
 
@@ -292,12 +315,14 @@ func runCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) int
 				log.Warnf("failed to send diagnostics: %s", err)
 			}
 
-			os.Exit(exitcode.ErrGeneric)
+			exitCode = exitcode.ErrGeneric
 		}
 	}()
 
+	var err error
+
 	// run command
-	exitCode, err := cmd(v)
+	exitCode, err = cmd(v)
 	// nolint:nestif
 	if err != nil {
 		if errwaka, ok := err.(wakaerror.Error); ok {