Merge pull request #2322 from jzelinskie/server-side-loop-repair

internal/datastore/pg: server-side loop for repair

Author: vroldanbet

internal/datastore/postgres/postgres.go

@@ -60,6 +60,13 @@ const (
 	gcBatchDeleteSize = 1000
 
 	primaryInstanceID = -1
+
+	// The number of loop iterations where pg_current_xact_id is called per
+	// query to the database during a repair.
+	//
+	// This value is ideally dependent on the underlying hardware, but 1M seems
+	// to be a reasonable starting place.
+	repairBatchSize = 1_000_000
 )
 
 var livingTupleConstraints = []string{"uq_relation_tuple_living_xid", "pk_relation_tuple"}
@@ -531,8 +538,6 @@ func (pgd *pgDatastore) Repair(ctx context.Context, operationName string, output
 	}
 }
 
-const batchSize = 10000
-
 func (pgd *pgDatastore) repairTransactionIDs(ctx context.Context, outputProgress bool) error {
 	conn, err := pgx.Connect(ctx, pgd.dburl)
 	if err != nil {
@@ -569,17 +574,8 @@ func (pgd *pgDatastore) repairTransactionIDs(ctx context.Context, outputProgress
 	}
 
 	for i := 0; i < counterDelta; i++ {
-		var batch pgx.Batch
-
-		batchCount := min(batchSize, counterDelta-i)
-		for j := 0; j < batchCount; j++ {
-			batch.Queue("begin;")
-			batch.Queue("select pg_current_xact_id();")
-			batch.Queue("rollback;")
-		}
-
-		br := conn.SendBatch(ctx, &batch)
-		if err := br.Close(); err != nil {
+		batchCount := min(repairBatchSize, counterDelta-i)
+		if _, err := conn.Exec(ctx, queryLoopXactID(batchCount)); err != nil {
 			return err
 		}
 
@@ -601,6 +597,17 @@ func (pgd *pgDatastore) repairTransactionIDs(ctx context.Context, outputProgress
 	return nil
 }
 
+// queryLoopXactID performs pg_current_xact_id() in a server-side loop in the
+// database in order to increment the xact_id.
+func queryLoopXactID(batchSize int) string {
+	return fmt.Sprintf(`DO $$
+BEGIN
+  FOR i IN 1..%d LOOP
+    PERFORM pg_current_xact_id(); ROLLBACK;
+  END LOOP;
+END $$;`, batchSize)
+}
+
 // RepairOperations returns the available repair operations for the datastore.
 func (pgd *pgDatastore) RepairOperations() []datastore.RepairOperation {
 	return []datastore.RepairOperation{