Fixing Timeout middleware Context propagation (#1910)

pafuent · web-flow · commit 02de901d7ef5 · 2021-07-09T23:36:03.000-03:00
This will let middlewares/handler later on the chain to properly handle the Timeout middleware Context cancellation. Fixes #1909
diff --git a/middleware/timeout.go b/middleware/timeout.go
@@ -2,9 +2,10 @@ package middleware
 
 import (
 	"context"
-	"github.com/labstack/echo/v4"
 	"net/http"
 	"time"
+
+	"github.com/labstack/echo/v4"
 )
 
 type (
@@ -87,6 +88,10 @@ type echoHandlerFuncWrapper struct {
 }
 
 func (t echoHandlerFuncWrapper) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+	// replace echo.Context Request with the one provided by TimeoutHandler to let later middlewares/handler on the chain
+	// handle properly it's cancellation
+	t.ctx.SetRequest(r)
+
 	// replace writer with TimeoutHandler custom one. This will guarantee that
 	// `writes by h to its ResponseWriter will return ErrHandlerTimeout.`
 	originalWriter := t.ctx.Response().Writer
diff --git a/middleware/timeout_test.go b/middleware/timeout_test.go
@@ -2,15 +2,16 @@ package middleware
 
 import (
 	"errors"
-	"github.com/labstack/echo/v4"
-	"github.com/stretchr/testify/assert"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"reflect"
 	"strings"
 	"testing"
 	"time"
+
+	"github.com/labstack/echo/v4"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestTimeoutSkipper(t *testing.T) {
@@ -273,3 +274,42 @@ func TestTimeoutWithDefaultErrorMessage(t *testing.T) {
 	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
 	assert.Equal(t, `<html><head><title>Timeout</title></head><body><h1>Timeout</h1></body></html>`, rec.Body.String())
 }
+
+func TestTimeoutCanHandleContextDeadlineOnNextHandler(t *testing.T) {
+	t.Parallel()
+
+	timeout := 1 * time.Millisecond
+	m := TimeoutWithConfig(TimeoutConfig{
+		Timeout:      timeout,
+		ErrorMessage: "Timeout! change me",
+	})
+
+	handlerFinishedExecution := make(chan bool)
+
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	rec := httptest.NewRecorder()
+
+	e := echo.New()
+	c := e.NewContext(req, rec)
+
+	stopChan := make(chan struct{})
+	err := m(func(c echo.Context) error {
+		// NOTE: when difference between timeout duration and handler execution time is almost the same (in range of 100microseconds)
+		// the result of timeout does not seem to be reliable - could respond timeout, could respond handler output
+		// difference over 500microseconds (0.5millisecond) response seems to be reliable
+		<-stopChan
+
+		// The Request Context should have a Deadline set by http.TimeoutHandler
+		if _, ok := c.Request().Context().Deadline(); !ok {
+			assert.Fail(t, "No timeout set on Request Context")
+		}
+		handlerFinishedExecution <- c.Request().Context().Err() == nil
+		return c.String(http.StatusOK, "Hello, World!")
+	})(c)
+	stopChan <- struct{}{}
+
+	assert.NoError(t, err)
+	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
+	assert.Equal(t, "Timeout! change me", rec.Body.String())
+	assert.False(t, <-handlerFinishedExecution)
+}
