Log limits file in throttle plugin (#890)

Author: timggggggg

cfg/config.go

@@ -564,6 +564,18 @@ func ParseFieldSelector(selector string) []string {
 	return result
 }
 
+func BuildFieldSelector(fields []string) string {
+	var builder strings.Builder
+	for i, field := range fields {
+		escaped := strings.ReplaceAll(field, ".", `\.`)
+		builder.WriteString(escaped)
+		if i != len(fields)-1 {
+			builder.WriteByte('.')
+		}
+	}
+	return builder.String()
+}
+
 // Parses several fields selectors and removes nested fields.
 // If there are empty field selectors or no selectors at all error returned.
 // For example:

cfg/config_test.go

@@ -864,3 +864,57 @@ func TestMergeYAMLs(t *testing.T) {
 		})
 	}
 }
+
+func TestBuildFieldSelector(t *testing.T) {
+	tests := []struct {
+		name     string
+		fields   []string
+		expected string
+	}{
+		{
+			name:     "Empty slice",
+			fields:   []string{},
+			expected: "",
+		},
+		{
+			name:     "Single field, no dot",
+			fields:   []string{"field"},
+			expected: "field",
+		},
+		{
+			name:     "Multiple fields, no dots",
+			fields:   []string{"a", "b", "c"},
+			expected: "a.b.c",
+		},
+		{
+			name:     "Field with dot",
+			fields:   []string{"a.b", "c"},
+			expected: "a\\.b.c",
+		},
+		{
+			name:     "All fields with dots",
+			fields:   []string{"a.b", "c.d"},
+			expected: "a\\.b.c\\.d",
+		},
+		{
+			name:     "Field with multiple dots",
+			fields:   []string{"a.b.c", "d"},
+			expected: "a\\.b\\.c.d",
+		},
+		{
+			name:     "Field with backslash",
+			fields:   []string{`a\b`, "c"},
+			expected: `a\b.c`,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			got := BuildFieldSelector(tt.fields)
+			require.Equal(t, tt.expected, got)
+		})
+	}
+}

plugin/action/throttle/README.md

@@ -227,5 +227,20 @@ Distribution object example:
 
 <br>
 
+**`limits_file`** *`string`* 
+
+The filename to store current log limits. Limits are loaded only on initialization
+> It's a `json` file. You can modify it manually. But the limit from the file will disappear if redis is available and it has a different value for this limit
+
+> ⚠ **Experimental feature**
+
+<br>
+
+**`limits_save_interval`** *`cfg.Duration`* *`default=3s`* 
+
+Defines the interval between each saving of log limits in file
+
+<br>
+
 
 <br>*Generated using [__insane-doc__](https://github.com/vitkovskii/insane-doc)*

plugin/action/throttle/distribution.go

@@ -136,3 +136,30 @@ func (ld *limitDistributions) copy() limitDistributions {
 		enabled:         ld.enabled,
 	}
 }
+
+func (ld *limitDistributions) getLimitDistributionsCfg() limitDistributionCfg {
+	field := cfg.BuildFieldSelector(ld.field)
+
+	ratioMap := make(map[float64][]string)
+	for value, idx := range ld.idxByKey {
+		if idx < 0 || idx >= len(ld.distributions) {
+			continue
+		}
+		ratio := ld.distributions[idx].ratio
+		ratioMap[ratio] = append(ratioMap[ratio], value)
+	}
+
+	ratios := make([]limitDistributionRatio, 0, len(ratioMap))
+	for ratio, values := range ratioMap {
+		ratios = append(ratios, limitDistributionRatio{
+			Ratio:  ratio,
+			Values: values,
+		})
+	}
+
+	return limitDistributionCfg{
+		Field:   field,
+		Ratios:  ratios,
+		Enabled: ld.enabled,
+	}
+}

plugin/action/throttle/distribution_test.go

@@ -3,6 +3,7 @@ package throttle
 import (
 	"fmt"
 	"slices"
+	"sort"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -196,3 +197,49 @@ func Test_parseLimitDistribution(t *testing.T) {
 		})
 	}
 }
+
+func TestGetLimitDistributionsCfg(t *testing.T) {
+	ld := &limitDistributions{
+		field: []string{"log", "level"},
+		idxByKey: map[string]int{
+			"A": 0,
+			"B": 1,
+			"C": 0,
+		},
+		distributions: []complexDistribution{
+			{ratio: 0.1, limit: 10},
+			{ratio: 0.2, limit: 20},
+		},
+		enabled: true,
+	}
+
+	expected := limitDistributionCfg{
+		Field: "log.level",
+		Ratios: []limitDistributionRatio{
+			{Ratio: 0.1, Values: []string{"A", "C"}},
+			{Ratio: 0.2, Values: []string{"B"}},
+		},
+		Enabled: true,
+	}
+
+	result := ld.getLimitDistributionsCfg()
+
+	require.Equal(t, expected.Field, result.Field, "wrong Field")
+	require.Equal(t, expected.Enabled, result.Enabled, "wrong Enabled")
+	require.True(t, isLimitDistributionCfgRatiosEqual(expected.Ratios, result.Ratios), "wrong Ratios")
+}
+
+func isLimitDistributionCfgRatiosEqual(expected, result []limitDistributionRatio) bool {
+	sort.Slice(result, func(i, j int) bool {
+		return result[i].Ratio < result[j].Ratio
+	})
+
+	for i, ldRatio := range result {
+		slices.Sort(ldRatio.Values)
+		if !slices.Equal(expected[i].Values, ldRatio.Values) {
+			return false
+		}
+	}
+
+	return true
+}

plugin/action/throttle/in_memory_limiter.go

@@ -1,6 +1,7 @@
 package throttle
 
 import (
+	"math"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -9,6 +10,8 @@ import (
 	"github.com/ozontech/file.d/pipeline"
 )
 
+const EPS = 1e-9
+
 type inMemoryLimiter struct {
 	limit   complexLimit
 	buckets buckets
@@ -236,6 +239,38 @@ func (l *inMemoryLimiter) bucketsMinID() int {
 	return l.buckets.getMinID()
 }
 
+func (l *inMemoryLimiter) getLimitCfg() limitCfg {
+	return limitCfg{}
+}
+
+func (l *inMemoryLimiter) getLimit() int64 {
+	return atomic.LoadInt64(&l.limit.value)
+}
+
+func (l *inMemoryLimiter) isLimitCfgChanged(curLimit int64, curDistribution []limitDistributionRatio) bool {
+	if l.getLimit() != curLimit {
+		return true
+	}
+
+	curDistributionsCount := 0
+
+	l.lock()
+	for _, ldRatio := range curDistribution {
+		curDistributionsCount += len(ldRatio.Values)
+		for _, fieldValue := range ldRatio.Values {
+			idx, has := l.limit.distributions.idxByKey[fieldValue]
+			if !has || math.Abs(l.limit.distributions.distributions[idx].ratio-ldRatio.Ratio) > EPS {
+				l.unlock()
+				return true
+			}
+		}
+	}
+	distributionsCount := len(l.limit.distributions.idxByKey)
+	l.unlock()
+
+	return distributionsCount != curDistributionsCount
+}
+
 func (l *inMemoryLimiter) setNowFn(fn func() time.Time) {
 	l.lock()
 	l.nowFn = fn

plugin/action/throttle/in_memory_limiter_test.go

@@ -0,0 +1,143 @@
+package throttle
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsLimitCfgChangedSame(t *testing.T) {
+	tests := []struct {
+		name            string
+		limiter         *inMemoryLimiter
+		curLimit        int64
+		curDistribution []limitDistributionRatio
+
+		want bool
+	}{
+		{
+			name: "same_limit_and_distribution",
+			limiter: &inMemoryLimiter{
+				limit: complexLimit{
+					value: 100,
+					distributions: limitDistributions{
+						idxByKey: map[string]int{
+							"A": 0,
+							"B": 1,
+						},
+						distributions: []complexDistribution{
+							{ratio: 0.1, limit: 10},
+							{ratio: 0.2, limit: 20},
+						},
+						enabled: true,
+					},
+				},
+			},
+			curLimit: 100,
+			curDistribution: []limitDistributionRatio{
+				{Ratio: 0.1, Values: []string{"A"}},
+				{Ratio: 0.2, Values: []string{"B"}},
+			},
+
+			want: false,
+		},
+		{
+			name: "limit_changed",
+			limiter: &inMemoryLimiter{
+				limit: complexLimit{
+					value: 200,
+				},
+			},
+			curLimit:        100,
+			curDistribution: []limitDistributionRatio{},
+
+			want: true,
+		},
+		{
+			name: "distribution_changed",
+			limiter: &inMemoryLimiter{
+				limit: complexLimit{
+					value: 100,
+					distributions: limitDistributions{
+						idxByKey: map[string]int{
+							"A": 0,
+							"B": 0,
+						},
+						distributions: []complexDistribution{
+							{ratio: 0.1, limit: 10},
+						},
+						enabled: true,
+					},
+				},
+			},
+			curLimit: 100,
+			curDistribution: []limitDistributionRatio{
+				{Ratio: 0.1, Values: []string{"A"}},
+				{Ratio: 0.2, Values: []string{"B"}},
+			},
+
+			want: true,
+		},
+		{
+			name: "increase_distributions_size",
+			limiter: &inMemoryLimiter{
+				limit: complexLimit{
+					value: 100,
+					distributions: limitDistributions{
+						idxByKey: map[string]int{
+							"A": 0,
+							"B": 1,
+							"C": 0,
+						},
+						distributions: []complexDistribution{
+							{ratio: 0.1, limit: 10},
+							{ratio: 0.2, limit: 20},
+						},
+						enabled: true,
+					},
+				},
+			},
+			curLimit: 100,
+			curDistribution: []limitDistributionRatio{
+				{Ratio: 0.1, Values: []string{"A"}},
+				{Ratio: 0.2, Values: []string{"B"}},
+			},
+
+			want: true,
+		},
+		{
+			name: "decrease_distributions_size",
+			limiter: &inMemoryLimiter{
+				limit: complexLimit{
+					value: 100,
+					distributions: limitDistributions{
+						idxByKey: map[string]int{
+							"A": 0,
+						},
+						distributions: []complexDistribution{
+							{ratio: 0.1, limit: 10},
+						},
+						enabled: true,
+					},
+				},
+			},
+			curLimit: 100,
+			curDistribution: []limitDistributionRatio{
+				{Ratio: 0.1, Values: []string{"A"}},
+				{Ratio: 0.2, Values: []string{"B"}},
+			},
+
+			want: true,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			got := tt.limiter.isLimitCfgChanged(tt.curLimit, tt.curDistribution)
+			require.Equal(t, tt.want, got)
+		})
+	}
+}