diff --git a/cmd/logfile/logfile.go b/cmd/logfile/logfile.go index a67a16b5..69d9d6d7 100644 --- a/cmd/logfile/logfile.go +++ b/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) diff --git a/cmd/logfile/logfile_test.go b/cmd/logfile/logfile_test.go index 12e2c8bf..1c0de61c 100644 --- a/cmd/logfile/logfile_test.go +++ b/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) diff --git a/cmd/root.go b/cmd/root.go index a5921dfe..48dede40 100644 --- a/cmd/root.go +++ b/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, diff --git a/cmd/run.go b/cmd/run.go index f878d3ec..772671f8 100644 --- a/cmd/run.go +++ b/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,28 +251,38 @@ 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) } @@ -267,7 +290,7 @@ func RunCmdWithOfflineSync(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, // 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 { diff --git a/cmd/run_test.go b/cmd/run_test.go index 29de4dfb..69e4b239 100644 --- a/cmd/run_test.go +++ b/cmd/run_test.go @@ -22,7 +22,7 @@ import ( bolt "go.etcd.io/bbolt" ) -func TestRunCmd_Error(t *testing.T) { +func TestRunCmd_Err(t *testing.T) { // this is exclusively run in subprocess if os.Getenv("TEST_RUN") == "1" { version.OS = "some os" @@ -46,9 +46,23 @@ func TestRunCmd_Error(t *testing.T) { v.Set("offline-queue-file", offlineQueueFile.Name()) v.Set("plugin", "vim") - cmd.RunCmd(v, false, false, func(v *viper.Viper) (int, error) { + var cmdNumCalls int + + cmdFn := func(v *viper.Viper) (int, error) { + cmdNumCalls++ return 42, errors.New("fail") - }) + } + + var shutdownNumCalls int + + shutdownFn := func() { + shutdownNumCalls++ + } + + cmd.RunCmd(v, false, false, cmdFn, shutdownFn) + + assert.Equal(t, 1, cmdNumCalls) + assert.Equal(t, 1, shutdownNumCalls) return } @@ -63,7 +77,7 @@ func TestRunCmd_Error(t *testing.T) { }) // run command in another runner, to effectively test os.Exit() - cmd := exec.Command(os.Args[0], "-test.run=TestRunCmd_Error") // nolint:gosec + cmd := exec.Command(os.Args[0], "-test.run=TestRunCmd_Err") // nolint:gosec cmd.Env = append(os.Environ(), "TEST_RUN=1") cmd.Env = append(cmd.Env, fmt.Sprintf("TEST_SERVER_URL=%s", testServerURL)) @@ -77,7 +91,7 @@ func TestRunCmd_Error(t *testing.T) { assert.Eventually(t, func() bool { return numCalls == 0 }, time.Second, 50*time.Millisecond) } -func TestRunCmd_Verbose_Error(t *testing.T) { +func TestRunCmd_Verbose_Err(t *testing.T) { // this is exclusively run in subprocess if os.Getenv("TEST_RUN") == "1" { version.OS = "some os" @@ -101,9 +115,23 @@ func TestRunCmd_Verbose_Error(t *testing.T) { v.Set("offline-queue-file", offlineQueueFile.Name()) v.Set("plugin", "vim") - cmd.RunCmd(v, true, false, func(v *viper.Viper) (int, error) { + var cmdNumCalls int + + cmdFn := func(v *viper.Viper) (int, error) { + cmdNumCalls++ return 42, errors.New("fail") - }) + } + + var shutdownNumCalls int + + shutdownFn := func() { + shutdownNumCalls++ + } + + cmd.RunCmd(v, true, false, cmdFn, shutdownFn) + + assert.Equal(t, 1, cmdNumCalls) + assert.Equal(t, 1, shutdownNumCalls) return } @@ -118,7 +146,7 @@ func TestRunCmd_Verbose_Error(t *testing.T) { }) // run command in another runner, to effectively test os.Exit() - cmd := exec.Command(os.Args[0], "-test.run=TestRunCmd_Verbose_Error") // nolint:gosec + cmd := exec.Command(os.Args[0], "-test.run=TestRunCmd_Verbose_Err") // nolint:gosec cmd.Env = append(os.Environ(), "TEST_RUN=1") cmd.Env = append(cmd.Env, fmt.Sprintf("TEST_SERVER_URL=%s", testServerURL)) @@ -132,7 +160,7 @@ func TestRunCmd_Verbose_Error(t *testing.T) { assert.Eventually(t, func() bool { return numCalls == 0 }, time.Second, 50*time.Millisecond) } -func TestRunCmd_SendDiagnostics_Error(t *testing.T) { +func TestRunCmd_SendDiagnostics_Err(t *testing.T) { // this is exclusively run in subprocess if os.Getenv("TEST_RUN") == "1" { version.OS = "some os" @@ -156,9 +184,23 @@ func TestRunCmd_SendDiagnostics_Error(t *testing.T) { v.Set("offline-queue-file", offlineQueueFile.Name()) v.Set("plugin", "vim") - cmd.RunCmd(v, true, true, func(v *viper.Viper) (int, error) { + var cmdNumCalls int + + cmdFn := func(v *viper.Viper) (int, error) { + cmdNumCalls++ return 42, errors.New("fail") - }) + } + + var shutdownNumCalls int + + shutdownFn := func() { + shutdownNumCalls++ + } + + cmd.RunCmd(v, true, true, cmdFn, shutdownFn) + + assert.Equal(t, 1, cmdNumCalls) + assert.Equal(t, 1, shutdownNumCalls) return } @@ -209,7 +251,7 @@ func TestRunCmd_SendDiagnostics_Error(t *testing.T) { }) // run command in another runner, to effectively test os.Exit() - cmd := exec.Command(os.Args[0], "-test.run=TestRunCmd_SendDiagnostics_Error") // nolint:gosec + cmd := exec.Command(os.Args[0], "-test.run=TestRunCmd_SendDiagnostics_Err") // nolint:gosec cmd.Env = append(os.Environ(), "TEST_RUN=1") cmd.Env = append(cmd.Env, fmt.Sprintf("TEST_SERVER_URL=%s", testServerURL)) @@ -247,9 +289,24 @@ func TestRunCmd_SendDiagnostics_Panic(t *testing.T) { v.Set("offline-queue-file", offlineQueueFile.Name()) v.Set("plugin", "vim") - cmd.RunCmd(v, true, false, func(v *viper.Viper) (int, error) { + var cmdNumCalls int + + cmdFn := func(v *viper.Viper) (int, error) { + cmdNumCalls++ + panic("fail") - }) + } + + var shutdownNumCalls int + + shutdownFn := func() { + shutdownNumCalls++ + } + + cmd.RunCmd(v, true, false, cmdFn, shutdownFn) + + assert.Equal(t, 1, cmdNumCalls) + assert.Equal(t, 1, shutdownNumCalls) return } @@ -339,9 +396,24 @@ func TestRunCmd_SendDiagnostics_NoLogs_Panic(t *testing.T) { v.Set("offline-queue-file", offlineQueueFile.Name()) v.Set("plugin", "vim") - cmd.RunCmd(v, false, false, func(v *viper.Viper) (int, error) { + var cmdNumCalls int + + cmdFn := func(v *viper.Viper) (int, error) { + cmdNumCalls++ + panic("fail") - }) + } + + var shutdownNumCalls int + + shutdownFn := func() { + shutdownNumCalls++ + } + + cmd.RunCmd(v, false, false, cmdFn, shutdownFn) + + assert.Equal(t, 1, cmdNumCalls) + assert.Equal(t, 1, shutdownNumCalls) return } @@ -429,9 +501,23 @@ func TestRunCmd_SendDiagnostics_WakaError(t *testing.T) { v.Set("offline-queue-file", offlineQueueFile.Name()) v.Set("plugin", "vim") - cmd.RunCmd(v, false, false, func(v *viper.Viper) (int, error) { + var cmdNumCalls int + + cmdFn := func(v *viper.Viper) (int, error) { + cmdNumCalls++ return 42, offline.ErrOpenDB{Err: errors.New("fail")} - }) + } + + var shutdownNumCalls int + + shutdownFn := func() { + shutdownNumCalls++ + } + + cmd.RunCmd(v, false, false, cmdFn, shutdownFn) + + assert.Equal(t, 1, cmdNumCalls) + assert.Equal(t, 1, shutdownNumCalls) return } @@ -518,9 +604,23 @@ func TestRunCmdWithOfflineSync(t *testing.T) { v.SetDefault("sync-offline-activity", 24) v.Set("plugin", "vim") - cmd.RunCmdWithOfflineSync(v, false, false, func(v *viper.Viper) (int, error) { + var cmdNumCalls int + + cmdFn := func(v *viper.Viper) (int, error) { + cmdNumCalls++ return 0, nil - }) + } + + var shutdownNumCalls int + + shutdownFn := func() { + shutdownNumCalls++ + } + + cmd.RunCmdWithOfflineSync(v, false, false, cmdFn, shutdownFn) + + assert.Equal(t, 1, cmdNumCalls) + assert.Equal(t, 1, shutdownNumCalls) return } diff --git a/pkg/metrics/metrics.go b/pkg/metrics/metrics.go new file mode 100644 index 00000000..c5a4b7de --- /dev/null +++ b/pkg/metrics/metrics.go @@ -0,0 +1,59 @@ +package metrics + +import ( + "fmt" + "os" + "path/filepath" + "runtime/pprof" + "strconv" + "time" + + "github.com/wakatime/wakatime-cli/pkg/ini" + "github.com/wakatime/wakatime-cli/pkg/log" +) + +// StartProfiling starts profiling cpu and memory. It returns a function that +// should be called to stop profiling and close the files. +func StartProfiling() (func(), error) { + folder, err := ini.WakaResourcesDir() + if err != nil { + return nil, fmt.Errorf("failed getting user's home directory: %s", err) + } + + metricsFolder := filepath.Join(folder, "metrics") + if err := os.MkdirAll(metricsFolder, 0750); err != nil { + return nil, fmt.Errorf("failed to create metrics folder: %s", err) + } + + now := strconv.FormatInt(time.Now().UTC().UnixNano(), 10) + + cpuf, err := os.Create(filepath.Join(metricsFolder, fmt.Sprintf("cpu_%s.profile", now))) // nolint:gosec + if err != nil { + return nil, fmt.Errorf("failed to create cpu profile file: %s", err) + } + + if err := pprof.StartCPUProfile(cpuf); err != nil { + log.Errorf("failed to start cpu profile: %s", err) + } + + memf, err := os.Create(filepath.Join(metricsFolder, fmt.Sprintf("mem_%s.profile", now))) // nolint:gosec + if err != nil { + return nil, fmt.Errorf("failed to create mem profile file: %s", err) + } + + if err := pprof.WriteHeapProfile(memf); err != nil { + log.Errorf("failed to write heap profile: %s", err) + } + + return func() { + pprof.StopCPUProfile() + + if err := cpuf.Close(); err != nil { + log.Errorf("failed to close cpu profile file: %s", err) + } + + if err := memf.Close(); err != nil { + log.Errorf("failed to close mem profile file: %s", err) + } + }, nil +}