Removing Benchmark & Updated implementation

src/mod/dynamicproxy/ratelimit.go

@@ -2,109 +2,55 @@ package dynamicproxy
 
 import (
 	"errors"
+	"log"
 	"net"
 	"net/http"
 	"sync"
+	"sync/atomic"
 	"time"
-
-	"log"
 )
 
-/*
+// IpTable is a rate limiter implementation using sync.Map with atomic int64
-	ratelimit.go
-
-	This file handles the ratelimit on proxy endpoints
-	if RateLimit is set to true
-*/
-
-// idk what this was for
-// func (h *ProxyHandler) handleRateLimitRouting(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
-// 	err := handleRateLimit(w, r, pe)
-// 	if err != nil {
-// 		h.logRequest(r, false, 429, "host", pe.Domain)
-// 	}
-// 	return err
-// }
-
 type IpTable struct {
-	sync.RWMutex
+	table sync.Map
-	table map[string]*IpTableValue
 }
 
-// Get the ip from the table
+// Increment the count of requests for a given IP
-func (t *IpTable) Get(ip string) (*IpTableValue, bool) {
-	t.RLock()
-	defer t.RUnlock()
-	v, ok := t.table[ip]
-	return v, ok
-}
-
-// Clear the ip from the table
-func (t *IpTable) Clear() {
-	t.Lock()
-	defer t.Unlock()
-	t.table = make(map[string]*IpTableValue)
-}
-
-// Increment the count of requests for a given ip
-// init ip in ipTable if not exists
 func (t *IpTable) Increment(ip string) {
-	t.Lock()
+	v, _ := t.table.LoadOrStore(ip, new(int64))
-	defer t.Unlock()
+	atomic.AddInt64(v.(*int64), 1)
-	v, ok := t.table[ip]
-	if !ok {
-		v = &IpTableValue{Count: 0, LastHit: time.Now()}
-	}
-	v.Count++
-	t.table[ip] = v
 }
 
-// Check if the ip is in the table and if it is, check if the count is less than the limit
+// Check if the IP is in the table and if it is, check if the count is less than the limit
 func (t *IpTable) Exceeded(ip string, limit int64) bool {
-	t.RLock()
+	v, ok := t.table.Load(ip)
-	defer t.RUnlock()
-	v, ok := t.table[ip]
 	if !ok {
 		return false
 	}
-	if v.Count < limit {
+	count := atomic.LoadInt64(v.(*int64))
-		return false
+	return count >= limit
-	}
-	return true
 }
 
-// Get the count of requests for a given ip
+// Get the count of requests for a given IP
-// returns 0 if ip is not in the table
 func (t *IpTable) GetCount(ip string) int64 {
-	t.RLock()
+	v, ok := t.table.Load(ip)
-	defer t.RUnlock()
-	v, ok := t.table[ip]
 	if !ok {
 		return 0
 	}
-	return v.Count
+	return atomic.LoadInt64(v.(*int64))
 }
 
-type IpTableValue struct {
+// Clear the IP table
-	Count   int64
+func (t *IpTable) Clear() {
-	LastHit time.Time
+	t.table.Range(func(key, value interface{}) bool {
+		t.table.Delete(key)
+		return true
+	})
 }
 
-var ipTable IpTable = IpTable{table: make(map[string]*IpTableValue)}
+var ipTable = IpTable{}
 
-// Handle rate limit logic
-// do not write to http.ResponseWriter if err return is not nil (already handled by this function)
 func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
-	// if len(pe.BasicAuthExceptionRules) > 0 {
-	// 	//Check if the current path matches the exception rules
-	// 	for _, exceptionRule := range pe.BasicAuthExceptionRules {
-	// 		if strings.HasPrefix(r.RequestURI, exceptionRule.PathPrefix) {
-	// 			//This path is excluded from basic auth
-	// 			return nil
-	// 		}
-	// 	}
-	// }
-
 	ip, _, err := net.SplitHostPort(r.RemoteAddr)
 	if err != nil {
 		w.WriteHeader(500)
@@ -114,7 +60,6 @@ func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint)
 
 	ipTable.Increment(ip)
 
-	// if ipTable.Exceeded(ip, pe.RateLimit) {
 	if ipTable.Exceeded(ip, 10) {
 		w.WriteHeader(429)
 		return errors.New("rate limit exceeded")

src/mod/dynamicproxy/ratelimit_benchmark_test.go

@@ -1,128 +0,0 @@
-package dynamicproxy
-
-import (
-	"net/http"
-	"net/http/httptest"
-	"sync"
-	"testing"
-)
-
-func BenchmarkRateLimitLockMapInt64(b *testing.B) {
-	go InitRateLimit()
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		handleRateLimit(w, r, &ProxyEndpoint{RateLimiting: 10})
-	})
-
-	request := httptest.NewRequest("GET", "/", nil)
-	request.RemoteAddr = "192.168.1.1:1234"
-
-	b.ResetTimer()
-
-	for i := 0; i < b.N; i++ {
-		recorder := httptest.NewRecorder()
-		handler.ServeHTTP(recorder, request)
-	}
-}
-
-func BenchmarkRateLimitSyncMapInt64(b *testing.B) {
-	go InitRateLimitSyncMapInt64()
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		handleRateLimitSyncMapInt64(w, r, &ProxyEndpoint{RateLimiting: 10})
-	})
-
-	request := httptest.NewRequest("GET", "/", nil)
-	request.RemoteAddr = "192.168.1.1:1234"
-
-	b.ResetTimer()
-
-	for i := 0; i < b.N; i++ {
-		recorder := httptest.NewRecorder()
-		handler.ServeHTTP(recorder, request)
-	}
-}
-
-func BenchmarkRateLimitSyncMapAtomicInt64(b *testing.B) {
-	go InitRateLimitSyncMapAtomicInt64()
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		handleRateLimitSyncMapAtomicInt64(w, r, &ProxyEndpoint{RateLimiting: 10})
-	})
-
-	request := httptest.NewRequest("GET", "/", nil)
-	request.RemoteAddr = "192.168.1.1:1234"
-
-	b.ResetTimer()
-
-	for i := 0; i < b.N; i++ {
-		recorder := httptest.NewRecorder()
-		handler.ServeHTTP(recorder, request)
-	}
-}
-
-func BenchmarkRateLimitLockMapInt64Concurrent(b *testing.B) {
-	go InitRateLimit()
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		handleRateLimit(w, r, &ProxyEndpoint{RateLimiting: 10})
-	})
-
-	request := httptest.NewRequest("GET", "/", nil)
-	request.RemoteAddr = "192.168.1.1:1234"
-
-	b.ResetTimer()
-
-	var wg sync.WaitGroup
-	for i := 0; i < b.N; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			recorder := httptest.NewRecorder()
-			handler.ServeHTTP(recorder, request)
-		}()
-	}
-	wg.Wait()
-}
-
-func BenchmarkRateLimitSyncMapInt64Concurrent(b *testing.B) {
-	go InitRateLimitSyncMapInt64()
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		handleRateLimitSyncMapInt64(w, r, &ProxyEndpoint{RateLimiting: 10})
-	})
-
-	request := httptest.NewRequest("GET", "/", nil)
-	request.RemoteAddr = "192.168.1.1:1234"
-
-	b.ResetTimer()
-
-	var wg sync.WaitGroup
-	for i := 0; i < b.N; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			recorder := httptest.NewRecorder()
-			handler.ServeHTTP(recorder, request)
-		}()
-	}
-	wg.Wait()
-}
-
-func BenchmarkRateLimitSyncMapAtomicInt64Concurrent(b *testing.B) {
-	go InitRateLimitSyncMapAtomicInt64()
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		handleRateLimitSyncMapAtomicInt64(w, r, &ProxyEndpoint{RateLimiting: 10})
-	})
-
-	request := httptest.NewRequest("GET", "/", nil)
-	request.RemoteAddr = "192.168.1.1:1234"
-
-	b.ResetTimer()
-
-	var wg sync.WaitGroup
-	for i := 0; i < b.N; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			recorder := httptest.NewRecorder()
-			handler.ServeHTTP(recorder, request)
-		}()
-	}
-	wg.Wait()
-}

src/mod/dynamicproxy/ratelimitb.go

@@ -1,79 +0,0 @@
-package dynamicproxy
-
-import (
-	"errors"
-	"log"
-	"net"
-	"net/http"
-	"sync"
-	"time"
-)
-
-// IpTableSyncMapInt64 is a rate limiter implementation using sync.Map with int64
-type IpTableSyncMapInt64 struct {
-	table sync.Map
-}
-
-// Increment the count of requests for a given IP
-func (t *IpTableSyncMapInt64) Increment(ip string) {
-	v, _ := t.table.LoadOrStore(ip, new(int64))
-	count := v.(*int64)
-	*count++
-}
-
-// Check if the IP is in the table and if it is, check if the count is less than the limit
-func (t *IpTableSyncMapInt64) Exceeded(ip string, limit int64) bool {
-	v, ok := t.table.Load(ip)
-	if !ok {
-		return false
-	}
-	count := v.(*int64)
-	return *count >= limit
-}
-
-// Get the count of requests for a given IP
-func (t *IpTableSyncMapInt64) GetCount(ip string) int64 {
-	v, ok := t.table.Load(ip)
-	if !ok {
-		return 0
-	}
-	count := v.(*int64)
-	return *count
-}
-
-// Clear the IP table
-func (t *IpTableSyncMapInt64) Clear() {
-	t.table.Range(func(key, value interface{}) bool {
-		t.table.Delete(key)
-		return true
-	})
-}
-
-var ipTableSyncMapInt64 = IpTableSyncMapInt64{}
-
-func handleRateLimitSyncMapInt64(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
-	ip, _, err := net.SplitHostPort(r.RemoteAddr)
-	if err != nil {
-		w.WriteHeader(500)
-		log.Println("Error resolving remote address", r.RemoteAddr, err)
-		return errors.New("internal server error")
-	}
-
-	ipTableSyncMapInt64.Increment(ip)
-
-	if ipTableSyncMapInt64.Exceeded(ip, 10) {
-		w.WriteHeader(429)
-		return errors.New("rate limit exceeded")
-	}
-
-	// log.Println("Rate limit check", ip, ipTableSyncMapInt64.GetCount(ip))
-
-	return nil
-}
-
-func InitRateLimitSyncMapInt64() {
-	for {
-		ipTableSyncMapInt64.Clear()
-		time.Sleep(time.Second)
-	}
-}

src/mod/dynamicproxy/ratelimitc.go

@@ -1,78 +0,0 @@
-package dynamicproxy
-
-import (
-	"errors"
-	"log"
-	"net"
-	"net/http"
-	"sync"
-	"sync/atomic"
-	"time"
-)
-
-// IpTableSyncMapAtomicInt64 is a rate limiter implementation using sync.Map with atomic int64
-type IpTableSyncMapAtomicInt64 struct {
-	table sync.Map
-}
-
-// Increment the count of requests for a given IP
-func (t *IpTableSyncMapAtomicInt64) Increment(ip string) {
-	v, _ := t.table.LoadOrStore(ip, new(int64))
-	atomic.AddInt64(v.(*int64), 1)
-}
-
-// Check if the IP is in the table and if it is, check if the count is less than the limit
-func (t *IpTableSyncMapAtomicInt64) Exceeded(ip string, limit int64) bool {
-	v, ok := t.table.Load(ip)
-	if !ok {
-		return false
-	}
-	count := atomic.LoadInt64(v.(*int64))
-	return count >= limit
-}
-
-// Get the count of requests for a given IP
-func (t *IpTableSyncMapAtomicInt64) GetCount(ip string) int64 {
-	v, ok := t.table.Load(ip)
-	if !ok {
-		return 0
-	}
-	return atomic.LoadInt64(v.(*int64))
-}
-
-// Clear the IP table
-func (t *IpTableSyncMapAtomicInt64) Clear() {
-	t.table.Range(func(key, value interface{}) bool {
-		t.table.Delete(key)
-		return true
-	})
-}
-
-var ipTableSyncMapAtomicInt64 = IpTableSyncMapAtomicInt64{}
-
-func handleRateLimitSyncMapAtomicInt64(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
-	ip, _, err := net.SplitHostPort(r.RemoteAddr)
-	if err != nil {
-		w.WriteHeader(500)
-		log.Println("Error resolving remote address", r.RemoteAddr, err)
-		return errors.New("internal server error")
-	}
-
-	ipTableSyncMapAtomicInt64.Increment(ip)
-
-	if ipTableSyncMapAtomicInt64.Exceeded(ip, 10) {
-		w.WriteHeader(429)
-		return errors.New("rate limit exceeded")
-	}
-
-	// log.Println("Rate limit check", ip, ipTableSyncMapAtomicInt64.GetCount(ip))
-
-	return nil
-}
-
-func InitRateLimitSyncMapAtomicInt64() {
-	for {
-		ipTableSyncMapAtomicInt64.Clear()
-		time.Sleep(time.Second)
-	}
-}