From be3942a83b1f2ced232e3a3d7aad9dd0b766e369 Mon Sep 17 00:00:00 2001
From: Geon Kim <geon0250@gmail.com>
Date: Sat, 30 Nov 2024 22:31:44 +0900
Subject: [PATCH] test(memlimit): add test for WithRefreshInterval

---
 memlimit/memlimit_common_test.go | 92 ++++++++++++++++++++++++++++++++
 memlimit/provider.go             |  3 ++
 2 files changed, 95 insertions(+)

diff --git a/memlimit/memlimit_common_test.go b/memlimit/memlimit_common_test.go
index d5310ac..4acc2ee 100644
--- a/memlimit/memlimit_common_test.go
+++ b/memlimit/memlimit_common_test.go
@@ -4,7 +4,9 @@ import (
 	"fmt"
 	"math"
 	"runtime/debug"
+	"sync/atomic"
 	"testing"
+	"time"
 )
 
 func TestLimit(t *testing.T) {
@@ -205,3 +207,93 @@ func TestSetGoMemLimitWithOpts(t *testing.T) {
 		})
 	}
 }
+
+func TestSetGoMemLimitWithOpts_rollbackOnPanic(t *testing.T) {
+	t.Cleanup(func() {
+		debug.SetMemoryLimit(math.MaxInt64)
+	})
+
+	limit := int64(987654321)
+	_ = debug.SetMemoryLimit(987654321)
+	_, err := SetGoMemLimitWithOpts(
+		WithProvider(func() (uint64, error) {
+			debug.SetMemoryLimit(123456789)
+			panic("panic")
+		}),
+		WithRatio(1),
+	)
+	if err == nil {
+		t.Error("SetGoMemLimtWithOpts() error = nil, want panic")
+	}
+
+	curr := debug.SetMemoryLimit(-1)
+	if curr != limit {
+		t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, limit)
+	}
+}
+
+func TestSetGoMemLimitWithOpts_WithRefreshInterval(t *testing.T) {
+	t.Cleanup(func() {
+		debug.SetMemoryLimit(math.MaxInt64)
+	})
+
+	var limit atomic.Int64
+	output, err := SetGoMemLimitWithOpts(
+		WithProvider(func() (uint64, error) {
+			l := limit.Load()
+			if l == 0 {
+				return 0, ErrNoLimit
+			}
+			return uint64(l), nil
+		}),
+		WithRatio(1),
+		WithRefreshInterval(10*time.Millisecond),
+	)
+	if err != nil {
+		t.Errorf("SetGoMemLimitWithOpts() error = %v", err)
+	} else if output != limit.Load() {
+		t.Errorf("SetGoMemLimitWithOpts() got = %v, want %v", output, limit.Load())
+	}
+
+	// 1. no limit
+	curr := debug.SetMemoryLimit(-1)
+	if curr != math.MaxInt64 {
+		t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, limit.Load())
+	}
+
+	// 2. max limit
+	limit.Add(math.MaxInt64)
+	time.Sleep(100 * time.Millisecond)
+
+	curr = debug.SetMemoryLimit(-1)
+	if curr != math.MaxInt64 {
+		t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, math.MaxInt64)
+	}
+
+	// 3. adjust limit
+	limit.Add(-1024)
+	time.Sleep(100 * time.Millisecond)
+
+	curr = debug.SetMemoryLimit(-1)
+	if curr != math.MaxInt64-1024 {
+		t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, math.MaxInt64-1024)
+	}
+
+	// 4. no limit again (don't change the limit)
+	limit.Store(0)
+	time.Sleep(100 * time.Millisecond)
+
+	curr = debug.SetMemoryLimit(-1)
+	if curr != math.MaxInt64-1024 {
+		t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, math.MaxInt64-1024)
+	}
+
+	// 5. new limit
+	limit.Store(math.MaxInt32)
+	time.Sleep(100 * time.Millisecond)
+
+	curr = debug.SetMemoryLimit(-1)
+	if curr != math.MaxInt32 {
+		t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, math.MaxInt32)
+	}
+}
diff --git a/memlimit/provider.go b/memlimit/provider.go
index 32cc1ee..4f83770 100644
--- a/memlimit/provider.go
+++ b/memlimit/provider.go
@@ -16,6 +16,9 @@ func Limit(limit uint64) func() (uint64, error) {
 
 // ApplyRationA is a helper Provider function that applies the given ratio to the given provider.
 func ApplyRatio(provider Provider, ratio float64) Provider {
+	if ratio == 1 {
+		return provider
+	}
 	return func() (uint64, error) {
 		if ratio <= 0 || ratio > 1 {
 			return 0, fmt.Errorf("invalid ratio: %f, ratio should be in the range (0.0,1.0]", ratio)