Repository: armadaproject/armada
Pull request #2391 (2023-04-24 09:49:47)

[Scheduler Ingester] Perform All Updates Inside A Single Transaction
Scheduler Ingester now applies a batch of operations inside a single transaction.  This has a couple of benefits:
- The operations may be reordered to improve ingestion performance.  By using a single transaction, we ensure that intermediate states are not visible to readers and so reordering only needs to ensure the final state is correct.
- If there is a failure, the entire batch will fail, therefore we won't end up in some sort of intermediate state.

As part of this PR I also implemented a lock around the actual database operations (implemented using a pg_advisory lock). This is needed because we rely on the scheduler seeing monotonically increasing sequence numbers and this will not be the case if two ingesters are writing to the dtaabase at the same time.  This was an issue previously, but would be much worse now that the individual transactions are much bigger.

┆Issue is synchronized with this [Jira Task](https://gr-oss.atlassian.net/browse/BATCH-192) by [Unito](https://www.unito.io)


codecov[bot]: ## [Codecov](https://codecov.io/gh/armadaproject/armada/pull/2391?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject) Report
Patch coverage: **`80.00`**% and project coverage change: **`+0.33`** :tada:
> Comparison is base [(`596fce2`)](https://codecov.io/gh/armadaproject/armada/commit/596fce2168bf0281e48ff6c682eae9d8166e7fd2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject) 58.95% compared to head [(`db5efed`)](https://codecov.io/gh/armadaproject/armada/pull/2391?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject) 59.29%.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #2391      +/-   ##
==========================================
+ Coverage   58.95%   59.29%   +0.33%     
==========================================
  Files         226      226              
  Lines       28809    28830      +21     
==========================================
+ Hits        16985    17095     +110     
+ Misses      10524    10424     -100     
- Partials     1300     1311      +11     
```

| Flag | Coverage Δ | |
|---|---|---|
| armada-server | `59.29% <80.00%> (+0.33%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/armadaproject/armada/pull/2391?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject) | Coverage Δ | |
|---|---|---|
| [internal/scheduleringester/ingester.go](https://codecov.io/gh/armadaproject/armada/pull/2391?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject#diff-aW50ZXJuYWwvc2NoZWR1bGVyaW5nZXN0ZXIvaW5nZXN0ZXIuZ28=) | `0.00% <0.00%> (ø)` | |
| [internal/common/database/upsert.go](https://codecov.io/gh/armadaproject/armada/pull/2391?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject#diff-aW50ZXJuYWwvY29tbW9uL2RhdGFiYXNlL3Vwc2VydC5nbw==) | `76.13% <75.00%> (ø)` | |
| [internal/scheduleringester/schedulerdb.go](https://codecov.io/gh/armadaproject/armada/pull/2391?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject#diff-aW50ZXJuYWwvc2NoZWR1bGVyaW5nZXN0ZXIvc2NoZWR1bGVyZGIuZ28=) | `72.16% <83.78%> (+6.20%)` | :arrow_up: |

... and [2 files with indirect coverage changes](https://codecov.io/gh/armadaproject/armada/pull/2391/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject)

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject)

</details>

[:umbrella: View full report in Codecov by Sentry](https://codecov.io/gh/armadaproject/armada/pull/2391?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject).


Patch:
diff --git a/internal/common/database/upsert.go b/internal/common/database/upsert.go
index 818e0bc5..34270b8f 100644
--- a/internal/common/database/upsert.go
+++ b/internal/common/database/upsert.go
@@ -6,12 +6,26 @@ import (
 	"reflect"
 	"strings"
 
+	"github.com/jackc/pgx/v4/pgxpool"
+
 	"github.com/google/uuid"
 	"github.com/jackc/pgx/v4"
-	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/pkg/errors"
 )
 
+func UpsertWithTransaction[T any](ctx context.Context, db *pgxpool.Pool, tableName string, records []T) error {
+	if len(records) == 0 {
+		return nil
+	}
+	return db.BeginTxFunc(ctx, pgx.TxOptions{
+		IsoLevel:       pgx.ReadCommitted,
+		AccessMode:     pgx.ReadWrite,
+		DeferrableMode: pgx.Deferrable,
+	}, func(tx pgx.Tx) error {
+		return Upsert(ctx, tx, tableName, records)
+	})
+}
+
 // Upsert is an optimised SQL call for bulk upserts.
 //
 // For efficiency, this function:
@@ -36,20 +50,7 @@ import (
 //
 // )
 // I.e., it should omit everything before and after the "(" and ")", respectively.
-func Upsert[T any](ctx context.Context, db *pgxpool.Pool, tableName string, records []T) error {
-	if len(records) == 0 {
-		return nil
-	}
-	return db.BeginTxFunc(ctx, pgx.TxOptions{
-		IsoLevel:       pgx.ReadCommitted,
-		AccessMode:     pgx.ReadWrite,
-		DeferrableMode: pgx.Deferrable,
-	}, func(tx pgx.Tx) error {
-		return CopyProtocolUpsert(ctx, tx, tableName, records)
-	})
-}
-
-func CopyProtocolUpsert[T any](ctx context.Context, tx pgx.Tx, tableName string, records []T) error {
+func Upsert[T any](ctx context.Context, tx pgx.Tx, tableName string, records []T) error {
 	if len(records) < 1 {
 		return nil
 	}
diff --git a/internal/scheduleringester/ingester.go b/internal/scheduleringester/ingester.go
index 90c37022..7d6e4ce7 100644
--- a/internal/scheduleringester/ingester.go
+++ b/internal/scheduleringester/ingester.go
@@ -25,7 +25,7 @@ func Run(config Configuration) {
 	if err != nil {
 		panic(errors.WithMessage(err, "Error opening connection to postgres"))
 	}
-	schedulerDb := NewSchedulerDb(db, svcMetrics, 100*time.Millisecond, 60*time.Second)
+	schedulerDb := NewSchedulerDb(db, svcMetrics, 100*time.Millisecond, 60*time.Second, 5*time.Second)
 
 	compressor, err := compress.NewZlibCompressor(1024)
 	if err != nil {
diff --git a/internal/scheduleringester/schedulerdb.go b/internal/scheduleringester/schedulerdb.go
index e687b3ac..2f180f32 100644
--- a/internal/scheduleringester/schedulerdb.go
+++ b/internal/scheduleringester/schedulerdb.go
@@ -6,13 +6,13 @@ import (
 	"strings"
 	"time"
 
+	"github.com/jackc/pgx/v4"
+
 	"github.com/google/uuid"
-	"github.com/hashicorp/go-multierror"
 	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/pkg/errors"
 	"golang.org/x/exp/maps"
 
-	"github.com/armadaproject/armada/internal/common/armadaerrors"
 	"github.com/armadaproject/armada/internal/common/database"
 	"github.com/armadaproject/armada/internal/common/ingest"
 	"github.com/armadaproject/armada/internal/common/ingest/metrics"
@@ -26,27 +26,70 @@ type SchedulerDb struct {
 	metrics        *metrics.Metrics
 	initialBackOff time.Duration
 	maxBackOff     time.Duration
+	lockTimeout    time.Duration
 }
 
-func NewSchedulerDb(db *pgxpool.Pool, metrics *metrics.Metrics, initialBackOff time.Duration, maxBackOff time.Duration) ingest.Sink[*DbOperationsWithMessageIds] {
-	return &SchedulerDb{db: db, metrics: metrics, initialBackOff: initialBackOff, maxBackOff: maxBackOff}
+func NewSchedulerDb(
+	db *pgxpool.Pool,
+	metrics *metrics.Metrics,
+	initialBackOff time.Duration,
+	maxBackOff time.Duration,
+	lockTimeout time.Duration,
+) ingest.Sink[*DbOperationsWithMessageIds] {
+	return &SchedulerDb{
+		db:             db,
+		metrics:        metrics,
+		initialBackOff: initialBackOff,
+		maxBackOff:     maxBackOff,
+		lockTimeout:    lockTimeout,
+	}
 }
 
+// Store persists all operations in the database.  Note that:
+//   - this function will retry until it either succeeds or a terminal error is encountered
+//   - this function will take out a postgres lock to ensure that other ingesters are not writing to the database
+//     at the same time (for details, see acquireLock())
 func (s *SchedulerDb) Store(ctx context.Context, instructions *DbOperationsWithMessageIds) error {
-	var result *multierror.Error = nil
-	for _, dbOp := range instructions.Ops {
-		err := ingest.WithRetry(func() (bool, error) {
-			err := s.WriteDbOp(ctx, dbOp)
-			shouldRetry := armadaerrors.IsNetworkError(err) || armadaerrors.IsRetryablePostgresError(err)
-			return shouldRetry, err
-		}, s.initialBackOff, s.maxBackOff)
-		result = multierror.Append(result, err)
-	}
-	return result.ErrorOrNil()
+	return ingest.WithRetry(func() (bool, error) {
+		err := s.db.BeginTxFunc(ctx, pgx.TxOptions{
+			IsoLevel:       pgx.ReadCommitted,
+			AccessMode:     pgx.ReadWrite,
+			DeferrableMode: pgx.Deferrable,
+		}, func(tx pgx.Tx) error {
+			// First acquire the write lock
+			lockCtx, cancel := context.WithTimeout(ctx, s.lockTimeout)
+			defer cancel()
+			err := s.acquireLock(lockCtx, tx)
+			if err != nil {
+				return err
+			}
+			// Now insert the ops
+			for _, dbOp := range instructions.Ops {
+				err := s.WriteDbOp(ctx, tx, dbOp)
+				if err != nil {
+					return err
+				}
+			}
+			return err
+		})
+		return true, err
+	}, s.initialBackOff, s.maxBackOff)
+}
+
+// acquireLock acquires the armada_scheduleringester_lock, which prevents two ingesters writing to the db at the same
+// time.  This is necessary because:
+// - when rows are inserted into the database they are stamped with a sequence number
+// - the scheduler relies on this sequence number increasing to ensure it has fetched all updated rows
+// - concurrent transactions will result in sequence numbers being interleaved across transactions.
+// - the interleaved sequences may result in the scheduler seeing sequence numbers that do not strictly increase over time.
+func (s *SchedulerDb) acquireLock(ctx context.Context, tx pgx.Tx) error {
+	const lockId = 8741339439634283896
+	_, err := tx.Exec(ctx, "SELECT pg_advisory_xact_lock($1)", lockId)
+	return errors.Wrapf(err, "Could not obtain lock")
 }
 
-func (s *SchedulerDb) WriteDbOp(ctx context.Context, op DbOperation) error {
-	queries := schedulerdb.New(s.db)
+func (s *SchedulerDb) WriteDbOp(ctx context.Context, tx pgx.Tx, op DbOperation) error {
+	queries := schedulerdb.New(tx)
 	switch o := op.(type) {
 	case InsertJobs:
 		records := make([]any, len(o))
@@ -55,7 +98,7 @@ func (s *SchedulerDb) WriteDbOp(ctx context.Context, op DbOperation) error {
 			records[i] = *v
 			i++
 		}
-		err := database.Upsert(ctx, s.db, "jobs", records)
+		err := database.Upsert(ctx, tx, "jobs", records)
 		if err != nil {
 			return err
 		}
@@ -66,7 +109,7 @@ func (s *SchedulerDb) WriteDbOp(ctx context.Context, op DbOperation) error {
 			records[i] = *v
 			i++
 		}
-		err := database.Upsert(ctx, s.db, "runs", records)
+		err := database.Upsert(ctx, tx, "runs", records)
 		if err != nil {
 			return err
 		}
@@ -214,7 +257,7 @@ func (s *SchedulerDb) WriteDbOp(ctx context.Context, op DbOperation) error {
 			records[i] = *v
 			i++
 		}
-		return database.Upsert(ctx, s.db, "job_run_errors", records)
+		return database.Upsert(ctx, tx, "job_run_errors", records)
 	case *InsertPartitionMarker:
 		for _, marker := range o.markers {
 			err := queries.InsertMarker(ctx, schedulerdb.InsertMarkerParams{


Test patch:
diff --git a/internal/common/database/upsert_test.go b/internal/common/database/upsert_test.go
index a16d1e59..d72f7309 100644
--- a/internal/common/database/upsert_test.go
+++ b/internal/common/database/upsert_test.go
@@ -75,13 +75,13 @@ func TestNamesValuesFromRecordPointer(t *testing.T) {
 	assert.Equal(t, []interface{}{r.Id, r.Value, r.Message, r.Serial}, values)
 }
 
-func TestUpsert(t *testing.T) {
+func TestUpsertWithTransaction(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Hour)
 	defer cancel()
 	err := withDb(func(db *pgxpool.Pool) error {
 		// Insert rows, read them back, and compare.
 		expected := makeRecords(10)
-		err := Upsert(ctx, db, TABLE_NAME, expected)
+		err := UpsertWithTransaction(ctx, db, TABLE_NAME, expected)
 		if !assert.NoError(t, err) {
 			return nil
 		}
@@ -96,7 +96,7 @@ func TestUpsert(t *testing.T) {
 
 		// Change one record, upsert, read back, and compare.
 		expected[0].Message = "foo"
-		err = Upsert(ctx, db, TABLE_NAME, expected)
+		err = UpsertWithTransaction(ctx, db, TABLE_NAME, expected)
 		if !assert.NoError(t, err) {
 			return nil
 		}
@@ -120,7 +120,7 @@ func TestConcurrency(t *testing.T) {
 			expected := makeRecords(10)
 			executor := fmt.Sprintf("executor-%d", i)
 			setMessageToExecutor(expected, executor)
-			err := Upsert(ctx, db, TABLE_NAME, expected)
+			err := UpsertWithTransaction(ctx, db, TABLE_NAME, expected)
 			if !assert.NoError(t, err) {
 				return nil
 			}
@@ -151,7 +151,7 @@ func TestAutoIncrement(t *testing.T) {
 	err := withDb(func(db *pgxpool.Pool) error {
 		// Insert two rows. These should automatically get auto-incrementing serial numbers.
 		records := makeRecords(2)
-		err := Upsert(ctx, db, TABLE_NAME, records)
+		err := UpsertWithTransaction(ctx, db, TABLE_NAME, records)
 		if !assert.NoError(t, err) {
 			return nil
 		}
@@ -172,7 +172,7 @@ func TestAutoIncrement(t *testing.T) {
 			records[1],
 		}
 
-		err = Upsert(ctx, db, TABLE_NAME, records)
+		err = UpsertWithTransaction(ctx, db, TABLE_NAME, records)
 		if !assert.NoError(t, err) {
 			return nil
 		}
@@ -256,7 +256,7 @@ func withDb(action func(db *pgxpool.Pool) error) error {
 				RETURN NEW;
 				END
 				$func$;
-			
+
 				CREATE TRIGGER next_serial
 				BEFORE INSERT or UPDATE ON %s
 				FOR EACH ROW
diff --git a/internal/scheduler/database/db_pruner_test.go b/internal/scheduler/database/db_pruner_test.go
index d7c38cc5..950b4ce9 100644
--- a/internal/scheduler/database/db_pruner_test.go
+++ b/internal/scheduler/database/db_pruner_test.go
@@ -116,11 +116,11 @@ func TestPruneDb_RemoveJobs(t *testing.T) {
 				jobsToInsert := commonutil.Map(tc.jobs, populateRequiredJobFields)
 				err := removeTriggers(ctx, db)
 				require.NoError(t, err)
-				err = database.Upsert(ctx, db, "jobs", jobsToInsert)
+				err = database.UpsertWithTransaction(ctx, db, "jobs", jobsToInsert)
 				require.NoError(t, err)
-				err = database.Upsert(ctx, db, "runs", tc.runs)
+				err = database.UpsertWithTransaction(ctx, db, "runs", tc.runs)
 				require.NoError(t, err)
-				err = database.Upsert(ctx, db, "job_run_errors", tc.errors)
+				err = database.UpsertWithTransaction(ctx, db, "job_run_errors", tc.errors)
 				require.NoError(t, err)
 				queries := New(db)
 				dbConn, err := db.Acquire(ctx)
diff --git a/internal/scheduler/database/job_repository_test.go b/internal/scheduler/database/job_repository_test.go
index caecce46..b5d0c9a7 100644
--- a/internal/scheduler/database/job_repository_test.go
+++ b/internal/scheduler/database/job_repository_test.go
@@ -87,9 +87,9 @@ func TestFetchJobUpdates(t *testing.T) {
 				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 
 				// Set up db
-				err := database.Upsert(ctx, repo.db, "jobs", tc.dbJobs)
+				err := database.UpsertWithTransaction(ctx, repo.db, "jobs", tc.dbJobs)
 				require.NoError(t, err)
-				err = database.Upsert(ctx, repo.db, "runs", tc.dbRuns)
+				err = database.UpsertWithTransaction(ctx, repo.db, "runs", tc.dbRuns)
 				require.NoError(t, err)
 
 				// Fetch updates
@@ -189,7 +189,7 @@ func TestFetchJobRunErrors(t *testing.T) {
 			err := withJobRepository(func(repo *PostgresJobRepository) error {
 				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 				// Set up db
-				err := database.Upsert(ctx, repo.db, "job_run_errors", tc.errorsInDb)
+				err := database.UpsertWithTransaction(ctx, repo.db, "job_run_errors", tc.errorsInDb)
 				require.NoError(t, err)
 
 				// Fetch updates
@@ -360,7 +360,7 @@ func TestFindInactiveRuns(t *testing.T) {
 				ctx, cancel := context.WithTimeout(context.Background(), 500*time.Second)
 
 				// Set up db
-				err := database.Upsert(ctx, repo.db, "runs", tc.dbRuns)
+				err := database.UpsertWithTransaction(ctx, repo.db, "runs", tc.dbRuns)
 				require.NoError(t, err)
 
 				inactive, err := repo.FindInactiveRuns(ctx, tc.runsToCheck)
@@ -490,9 +490,9 @@ func TestFetchJobRunLeases(t *testing.T) {
 				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 
 				// Set up db
-				err := database.Upsert(ctx, repo.db, "jobs", tc.dbJobs)
+				err := database.UpsertWithTransaction(ctx, repo.db, "jobs", tc.dbJobs)
 				require.NoError(t, err)
-				err = database.Upsert(ctx, repo.db, "runs", tc.dbRuns)
+				err = database.UpsertWithTransaction(ctx, repo.db, "runs", tc.dbRuns)
 				require.NoError(t, err)
 
 				leases, err := repo.FetchJobRunLeases(ctx, tc.executor, tc.maxRowsToFetch, tc.excludedRuns)
diff --git a/internal/scheduleringester/schedulerdb_test.go b/internal/scheduleringester/schedulerdb_test.go
index 1eef9b01..49bbd15f 100644
--- a/internal/scheduleringester/schedulerdb_test.go
+++ b/internal/scheduleringester/schedulerdb_test.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/jackc/pgx/v4"
 	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
@@ -13,6 +14,7 @@ import (
 	"golang.org/x/exp/constraints"
 	"golang.org/x/exp/maps"
 
+	"github.com/armadaproject/armada/internal/common/ingest/metrics"
 	"github.com/armadaproject/armada/internal/common/util"
 	schedulerdb "github.com/armadaproject/armada/internal/scheduler/database"
 )
@@ -288,7 +290,13 @@ func assertOpSuccess(t *testing.T, schedulerDb *SchedulerDb, serials map[string]
 	defer cancel()
 
 	// Apply the op to the database.
-	err := schedulerDb.WriteDbOp(ctx, op)
+	err := schedulerDb.db.BeginTxFunc(ctx, pgx.TxOptions{
+		IsoLevel:       pgx.ReadCommitted,
+		AccessMode:     pgx.ReadWrite,
+		DeferrableMode: pgx.Deferrable,
+	}, func(tx pgx.Tx) error {
+		return schedulerDb.WriteDbOp(ctx, tx, op)
+	})
 	if err != nil {
 		return err
 	}
@@ -577,6 +585,43 @@ func assertOpSuccess(t *testing.T, schedulerDb *SchedulerDb, serials map[string]
 	return nil
 }
 
+func TestStore(t *testing.T) {
+	jobId := util.ULID().String()
+	runId := uuid.New()
+	ops := []DbOperation{
+		InsertJobs{
+			jobId: &schedulerdb.Job{
+				JobID:          jobId,
+				JobSet:         "set1",
+				Groups:         make([]byte, 0),
+				SubmitMessage:  make([]byte, 0),
+				SchedulingInfo: make([]byte, 0),
+			},
+		},
+		InsertRuns{
+			runId: &schedulerdb.Run{JobID: jobId, RunID: runId},
+		},
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	err := schedulerdb.WithTestDb(func(q *schedulerdb.Queries, db *pgxpool.Pool) error {
+		schedulerDb := NewSchedulerDb(db, metrics.NewMetrics("test"), time.Second, time.Second, 10*time.Second)
+		err := schedulerDb.Store(ctx, &DbOperationsWithMessageIds{Ops: ops})
+		require.NoError(t, err)
+
+		jobIds, err := q.SelectAllJobIds(ctx)
+		require.NoError(t, err)
+		require.Equal(t, []string{jobId}, jobIds)
+
+		runIds, err := q.SelectAllRunIds(ctx)
+		require.NoError(t, err)
+		require.Equal(t, []uuid.UUID{runId}, runIds)
+
+		return nil
+	})
+	require.NoError(t, err)
+}
+
 func max[E constraints.Ordered](a, b E) E {
 	if a > b {
 		return a