Merge #144458

144458: workload: use multi region partitioner in auditor r=srosenberg a=csgourav

Audtior should user multi region partitioner in multi region run of tpcc. Otherwise the audit checks fail for some warehouses.

Fixes: [CRDB-49260](https://cockroachlabs.atlassian.net/browse/CRDB-49260)
Epic: none

Co-authored-by: Gourav Kumar <[email protected]>

Author: craig[bot]

pkg/workload/tpcc/new_order.go

@@ -187,9 +187,17 @@ func (n *newOrder) run(ctx context.Context, wID int) (interface{}, time.Duration
 			allLocal = 0
 			// To avoid picking the local warehouse again, randomly choose among n-1
 			// warehouses and swap in the nth if necessary.
-			item.olSupplyWID = n.config.wPart.randActive(rng)
-			for item.olSupplyWID == wID {
+			if len(n.config.multiRegionCfg.regions) > 0 {
+				// For multi-region configurations, use the multi-region partitioner
+				item.olSupplyWID = n.config.wMRPart.randActive(rng)
+				for item.olSupplyWID == wID {
+					item.olSupplyWID = n.config.wMRPart.randActive(rng)
+				}
+			} else {
 				item.olSupplyWID = n.config.wPart.randActive(rng)
+				for item.olSupplyWID == wID {
+					item.olSupplyWID = n.config.wPart.randActive(rng)
+				}
 			}
 			n.config.auditor.Lock()
 			n.config.auditor.orderLineRemoteWarehouseFreq[item.olSupplyWID]++

pkg/workload/tpcc/payment.go

@@ -223,10 +223,19 @@ func (p *payment) run(ctx context.Context, wID int) (interface{}, time.Duration,
 		d.cWID = wID
 		d.cDID = d.dID
 	} else {
-		d.cWID = p.config.wPart.randActive(rng)
-		// Find a cWID != w_id if there's more than 1 configured warehouse.
-		for d.cWID == wID && p.config.activeWarehouses > 1 {
+		if len(p.config.multiRegionCfg.regions) > 0 {
+			// For multi-region configurations, use the multi-region partitioner
+			d.cWID = p.config.wMRPart.randActive(rng)
+			// Find a cWID != w_id if there's more than 1 configured warehouse.
+			for d.cWID == wID && p.config.activeWarehouses > 1 {
+				d.cWID = p.config.wMRPart.randActive(rng)
+			}
+		} else {
 			d.cWID = p.config.wPart.randActive(rng)
+			// Find a cWID != w_id if there's more than 1 configured warehouse.
+			for d.cWID == wID && p.config.activeWarehouses > 1 {
+				d.cWID = p.config.wPart.randActive(rng)
+			}
 		}
 		p.config.auditor.Lock()
 		p.config.auditor.paymentRemoteWarehouseFreq[d.cWID]++

pkg/workload/tpcc/tpcc.go

@@ -484,9 +484,11 @@ func (w *tpcc) Hooks() workload.Hooks {
 				if err != nil {
 					return errors.Wrap(err, "error creating multi-region partitioner")
 				}
+				w.auditor = newAuditor(w.activeWarehouses, w.wMRPart, w.affinityPartitions)
+			} else {
+				w.auditor = newAuditor(w.activeWarehouses, w.wPart, w.affinityPartitions)
 			}
 
-			w.auditor = newAuditor(w.activeWarehouses, w.wPart, w.affinityPartitions)
 			return initializeMix(w)
 		},
 		PreCreate: func(db *gosql.DB) error {