From 32a167b81356e19e9e173bb58a0503eea5e80e3d Mon Sep 17 00:00:00 2001
From: Marc Lopez Rubio <marc5.12@outlook.com>
Date: Tue, 4 Apr 2023 11:45:57 +0800
Subject: [PATCH] metrics: Count context.DeadlineExceeded as timeout (#10594)

Fixes the response.errors.timeout counter.

---------

Signed-off-by: Marc Lopez Rubio <marc5.12@outlook.com>
---
 changelogs/8.7.asciidoc                       |  1 +
 .../beater/middleware/timeout_middleware.go   |  2 +-
 .../middleware/timeout_middleware_test.go     | 54 ++++++++++++-------
 3 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/changelogs/8.7.asciidoc b/changelogs/8.7.asciidoc
index 7224dfed5c9..ea16e29c205 100644
--- a/changelogs/8.7.asciidoc
+++ b/changelogs/8.7.asciidoc
@@ -16,6 +16,7 @@ https://github.com/elastic/apm-server/compare/v8.7.0\...v8.7.1[View commits]
 ==== Bug fixes
 - Fix indexing failures for metric, error and log documents when `agent.activation_method` is set {pull}10552[10552]
 - Use droppedspan outcome in service destination aggregation {pull}10592[10592]
+- metrics: Count context.DeadlineExceeded as timeout {pull}10594[10594]
 
 [float]
 [[release-notes-8.7.0]]
diff --git a/internal/beater/middleware/timeout_middleware.go b/internal/beater/middleware/timeout_middleware.go
index 3bbb952e978..abd91a301e9 100644
--- a/internal/beater/middleware/timeout_middleware.go
+++ b/internal/beater/middleware/timeout_middleware.go
@@ -35,7 +35,7 @@ func TimeoutMiddleware() Middleware {
 			h(c)
 
 			err := c.Request.Context().Err()
-			if errors.Is(err, context.Canceled) {
+			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
 				c.Result.SetDefault(request.IDResponseErrorsTimeout)
 				c.Result.Err = tErr
 				c.Result.Body = tErr.Error()
diff --git a/internal/beater/middleware/timeout_middleware_test.go b/internal/beater/middleware/timeout_middleware_test.go
index 09020c19178..d98f19994e4 100644
--- a/internal/beater/middleware/timeout_middleware_test.go
+++ b/internal/beater/middleware/timeout_middleware_test.go
@@ -22,6 +22,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 
@@ -29,24 +30,41 @@ import (
 )
 
 func TestTimeoutMiddleware(t *testing.T) {
-	var err error
-	m := TimeoutMiddleware()
-	h := request.Handler(func(c *request.Context) {
-		ctx := c.Request.Context()
-		ctx, cancel := context.WithCancel(ctx)
-		r := c.Request.WithContext(ctx)
-		c.Request = r
-		cancel()
-	})
-
-	h, err = m(h)
-	assert.NoError(t, err)
+	test := func(t *testing.T, handler request.Handler) {
+		var err error
+		h, err := TimeoutMiddleware()(handler)
+		assert.NoError(t, err)
 
-	c := request.NewContext()
-	r, err := http.NewRequest("GET", "/", nil)
-	assert.NoError(t, err)
-	c.Reset(httptest.NewRecorder(), r)
-	h(c)
+		c := request.NewContext()
+		r, err := http.NewRequest("GET", "/", nil)
+		assert.NoError(t, err)
+		c.Reset(httptest.NewRecorder(), r)
+		h(c)
 
-	assert.Equal(t, http.StatusServiceUnavailable, c.Result.StatusCode)
+		assert.Equal(t, http.StatusServiceUnavailable, c.Result.StatusCode)
+	}
+	t.Run("Cancelled", func(t *testing.T) {
+		test(t, request.Handler(func(c *request.Context) {
+			ctx := c.Request.Context()
+			ctx, cancel := context.WithCancel(ctx)
+			r := c.Request.WithContext(ctx)
+			c.Request = r
+			cancel()
+		}))
+	})
+	t.Run("DeadlineExceeded", func(t *testing.T) {
+		var cancel func()
+		defer func() {
+			if cancel != nil {
+				cancel()
+			}
+		}()
+		test(t, request.Handler(func(c *request.Context) {
+			ctx := c.Request.Context()
+			ctx, cancel = context.WithTimeout(ctx, time.Nanosecond)
+			r := c.Request.WithContext(ctx)
+			c.Request = r
+			time.Sleep(time.Second)
+		}))
+	})
 }