github.com/m-lab/locate@v0.17.6/limits/ratelimiter_bench_test.go (about)

     1  package limits
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/gomodule/redigo/redis"
    10  )
    11  
    12  // getRedisMemory returns the current memory usage in bytes
    13  func getRedisMemory(pool *redis.Pool) (int64, error) {
    14  	conn := pool.Get()
    15  	defer conn.Close()
    16  
    17  	info, err := redis.String(conn.Do("INFO", "memory"))
    18  	if err != nil {
    19  		return 0, err
    20  	}
    21  
    22  	var used int64
    23  	for _, line := range strings.Split(info, "\r\n") {
    24  		if strings.HasPrefix(line, "used_memory:") {
    25  			fmt.Sscanf(line, "used_memory:%d", &used)
    26  			break
    27  		}
    28  	}
    29  	return used, nil
    30  }
    31  
    32  func BenchmarkRateLimiter_RealWorld(b *testing.B) {
    33  	pool := &redis.Pool{
    34  		MaxIdle:     3,
    35  		IdleTimeout: 240 * time.Second,
    36  		Dial: func() (redis.Conn, error) {
    37  			return redis.Dial("tcp", "localhost:6379")
    38  		},
    39  	}
    40  
    41  	// Create rate limiter with different limits for IP and IP+UA
    42  	limiter := NewRateLimiter(pool, RateLimitConfig{
    43  		IPConfig: LimitConfig{
    44  			Interval:  time.Hour,
    45  			MaxEvents: 120, // More permissive IP-only limit
    46  		},
    47  		IPUAConfig: LimitConfig{
    48  			Interval:  time.Hour,
    49  			MaxEvents: 60, // Stricter IP+UA limit
    50  		},
    51  		KeyPrefix: "benchmark:",
    52  	})
    53  
    54  	// Clean up before and after benchmark
    55  	conn := pool.Get()
    56  	conn.Do("FLUSHDB")
    57  	conn.Close()
    58  
    59  	defer func() {
    60  		conn := pool.Get()
    61  		defer conn.Close()
    62  		conn.Do("FLUSHDB")
    63  	}()
    64  
    65  	tests := []struct {
    66  		name     string
    67  		duration time.Duration
    68  		ipRange  int  // Number of unique IPs
    69  		uaRange  int  // Number of unique UAs
    70  		ipOnly   bool // Whether to test IP-only limiting
    71  
    72  	}{
    73  		{
    74  			name:    "SingleIPUA_Long",
    75  			ipRange: 1,
    76  			uaRange: 1,
    77  		},
    78  		{
    79  			name:    "ManyIPs_OneUA",
    80  			ipRange: 100000,
    81  			uaRange: 1,
    82  		},
    83  		{
    84  			name:    "OneIP_ManyUAs",
    85  			ipRange: 1,
    86  			uaRange: 100000,
    87  		},
    88  		{
    89  			name:    "ManyIPs_ManyUAs",
    90  			ipRange: 100,
    91  			uaRange: 10000,
    92  		},
    93  	}
    94  
    95  	for _, tt := range tests {
    96  		b.Run(tt.name, func(b *testing.B) {
    97  			// Get initial memory usage
    98  			memBefore, err := getRedisMemory(pool)
    99  			if err != nil {
   100  				b.Fatalf("Failed to get initial memory: %v", err)
   101  			}
   102  
   103  			// Reset timer for the actual benchmark
   104  			b.ResetTimer()
   105  
   106  			// Run b.N iterations
   107  			for i := 0; i < b.N; i++ {
   108  				ip := fmt.Sprintf("%d", i%tt.ipRange)
   109  				ua := fmt.Sprintf("%d", i%tt.uaRange)
   110  
   111  				_, err := limiter.IsLimited(ip, ua)
   112  				if err != nil {
   113  					b.Fatalf("IsLimited failed: %v", err)
   114  				}
   115  			}
   116  
   117  			b.StopTimer()
   118  
   119  			// Get final memory usage
   120  			memAfter, err := getRedisMemory(pool)
   121  			if err != nil {
   122  				b.Fatalf("Failed to get final memory: %v", err)
   123  			}
   124  
   125  			memoryUsed := memAfter - memBefore
   126  			b.ReportMetric(float64(memoryUsed)/float64(b.N), "bytes/op")
   127  
   128  			b.Logf("Iterations: %d", b.N)
   129  			b.Logf("  Memory used: %.2f MB", float64(memoryUsed)/(1024*1024))
   130  			b.Logf("  Memory per operation: %.2f bytes", float64(memoryUsed)/float64(b.N))
   131  		})
   132  
   133  		// Clean up between tests
   134  		conn := pool.Get()
   135  		conn.Do("FLUSHDB")
   136  		conn.Close()
   137  	}
   138  }