internal/database: support serializable transactions

Add DB.TransactSerializable, which executes a transaction with
serializable isolation. Although serializable transactions are more
expensive, they reduce the risk for anomalies, like the constraint
violations we sometimes see.

This CL does not change any existing behavior. In particular, the
DB.Transact method, despite calling sql.DB.BeginTx instead of
sql.DB.Begin, uses the background context exactly as Begin does
(source:
https://github.com/golang/go/blob/go1.14.3/src/database/sql/sql.go#L1689).

Change-Id: Iab9c99ca35de9ef884b149a3b39ecac3300ecac0
Reviewed-on: https://team-review.git.corp.google.com/c/golang/discovery/+/750952
Reviewed-by: Julie Qiu <julieqiu@google.com>

1 files changed, 80 insertions, 0 deletions

diff --git a/internal/database/database_test.go b/internal/database/database_test.go
index 589acf17..7d95bd83 100644
--- a/internal/database/database_test.go
+++ b/internal/database/database_test.go
@@ -308,3 +308,83 @@ func TestBulkUpdate(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestTransactSerializable(t *testing.T) {
+	// Test that serializable transactions retry until success.
+	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	defer cancel()
+
+	// This test was taken from the example at https://www.postgresql.org/docs/11/transaction-iso.html,
+	// section 13.2.3.
+	for _, stmt := range []string{
+		`DROP TABLE IF EXISTS ser`,
+		`CREATE TABLE ser (id INTEGER GENERATED ALWAYS AS IDENTITY PRIMARY KEY, class INTEGER, value INTEGER)`,
+		`INSERT INTO ser (class, value) VALUES (1, 10), (1, 20), (2, 100), (2, 200)`,
+	} {
+		if _, err := testDB.Exec(ctx, stmt); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// A transaction that sums values in class 1 and inserts that sum into class 2,
+	// or vice versa.
+	insertSum := func(tx *DB, queryClass int) error {
+		var sum int
+		err := tx.QueryRow(ctx, `SELECT SUM(value) FROM ser WHERE class = $1`, queryClass).Scan(&sum)
+		if err != nil {
+			return err
+		}
+		insertClass := 3 - queryClass
+		_, err = tx.Exec(ctx, `INSERT INTO ser (class, value) VALUES ($1, $2)`, insertClass, sum)
+		return err
+	}
+
+	// Run the following two transactions multiple times concurrently:
+	//   sum rows with class = 1 and insert as a row with class 2
+	//   sum rows with class = 2 and insert as a row with class 1
+	// We determined empirically that this number of transactions produces a serialization conflict
+	// 100 times out of 100.
+	const numTransactions = 10
+	errc := make(chan error, numTransactions)
+	for i := 0; i < numTransactions; i++ {
+		i := i
+		go func() {
+			errc <- testDB.TransactSerializable(ctx, func(tx *DB) error { return insertSum(tx, 1+i%2) })
+		}()
+	}
+	// None of the transactions should fail.
+	for i := 0; i < numTransactions; i++ {
+		if err := <-errc; err != nil {
+			t.Fatal(err)
+		}
+	}
+	t.Logf("max retries: %d", testDB.MaxRetries())
+	// If nothing got retried, this test isn't exercising some important behavior.
+	if testDB.MaxRetries() == 0 {
+		t.Fatal("did not see any retries")
+	}
+
+	// Demonstrate serializability: there should be numTransactions new rows in
+	// addition to the 4 we started with, and viewing the rows in insertion
+	// order, each of the new rows should have the sum of the other class's rows
+	// so far.
+	type row struct {
+		Class, Value int
+	}
+	var rows []row
+	if err := testDB.CollectStructs(ctx, &rows, `SELECT class, value FROM ser ORDER BY id`); err != nil {
+		t.Fatal(err)
+	}
+	const initialRows = 4
+	if got, want := len(rows), initialRows+numTransactions; got != want {
+		t.Fatalf("got %d rows, want %d", got, want)
+	}
+	sum := make([]int, 2)
+	for i, r := range rows {
+		if got, want := r.Value, sum[2-r.Class]; got != want && i >= initialRows {
+			t.Fatalf("row #%d: got %d, want %d", i, got, want)
+		}
+		sum[r.Class-1] += r.Value
+	}
+
+}