From e4ad352addfe5d3bc6e63b62e009f394b2b0ce74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Henrique=20Guard=C3=A3o=20Gandarez?= Date: Wed, 24 Jul 2024 13:41:07 -0300 Subject: [PATCH] Read extra heartbeats once --- cmd/heartbeat/heartbeat.go | 12 ++++-- cmd/heartbeat/heartbeat_test.go | 10 ----- cmd/offline/offline.go | 6 +++ cmd/params/params.go | 37 +++++++---------- cmd/params/params_test.go | 9 ----- main_test.go | 72 +++++++++++++++++++++++++++++++++ 6 files changed, 101 insertions(+), 45 deletions(-) diff --git a/cmd/heartbeat/heartbeat.go b/cmd/heartbeat/heartbeat.go index 8799a8e2..8f3661bf 100644 --- a/cmd/heartbeat/heartbeat.go +++ b/cmd/heartbeat/heartbeat.go @@ -70,7 +70,7 @@ func Run(v *viper.Viper) (int, error) { // heartbeats from the offline queue, if available and offline sync is not // explicitly disabled. func SendHeartbeats(v *viper.Viper, queueFilepath string) error { - params, err := LoadParams(v) + params, err := loadParams(v) if err != nil { return fmt.Errorf("failed to load command parameters: %w", err) } @@ -83,7 +83,11 @@ func SendHeartbeats(v *viper.Viper, queueFilepath string) error { LastSentAt: params.Offline.LastSentAt, Timeout: params.Offline.RateLimit, }) { - if err = offlinecmd.SaveHeartbeats(v, nil, queueFilepath); err == nil { + // it prevents reading extra heartbeats again from stdin when we're rate limited. + // Otherwise, it would fail to read. + v.Set("extra-heartbeats", false) + + if err = offlinecmd.SaveHeartbeats(v, params.Heartbeat.ExtraHeartbeats, queueFilepath); err == nil { return nil } @@ -165,9 +169,9 @@ func SendHeartbeats(v *viper.Viper, queueFilepath string) error { return nil } -// LoadParams loads params from viper.Viper instance. Returns ErrAuth +// loadParams loads params from viper.Viper instance. Returns ErrAuth // if failed to retrieve api key. -func LoadParams(v *viper.Viper) (paramscmd.Params, error) { +func loadParams(v *viper.Viper) (paramscmd.Params, error) { if v == nil { return paramscmd.Params{}, errors.New("viper instance unset") } diff --git a/cmd/heartbeat/heartbeat_test.go b/cmd/heartbeat/heartbeat_test.go index d6d09e63..9e34a203 100644 --- a/cmd/heartbeat/heartbeat_test.go +++ b/cmd/heartbeat/heartbeat_test.go @@ -10,13 +10,11 @@ import ( "path/filepath" "runtime" "strings" - "sync" "testing" "time" "github.com/wakatime/wakatime-cli/cmd" cmdheartbeat "github.com/wakatime/wakatime-cli/cmd/heartbeat" - cmdparams "github.com/wakatime/wakatime-cli/cmd/params" "github.com/wakatime/wakatime-cli/pkg/api" "github.com/wakatime/wakatime-cli/pkg/heartbeat" "github.com/wakatime/wakatime-cli/pkg/ini" @@ -292,8 +290,6 @@ func TestSendHeartbeats_ExtraHeartbeats(t *testing.T) { os.Stdin = r - cmdparams.Once = sync.Once{} - data, err := os.ReadFile("testdata/extra_heartbeats.json") require.NoError(t, err) @@ -376,8 +372,6 @@ func TestSendHeartbeats_ExtraHeartbeats_Sanitize(t *testing.T) { os.Stdin = r - cmdparams.Once = sync.Once{} - data, err := os.ReadFile("testdata/extra_heartbeats.json") require.NoError(t, err) @@ -586,8 +580,6 @@ func TestSendHeartbeats_IsUnsavedEntity(t *testing.T) { os.Stdin = inr - cmdparams.Once = sync.Once{} - data, err := os.ReadFile("testdata/extra_heartbeats_is_unsaved_entity.json") require.NoError(t, err) @@ -718,8 +710,6 @@ func TestSendHeartbeats_NonExistingExtraHeartbeatsEntity(t *testing.T) { os.Stdin = inr - cmdparams.Once = sync.Once{} - data, err := os.ReadFile("testdata/extra_heartbeats_nonexisting_entity.json") require.NoError(t, err) diff --git a/cmd/offline/offline.go b/cmd/offline/offline.go index d84ec422..165a5448 100644 --- a/cmd/offline/offline.go +++ b/cmd/offline/offline.go @@ -57,7 +57,13 @@ func SaveHeartbeats(v *viper.Viper, heartbeats []heartbeat.Heartbeat, queueFilep return nil } +// loadParams loads params from viper.Viper instance. Returns ErrAuth +// if failed to retrieve api key. func loadParams(v *viper.Viper) (paramscmd.Params, error) { + if v == nil { + return paramscmd.Params{}, errors.New("viper instance unset") + } + paramAPI, err := paramscmd.LoadAPIParams(v) if err != nil { log.Warnf("failed to load API parameters: %s", err) diff --git a/cmd/params/params.go b/cmd/params/params.go index 72a63ae2..d09f16d5 100644 --- a/cmd/params/params.go +++ b/cmd/params/params.go @@ -13,7 +13,6 @@ import ( "regexp" "strconv" "strings" - "sync" "time" "github.com/wakatime/wakatime-cli/pkg/api" @@ -399,7 +398,10 @@ func LoadHeartbeatParams(v *viper.Viper) (Heartbeat, error) { var extraHeartbeats []heartbeat.Heartbeat if v.GetBool("extra-heartbeats") { - extraHeartbeats = readExtraHeartbeats() + extraHeartbeats, err = readExtraHeartbeats() + if err != nil { + log.Errorf("failed to read extra heartbeats: %s", err) + } } var isWrite *bool @@ -750,29 +752,20 @@ func readAPIKeyFromCommand(cmdStr string) (string, error) { return strings.TrimSpace(string(out)), nil } -var extraHeartbeatsCache *[]heartbeat.Heartbeat // nolint:gochecknoglobals - -// Once prevents reading from stdin twice. -var Once sync.Once // nolint:gochecknoglobals - -func readExtraHeartbeats() []heartbeat.Heartbeat { - Once.Do(func() { - in := bufio.NewReader(os.Stdin) - - input, err := in.ReadString('\n') - if err != nil && err != io.EOF { - log.Debugf("failed to read data from stdin: %s", err) - } +func readExtraHeartbeats() ([]heartbeat.Heartbeat, error) { + in := bufio.NewReader(os.Stdin) - heartbeats, err := parseExtraHeartbeats(input) - if err != nil { - log.Errorf("failed parsing: %s", err) - } + input, err := in.ReadString('\n') + if err != nil && err != io.EOF { + log.Debugf("failed to read data from stdin: %s", err) + } - extraHeartbeatsCache = &heartbeats - }) + heartbeats, err := parseExtraHeartbeats(input) + if err != nil { + log.Errorf("failed parsing: %s", err) + } - return *extraHeartbeatsCache + return heartbeats, nil } func parseExtraHeartbeats(data string) ([]heartbeat.Heartbeat, error) { diff --git a/cmd/params/params_test.go b/cmd/params/params_test.go index 4738f53f..40c1c249 100644 --- a/cmd/params/params_test.go +++ b/cmd/params/params_test.go @@ -9,7 +9,6 @@ import ( "regexp" "runtime" "strings" - "sync" "testing" "time" @@ -233,8 +232,6 @@ func TestLoadHeartbeatParams_ExtraHeartbeats(t *testing.T) { os.Stdin = r - cmdparams.Once = sync.Once{} - data, err := os.ReadFile("testdata/extra_heartbeats.json") require.NoError(t, err) @@ -307,8 +304,6 @@ func TestLoadHeartbeatParams_ExtraHeartbeats_WithStringValues(t *testing.T) { os.Stdin = r - cmdparams.Once = sync.Once{} - data, err := os.ReadFile("testdata/extra_heartbeats_with_string_values.json") require.NoError(t, err) @@ -376,8 +371,6 @@ func TestLoadHeartbeatParams_ExtraHeartbeats_WithEOF(t *testing.T) { os.Stdin = r - cmdparams.Once = sync.Once{} - data, err := os.ReadFile("testdata/extra_heartbeats.json") require.NoError(t, err) @@ -456,8 +449,6 @@ func TestLoadHeartbeatParams_ExtraHeartbeats_NoData(t *testing.T) { os.Stdin = r - cmdparams.Once = sync.Once{} - go func() { _, err := w.Write([]byte{}) require.NoError(t, err) diff --git a/main_test.go b/main_test.go index 422c8a2c..0ab77fbf 100644 --- a/main_test.go +++ b/main_test.go @@ -332,6 +332,78 @@ func TestSendHeartbeats_ExtraHeartbeats(t *testing.T) { assert.Eventually(t, func() bool { return numCalls == 2 }, time.Second, 50*time.Millisecond) } +func TestSendHeartbeats_ExtraHeartbeats_RateLimited(t *testing.T) { + apiURL, router, close := setupTestServer() + defer close() + + var numCalls int + + router.HandleFunc("/users/current/heartbeats.bulk", func(_ http.ResponseWriter, _ *http.Request) { + numCalls++ + }) + + tmpDir := t.TempDir() + + offlineQueueFile, err := os.CreateTemp(tmpDir, "") + require.NoError(t, err) + + defer offlineQueueFile.Close() + + offlineQueueFileLegacy, err := os.CreateTemp(tmpDir, "") + require.NoError(t, err) + + defer offlineQueueFileLegacy.Close() + + tmpConfigFile, err := os.CreateTemp(tmpDir, "wakatime.cfg") + require.NoError(t, err) + + defer tmpConfigFile.Close() + + tmpInternalConfigFile, err := os.CreateTemp(tmpDir, "wakatime-internal.cfg") + require.NoError(t, err) + + _, err = tmpInternalConfigFile.WriteString( + fmt.Sprintf("[internal]\nheartbeats_last_sent_at = %s\n", time.Now().Format(time.RFC3339)), + ) + require.NoError(t, err) + + defer tmpInternalConfigFile.Close() + + data, err := os.ReadFile("testdata/extra_heartbeats.json") + require.NoError(t, err) + + buffer := bytes.NewBuffer(data) + + runWakatimeCli( + t, + buffer, + "--api-url", apiURL, + "--key", "00000000-0000-4000-8000-000000000000", + "--config", tmpConfigFile.Name(), + "--heartbeat-rate-limit-seconds", "100", + "--internal-config", tmpInternalConfigFile.Name(), + "--entity", "testdata/main.go", + "--extra-heartbeats", "true", + "--cursorpos", "12", + "--sync-offline-activity", "2", + "--offline-queue-file", offlineQueueFile.Name(), + "--offline-queue-file-legacy", offlineQueueFileLegacy.Name(), + "--lineno", "42", + "--lines-in-file", "100", + "--time", "1585598059", + "--hide-branch-names", ".*", + "--write", + "--verbose", + ) + + offlineCount, err := offline.CountHeartbeats(offlineQueueFile.Name()) + require.NoError(t, err) + + assert.Equal(t, 26, offlineCount) + + assert.Zero(t, numCalls) +} + func TestSendHeartbeats_ExtraHeartbeats_SyncLegacyOfflineActivity(t *testing.T) { apiURL, router, close := setupTestServer() defer close()