adds throttle grace

rakutentech · Oct 24, 2023 · 759a8c7 · 759a8c7
1 parent 5c251e0
commit 759a8c7
Show file tree

Hide file tree

Showing 3 changed files with 133 additions and 6 deletions.
diff --git a/README.md b/README.md
@@ -46,12 +46,12 @@ go install github.com/rakutentech/go-alertnotification@latest
 
 ### Throttling Configs
 
-| Env Variable           | default                                      | Explanation                    |
-| :--------------------- | :------------------------------------------- | :----------------------------- |
-| THROTTLE_DURATION      | 7                                            | throttling duration in minutes |
-| THROTTLE_GRACE_SECONDS | 0                                            | throttling grace in seconds    |
-| THROTTLE_DISKCACHE_DIR | `/tmp/cache/{APP_NAME}_throttler_disk_cache` | disk location for throttling   |
-| THROTTLE_ENABLED       | true                                         | Disable all together           |
+| Env Variable           | default                                      | Explanation                      |
+| :--------------------- | :------------------------------------------- | :------------------------------- |
+| THROTTLE_DURATION      | 7                                            | throttling duration in (minutes) |
+| THROTTLE_GRACE_SECONDS | 0                                            | throttling grace in (seconds)    |
+| THROTTLE_DISKCACHE_DIR | `/tmp/cache/{APP_NAME}_throttler_disk_cache` | disk location for throttling     |
+| THROTTLE_ENABLED       | true                                         | Disable all together             |
 
 ## Usage
 

diff --git a/throttler.go b/throttler.go
@@ -49,6 +49,14 @@ func (t *Throttler) IsThrottled(ocError error) bool {
 	if err != nil {
 		return false
 	}
+
+	gracedTime, gracedSeconds := getGracePeriod(dc, ocError)
+
+	if isWithinGracePeriod(string(gracedTime), gracedSeconds) {
+		// already graced and not over throttling duration, do nothing
+		return true
+	}
+
 	cachedTime, throttled := dc.Get(ocError.Error())
 
 	if throttled && !isOverThrottleDuration(string(cachedTime), t.ThrottleDuration) {
@@ -63,6 +71,35 @@ func (t *Throttler) IsThrottled(ocError error) bool {
 	return false
 }
 
+func getGracePeriod(dc *diskache.Diskache, ocError error) ([]byte, int) {
+	graceKeyPrefix := "GRACE-"
+	now := time.Now().Format(time.RFC3339)
+	gracedTime, gracedCached := dc.Get(graceKeyPrefix + ocError.Error())
+	gracedSeconds, err := strconv.Atoi(os.Getenv("THROTTLE_GRACE_SECONDS"))
+	if err != nil {
+		return []byte{}, 0 // no grace period
+	}
+
+	if !gracedCached {
+		gracedTime = []byte(now)
+		err = dc.Set(graceKeyPrefix+ocError.Error(), gracedTime)
+		if err != nil {
+			return []byte{}, 0 // no grace period
+		}
+	}
+	return gracedTime, gracedSeconds
+}
+
+func isWithinGracePeriod(cachedTime string, gracedSeconds int) bool {
+	gracedTime, err := time.Parse(time.RFC3339, string(cachedTime))
+	if err != nil {
+		return false
+	}
+	now := time.Now()
+	diff := int(now.Sub(gracedTime).Seconds())
+	return diff < gracedSeconds
+}
+
 func isOverThrottleDuration(cachedTime string, throttleDuration int) bool {
 	throttledTime, err := time.Parse(time.RFC3339, string(cachedTime))
 	if err != nil {

diff --git a/throttler_test.go b/throttler_test.go
@@ -113,6 +113,96 @@ func TestThrottler_IsThrottled(t *testing.T) {
 		})
 	}
 }
+func TestThrottler_IsThrottledGraced(t *testing.T) {
+	type fields struct {
+		CacheOpt         string
+		ThrottleDuration int
+	}
+	type args struct {
+		ocError error
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   bool
+	}{
+		{
+			name: "default",
+			fields: fields{
+				CacheOpt:         fmt.Sprintf("/tmp/cache/%v_throttler_disk_cache", os.Getenv("APP_NAME")),
+				ThrottleDuration: 5,
+			},
+			args: args{
+				ocError: errors.New("test_throttling"),
+			},
+			want: true,
+		},
+	}
+
+	os.Setenv("THROTTLE_GRACE_SECONDS", "10")
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			th := &Throttler{
+				CacheOpt:         tt.fields.CacheOpt,
+				ThrottleDuration: tt.fields.ThrottleDuration,
+			}
+			if got := th.IsThrottled(tt.args.ocError); got != tt.want {
+				t.Errorf("Throttler.IsThrottled() = %v, want %v", got, tt.want)
+			}
+			err := th.CleanThrottlingCache()
+			if err != nil {
+				t.Errorf("Cannot clean after test. %+v", err)
+			}
+
+		})
+	}
+}
+func TestThrottler_IsThrottledOverGraced(t *testing.T) {
+	type fields struct {
+		CacheOpt         string
+		ThrottleDuration int
+	}
+	type args struct {
+		ocError error
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   bool
+	}{
+		{
+			name: "default",
+			fields: fields{
+				CacheOpt:         fmt.Sprintf("/tmp/cache/%v_throttler_disk_cache", os.Getenv("APP_NAME")),
+				ThrottleDuration: 5,
+			},
+			args: args{
+				ocError: errors.New("test_throttling"),
+			},
+			want: false,
+		},
+	}
+
+	os.Setenv("THROTTLE_GRACE_SECONDS", "0")
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			th := &Throttler{
+				CacheOpt:         tt.fields.CacheOpt,
+				ThrottleDuration: tt.fields.ThrottleDuration,
+			}
+			if got := th.IsThrottled(tt.args.ocError); got != tt.want {
+				t.Errorf("Throttler.IsThrottled() = %v, want %v", got, tt.want)
+			}
+			err := th.CleanThrottlingCache()
+			if err != nil {
+				t.Errorf("Cannot clean after test. %+v", err)
+			}
+
+		})
+	}
+}
 
 func TestThrottler_ThrottleError(t *testing.T) {
 	type fields struct {