Update RecordContractSpending (#1417)

stores/metadata.go

@@ -52,9 +52,7 @@ var (
 	pruneDirsAlertID  = frand.Entropy256()
 )
 
-var (
-	objectDeleteBatchSizes = []int64{10, 50, 100, 200, 500, 1000, 5000, 10000, 50000, 100000}
-)
+var objectDeleteBatchSizes = []int64{10, 50, 100, 200, 500, 1000, 5000, 10000, 50000, 100000}
 
 type (
 	contractState uint8
@@ -715,16 +713,11 @@ func (s *SQLStore) RecordContractSpending(ctx context.Context, records []api.Con
 	}
 	metrics := make([]api.ContractMetric, 0, len(squashedRecords))
 	for fcid, newSpending := range squashedRecords {
-		err := s.retryTransaction(ctx, func(tx *gorm.DB) error {
-			var contract dbContract
-			err := tx.Model(&dbContract{}).
-				Where("fcid", fileContractID(fcid)).
-				Joins("Host").
-				Take(&contract).Error
-			if errors.Is(err, gorm.ErrRecordNotFound) {
-				return nil // contract not found, continue with next one
+		err := s.db.Transaction(ctx, func(tx sql.DatabaseTx) error {
+			contract, err := tx.Contract(ctx, fcid)
+			if errors.Is(err, api.ErrContractNotFound) {
 			} else if err != nil {
-				return err
+				return fmt.Errorf("failed to fetch contract: %w", err)
 			}
 
 			remainingCollateral := types.ZeroCurrency
@@ -734,38 +727,35 @@ func (s *SQLStore) RecordContractSpending(ctx context.Context, records []api.Con
 			m := api.ContractMetric{
 				Timestamp:           api.TimeNow(),
 				ContractID:          fcid,
-				HostKey:             types.PublicKey(contract.Host.PublicKey),
+				HostKey:             contract.HostKey,
 				RemainingCollateral: remainingCollateral,
 				RemainingFunds:      latestValues[fcid].validRenterPayout,
 				RevisionNumber:      latestValues[fcid].revision,
-				UploadSpending:      types.Currency(contract.UploadSpending).Add(newSpending.Uploads),
-				DownloadSpending:    types.Currency(contract.DownloadSpending).Add(newSpending.Downloads),
-				FundAccountSpending: types.Currency(contract.FundAccountSpending).Add(newSpending.FundAccount),
-				DeleteSpending:      types.Currency(contract.DeleteSpending).Add(newSpending.Deletions),
-				ListSpending:        types.Currency(contract.ListSpending).Add(newSpending.SectorRoots),
+				UploadSpending:      contract.Spending.Uploads.Add(newSpending.Uploads),
+				DownloadSpending:    contract.Spending.Downloads.Add(newSpending.Downloads),
+				FundAccountSpending: contract.Spending.FundAccount.Add(newSpending.FundAccount),
+				DeleteSpending:      contract.Spending.Deletions.Add(newSpending.Deletions),
+				ListSpending:        contract.Spending.SectorRoots.Add(newSpending.SectorRoots),
 			}
 			metrics = append(metrics, m)
 
-			updates := make(map[string]interface{})
+			var updates api.ContractSpending
 			if !newSpending.Uploads.IsZero() {
-				updates["upload_spending"] = currency(m.UploadSpending)
+				updates.Uploads = m.UploadSpending
 			}
 			if !newSpending.Downloads.IsZero() {
-				updates["download_spending"] = currency(m.DownloadSpending)
+				updates.Downloads = m.DownloadSpending
 			}
 			if !newSpending.FundAccount.IsZero() {
-				updates["fund_account_spending"] = currency(m.FundAccountSpending)
+				updates.FundAccount = m.FundAccountSpending
 			}
 			if !newSpending.Deletions.IsZero() {
-				updates["delete_spending"] = currency(m.DeleteSpending)
+				updates.Deletions = m.DeleteSpending
 			}
 			if !newSpending.SectorRoots.IsZero() {
-				updates["list_spending"] = currency(m.ListSpending)
+				updates.SectorRoots = m.ListSpending
 			}
-			updates["revision_number"] = latestValues[fcid].revision
-			updates["size"] = latestValues[fcid].size
-			err = tx.Model(&contract).Updates(updates).Error
-			return err
+			return tx.RecordContractSpending(ctx, fcid, latestValues[fcid].revision, latestValues[fcid].size, updates)
 		})
 		if err != nil {
 			return err

stores/metadata_test.go

@@ -2328,9 +2328,10 @@ func TestRecordContractSpending(t *testing.T) {
 	cm, err := ss.addTestContract(fcid, hk)
 	if err != nil {
 		t.Fatal(err)
-	}
-	if cm.Spending != (api.ContractSpending{}) {
+	} else if cm.Spending != (api.ContractSpending{}) {
 		t.Fatal("spending should be all 0")
+	} else if cm.Size != 0 && cm.RevisionNumber != 0 {
+		t.Fatalf("unexpected size or revision number, %v %v", cm.Size, cm.RevisionNumber)
 	}
 
 	// Record some spending.
@@ -2350,6 +2351,8 @@ func TestRecordContractSpending(t *testing.T) {
 		{
 			ContractID:       fcid,
 			ContractSpending: expectedSpending,
+			RevisionNumber:   100,
+			Size:             200,
 		},
 	})
 	if err != nil {
@@ -2358,16 +2361,20 @@ func TestRecordContractSpending(t *testing.T) {
 	cm2, err := ss.Contract(context.Background(), fcid)
 	if err != nil {
 		t.Fatal(err)
-	}
-	if cm2.Spending != expectedSpending {
+	} else if cm2.Spending != expectedSpending {
 		t.Fatal("invalid spending", cm2.Spending, expectedSpending)
+	} else if cm2.Size != 200 && cm2.RevisionNumber != 100 {
+		t.Fatalf("unexpected size or revision number, %v %v", cm2.Size, cm2.RevisionNumber)
 	}
 
-	// Record the same spending again.
+	// Record the same spending again but with a lower revision number. This
+	// shouldn't update the size.
 	err = ss.RecordContractSpending(context.Background(), []api.ContractSpendingRecord{
 		{
 			ContractID:       fcid,
 			ContractSpending: expectedSpending,
+			RevisionNumber:   100,
+			Size:             200,
 		},
 	})
 	if err != nil {
@@ -2377,9 +2384,10 @@ func TestRecordContractSpending(t *testing.T) {
 	cm3, err := ss.Contract(context.Background(), fcid)
 	if err != nil {
 		t.Fatal(err)
-	}
-	if cm3.Spending != expectedSpending {
+	} else if cm3.Spending != expectedSpending {
 		t.Fatal("invalid spending")
+	} else if cm2.Size != 200 && cm2.RevisionNumber != 100 {
+		t.Fatalf("unexpected size or revision number, %v %v", cm2.Size, cm2.RevisionNumber)
 	}
 }
 

stores/sql/database.go

@@ -235,6 +235,9 @@ type (
 		// or slab buffer.
 		PruneSlabs(ctx context.Context, limit int64) (int64, error)
 
+		// RecordContractSpending records new spending for a contract
+		RecordContractSpending(ctx context.Context, fcid types.FileContractID, revisionNumber, size uint64, newSpending api.ContractSpending) error
+
 		// RecordHostScans records the results of host scans in the database
 		// such as recording the settings and price table of a host in case of
 		// success and updating the uptime and downtime of a host.

stores/sql/main.go

@@ -2763,3 +2763,42 @@ func MarkPackedSlabUploaded(ctx context.Context, tx Tx, slab api.UploadedPackedS
 	}
 	return bufferFileName, nil
 }
+
+func RecordContractSpending(ctx context.Context, tx Tx, fcid types.FileContractID, revisionNumber, size uint64, newSpending api.ContractSpending) error {
+	var updateKeys []string
+	var updateValues []interface{}
+
+	if !newSpending.Uploads.IsZero() {
+		updateKeys = append(updateKeys, "upload_spending = ?")
+		updateValues = append(updateValues, Currency(newSpending.Uploads))
+	}
+	if !newSpending.Downloads.IsZero() {
+		updateKeys = append(updateKeys, "download_spending = ?")
+		updateValues = append(updateValues, Currency(newSpending.Downloads))
+	}
+	if !newSpending.FundAccount.IsZero() {
+		updateKeys = append(updateKeys, "fund_account_spending = ?")
+		updateValues = append(updateValues, Currency(newSpending.FundAccount))
+	}
+	if !newSpending.Deletions.IsZero() {
+		updateKeys = append(updateKeys, "delete_spending = ?")
+		updateValues = append(updateValues, Currency(newSpending.Deletions))
+	}
+	if !newSpending.SectorRoots.IsZero() {
+		updateKeys = append(updateKeys, "list_spending = ?")
+		updateValues = append(updateValues, Currency(newSpending.SectorRoots))
+	}
+	updateKeys = append(updateKeys, "revision_number = ?", "size = ?")
+	updateValues = append(updateValues, revisionNumber, size)
+
+	updateValues = append(updateValues, FileContractID(fcid))
+	_, err := tx.Exec(ctx, fmt.Sprintf(`
+    UPDATE contracts
+    SET %s
+    WHERE fcid = ?
+  `, strings.Join(updateKeys, ",")), updateValues...)
+	if err != nil {
+		return fmt.Errorf("failed to record contract spending: %w", err)
+	}
+	return nil
+}

stores/sql/mysql/main.go

@@ -611,6 +611,10 @@ func (tx *MainDatabaseTx) PruneSlabs(ctx context.Context, limit int64) (int64, e
 	return res.RowsAffected()
 }
 
+func (tx *MainDatabaseTx) RecordContractSpending(ctx context.Context, fcid types.FileContractID, revisionNumber, size uint64, newSpending api.ContractSpending) error {
+	return ssql.RecordContractSpending(ctx, tx, fcid, revisionNumber, size, newSpending)
+}
+
 func (tx *MainDatabaseTx) RecordHostScans(ctx context.Context, scans []api.HostScan) error {
 	return ssql.RecordHostScans(ctx, tx, scans)
 }

stores/sql/sqlite/main.go

@@ -608,6 +608,10 @@ func (tx *MainDatabaseTx) PruneSlabs(ctx context.Context, limit int64) (int64, e
 	return res.RowsAffected()
 }
 
+func (tx *MainDatabaseTx) RecordContractSpending(ctx context.Context, fcid types.FileContractID, revisionNumber, size uint64, newSpending api.ContractSpending) error {
+	return ssql.RecordContractSpending(ctx, tx, fcid, revisionNumber, size, newSpending)
+}
+
 func (tx *MainDatabaseTx) RecordHostScans(ctx context.Context, scans []api.HostScan) error {
 	return ssql.RecordHostScans(ctx, tx, scans)
 }