Merge pull request #1071 from wakatime/feature/rate-limiting

Rate limiting by default

.golangci.yml

@@ -67,3 +67,4 @@ issues:
         - gosec
   include:
     - EXC0002
+  fix: true

USAGE.md

@@ -91,6 +91,7 @@ some/submodule/name = new project name
 | api_key                        | Your wakatime api key. | _string_ | |
 | api_key_vault_cmd              | A command to get your api key, perhaps from some sort of secure vault. Actually a space-separated list of an executable and its arguments. Executables in PATH can be referred to by their basenames. Shell syntax not supported. | _string_ | |
 | api_url                        | The WakaTime API base url. | _string_ | <https://api.wakatime.com/api/v1> |
+| heartbeat_rate_limit_seconds   | Rate limit sending heartbeats to the API once per duration. Set to 0 to disable rate limiting. | _int_ | `120` |
 | hide_file_names                | Obfuscate filenames. Will not send file names to api. | _bool_;_list_ | `false` |
 | hide_project_names             | Obfuscate project names. When a project folder is detected instead of using the folder name as the project, a `.wakatime-project file` is created with a random project name. | _bool_;_list_ | `false` |
 | hide_branch_names              | Obfuscate branch names. Will not send revision control branch names to api. | _bool_;_list_ | `false` |

cmd/configwrite/configwrite_test.go

@@ -132,7 +132,7 @@ func TestWriteErr(t *testing.T) {
 
 func TestWriteSaveErr(t *testing.T) {
 	v := viper.New()
-	w := &writerMock{
+	w := &mockWriter{
 		WriteFn: func(section string, keyValue map[string]string) error {
 			assert.Equal(t, "settings", section)
 			assert.Equal(t, map[string]string{"debug": "false"}, keyValue)
@@ -148,10 +148,10 @@ func TestWriteSaveErr(t *testing.T) {
 	assert.Error(t, err)
 }
 
-type writerMock struct {
+type mockWriter struct {
 	WriteFn func(section string, keyValue map[string]string) error
 }
 
-func (m *writerMock) Write(section string, keyValue map[string]string) error {
+func (m *mockWriter) Write(section string, keyValue map[string]string) error {
 	return m.WriteFn(section, keyValue)
 }

cmd/heartbeat/heartbeat.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"time"
 
 	apicmd "github.com/wakatime/wakatime-cli/cmd/api"
 	offlinecmd "github.com/wakatime/wakatime-cli/cmd/offline"
@@ -16,6 +17,7 @@ import (
 	"github.com/wakatime/wakatime-cli/pkg/filestats"
 	"github.com/wakatime/wakatime-cli/pkg/filter"
 	"github.com/wakatime/wakatime-cli/pkg/heartbeat"
+	"github.com/wakatime/wakatime-cli/pkg/ini"
 	"github.com/wakatime/wakatime-cli/pkg/language"
 	_ "github.com/wakatime/wakatime-cli/pkg/lexer" // force to load all lexers
 	"github.com/wakatime/wakatime-cli/pkg/log"
@@ -76,6 +78,19 @@ func SendHeartbeats(v *viper.Viper, queueFilepath string) error {
 	setLogFields(params)
 	log.Debugf("params: %s", params)
 
+	if RateLimited(RateLimitParams{
+		Disabled:   params.Offline.Disabled,
+		LastSentAt: params.Offline.LastSentAt,
+		Timeout:    params.Offline.RateLimit,
+	}) {
+		if err = offlinecmd.SaveHeartbeats(v, nil, queueFilepath); err == nil {
+			return nil
+		}
+
+		// log offline db error then try to send heartbeats to API so they're not lost
+		log.Errorf("failed to save rate limited heartbeats: %s", err)
+	}
+
 	heartbeats := buildHeartbeats(params)
 
 	var chOfflineSave = make(chan bool)
@@ -143,6 +158,10 @@ func SendHeartbeats(v *viper.Viper, queueFilepath string) error {
 		}
 	}
 
+	if err := ResetRateLimit(v); err != nil {
+		log.Errorf("failed to reset rate limit: %s", err)
+	}
+
 	return nil
 }
 
@@ -170,6 +189,48 @@ func LoadParams(v *viper.Viper) (paramscmd.Params, error) {
 	}, nil
 }
 
+// RateLimitParams contains params for the RateLimited function.
+type RateLimitParams struct {
+	Disabled   bool
+	LastSentAt time.Time
+	Timeout    time.Duration
+}
+
+// RateLimited determines if we should send heartbeats to the API or save to the offline db.
+func RateLimited(params RateLimitParams) bool {
+	if params.Disabled {
+		return false
+	}
+
+	if params.Timeout == 0 {
+		return false
+	}
+
+	if params.LastSentAt.IsZero() {
+		return false
+	}
+
+	return time.Since(params.LastSentAt) < params.Timeout
+}
+
+// ResetRateLimit updates the internal.heartbeats_last_sent_at timestamp.
+func ResetRateLimit(v *viper.Viper) error {
+	w, err := ini.NewWriter(v, ini.InternalFilePath)
+	if err != nil {
+		return fmt.Errorf("failed to parse config file: %s", err)
+	}
+
+	keyValue := map[string]string{
+		"heartbeats_last_sent_at": time.Now().Format(ini.DateFormat),
+	}
+
+	if err := w.Write("internal", keyValue); err != nil {
+		return fmt.Errorf("failed to write to internal config file: %s", err)
+	}
+
+	return nil
+}
+
 func buildHeartbeats(params paramscmd.Params) []heartbeat.Heartbeat {
 	heartbeats := []heartbeat.Heartbeat{}
 

cmd/heartbeat/heartbeat_test.go

@@ -47,6 +47,8 @@ func TestSendHeartbeats(t *testing.T) {
 	subfolders := project.CountSlashesInProjectFolder(projectFolder)
 
 	router.HandleFunc("/users/current/heartbeats.bulk", func(w http.ResponseWriter, req *http.Request) {
+		numCalls++
+
 		// check request
 		assert.Equal(t, http.MethodPost, req.Method)
 		assert.Equal(t, []string{"application/json"}, req.Header["Accept"])
@@ -85,7 +87,47 @@ func TestSendHeartbeats(t *testing.T) {
 
 		_, err = io.Copy(w, f)
 		require.NoError(t, err)
+	})
+
+	v := viper.New()
+	v.SetDefault("sync-offline-activity", 1000)
+	v.Set("api-url", testServerURL)
+	v.Set("category", "debugging")
+	v.Set("cursorpos", 42)
+	v.Set("entity", "testdata/main.go")
+	v.Set("entity-type", "file")
+	v.Set("key", "00000000-0000-4000-8000-000000000000")
+	v.Set("language", "Go")
+	v.Set("alternate-language", "Golang")
+	v.Set("hide-branch-names", true)
+	v.Set("project", "wakatime-cli")
+	v.Set("lineno", 13)
+	v.Set("local-file", "testdata/localfile.go")
+	v.Set("plugin", plugin)
+	v.Set("time", 1585598059.1)
+	v.Set("timeout", 5)
+	v.Set("write", true)
+
+	offlineQueueFile, err := os.CreateTemp(t.TempDir(), "")
+	require.NoError(t, err)
+
+	err = cmdheartbeat.SendHeartbeats(v, offlineQueueFile.Name())
+	require.NoError(t, err)
+
+	assert.Eventually(t, func() bool { return numCalls == 1 }, time.Second, 50*time.Millisecond)
+}
+
+func TestSendHeartbeats_RateLimited(t *testing.T) {
+	testServerURL, router, tearDown := setupTestServer()
+	defer tearDown()
+
+	var (
+		plugin   = "plugin/0.0.1"
+		numCalls int
+	)
 
+	router.HandleFunc("/users/current/heartbeats.bulk", func(_ http.ResponseWriter, _ *http.Request) {
+		// Should not be called
 		numCalls++
 	})
 
@@ -107,14 +149,16 @@ func TestSendHeartbeats(t *testing.T) {
 	v.Set("time", 1585598059.1)
 	v.Set("timeout", 5)
 	v.Set("write", true)
+	v.Set("heartbeat-rate-limit-seconds", 500)
+	v.Set("internal.heartbeats_last_sent_at", time.Now().Add(-time.Minute).Format(time.RFC3339))
 
 	offlineQueueFile, err := os.CreateTemp(t.TempDir(), "")
 	require.NoError(t, err)
 
 	err = cmdheartbeat.SendHeartbeats(v, offlineQueueFile.Name())
 	require.NoError(t, err)
 
-	assert.Eventually(t, func() bool { return numCalls == 1 }, time.Second, 50*time.Millisecond)
+	assert.Zero(t, numCalls)
 }
 
 func TestSendHeartbeats_WithFiltering_Exclude(t *testing.T) {
@@ -1052,6 +1096,75 @@ func TestSendHeartbeats_ObfuscateProjectNotBranch(t *testing.T) {
 	assert.Eventually(t, func() bool { return numCalls == 1 }, time.Second, 50*time.Millisecond)
 }
 
+func TestRateLimited(t *testing.T) {
+	p := cmdheartbeat.RateLimitParams{
+		Timeout:    time.Duration(offline.RateLimitDefaultSeconds) * time.Second,
+		LastSentAt: time.Now(),
+	}
+
+	assert.True(t, cmdheartbeat.RateLimited(p))
+}
+
+func TestRateLimited_NotLimited(t *testing.T) {
+	p := cmdheartbeat.RateLimitParams{
+		LastSentAt: time.Now().Add(time.Duration(-offline.RateLimitDefaultSeconds*2) * time.Second),
+		Timeout:    time.Duration(offline.RateLimitDefaultSeconds) * time.Second,
+	}
+
+	assert.False(t, cmdheartbeat.RateLimited(p))
+}
+
+func TestRateLimited_Disabled(t *testing.T) {
+	p := cmdheartbeat.RateLimitParams{
+		Disabled: true,
+	}
+
+	assert.False(t, cmdheartbeat.RateLimited(p))
+}
+
+func TestRateLimited_TimeoutZero(t *testing.T) {
+	p := cmdheartbeat.RateLimitParams{
+		LastSentAt: time.Time{},
+	}
+
+	assert.False(t, cmdheartbeat.RateLimited(p))
+}
+
+func TestRateLimited_LastSentAtZero(t *testing.T) {
+	p := cmdheartbeat.RateLimitParams{
+		Timeout: 0,
+	}
+
+	assert.False(t, cmdheartbeat.RateLimited(p))
+}
+
+func TestResetRateLimit(t *testing.T) {
+	tmpFile, err := os.CreateTemp(t.TempDir(), "wakatime")
+	require.NoError(t, err)
+
+	defer tmpFile.Close()
+
+	v := viper.New()
+	v.Set("config", tmpFile.Name())
+	v.Set("internal-config", tmpFile.Name())
+
+	writer, err := ini.NewWriter(v, func(vp *viper.Viper) (string, error) {
+		assert.Equal(t, v, vp)
+		return tmpFile.Name(), nil
+	})
+	require.NoError(t, err)
+
+	err = cmdheartbeat.ResetRateLimit(v)
+	require.NoError(t, err)
+
+	err = writer.File.Reload()
+	require.NoError(t, err)
+
+	lastSentAt := writer.File.Section("internal").Key("heartbeats_last_sent_at").MustTimeFormat(ini.DateFormat)
+
+	assert.WithinDuration(t, time.Now(), lastSentAt, 1*time.Second)
+}
+
 func setupTestServer() (string, *http.ServeMux, func()) {
 	router := http.NewServeMux()
 	srv := httptest.NewServer(router)

cmd/offlinesync/offlinesync.go

@@ -5,6 +5,7 @@ import (
 	"os"
 
 	cmdapi "github.com/wakatime/wakatime-cli/cmd/api"
+	cmdheartbeat "github.com/wakatime/wakatime-cli/cmd/heartbeat"
 	"github.com/wakatime/wakatime-cli/cmd/params"
 	"github.com/wakatime/wakatime-cli/pkg/apikey"
 	"github.com/wakatime/wakatime-cli/pkg/exitcode"
@@ -16,8 +17,33 @@ import (
 	"github.com/spf13/viper"
 )
 
-// Run executes the sync-offline-activity command.
-func Run(v *viper.Viper) (int, error) {
+// RunWithoutRateLimiting executes the sync-offline-activity command without rate limiting.
+func RunWithoutRateLimiting(v *viper.Viper) (int, error) {
+	return run(v)
+}
+
+// RunWithRateLimiting executes sync-offline-activity command with rate limiting enabled.
+func RunWithRateLimiting(v *viper.Viper) (int, error) {
+	paramOffline := params.LoadOfflineParams(v)
+
+	if cmdheartbeat.RateLimited(cmdheartbeat.RateLimitParams{
+		Disabled:   paramOffline.Disabled,
+		LastSentAt: paramOffline.LastSentAt,
+		Timeout:    paramOffline.RateLimit,
+	}) {
+		log.Debugln("skip syncing offline activity to respect rate limit")
+		return exitcode.Success, nil
+	}
+
+	return run(v)
+}
+
+func run(v *viper.Viper) (int, error) {
+	paramOffline := params.LoadOfflineParams(v)
+	if paramOffline.Disabled {
+		return exitcode.Success, nil
+	}
+
 	queueFilepath, err := offline.QueueFilepath()
 	if err != nil {
 		return exitcode.ErrGeneric, fmt.Errorf(
@@ -97,8 +123,6 @@ func syncOfflineActivityLegacy(v *viper.Viper, queueFilepath string) error {
 // SyncOfflineActivity syncs offline activity by sending heartbeats
 // from the offline queue to the WakaTime API.
 func SyncOfflineActivity(v *viper.Viper, queueFilepath string) error {
-	paramOffline := params.LoadOfflineParams(v)
-
 	paramAPI, err := params.LoadAPIParams(v)
 	if err != nil {
 		return fmt.Errorf("failed to load API parameters: %w", err)
@@ -109,6 +133,8 @@ func SyncOfflineActivity(v *viper.Viper, queueFilepath string) error {
 		return fmt.Errorf("failed to initialize api client: %w", err)
 	}
 
+	paramOffline := params.LoadOfflineParams(v)
+
 	if paramOffline.QueueFile != "" {
 		queueFilepath = paramOffline.QueueFile
 	}
@@ -126,5 +152,9 @@ func SyncOfflineActivity(v *viper.Viper, queueFilepath string) error {
 		return err
 	}
 
+	if err := cmdheartbeat.ResetRateLimit(v); err != nil {
+		log.Errorf("failed to reset rate limit: %s", err)
+	}
+
 	return nil
 }