updated the Handle function

ziti-agent/cmd/webhook/mux.go

@@ -1,9 +1,7 @@
 package webhook
 
 import (
-	"context"
 	"net/http"
-	"sync"
 
 	"golang.org/x/time/rate"
 )
@@ -14,45 +12,47 @@ var rateLimiter = rate.NewLimiter(rate.Limit(1), 5)
 // CustomServeMux implements a custom ServeMux with a rate limiter and channel queue
 type customMux struct {
 	*http.ServeMux
-	queue chan *http.Request
-	wg    sync.WaitGroup
+	// queue chan *http.Request
+	// wg    sync.WaitGroup
 }
 
 func NewCustomMux() *customMux {
 	return &customMux{
 		http.NewServeMux(),
-		make(chan *http.Request, 10),
-		sync.WaitGroup{},
+		// make(chan *http.Request, 10),
+		// sync.WaitGroup{},
 	}
 }
 
-func (m *customMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	// Check rate limit
-	if rateLimiter.Allow() == false {
-		http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
-		return
-	}
+func (m *customMux) Handle(pattern string, handler http.HandlerFunc) {
+	m.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
+		// Check rate limit
+		if rateLimiter.Allow() == false {
+			http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
+			return
+		}
 
-	m.queue <- r
-	m.wg.Add(1)
+		// m.queue <- r
+		// m.wg.Add(1)
 
-	defer m.wg.Done()
+		// defer m.wg.Done()
 
-	for {
-		select {
-		case req := <-m.queue:
-			m.ServeMux.ServeHTTP(w, req)
-		case <-context.Background().Done():
-			return
-		}
-	}
+		// for {
+		// 	select {
+		// 	case req := <-m.queue:
+		// 		m.ServeMux.ServeHTTP(w, req)
+		// 	case <-context.Background().Done():
+		// 		return
+		// 	}
+		// }
 
-	// // Serve next request
-	// m.ServeMux.ServeHTTP(w, r)
+		// Serve next request
+		m.ServeMux.ServeHTTP(w, r)
+	})
 }
 
-func (m *customMux) Shutdown(ctx context.Context) error {
-	close(m.queue)
-	m.wg.Wait()
-	return nil
-}
+// func (m *customMux) Shutdown(ctx context.Context) error {
+// 	close(m.queue)
+// 	m.wg.Wait()
+// 	return nil
+// }