Add timeout on horizon db ingestion statements

stellar · May 8, 2020 · 3ccaf33 · 3ccaf33
1 parent ebd7214
commit 3ccaf33
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 6 deletions.
diff --git a/services/horizon/internal/expingest/main.go b/services/horizon/internal/expingest/main.go
@@ -47,6 +47,8 @@ const (
 	// MaxDBConnections is the size of the postgres connection pool dedicated to Horizon ingestion
 	MaxDBConnections = 2
 
+	dbSessionTimeout = 20 * time.Second
+
 	defaultCoreCursorName           = "HORIZON"
 	stateVerificationErrorThreshold = 3
 )
@@ -146,6 +148,7 @@ func NewSystem(config Config) (*System, error) {
 
 	coreSession := config.CoreSession.Clone()
 	coreSession.Ctx = ctx
+	coreSession.Timeout = dbSessionTimeout
 
 	ledgerBackend, err := ledgerbackend.NewDatabaseBackendFromSession(coreSession)
 	if err != nil {
@@ -155,6 +158,7 @@ func NewSystem(config Config) (*System, error) {
 
 	historyQ := &history.Q{config.HistorySession.Clone()}
 	historyQ.Ctx = ctx
+	historyQ.Timeout = dbSessionTimeout
 
 	historyAdapter := adapters.MakeHistoryArchiveAdapter(archive)
 

diff --git a/support/db/main.go b/support/db/main.go
@@ -14,6 +14,7 @@ package db
 import (
 	"context"
 	"database/sql"
+	"time"
 
 	"github.com/Masterminds/squirrel"
 	"github.com/jmoiron/sqlx"
@@ -101,6 +102,11 @@ type Session struct {
 	// Ctx is the context in which the repo is operating under.
 	Ctx context.Context
 
+	// Timeout is a timeout applied on all SQL Select queries and Exec statements.
+	// If Timeout is 0 then no timeout is applied. Note that the timeout is not
+	// applied on BEGIN, COMMIT, or ROLLBACK statements.
+	Timeout time.Duration
+
 	tx *sqlx.Tx
 }
 

diff --git a/support/db/session.go b/support/db/session.go
@@ -65,8 +65,9 @@ func (s *Session) GetTx() *sqlx.Tx {
 // source is currently within.
 func (s *Session) Clone() *Session {
 	return &Session{
-		DB:  s.DB,
-		Ctx: s.Ctx,
+		DB:      s.DB,
+		Ctx:     s.Ctx,
+		Timeout: s.Timeout,
 	}
 }
 
@@ -120,6 +121,17 @@ func (s *Session) Get(dest interface{}, query sq.Sqlizer) error {
 	return s.GetRaw(dest, sql, args...)
 }
 
+func noop() {
+
+}
+
+func (s *Session) contextWithTimeout() (context.Context, context.CancelFunc) {
+	if s.Timeout == 0 {
+		return s.Ctx, noop
+	}
+	return context.WithTimeout(s.Ctx, s.Timeout)
+}
+
 // GetRaw runs `query` with `args`, setting the first result found on
 // `dest`, if any.
 func (s *Session) GetRaw(dest interface{}, query string, args ...interface{}) error {
@@ -129,7 +141,9 @@ func (s *Session) GetRaw(dest interface{}, query string, args ...interface{}) er
 	}
 
 	start := time.Now()
-	err = s.conn().GetContext(s.Ctx, dest, query, args...)
+	ctx, cancel := s.contextWithTimeout()
+	defer cancel()
+	err = s.conn().GetContext(ctx, dest, query, args...)
 	s.log("get", start, query, args)
 
 	if err == nil {
@@ -198,7 +212,9 @@ func (s *Session) ExecRaw(query string, args ...interface{}) (sql.Result, error)
 	}
 
 	start := time.Now()
-	result, err := s.conn().ExecContext(s.Ctx, query, args...)
+	ctx, cancel := s.contextWithTimeout()
+	defer cancel()
+	result, err := s.conn().ExecContext(ctx, query, args...)
 	s.log("exec", start, query, args)
 
 	if err == nil {
@@ -244,7 +260,9 @@ func (s *Session) QueryRaw(query string, args ...interface{}) (*sqlx.Rows, error
 	}
 
 	start := time.Now()
-	result, err := s.conn().QueryxContext(s.Ctx, query, args...)
+	ctx, cancel := s.contextWithTimeout()
+	defer cancel()
+	result, err := s.conn().QueryxContext(ctx, query, args...)
 	s.log("query", start, query, args)
 
 	if err == nil {
@@ -308,7 +326,9 @@ func (s *Session) SelectRaw(
 	}
 
 	start := time.Now()
-	err = s.conn().SelectContext(s.Ctx, dest, query, args...)
+	ctx, cancel := s.contextWithTimeout()
+	defer cancel()
+	err = s.conn().SelectContext(ctx, dest, query, args...)
 	s.log("select", start, query, args)
 
 	if err == nil {

diff --git a/support/db/session_test.go b/support/db/session_test.go
@@ -3,12 +3,25 @@ package db
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/stellar/go/support/db/dbtest"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func TestClone(t *testing.T) {
+	session := &Session{Ctx: context.Background()}
+	sessionClone := session.Clone()
+	assert.Equal(t, time.Duration(0), sessionClone.Timeout)
+	assert.Equal(t, session.Ctx, sessionClone.Ctx)
+
+	session.Timeout = time.Second
+	sessionClone = session.Clone()
+	assert.Equal(t, session.Timeout, sessionClone.Timeout)
+	assert.Equal(t, session.Ctx, sessionClone.Ctx)
+}
+
 func TestSession(t *testing.T) {
 	db := dbtest.Postgres(t).Load(testSchema)
 	defer db.Close()