feat: add global api timeout

formancehq · Oct 22, 2024 · 0d5c5b9 · 0d5c5b9
1 parent 084d78d
commit 0d5c5b9
Show file tree

Hide file tree

Showing 17 changed files with 258 additions and 29 deletions.
diff --git a/cmd/serve.go b/cmd/serve.go
@@ -5,6 +5,7 @@ import (
 	"github.com/formancehq/go-libs/v2/health"
 	"github.com/formancehq/go-libs/v2/httpserver"
 	"github.com/formancehq/go-libs/v2/otlp"
+	"github.com/formancehq/ledger/internal/api/common"
 	"github.com/formancehq/ledger/internal/storage/driver"
 	"github.com/go-chi/chi/v5"
 	"go.opentelemetry.io/otel/sdk/metric"
@@ -31,11 +32,13 @@ import (
 )
 
 const (
-	BindFlag                   = "bind"
-	BallastSizeInBytesFlag     = "ballast-size"
-	NumscriptCacheMaxCountFlag = "numscript-cache-max-count"
-	AutoUpgradeFlag            = "auto-upgrade"
-	ExperimentalFeaturesFlag   = "experimental-features"
+	BindFlag                         = "bind"
+	BallastSizeInBytesFlag           = "ballast-size"
+	NumscriptCacheMaxCountFlag       = "numscript-cache-max-count"
+	AutoUpgradeFlag                  = "auto-upgrade"
+	APIResponseTimeoutDelayFlag      = "api-response-timeout-delay"
+	APIResponseTimeoutStatusCodeFlag = "api-response-timeout-status-code"
+	ExperimentalFeaturesFlag         = "experimental-features"
 )
 
 func NewServeCommand() *cobra.Command {
@@ -55,6 +58,16 @@ func NewServeCommand() *cobra.Command {
 				return err
 			}
 
+			apiResponseTimeoutDelay, err := cmd.Flags().GetDuration(APIResponseTimeoutDelayFlag)
+			if err != nil {
+				return err
+			}
+
+			apiResponseTimeoutStatusCode, err := cmd.Flags().GetInt(APIResponseTimeoutStatusCodeFlag)
+			if err != nil {
+				return err
+			}
+
 			options := []fx.Option{
 				fx.NopLogger,
 				otlp.FXModuleFromFlags(cmd),
@@ -79,6 +92,10 @@ func NewServeCommand() *cobra.Command {
 				api.Module(api.Config{
 					Version: Version,
 					Debug:   service.IsDebug(cmd),
+					Timeout: common.TimeoutConfiguration{
+						Timeout:    apiResponseTimeoutDelay,
+						StatusCode: apiResponseTimeoutStatusCode,
+					},
 				}),
 				fx.Decorate(func(
 					params struct {
@@ -112,6 +129,8 @@ func NewServeCommand() *cobra.Command {
 	cmd.Flags().Bool(AutoUpgradeFlag, false, "Automatically upgrade all schemas")
 	cmd.Flags().String(BindFlag, "0.0.0.0:3068", "API bind address")
 	cmd.Flags().Bool(ExperimentalFeaturesFlag, false, "Enable features configurability")
+	cmd.Flags().Duration(APIResponseTimeoutDelayFlag, 0, "API response timeout delay")
+	cmd.Flags().Int(APIResponseTimeoutStatusCodeFlag, http.StatusGatewayTimeout, "API response timeout status code")
 
 	service.AddFlags(cmd.Flags())
 	bunconnect.AddFlags(cmd.Flags())

diff --git a/internal/api/common/middlewares_timeout.go b/internal/api/common/middlewares_timeout.go
@@ -0,0 +1,89 @@
+package common
+
+import (
+	"context"
+	"fmt"
+	"github.com/formancehq/go-libs/v2/api"
+	"net/http"
+	"sync"
+	"time"
+)
+
+type responseWriter struct {
+	http.ResponseWriter
+	sync.Mutex
+	contextWithTimeout context.Context
+	headerWritten      bool
+}
+
+func (r *responseWriter) Write(bytes []byte) (int, error) {
+	r.Lock()
+	defer r.Unlock()
+
+	// As the http status code has already been sent, we ignore the context canceled in this case.
+	// The timeout is usually shorter than the allowed window.
+	// We let a chance to the client to get their data.
+	if r.headerWritten {
+		return r.ResponseWriter.Write(bytes)
+	}
+
+	select {
+	case <-r.contextWithTimeout.Done():
+		return 0, nil
+	default:
+		return r.ResponseWriter.Write(bytes)
+	}
+}
+
+func (r *responseWriter) WriteHeader(statusCode int) {
+	r.Lock()
+	defer r.Unlock()
+
+	select {
+	case <-r.contextWithTimeout.Done():
+	default:
+		r.ResponseWriter.WriteHeader(statusCode)
+		r.headerWritten = true
+	}
+}
+
+var _ http.ResponseWriter = &responseWriter{}
+
+type TimeoutConfiguration struct {
+	Timeout    time.Duration
+	StatusCode int
+}
+
+func Timeout(configuration TimeoutConfiguration) func(h http.Handler) http.Handler {
+	return func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			newRequestContext, cancelRequestContext := context.WithCancel(r.Context())
+			defer cancelRequestContext()
+
+			contextWithTimeout, cancelContextWithTimeout := context.WithTimeout(context.Background(), configuration.Timeout)
+			defer cancelContextWithTimeout()
+
+			rw := &responseWriter{
+				ResponseWriter:     w,
+				contextWithTimeout: contextWithTimeout,
+			}
+			done := make(chan struct{})
+			go func() {
+				select {
+				case <-done:
+				case <-contextWithTimeout.Done():
+					rw.Lock()
+					defer rw.Unlock()
+
+					if !rw.headerWritten {
+						cancelRequestContext()
+						api.WriteErrorResponse(w, configuration.StatusCode, "TIMEOUT", fmt.Errorf("request timed out"))
+					}
+				}
+			}()
+			defer close(done)
+
+			h.ServeHTTP(w, r.WithContext(newRequestContext))
+		})
+	}
+}
diff --git a/internal/api/module.go b/internal/api/module.go
@@ -2,19 +2,20 @@ package api
 
 import (
 	_ "embed"
+	"github.com/formancehq/go-libs/v2/auth"
+	"github.com/formancehq/go-libs/v2/health"
 	"github.com/formancehq/go-libs/v2/logging"
+	"github.com/formancehq/ledger/internal/api/common"
 	"github.com/formancehq/ledger/internal/controller/system"
 	"github.com/go-chi/chi/v5"
 	"go.opentelemetry.io/otel/trace"
-
-	"github.com/formancehq/go-libs/v2/auth"
-	"github.com/formancehq/go-libs/v2/health"
 	"go.uber.org/fx"
 )
 
 type Config struct {
 	Version string
 	Debug   bool
+	Timeout common.TimeoutConfiguration
 }
 
 func Module(cfg Config) fx.Option {
@@ -32,6 +33,7 @@ func Module(cfg Config) fx.Option {
 				"develop",
 				cfg.Debug,
 				WithTracer(tracer.Tracer("api")),
+				WithTimeout(cfg.Timeout),
 			)
 		}),
 		health.Module(),

diff --git a/internal/api/router.go b/internal/api/router.go
@@ -1,12 +1,12 @@
 package api
 
 import (
+	"github.com/formancehq/go-libs/v2/api"
+	"github.com/formancehq/ledger/internal/controller/system"
 	"go.opentelemetry.io/otel/trace"
 	nooptracer "go.opentelemetry.io/otel/trace/noop"
 	"net/http"
 
-	"github.com/formancehq/ledger/internal/controller/system"
-
 	"github.com/formancehq/go-libs/v2/logging"
 	"github.com/go-chi/chi/v5/middleware"
 	"github.com/go-chi/cors"
@@ -53,11 +53,19 @@ func NewRouter(
 		common.LogID(),
 	)
 
+	commonMiddlewares := []func(http.Handler) http.Handler{
+		middleware.RequestLogger(api.NewLogFormatter()),
+	}
+	if routerOptions.timeoutConfiguration.Timeout != 0 {
+		commonMiddlewares = append(commonMiddlewares, common.Timeout(routerOptions.timeoutConfiguration))
+	}
+
 	v2Router := v2.NewRouter(
 		systemController,
 		authenticator,
 		debug,
 		v2.WithTracer(routerOptions.tracer),
+		v2.WithMiddlewares(commonMiddlewares...),
 	)
 	mux.Handle("/v2*", http.StripPrefix("/v2", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		chi.RouteContext(r.Context()).Reset()
@@ -69,13 +77,15 @@ func NewRouter(
 		version,
 		debug,
 		v1.WithTracer(routerOptions.tracer),
+		v1.WithMiddlewares(commonMiddlewares...),
 	))
 
 	return mux
 }
 
 type routerOptions struct {
-	tracer trace.Tracer
+	tracer               trace.Tracer
+	timeoutConfiguration common.TimeoutConfiguration
 }
 
 type RouterOption func(ro *routerOptions)
@@ -86,6 +96,12 @@ func WithTracer(tracer trace.Tracer) RouterOption {
 	}
 }
 
+func WithTimeout(timeoutConfiguration common.TimeoutConfiguration) RouterOption {
+	return func(ro *routerOptions) {
+		ro.timeoutConfiguration = timeoutConfiguration
+	}
+}
+
 var defaultRouterOptions = []RouterOption{
 	WithTracer(nooptracer.Tracer{}),
 }
diff --git a/internal/api/v1/routes.go b/internal/api/v1/routes.go
@@ -7,9 +7,6 @@ import (
 
 	"github.com/formancehq/ledger/internal/controller/system"
 
-	"github.com/formancehq/go-libs/v2/api"
-	"github.com/go-chi/chi/v5/middleware"
-
 	"github.com/formancehq/go-libs/v2/service"
 
 	"github.com/formancehq/go-libs/v2/auth"
@@ -35,7 +32,7 @@ func NewRouter(
 	router.Get("/_info", getInfo(systemController, version))
 
 	router.Group(func(router chi.Router) {
-		router.Use(middleware.RequestLogger(api.NewLogFormatter()))
+		router.Use(routerOptions.middlewares...)
 		router.Use(auth.Middleware(authenticator))
 		router.Use(service.OTLPMiddleware("ledger", debug))
 
@@ -85,7 +82,8 @@ func NewRouter(
 }
 
 type routerOptions struct {
-	tracer trace.Tracer
+	tracer      trace.Tracer
+	middlewares []func(handler http.Handler) http.Handler
 }
 
 type RouterOption func(ro *routerOptions)
@@ -96,6 +94,12 @@ func WithTracer(tracer trace.Tracer) RouterOption {
 	}
 }
 
+func WithMiddlewares(handlers ...func(http.Handler) http.Handler) RouterOption {
+	return func(ro *routerOptions) {
+		ro.middlewares = append(ro.middlewares, handlers...)
+	}
+}
+
 var defaultRouterOptions = []RouterOption{
 	WithTracer(nooptracer.Tracer{}),
 }
diff --git a/internal/api/v2/routes.go b/internal/api/v2/routes.go
@@ -6,9 +6,6 @@ import (
 
 	"github.com/formancehq/ledger/internal/controller/system"
 
-	"github.com/formancehq/go-libs/v2/api"
-	"github.com/go-chi/chi/v5/middleware"
-
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 
@@ -33,7 +30,7 @@ func NewRouter(
 	router := chi.NewMux()
 
 	router.Group(func(router chi.Router) {
-		router.Use(middleware.RequestLogger(api.NewLogFormatter()))
+		router.Use(routerOptions.middlewares...)
 		router.Use(auth.Middleware(authenticator))
 		router.Use(service.OTLPMiddleware("ledger", debug))
 
@@ -93,7 +90,8 @@ func NewRouter(
 }
 
 type routerOptions struct {
-	tracer trace.Tracer
+	tracer      trace.Tracer
+	middlewares []func(http.Handler) http.Handler
 }
 
 type RouterOption func(ro *routerOptions)
@@ -104,6 +102,12 @@ func WithTracer(tracer trace.Tracer) RouterOption {
 	}
 }
 
+func WithMiddlewares(middlewares ...func(http.Handler) http.Handler) RouterOption {
+	return func(ro *routerOptions) {
+		ro.middlewares = append(ro.middlewares, middlewares...)
+	}
+}
+
 var defaultRouterOptions = []RouterOption{
 	WithTracer(nooptracer.Tracer{}),
 }
diff --git a/openapi/v1.yaml b/openapi/v1.yaml
@@ -1852,6 +1852,7 @@ components:
         - COMPILATION_FAILED
         - METADATA_OVERRIDE
         - NOT_FOUND
+        - TIMEOUT
       example: INSUFFICIENT_FUND
     LedgerInfoResponse:
       properties:

diff --git a/openapi/v2.yaml b/openapi/v2.yaml
@@ -1722,6 +1722,7 @@ components:
         - NO_POSTINGS
         - LEDGER_NOT_FOUND
         - IMPORT
+        - TIMEOUT
       example: VALIDATION
     V2LedgerInfoResponse:
       type: object

diff --git a/pkg/client/.speakeasy/gen.lock b/pkg/client/.speakeasy/gen.lock
@@ -1,12 +1,12 @@
 lockVersion: 2.0.0
 id: a9ac79e1-e429-4ee3-96c4-ec973f19bec3
 management:
-  docChecksum: 0558cb1f36313c824f51744f6f8c355d
+  docChecksum: 6d88216daa0351f4aa07aaa5646d4d8b
   docVersion: v1
   speakeasyVersion: 1.351.0
   generationVersion: 2.384.1
-  releaseVersion: 0.4.13
-  configChecksum: a8fb416ba04bb07465d883f3e4ed1f7a
+  releaseVersion: 0.4.14
+  configChecksum: 583ed6e25e9905dc157e817350d207e0
 features:
   go:
     additionalDependencies: 0.1.0

diff --git a/pkg/client/.speakeasy/gen.yaml b/pkg/client/.speakeasy/gen.yaml
@@ -15,7 +15,7 @@ generation:
   auth:
     oAuth2ClientCredentialsEnabled: true
 go:
-  version: 0.4.13
+  version: 0.4.14
   additionalDependencies: {}
   allowUnknownFieldsInWeakUnions: false
   clientServerStatusCodesAsErrors: true

diff --git a/pkg/client/docs/models/components/errorsenum.md b/pkg/client/docs/models/components/errorsenum.md
@@ -12,4 +12,5 @@
 | `ErrorsEnumNoScript`          | NO_SCRIPT                     |
 | `ErrorsEnumCompilationFailed` | COMPILATION_FAILED            |
 | `ErrorsEnumMetadataOverride`  | METADATA_OVERRIDE             |
-| `ErrorsEnumNotFound`          | NOT_FOUND                     |
+| `ErrorsEnumNotFound`          | NOT_FOUND                     |
+| `ErrorsEnumTimeout`           | TIMEOUT                       |
diff --git a/pkg/client/docs/models/components/v2errorsenum.md b/pkg/client/docs/models/components/v2errorsenum.md
@@ -16,4 +16,5 @@
 | `V2ErrorsEnumAlreadyRevert`     | ALREADY_REVERT                  |
 | `V2ErrorsEnumNoPostings`        | NO_POSTINGS                     |
 | `V2ErrorsEnumLedgerNotFound`    | LEDGER_NOT_FOUND                |
-| `V2ErrorsEnumImport`            | IMPORT                          |
+| `V2ErrorsEnumImport`            | IMPORT                          |
+| `V2ErrorsEnumTimeout`           | TIMEOUT                         |