diff --git a/backend/http.go b/backend/http.go index a3c4396..a3f4172 100644 --- a/backend/http.go +++ b/backend/http.go @@ -1,7 +1,8 @@ package backend import ( - "bytes" + "fmt" + "io" "net/http" "time" @@ -9,6 +10,7 @@ import ( ) const defaultTimeout = 30 * time.Second +const defaultMaxReqPerHost = 50 // HTTPBackend represents an HTTP backend for handling requests. type HTTPBackend struct { @@ -18,6 +20,7 @@ type HTTPBackend struct { type httpBackendConfig struct { defaultTimeout time.Duration + maxReqPerHost int } type HTTPBackendOption func(*httpBackendConfig) @@ -26,6 +29,7 @@ type HTTPBackendOption func(*httpBackendConfig) func NewBackend(factory RequestFactory, options ...HTTPBackendOption) *HTTPBackend { httpBackendConfig := &httpBackendConfig{ defaultTimeout: defaultTimeout, + maxReqPerHost: defaultMaxReqPerHost, } for _, option := range options { @@ -36,10 +40,15 @@ func NewBackend(factory RequestFactory, options ...HTTPBackendOption) *HTTPBacke factory: factory, client: &http.Client{ Timeout: httpBackendConfig.defaultTimeout, + Transport: &http.Transport{ + MaxConnsPerHost: httpBackendConfig.maxReqPerHost, + }, }, } } +// Handle handles the incoming request by sending it to the HTTP server and returning the response. +// It takes a connection and a request as parameters and returns an error if any. func (b *HTTPBackend) Handle(conn wasabi.Connection, r wasabi.Request) error { httpReq, err := b.factory(r) if err != nil { @@ -55,18 +64,33 @@ func (b *HTTPBackend) Handle(conn wasabi.Connection, r wasabi.Request) error { defer resp.Body.Close() - respBody := bytes.NewBuffer(make([]byte, 0)) + length := 0 + if resp.ContentLength > 0 { + length = int(resp.ContentLength) + } else { + return fmt.Errorf("response content length is unknown") + } - _, err = respBody.ReadFrom(resp.Body) - if err != nil { + body := make([]byte, length) + _, err = resp.Body.Read(body) + + if err != nil && err != io.EOF { return err } - return conn.Send(wasabi.MsgTypeText, respBody.Bytes()) + return conn.Send(wasabi.MsgTypeText, body) } +// WithTimeout sets the default timeout for the HTTP client. func WithTimeout(timeout time.Duration) HTTPBackendOption { return func(cfg *httpBackendConfig) { cfg.defaultTimeout = timeout } } + +// WithMaxRequestsPerHost sets the maximum number of requests per host. +func WithMaxRequestsPerHost(maxReqPerHost int) HTTPBackendOption { + return func(cfg *httpBackendConfig) { + cfg.maxReqPerHost = maxReqPerHost + } +} diff --git a/backend/http_test.go b/backend/http_test.go index d44d40a..b1d6805 100644 --- a/backend/http_test.go +++ b/backend/http_test.go @@ -106,3 +106,20 @@ func TestHTTPBackend_Handle_TimeoutRequestByContext(t *testing.T) { t.Errorf("Expected error to be %v, but got %v", context.DeadlineExceeded, err) } } +func TestWithMaxRequestsPerHost(t *testing.T) { + maxReqPerHost := 100 + + backend := NewBackend(nil, WithMaxRequestsPerHost(maxReqPerHost)) + + if backend.client.Transport.(*http.Transport).MaxConnsPerHost != maxReqPerHost { + t.Errorf("Expected MaxConnsPerHost to be %v, but got %v", maxReqPerHost, backend.client.Transport.(*http.Transport).MaxConnsPerHost) + } +} + +func TestWithMaxRequestsPerHost_DefaultValue(t *testing.T) { + backend := NewBackend(nil) + + if backend.client.Transport.(*http.Transport).MaxConnsPerHost != defaultMaxReqPerHost { + t.Errorf("Expected MaxConnsPerHost to be %v, but got %v", defaultMaxReqPerHost, backend.client.Transport.(*http.Transport).MaxConnsPerHost) + } +} diff --git a/examples/http_backend/main.go b/examples/http_backend/main.go index b0997d4..d71bd14 100644 --- a/examples/http_backend/main.go +++ b/examples/http_backend/main.go @@ -11,7 +11,6 @@ import ( "github.com/ksysoev/wasabi/backend" "github.com/ksysoev/wasabi/channel" "github.com/ksysoev/wasabi/dispatch" - "github.com/ksysoev/wasabi/middleware/request" "github.com/ksysoev/wasabi/server" ) @@ -31,16 +30,9 @@ func main() { return httpReq, nil }) - ErrHandler := request.NewErrorHandlingMiddleware(func(conn wasabi.Connection, req wasabi.Request, err error) error { - conn.Send(wasabi.MsgTypeText, []byte("Failed to process request: "+err.Error())) - return nil - }) - dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, msgType wasabi.MessageType, data []byte) wasabi.Request { return dispatch.NewRawRequest(conn.Context(), msgType, data) }) - dispatcher.Use(ErrHandler) - dispatcher.Use(request.NewTrottlerMiddleware(100)) channel := channel.NewChannel("/", dispatcher, channel.NewConnectionRegistry(), channel.WithOriginPatterns("*"))