Read extra heartbeats once

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")
 	}

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)
 

cmd/offline/offline.go

@@ -57,7 +57,13 @@
 	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)

cmd/params/params.go

@@ -13,7 +13,6 @@
 	"regexp"
 	"strconv"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/wakatime/wakatime-cli/pkg/api"
@@ -399,7 +398,10 @@
 	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 @@
 	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) {

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)

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()