IG-15850: Add connection semaphore per context. (#66)

* Add connection semaphore per context. * Default to no semaphore. * int64 -> int Co-authored-by: Gal Topper <[email protected]>
v3io · Jun 2, 2020 · 9a94144 · 9a94144
1 parent 3bbe8d8
commit 9a94144
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 4 deletions.
diff --git a/go.mod b/go.mod
@@ -18,5 +18,6 @@ require (
 	github.com/valyala/fasthttp v1.2.0
 	go.uber.org/atomic v1.3.2 // indirect
 	go.uber.org/multierr v1.1.0 // indirect
+	golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a
 	zombiezen.com/go/capnproto2 v2.17.0+incompatible
 )
diff --git a/go.sum b/go.sum
@@ -44,6 +44,8 @@ go.uber.org/multierr v1.1.0 h1:HoEmRHQPVSqub6w2z2d2EOVs2fjyFRGyofhKuyDq0QI=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 golang.org/x/net v0.0.0-20180911220305-26e67e76b6c3 h1:czFLhve3vsQetD6JOJ8NZZvGQIXlnN3/yXxbT6/awxI=
 golang.org/x/net v0.0.0-20180911220305-26e67e76b6c3/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
+golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a h1:WXEvlFVvvGxCJLG6REjsT03iWnKLEWinaScsxF2Vm2o=
+golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223 h1:DH4skfRX4EBpamg7iV4ZlCpblAHI6s6TDM39bFZumv8=
 golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 zombiezen.com/go/capnproto2 v2.17.0+incompatible h1:sIoKPFGNlM38Qh+PBLa9Wzg1j99oInS/Qlk+5N/CHa4=

diff --git a/pkg/dataplane/http/context.go b/pkg/dataplane/http/context.go
@@ -2,6 +2,7 @@ package v3iohttp
 
 import (
 	"bytes"
+	goctx "context"
 	"crypto/tls"
 	"encoding/base64"
 	"encoding/json"
@@ -26,17 +27,19 @@ import (
 	"github.com/nuclio/errors"
 	"github.com/nuclio/logger"
 	"github.com/valyala/fasthttp"
+	"golang.org/x/sync/semaphore"
 	"zombiezen.com/go/capnproto2"
 )
 
 // TODO: Request should have a global pool
 var requestID uint64
 
 type context struct {
-	logger           logger.Logger
-	requestChan      chan *v3io.Request
-	httpClient       *fasthttp.Client
-	numWorkers       int
+	logger        logger.Logger
+	requestChan   chan *v3io.Request
+	httpClient    *fasthttp.Client
+	numWorkers    int
+	connSemaphore *semaphore.Weighted
 }
 
 type NewClientInput struct {
@@ -89,6 +92,10 @@ func NewContext(parentLogger logger.Logger, newContextInput *NewContextInput) (v
 		numWorkers:  numWorkers,
 	}
 
+	if newContextInput.MaxConns > 0 {
+		newContext.connSemaphore = semaphore.NewWeighted(int64(newContextInput.MaxConns))
+	}
+
 	for workerIndex := 0; workerIndex < numWorkers; workerIndex++ {
 		go newContext.workerEntry(workerIndex)
 	}
@@ -997,11 +1004,17 @@ func (c *context) sendRequest(dataPlaneInput *v3io.DataPlaneInput,
 		"method", method,
 		"body-length", len(body))
 
+	if c.connSemaphore != nil {
+		c.connSemaphore.Acquire(goctx.TODO(), 1)
+	}
 	if dataPlaneInput.Timeout <= 0 {
 		err = c.httpClient.Do(request, response.HTTPResponse)
 	} else {
 		err = c.httpClient.DoTimeout(request, response.HTTPResponse, dataPlaneInput.Timeout)
 	}
+	if c.connSemaphore != nil {
+		c.connSemaphore.Release(1)
+	}
 
 	if err != nil {
 		goto cleanup

diff --git a/pkg/dataplane/http/types.go b/pkg/dataplane/http/types.go
@@ -6,4 +6,5 @@ type NewContextInput struct {
 	HTTPClient     *fasthttp.Client
 	NumWorkers     int
 	RequestChanLen int
+	MaxConns       int
 }