Adds timeout middleware

middleware/request/timeout.go

@@ -0,0 +1,27 @@
+package request
+
+import (
+	"context"
+	"time"
+
+	"github.com/ksysoev/wasabi"
+	"github.com/ksysoev/wasabi/dispatch"
+)
+
+// NewSetTimeoutMiddleware returns a middleware that sets a timeout for each request.
+// The timeout duration is specified by the 'timeout' parameter.
+// The returned middleware is a function that takes a 'next' request handler as input
+// and returns a new request handler that applies the timeout to the incoming request.
+func NewSetTimeoutMiddleware(timeout time.Duration) func(next wasabi.RequestHandler) wasabi.RequestHandler {
+	return func(next wasabi.RequestHandler) wasabi.RequestHandler {
+		return dispatch.RequestHandlerFunc(func(conn wasabi.Connection, req wasabi.Request) error {
+			ctx := req.Context()
+			ctx, cancel := context.WithTimeout(ctx, timeout)
+			defer cancel()
+
+			req = req.WithContext(ctx)
+
+			return next.Handle(conn, req)
+		})
+	}
+}

middleware/request/timeout_test.go

@@ -0,0 +1,36 @@
+package request
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/ksysoev/wasabi"
+	"github.com/ksysoev/wasabi/dispatch"
+	"github.com/ksysoev/wasabi/mocks"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestNewSetTimeoutMiddleware(t *testing.T) {
+	timeout := 5 * time.Second
+	handler := dispatch.RequestHandlerFunc(func(conn wasabi.Connection, req wasabi.Request) error {
+		// Your custom handler logic here
+		return nil
+	})
+
+	middleware := NewSetTimeoutMiddleware(timeout)(handler)
+
+	ctx := context.Background()
+
+	conn := mocks.NewMockConnection(t) // Create a mock connection
+	req := mocks.NewMockRequest(t)     // Create a mock request
+
+	req.EXPECT().Context().Return(ctx)
+	req.EXPECT().WithContext(mock.AnythingOfType("*context.timerCtx")).Return(req)
+
+	err := middleware.Handle(conn, req)
+
+	if err != nil {
+		t.Errorf("Expected no error, but got %v", err)
+	}
+}