From d3dc094dbc68a3bd957e1fcf45d9bd0861cf8ced Mon Sep 17 00:00:00 2001
From: Artem Poltorzhitskiy <aopoltorzhicky@gmail.com>
Date: Wed, 3 Jul 2024 11:37:12 +0200
Subject: [PATCH] Fix: timeout middleware (#234)

---
 cmd/api/handler/error.go |  7 ++++++-
 cmd/api/init.go          |  8 ++------
 cmd/api/timeout.go       | 19 +++++++++++++++++++
 3 files changed, 27 insertions(+), 7 deletions(-)
 create mode 100644 cmd/api/timeout.go

diff --git a/cmd/api/handler/error.go b/cmd/api/handler/error.go
index a769e266..c95e60ab 100644
--- a/cmd/api/handler/error.go
+++ b/cmd/api/handler/error.go
@@ -48,7 +48,12 @@ func handleError(c echo.Context, err error, noRows NoRows) error {
 	if err.Error() == errCancelRequest {
 		return nil
 	}
-	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+	if errors.Is(err, context.DeadlineExceeded) {
+		return c.JSON(http.StatusRequestTimeout, Error{
+			Message: "timeout",
+		})
+	}
+	if errors.Is(err, context.Canceled) {
 		return c.JSON(http.StatusBadGateway, Error{
 			Message: err.Error(),
 		})
diff --git a/cmd/api/init.go b/cmd/api/init.go
index 164ef0e4..e3ffd984 100644
--- a/cmd/api/init.go
+++ b/cmd/api/init.go
@@ -202,11 +202,7 @@ func initEcho(cfg ApiConfig, db postgres.Storage, env string) *echo.Echo {
 	if cfg.RequestTimeout > 0 {
 		timeout = time.Duration(cfg.RequestTimeout) * time.Second
 	}
-	e.Use(middleware.TimeoutWithConfig(middleware.TimeoutConfig{
-		Skipper:      websocketSkipper,
-		Timeout:      timeout,
-		ErrorMessage: `{"message":"timeout"}`,
-	}))
+	e.Use(RequestTimeout(timeout))
 
 	e.Use(middleware.GzipWithConfig(middleware.GzipConfig{
 		Skipper: gzipSkipper,
@@ -281,7 +277,7 @@ func initHandlers(ctx context.Context, e *echo.Echo, cfg Config, db postgres.Sto
 	searchHandler := handler.NewSearchHandler(db.Search, db.Address, db.Blocks, db.Tx, db.Namespace, db.Validator, db.Rollup)
 	v1.GET("/search", searchHandler.Search)
 
-	ttlCache := cache.NewTTLCache(cache.Config{MaxEntitiesCount: 1000}, time.Minute*15)
+	ttlCache := cache.NewTTLCache(cache.Config{MaxEntitiesCount: 1000}, time.Minute*30)
 	ttlCacheMiddleware := cache.Middleware(ttlCache, nil)
 
 	addressHandlers := handler.NewAddressHandler(db.Address, db.Tx, db.BlobLogs, db.Message, db.Delegation, db.Undelegation, db.Redelegation, db.VestingAccounts, db.Grants, db.State, cfg.Indexer.Name)
diff --git a/cmd/api/timeout.go b/cmd/api/timeout.go
new file mode 100644
index 00000000..a9d861e5
--- /dev/null
+++ b/cmd/api/timeout.go
@@ -0,0 +1,19 @@
+package main
+
+import (
+	"context"
+	"time"
+
+	"github.com/labstack/echo/v4"
+)
+
+func RequestTimeout(timeout time.Duration) echo.MiddlewareFunc {
+	return func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) error {
+			timeoutCtx, cancel := context.WithTimeout(c.Request().Context(), timeout)
+			c.SetRequest(c.Request().WithContext(timeoutCtx))
+			defer cancel()
+			return next(c)
+		}
+	}
+}