support/db: Delay canceling queries from client side when there's a statement / transaction timeout configured in postgres

services/horizon/internal/app.go

@@ -532,6 +532,7 @@ func (a *App) init() error {
 		SSEUpdateFrequency:       a.config.SSEUpdateFrequency,
 		StaleThreshold:           a.config.StaleThreshold,
 		ConnectionTimeout:        a.config.ConnectionTimeout,
+		CancelDBQueryTimeout:     a.config.CancelDBQueryTimeout,
 		MaxHTTPRequestSize:       a.config.MaxHTTPRequestSize,
 		NetworkPassphrase:        a.config.NetworkPassphrase,
 		MaxPathLength:            a.config.MaxPathLength,

services/horizon/internal/config.go

@@ -36,8 +36,9 @@ type Config struct {
 	HorizonDBMaxOpenConnections int
 	HorizonDBMaxIdleConnections int
 
-	SSEUpdateFrequency time.Duration
-	ConnectionTimeout  time.Duration
+	SSEUpdateFrequency   time.Duration
+	ConnectionTimeout    time.Duration
+	CancelDBQueryTimeout time.Duration
 	// MaxHTTPRequestSize is the maximum allowed request payload size
 	MaxHTTPRequestSize uint
 	RateQuota          *throttled.RateQuota

services/horizon/internal/flags.go

@@ -437,6 +437,18 @@ func Flags() (*Config, support.ConfigOptions) {
 			Usage:          "defines the timeout of connection after which 504 response will be sent or stream will be closed, if Horizon is behind a load balancer with idle connection timeout, this should be set to a few seconds less that idle timeout, does not apply to POST /transactions",
 			UsedInCommands: ApiServerCommands,
 		},
+		&support.ConfigOption{
+			Name:           "cancel-db-query-timeout",
+			ConfigKey:      &config.CancelDBQueryTimeout,
+			OptType:        types.Int,
+			FlagDefault:    0,
+			CustomSetValue: support.SetDuration,
+			Usage: "defines the timeout for when horizon will cancel all postgres queries connected to an HTTP request. The timeout is measured in seconds since the start of the HTTP request. Note, this timeout does not apply to POST /transactions. " +
+				"The difference between cancel-db-query-timeout and connection-timeout is that connection-timeout applies a postgres statement timeout whereas cancel-db-query-timeout will send an additional request to postgres to cancel the ongoing query. " +
+				"Generally, cancel-db-query-timeout should be configured to be higher than connection-timeout to allow the postgres statement timeout to kill long running queries without having to send the additional cancel request to postgres. " +
+				"By default, cancel-db-query-timeout will be set to twice the connection-timeout.",
+			UsedInCommands: ApiServerCommands,
+		},
 		&support.ConfigOption{
 			Name:           "max-http-request-size",
 			ConfigKey:      &config.MaxHTTPRequestSize,
@@ -983,5 +995,10 @@ func ApplyFlags(config *Config, flags support.ConfigOptions, options ApplyOption
 			" If Horizon is behind both, use --behind-cloudflare only")
 	}
 
+	if config.CancelDBQueryTimeout == 0 {
+		// the default value for cancel-db-query-timeout is twice the connection-timeout
+		config.CancelDBQueryTimeout = config.ConnectionTimeout * 2
+	}
+
 	return nil
 }

services/horizon/internal/httpx/middleware.go

@@ -197,14 +197,15 @@ func recoverMiddleware(h http.Handler) http.Handler {
 // NewHistoryMiddleware adds session to the request context and ensures Horizon
 // is not in a stale state, which is when the difference between latest core
 // ledger and latest history ledger is higher than the given threshold
-func NewHistoryMiddleware(ledgerState *ledger.State, staleThreshold int32, session db.SessionInterface) func(http.Handler) http.Handler {
+func NewHistoryMiddleware(ledgerState *ledger.State, staleThreshold int32, session db.SessionInterface, contextDBTimeout time.Duration) func(http.Handler) http.Handler {
 	return func(h http.Handler) http.Handler {
 
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			ctx := r.Context()
 			if routePattern := supportHttp.GetChiRoutePattern(r); routePattern != "" {
 				ctx = context.WithValue(ctx, &db.RouteContextKey, routePattern)
 			}
+			ctx = setContextDBTimeout(contextDBTimeout, ctx)
 			if staleThreshold > 0 {
 				ls := ledgerState.CurrentStatus()
 				isStale := (ls.CoreLatest - ls.HistoryLatest) > int32(staleThreshold)
@@ -237,8 +238,9 @@ func NewHistoryMiddleware(ledgerState *ledger.State, staleThreshold int32, sessi
 // has been verified and is correct (Otherwise returns `500 Internal Server Error` to prevent
 // returning invalid data to the user)
 type StateMiddleware struct {
-	HorizonSession      db.SessionInterface
-	NoStateVerification bool
+	HorizonSession       db.SessionInterface
+	CancelDBQueryTimeout time.Duration
+	NoStateVerification  bool
 }
 
 func ingestionStatus(ctx context.Context, q *history.Q) (uint32, bool, error) {
@@ -276,6 +278,7 @@ func (m *StateMiddleware) WrapFunc(h http.HandlerFunc) http.HandlerFunc {
 		if routePattern := supportHttp.GetChiRoutePattern(r); routePattern != "" {
 			ctx = context.WithValue(ctx, &db.RouteContextKey, routePattern)
 		}
+		ctx = setContextDBTimeout(m.CancelDBQueryTimeout, ctx)
 		session := m.HorizonSession.Clone()
 		q := &history.Q{session}
 		sseRequest := render.Negotiate(r) == render.MimeEventStream
@@ -344,6 +347,13 @@ func (m *StateMiddleware) WrapFunc(h http.HandlerFunc) http.HandlerFunc {
 	}
 }
 
+func setContextDBTimeout(timeout time.Duration, ctx context.Context) context.Context {
+	if timeout > 0 {
+		ctx = context.WithValue(ctx, &db.DeadlineCtxKey, time.Now().Add(timeout))
+	}
+	return ctx
+}
+
 // WrapFunc executes the middleware on a given HTTP handler function
 func (m *StateMiddleware) Wrap(h http.Handler) http.Handler {
 	return m.WrapFunc(h.ServeHTTP)

services/horizon/internal/httpx/router.go

@@ -38,6 +38,7 @@ type RouterConfig struct {
 	SSEUpdateFrequency       time.Duration
 	StaleThreshold           uint
 	ConnectionTimeout        time.Duration
+	CancelDBQueryTimeout     time.Duration
 	MaxHTTPRequestSize       uint
 	NetworkPassphrase        string
 	MaxPathLength            uint
@@ -139,7 +140,8 @@ func (r *Router) addMiddleware(config *RouterConfig,
 
 func (r *Router) addRoutes(config *RouterConfig, rateLimiter *throttled.HTTPRateLimiter, ledgerState *ledger.State) {
 	stateMiddleware := StateMiddleware{
-		HorizonSession: config.DBSession,
+		HorizonSession:       config.DBSession,
+		CancelDBQueryTimeout: config.CancelDBQueryTimeout,
 	}
 
 	r.Method(http.MethodGet, "/health", config.HealthCheck)
@@ -157,7 +159,7 @@ func (r *Router) addRoutes(config *RouterConfig, rateLimiter *throttled.HTTPRate
 		LedgerSourceFactory: historyLedgerSourceFactory{ledgerState: ledgerState, updateFrequency: config.SSEUpdateFrequency},
 	}
 
-	historyMiddleware := NewHistoryMiddleware(ledgerState, int32(config.StaleThreshold), config.DBSession)
+	historyMiddleware := NewHistoryMiddleware(ledgerState, int32(config.StaleThreshold), config.DBSession, config.CancelDBQueryTimeout)
 	// State endpoints behind stateMiddleware
 	r.Group(func(r chi.Router) {
 		r.Route("/accounts", func(r chi.Router) {

services/horizon/internal/init.go

@@ -45,38 +45,29 @@ func mustInitHorizonDB(app *App) {
 			log.Fatalf("max open connections to horizon db must be greater than %d", ingest.MaxDBConnections)
 		}
 	}
+	serverSidePGTimeoutConfigs := []db.ClientConfig{
+		db.StatementTimeout(app.config.ConnectionTimeout),
+		db.IdleTransactionTimeout(app.config.ConnectionTimeout),
+	}
 
 	if app.config.RoDatabaseURL == "" {
-		var clientConfigs []db.ClientConfig
-		if !app.config.Ingest {
-			// if we are not ingesting then we don't expect to have long db queries / transactions
-			clientConfigs = append(
-				clientConfigs,
-				db.StatementTimeout(app.config.ConnectionTimeout),
-				db.IdleTransactionTimeout(app.config.ConnectionTimeout),
-			)
-		}
 		app.historyQ = &history.Q{mustNewDBSession(
 			db.HistorySubservice,
 			app.config.DatabaseURL,
 			maxIdle,
 			maxOpen,
 			app.prometheusRegistry,
-			clientConfigs...,
+			serverSidePGTimeoutConfigs...,
 		)}
 	} else {
 		// If RO set, use it for all DB queries
-		roClientConfigs := []db.ClientConfig{
-			db.StatementTimeout(app.config.ConnectionTimeout),
-			db.IdleTransactionTimeout(app.config.ConnectionTimeout),
-		}
 		app.historyQ = &history.Q{mustNewDBSession(
 			db.HistorySubservice,
 			app.config.RoDatabaseURL,
 			maxIdle,
 			maxOpen,
 			app.prometheusRegistry,
-			roClientConfigs...,
+			serverSidePGTimeoutConfigs...,
 		)}
 
 		app.primaryHistoryQ = &history.Q{mustNewDBSession(
@@ -85,6 +76,7 @@ func mustInitHorizonDB(app *App) {
 			maxIdle,
 			maxOpen,
 			app.prometheusRegistry,
+			serverSidePGTimeoutConfigs...,
 		)}
 	}
 }

services/horizon/internal/middleware_test.go

@@ -402,7 +402,7 @@ func TestCheckHistoryStaleMiddleware(t *testing.T) {
 			}
 			ledgerState := &ledger.State{}
 			ledgerState.SetStatus(state)
-			historyMiddleware := httpx.NewHistoryMiddleware(ledgerState, testCase.staleThreshold, tt.HorizonSession())
+			historyMiddleware := httpx.NewHistoryMiddleware(ledgerState, testCase.staleThreshold, tt.HorizonSession(), 0)
 			handler := chi.NewRouter()
 			handler.With(historyMiddleware).MethodFunc("GET", "/", endpoint)
 			w := httptest.NewRecorder()

staticcheck.sh

@@ -1,7 +1,7 @@
 #! /bin/bash
 set -e
 
-version='2023.1.1'
+version='2023.1.7'
 
 staticcheck='go run honnef.co/go/tools/cmd/staticcheck@'"$version"
 

support/db/main.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/Masterminds/squirrel"
 	"github.com/jmoiron/sqlx"
+
 	"github.com/stellar/go/support/errors"
 
 	// Enable postgres
@@ -119,6 +120,7 @@ type Session struct {
 	DB *sqlx.DB
 
 	tx            *sqlx.Tx
+	txCancel      context.CancelFunc
 	txOptions     *sql.TxOptions
 	errorHandlers []ErrorHandlerFunc
 }

support/db/session.go

@@ -12,28 +12,60 @@ import (
 	sq "github.com/Masterminds/squirrel"
 	"github.com/jmoiron/sqlx"
 	"github.com/lib/pq"
+
 	"github.com/stellar/go/support/db/sqlutils"
 	"github.com/stellar/go/support/errors"
 	"github.com/stellar/go/support/log"
 )
 
+var DeadlineCtxKey = CtxKey("deadline")
+
+func noop() {}
+
+// context() checks if there is a override on the context timeout which is configured using DeadlineCtxKey.
+// If the override exists, we return a new context with the desired deadline. Otherwise, we return the
+// original context.
+// Note that the override will not be applied if requestCtx has already been terminated.
+func (s *Session) context(requestCtx context.Context) (context.Context, context.CancelFunc, error) {
+	deadline, ok := requestCtx.Value(&DeadlineCtxKey).(time.Time)
+	if !ok {
+		return requestCtx, noop, nil
+	}
+
+	// if requestCtx is already terminated don't proceed with the db statement
+	if requestCtx.Err() != nil {
+		return requestCtx, noop, requestCtx.Err()
+	}
+
+	ctx, cancel := context.WithDeadline(context.Background(), deadline)
+	return ctx, cancel, nil
+}
+
 // Begin binds this session to a new transaction.
 func (s *Session) Begin(ctx context.Context) error {
 	if s.tx != nil {
 		return errors.New("already in transaction")
 	}
+	ctx, cancel, err := s.context(ctx)
+	if err != nil {
+		cancel()
+		return err
+	}
 
 	tx, err := s.DB.BeginTxx(ctx, nil)
 	if err != nil {
 		if knownErr := s.handleError(err, ctx); knownErr != nil {
+			cancel()
 			return knownErr
 		}
 
+		cancel()
 		return errors.Wrap(err, "beginx failed")
 	}
 	log.Debug("sql: begin")
 	s.tx = tx
 	s.txOptions = nil
+	s.txCancel = cancel
 	return nil
 }
 
@@ -43,19 +75,27 @@ func (s *Session) BeginTx(ctx context.Context, opts *sql.TxOptions) error {
 	if s.tx != nil {
 		return errors.New("already in transaction")
 	}
+	ctx, cancel, err := s.context(ctx)
+	if err != nil {
+		cancel()
+		return err
+	}
 
 	tx, err := s.DB.BeginTxx(ctx, opts)
 	if err != nil {
 		if knownErr := s.handleError(err, ctx); knownErr != nil {
+			cancel()
 			return knownErr
 		}
 
+		cancel()
 		return errors.Wrap(err, "beginTx failed")
 	}
 	log.Debug("sql: begin")
 
 	s.tx = tx
 	s.txOptions = opts
+	s.txCancel = cancel
 	return nil
 }
 
@@ -93,6 +133,8 @@ func (s *Session) Commit() error {
 	log.Debug("sql: commit")
 	s.tx = nil
 	s.txOptions = nil
+	s.txCancel()
+	s.txCancel = nil
 
 	if knownErr := s.handleError(err, context.Background()); knownErr != nil {
 		return knownErr
@@ -135,7 +177,13 @@ func (s *Session) Get(ctx context.Context, dest interface{}, query sq.Sqlizer) e
 // GetRaw runs `query` with `args`, setting the first result found on
 // `dest`, if any.
 func (s *Session) GetRaw(ctx context.Context, dest interface{}, query string, args ...interface{}) error {
-	query, err := s.ReplacePlaceholders(query)
+	ctx, cancel, err := s.context(ctx)
+	defer cancel()
+	if err != nil {
+		return err
+	}
+
+	query, err = s.ReplacePlaceholders(query)
 	if err != nil {
 		return errors.Wrap(err, "replace placeholders failed")
 	}
@@ -204,7 +252,13 @@ func (s *Session) ExecAll(ctx context.Context, script string) error {
 
 // ExecRaw runs `query` with `args`
 func (s *Session) ExecRaw(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
-	query, err := s.ReplacePlaceholders(query)
+	ctx, cancel, err := s.context(ctx)
+	defer cancel()
+	if err != nil {
+		return nil, err
+	}
+
+	query, err = s.ReplacePlaceholders(query)
 	if err != nil {
 		return nil, errors.Wrap(err, "replace placeholders failed")
 	}
@@ -304,7 +358,13 @@ func (s *Session) Query(ctx context.Context, query sq.Sqlizer) (*sqlx.Rows, erro
 
 // QueryRaw runs `query` with `args`
 func (s *Session) QueryRaw(ctx context.Context, query string, args ...interface{}) (*sqlx.Rows, error) {
-	query, err := s.ReplacePlaceholders(query)
+	ctx, cancel, err := s.context(ctx)
+	defer cancel()
+	if err != nil {
+		return nil, err
+	}
+
+	query, err = s.ReplacePlaceholders(query)
 	if err != nil {
 		return nil, errors.Wrap(err, "replace placeholders failed")
 	}
@@ -350,6 +410,8 @@ func (s *Session) Rollback() error {
 	log.Debug("sql: rollback")
 	s.tx = nil
 	s.txOptions = nil
+	s.txCancel()
+	s.txCancel = nil
 
 	if knownErr := s.handleError(err, context.Background()); knownErr != nil {
 		return knownErr
@@ -381,8 +443,14 @@ func (s *Session) SelectRaw(
 	query string,
 	args ...interface{},
 ) error {
+	ctx, cancel, err := s.context(ctx)
+	defer cancel()
+	if err != nil {
+		return err
+	}
+
 	s.clearSliceIfPossible(dest)
-	query, err := s.ReplacePlaceholders(query)
+	query, err = s.ReplacePlaceholders(query)
 	if err != nil {
 		return errors.Wrap(err, "replace placeholders failed")
 	}