Add concurrency limitation feature

examples/go.mod

@@ -1,8 +1,18 @@
 module github.com/elazarl/goproxy/examples/goproxy-transparent
 
+go 1.20
+
+require (
+	github.com/elazarl/goproxy v0.0.0-20241217120900-7711dfa3811c
+	github.com/elazarl/goproxy/ext v0.0.0-20241217120900-7711dfa3811c
+	github.com/gorilla/websocket v1.5.3
+	github.com/inconshreveable/go-vhost v1.0.0
+)
+
 require (
-	github.com/elazarl/goproxy v0.0.0-20181111060418-2ce16c963a8a
-	github.com/inconshreveable/go-vhost v0.0.0-20160627193104-06d84117953b
+	github.com/rogpeppe/go-charset v0.0.0-20190617161244-0dc95cdf6f31 // indirect
+	golang.org/x/net v0.32.0 // indirect
+	golang.org/x/text v0.21.0 // indirect
 )
 
 replace github.com/elazarl/goproxy => ../

examples/go.sum

@@ -1,4 +1,13 @@
-github.com/elazarl/goproxy/ext v0.0.0-20190711103511-473e67f1d7d2/go.mod h1:gNh8nYJoAm43RfaxurUnxr+N1PwuFV3ZMl/efxlIlY8=
-github.com/inconshreveable/go-vhost v0.0.0-20160627193104-06d84117953b h1:IpLPmn6Re21F0MaV6Zsc5RdSE6KuoFpWmHiUSEs3PrE=
-github.com/inconshreveable/go-vhost v0.0.0-20160627193104-06d84117953b/go.mod h1:aA6DnFhALT3zH0y+A39we+zbrdMC2N0X/q21e6FI0LU=
+github.com/elazarl/goproxy/ext v0.0.0-20241217120900-7711dfa3811c h1:R+i10jtNSzKJKqEZAYJnR9M8y14k0zrNHqD1xkv/A2M=
+github.com/elazarl/goproxy/ext v0.0.0-20241217120900-7711dfa3811c/go.mod h1:gNh8nYJoAm43RfaxurUnxr+N1PwuFV3ZMl/efxlIlY8=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/inconshreveable/go-vhost v1.0.0 h1:IK4VZTlXL4l9vz2IZoiSFbYaaqUW7dXJAiPriUN5Ur8=
+github.com/inconshreveable/go-vhost v1.0.0/go.mod h1:aA6DnFhALT3zH0y+A39we+zbrdMC2N0X/q21e6FI0LU=
 github.com/rogpeppe/go-charset v0.0.0-20180617210344-2471d30d28b4/go.mod h1:qgYeAmZ5ZIpBWTGllZSQnw97Dj+woV0toclVaRGI8pc=
+github.com/rogpeppe/go-charset v0.0.0-20190617161244-0dc95cdf6f31 h1:DE4LcMKyqAVa6a0CGmVxANbnVb7stzMmPkQiieyNmfQ=
+github.com/rogpeppe/go-charset v0.0.0-20190617161244-0dc95cdf6f31/go.mod h1:qgYeAmZ5ZIpBWTGllZSQnw97Dj+woV0toclVaRGI8pc=
+golang.org/x/net v0.32.0 h1:ZqPmj8Kzc+Y6e0+skZsuACbx+wzMgo5MQsJh9Qd6aYI=
+golang.org/x/net v0.32.0/go.mod h1:CwU0IoeOlnQQWJ6ioyFrfRuomB8GKF6KbYXZVyeXNfs=
+golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
+golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=

ext/go.mod

@@ -1,3 +1,13 @@
 module github.com/elazarl/goproxy/ext
 
-require github.com/rogpeppe/go-charset v0.0.0-20180617210344-2471d30d28b4 // indirect
+go 1.20
+
+require (
+	github.com/elazarl/goproxy v0.0.0-20241217120900-7711dfa3811c
+	github.com/rogpeppe/go-charset v0.0.0-20190617161244-0dc95cdf6f31
+)
+
+require (
+	golang.org/x/net v0.32.0 // indirect
+	golang.org/x/text v0.21.0 // indirect
+)

ext/go.sum

@@ -1,2 +1,10 @@
+github.com/elazarl/goproxy v0.0.0-20241217120900-7711dfa3811c h1:yWAGp1CjD1mQGLUsADqPn5s1n2AkGAX33XLDUgoXzyo=
+github.com/elazarl/goproxy v0.0.0-20241217120900-7711dfa3811c/go.mod h1:P73liMk9TZCyF9fXG/RyMeSizmATvpvy3ZS61/1eXn4=
 github.com/rogpeppe/go-charset v0.0.0-20180617210344-2471d30d28b4 h1:BN/Nyn2nWMoqGRA7G7paDNDqTXE30mXGqzzybrfo05w=
 github.com/rogpeppe/go-charset v0.0.0-20180617210344-2471d30d28b4/go.mod h1:qgYeAmZ5ZIpBWTGllZSQnw97Dj+woV0toclVaRGI8pc=
+github.com/rogpeppe/go-charset v0.0.0-20190617161244-0dc95cdf6f31 h1:DE4LcMKyqAVa6a0CGmVxANbnVb7stzMmPkQiieyNmfQ=
+github.com/rogpeppe/go-charset v0.0.0-20190617161244-0dc95cdf6f31/go.mod h1:qgYeAmZ5ZIpBWTGllZSQnw97Dj+woV0toclVaRGI8pc=
+golang.org/x/net v0.32.0 h1:ZqPmj8Kzc+Y6e0+skZsuACbx+wzMgo5MQsJh9Qd6aYI=
+golang.org/x/net v0.32.0/go.mod h1:CwU0IoeOlnQQWJ6ioyFrfRuomB8GKF6KbYXZVyeXNfs=
+golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
+golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=

ext/limitation/concurrency.go

@@ -0,0 +1,32 @@
+package limitation
+
+import (
+	"net/http"
+
+	"github.com/elazarl/goproxy"
+)
+
+// ConcurrentRequests implements a mechanism to limit the number of
+// concurrently handled HTTP requests, configurable by the user.
+// The ReqHandler can simply be added to the server with OnRequest().
+func ConcurrentRequests(limit int) goproxy.ReqHandler {
+	// Do nothing when the specified limit is invalid
+	if limit <= 0 {
+		return goproxy.FuncReqHandler(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
+			return req, nil
+		})
+	}
+
+	limitation := make(chan struct{}, limit)
+	return goproxy.FuncReqHandler(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
+		limitation <- struct{}{}
+
+		// Release semaphore when request finishes
+		go func() {
+			<-req.Context().Done()
+			<-limitation
+		}()
+
+		return req, nil
+	})
+}

ext/limitation/concurrency_test.go

@@ -0,0 +1,96 @@
+package limitation_test
+
+import (
+	"context"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/elazarl/goproxy"
+	"github.com/elazarl/goproxy/ext/limitation"
+)
+
+func TestConcurrentRequests(t *testing.T) {
+	mockRequest := &http.Request{Host: "test.com"}
+	ctx := &goproxy.ProxyCtx{}
+	maximumDuration := 100 * time.Millisecond
+
+	t.Run("empty limitation", func(t *testing.T) {
+		timer := time.NewTimer(maximumDuration)
+		defer timer.Stop()
+		done := make(chan struct{})
+
+		go func() {
+			zeroLimiter := limitation.ConcurrentRequests(0)
+			zeroLimiter.Handle(mockRequest, ctx)
+			done <- struct{}{}
+		}()
+
+		select {
+		case <-timer.C:
+			t.Error("Limiter took too long")
+		case <-done:
+		}
+	})
+
+	t.Run("normal limitation", func(t *testing.T) {
+		timer := time.NewTimer(maximumDuration)
+		defer timer.Stop()
+		done := make(chan struct{})
+
+		go func() {
+			oneLimiter := limitation.ConcurrentRequests(1)
+			oneLimiter.Handle(mockRequest, ctx)
+			done <- struct{}{}
+		}()
+
+		select {
+		case <-timer.C:
+			t.Error("Limiter took too long")
+		case <-done:
+		}
+	})
+
+	t.Run("more than the limit", func(t *testing.T) {
+		timer := time.NewTimer(maximumDuration)
+		defer timer.Stop()
+		done := make(chan struct{})
+
+		go func() {
+			oneLimiter := limitation.ConcurrentRequests(1)
+			oneLimiter.Handle(mockRequest, ctx)
+			oneLimiter.Handle(mockRequest, ctx)
+			done <- struct{}{}
+		}()
+
+		select {
+		case <-timer.C:
+			// Do nothing, we expect to reach the timeout
+		case <-done:
+			t.Error("Limiter was too fast")
+		}
+	})
+
+	t.Run("more than the limit but one request finishes", func(t *testing.T) {
+		timer := time.NewTimer(maximumDuration)
+		defer timer.Stop()
+		done := make(chan struct{})
+
+		timeoutCtx, cancel := context.WithCancel(mockRequest.Context())
+		mockRequestWithCancel := mockRequest.WithContext(timeoutCtx)
+
+		go func() {
+			oneLimiter := limitation.ConcurrentRequests(1)
+			oneLimiter.Handle(mockRequestWithCancel, ctx)
+			cancel()
+			oneLimiter.Handle(mockRequest, ctx)
+			done <- struct{}{}
+		}()
+
+		select {
+		case <-timer.C:
+			t.Error("Limiter took too long")
+		case <-done:
+		}
+	})
+}

go.mod

@@ -3,7 +3,7 @@ module github.com/elazarl/goproxy
 go 1.20
 
 require (
-	github.com/elazarl/goproxy/ext v0.0.0-20241216102027-e85c60b37433
+	github.com/elazarl/goproxy/ext v0.0.0-20241217120900-7711dfa3811c
 	golang.org/x/net v0.32.0
 )
 

go.sum

@@ -1,5 +1,5 @@
-github.com/elazarl/goproxy/ext v0.0.0-20241216102027-e85c60b37433 h1:zezqs+UN/8nYMOm1PobfrT/FxliWYq5Um1DLxyHA8d0=
-github.com/elazarl/goproxy/ext v0.0.0-20241216102027-e85c60b37433/go.mod h1:gNh8nYJoAm43RfaxurUnxr+N1PwuFV3ZMl/efxlIlY8=
+github.com/elazarl/goproxy/ext v0.0.0-20241217120900-7711dfa3811c h1:R+i10jtNSzKJKqEZAYJnR9M8y14k0zrNHqD1xkv/A2M=
+github.com/elazarl/goproxy/ext v0.0.0-20241217120900-7711dfa3811c/go.mod h1:gNh8nYJoAm43RfaxurUnxr+N1PwuFV3ZMl/efxlIlY8=
 github.com/rogpeppe/go-charset v0.0.0-20180617210344-2471d30d28b4/go.mod h1:qgYeAmZ5ZIpBWTGllZSQnw97Dj+woV0toclVaRGI8pc=
 golang.org/x/net v0.32.0 h1:ZqPmj8Kzc+Y6e0+skZsuACbx+wzMgo5MQsJh9Qd6aYI=
 golang.org/x/net v0.32.0/go.mod h1:CwU0IoeOlnQQWJ6ioyFrfRuomB8GKF6KbYXZVyeXNfs=