support/db: Support concurrent queries in a transaction (#2024)

This commit adds `Synchronized` flag to `support/db.Session`. When set
to `true` and `Session` runs a transaction all `Exec*` and `Select*`
methods are protected by mutex what allows running them in multiple
goroutines.

This is an experimental feature (see below) and not a breaking change
(default is `false`).

Postgres protocol does not allow sending Exec query results if
previously sent Query haven't been fully read. This issue manifested
itself when a PR that's sending read and write queries in multiple
goroutines was merged.

More info: lib/pq#81 lib/pq#635

Known limitations:

* It's possible that it will not be needed if we decide to remove
concurrency from ingestion pipeline (see #1983). We are adding this to
unblock development of new ingestion processors with a more readable
code (currently all database processors are merged into one big
processor that is hard to read and test). Splitting `DatabaseProcessor`
will be done in a separate PR/PRs.
* During Horizon Team meeting we agreed to create a new
`SynchronizedSession` struct embedding `db.Session` inside that would
not be placed in a shared `support` package. The problem is that
`history.Q` embeds `db.Session` inside. Changing this to
`db.SessionInterface` requires updating a lot of files that create
`history.Q` objects, tests and mocks. Given that we may want to revert
this change in the future, the change in this commit seems to be
simpler.

Author: bartekn

services/horizon/internal/expingest/main.go

@@ -137,7 +137,12 @@ func NewSystem(config Config) (*System, error) {
 		return nil, errors.Wrap(err, "error creating ledger backend")
 	}
 
-	historyQ := &history.Q{config.HistorySession}
+	// Make historySession synchronized so it can be used in the pipeline
+	// (saving to DB in multiple goroutines at the same time).
+	historySession := config.HistorySession.Clone()
+	historySession.Synchronized = true
+
+	historyQ := &history.Q{historySession}
 
 	session := &ingest.LiveSession{
 		Archive:          archive,
@@ -157,7 +162,7 @@ func NewSystem(config Config) (*System, error) {
 
 	system := &System{
 		session:                  session,
-		historySession:           config.HistorySession,
+		historySession:           historySession,
 		historyQ:                 historyQ,
 		graph:                    config.OrderBookGraph,
 		retry:                    alwaysRetry{time.Second},

support/db/main.go

@@ -14,6 +14,7 @@ package db
 import (
 	"context"
 	"database/sql"
+	"sync"
 
 	"github.com/Masterminds/squirrel"
 	"github.com/jmoiron/sqlx"
@@ -117,7 +118,16 @@ type Session struct {
 	// Ctx is the context in which the repo is operating under.
 	Ctx context.Context
 
-	tx *sqlx.Tx
+	// Synchronized is an EXPERIMENTAL flag that allows sending queries
+	// concurrently in a DB tx. When set to `true` all Exec and Select methods
+	// sent in a transaction are protected by mutex. It is not needed outside
+	// transaction context because then all queries are sent in a separate DB
+	// connections.
+	// Please note that Begin, Commit and Rollback must not be run in parallel.
+	// Also, Query methods are not available when in Synchronized transaction.
+	Synchronized bool
+	txMutex      sync.Mutex
+	tx           *sqlx.Tx
 }
 
 type SessionInterface interface {

support/db/session.go

@@ -67,6 +67,8 @@ func (s *Session) Clone() *Session {
 	return &Session{
 		DB:  s.DB,
 		Ctx: s.Ctx,
+
+		Synchronized: s.Synchronized,
 	}
 }
 
@@ -128,6 +130,11 @@ func (s *Session) GetRaw(dest interface{}, query string, args ...interface{}) er
 		return errors.Wrap(err, "replace placeholders failed")
 	}
 
+	if s.Synchronized && s.tx != nil {
+		s.txMutex.Lock()
+		defer s.txMutex.Unlock()
+	}
+
 	start := time.Now()
 	err = s.conn().GetContext(s.Ctx, dest, query, args...)
 	s.log("get", start, query, args)
@@ -197,6 +204,11 @@ func (s *Session) ExecRaw(query string, args ...interface{}) (sql.Result, error)
 		return nil, errors.Wrap(err, "replace placeholders failed")
 	}
 
+	if s.Synchronized && s.tx != nil {
+		s.txMutex.Lock()
+		defer s.txMutex.Unlock()
+	}
+
 	start := time.Now()
 	result, err := s.conn().ExecContext(s.Ctx, query, args...)
 	s.log("exec", start, query, args)
@@ -238,6 +250,10 @@ func (s *Session) Query(query sq.Sqlizer) (*sqlx.Rows, error) {
 
 // QueryRaw runs `query` with `args`
 func (s *Session) QueryRaw(query string, args ...interface{}) (*sqlx.Rows, error) {
+	if s.Synchronized && s.tx != nil {
+		return nil, errors.New("QueryRaw cannot be run in transaction when Synchronized is true")
+	}
+
 	query, err := s.ReplacePlaceholders(query)
 	if err != nil {
 		return nil, errors.Wrap(err, "replace placeholders failed")
@@ -307,6 +323,11 @@ func (s *Session) SelectRaw(
 		return errors.Wrap(err, "replace placeholders failed")
 	}
 
+	if s.Synchronized && s.tx != nil {
+		s.txMutex.Lock()
+		defer s.txMutex.Unlock()
+	}
+
 	start := time.Now()
 	err = s.conn().SelectContext(s.Ctx, dest, query, args...)
 	s.log("select", start, query, args)

support/db/session_test.go

@@ -2,13 +2,57 @@ package db
 
 import (
 	"context"
+	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/stellar/go/support/db/dbtest"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func TestConcurrentQueriesTransaction(t *testing.T) {
+	db := dbtest.Postgres(t).Load(testSchema)
+	defer db.Close()
+
+	sess := &Session{
+		Ctx: context.Background(),
+		DB:  db.Open(),
+		// This test would fail for `Synchronized: false`
+		Synchronized: true,
+	}
+	defer sess.DB.Close()
+
+	err := sess.Begin()
+	assert.NoError(t, err)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 1000; i++ {
+		wg.Add(1)
+		go func(i int) {
+			istr := fmt.Sprintf("%d", i)
+			var err2 error
+			if i%3 == 0 {
+				var names []string
+				err2 = sess.SelectRaw(&names, "SELECT name FROM people")
+			} else if i%3 == 1 {
+				var name string
+				err2 = sess.GetRaw(&name, "SELECT name FROM people LIMIT 1")
+			} else {
+				_, err2 = sess.ExecRaw(
+					"INSERT INTO people (name, hunger_level) VALUES ('bartek" + istr + "', " + istr + ")",
+				)
+			}
+			assert.NoError(t, err2)
+			wg.Done()
+		}(i)
+	}
+
+	wg.Wait()
+	err = sess.Rollback()
+	assert.NoError(t, err)
+}
+
 func TestSession(t *testing.T) {
 	db := dbtest.Postgres(t).Load(testSchema)
 	defer db.Close()