From ee03f4127e16d545b9c99d39226d591b97074d8c Mon Sep 17 00:00:00 2001
From: Kirill Sysoev <k.sysoev@me.com>
Date: Sat, 13 Apr 2024 14:26:55 +0800
Subject: [PATCH] Improves timeout handling at http backend

---
 backend/http.go      | 33 +++++++++++++++++++++++++++++----
 backend/http_test.go |  3 +++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/backend/http.go b/backend/http.go
index c16397f..2e734df 100644
--- a/backend/http.go
+++ b/backend/http.go
@@ -3,32 +3,51 @@ package backend
 import (
 	"bytes"
 	"net/http"
+	"time"
 
 	"github.com/ksysoev/wasabi"
 )
 
+const defaultTimeout = 30 * time.Second
+
 // HTTPBackend represents an HTTP backend for handling requests.
 type HTTPBackend struct {
 	factory RequestFactory
 	client  *http.Client
 }
 
+type httpBackendConfig struct {
+	defaultTimeout time.Duration
+}
+
+type HTTPBackendOption func(*httpBackendConfig)
+
 // NewBackend creates a new instance of HTTPBackend with the given RequestFactory.
-func NewBackend(factory RequestFactory) *HTTPBackend {
+func NewBackend(factory RequestFactory, options ...HTTPBackendOption) *HTTPBackend {
+	httpBackendConfig := &httpBackendConfig{
+		defaultTimeout: defaultTimeout,
+	}
+
+	for _, option := range options {
+		option(httpBackendConfig)
+	}
+
 	return &HTTPBackend{
 		factory: factory,
-		client:  &http.Client{},
+		client: &http.Client{
+			Timeout: httpBackendConfig.defaultTimeout,
+		},
 	}
 }
 
-// Handle handles the incoming connection and request.
-// It sends the request to the backend server and returns the response to the connection.
 func (b *HTTPBackend) Handle(conn wasabi.Connection, r wasabi.Request) error {
 	httpReq, err := b.factory(r)
 	if err != nil {
 		return err
 	}
 
+	httpReq = httpReq.WithContext(r.Context())
+
 	resp, err := b.client.Do(httpReq)
 	if err != nil {
 		return err
@@ -45,3 +64,9 @@ func (b *HTTPBackend) Handle(conn wasabi.Connection, r wasabi.Request) error {
 
 	return conn.Send(respBody.String())
 }
+
+func WithDefaultTimeout(timeout time.Duration) HTTPBackendOption {
+	return func(cfg *httpBackendConfig) {
+		cfg.defaultTimeout = timeout
+	}
+}
diff --git a/backend/http_test.go b/backend/http_test.go
index efa717b..33259d2 100644
--- a/backend/http_test.go
+++ b/backend/http_test.go
@@ -2,6 +2,7 @@ package backend
 
 import (
 	"bytes"
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -19,6 +20,8 @@ func TestHTTPBackend_Handle(t *testing.T) {
 	mockConn := mocks.NewMockConnection(t)
 	mockReq := mocks.NewMockRequest(t)
 
+	mockReq.EXPECT().Context().Return(context.Background())
+
 	mockConn.EXPECT().Send("OK").Return(nil)
 	mockReq.EXPECT().Data().Return([]byte("test request"))