feat: Add connection limit option to the server

freepik-company · Oct 21, 2024 · ff17a70 · ff17a70
1 parent c3612c7
commit ff17a70
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 3 deletions.
diff --git a/api/config_types.go b/api/config_types.go
@@ -37,8 +37,9 @@ type ListenerT struct {
 }
 
 type ListenerOptionsT struct {
-	ReadTimeout  string `yaml:"readTimeout"`
-	WriteTimeout string `yaml:"writeTimeout"`
+	ReadTimeout              string `yaml:"readTimeout"`
+	WriteTimeout             string `yaml:"writeTimeout"`
+	MaxConcurrentConnections int    `yaml:"maxConcurrentConnections"`
 
 	// Carry stuff
 	ReadTimeoutDuration  *time.Duration

diff --git a/docs/samples/bifrost.yaml b/docs/samples/bifrost.yaml
@@ -10,6 +10,7 @@ listener:
   options:
     readTimeout: 10s
     writeTimeout: 10s
+    maxConcurrentConnections: 1000
 
 # (Optional) Authentication configuration
 authentication:

diff --git a/go.mod b/go.mod
@@ -24,4 +24,5 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/sts v1.32.2 // indirect
 	github.com/aws/smithy-go v1.22.0 // indirect
 	go.uber.org/multierr v1.10.0 // indirect
+	golang.org/x/net v0.30.0 // indirect
 )
diff --git a/go.sum b/go.sum
@@ -36,6 +36,8 @@ go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ=
 go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
 go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
 go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
+golang.org/x/net v0.30.0 h1:AcW1SDZMkb8IpzCdQUaIq2sP4sZ4zw+55h6ynffypl4=
+golang.org/x/net v0.30.0/go.mod h1:2wGyMJ5iFasEhkwi13ChkO/t1ECNC4X4eBKkVFyYFlU=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

diff --git a/internal/httpserver/httpserver.go b/internal/httpserver/httpserver.go
@@ -14,7 +14,9 @@ import (
 	"bifrost/api"
 	"bifrost/internal/globals"
 	"bifrost/internal/signature"
+
 	//
+	"golang.org/x/net/netutil"
 )
 
 type HttpServer struct {
@@ -34,6 +36,16 @@ func NewHttpServer() (server *HttpServer) {
 	return server
 }
 
+// TODO
+func (s *HttpServer) SetAddr(addr string) {
+	s.Server.Addr = addr
+}
+
+// TODO
+func (s *HttpServer) SetHandler(handler http.Handler) {
+	s.Server.Handler = handler
+}
+
 // getRequestAuthParam extracts a parameter from the Authorization header of the request
 // The header value is expected to be in the format: <Auth type> <Param1=value1,Param2=value2>
 func getRequestAuthParam(request *http.Request, param string) (string, error) {
@@ -371,10 +383,25 @@ func (s *HttpServer) Run(httpAddr string) {
 	globals.Application.Logger.Infof("Starting HTTP server on %s", httpAddr)
 
 	// TODO: Configure and use the server previously crafted
-	err := http.ListenAndServe(httpAddr, mux)
+	s.SetAddr(httpAddr)
+	s.SetHandler(mux)
+
+	//
+	listener, err := net.Listen("tcp", s.Server.Addr)
+	if err != nil {
+		globals.Application.Logger.Errorf("Server failed. Reason: %s", err.Error())
+	}
+
+	limitedListener := listener
+	if globals.Application.Config.Listener.Options.MaxConcurrentConnections > 0 {
+		limitedListener = netutil.LimitListener(listener, globals.Application.Config.Listener.Options.MaxConcurrentConnections)
+	}
+
+	err = s.Server.Serve(limitedListener)
 	if err != nil {
 		globals.Application.Logger.Errorf("Server failed. Reason: %s", err.Error())
 	}
+
 }
 
 func (s *HttpServer) Stop() {