From 5bec2d2478e5fed41e00f18b27d2018c3b1616c1 Mon Sep 17 00:00:00 2001
From: PJ <peterjan.brone@gmail.com>
Date: Thu, 12 Oct 2023 12:32:33 +0200
Subject: [PATCH] worker: record spending for every batch

---
 worker/rhpv2.go | 28 +++++++++++-----------------
 worker/rhpv3.go |  2 +-
 2 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/worker/rhpv2.go b/worker/rhpv2.go
index 553f67ba3..6b849fbb1 100644
--- a/worker/rhpv2.go
+++ b/worker/rhpv2.go
@@ -309,6 +309,10 @@ func (w *worker) PruneContract(ctx context.Context, hostIP string, hostKey types
 				if deleted < uint64(len(indices)) {
 					remaining = uint64(len(indices)) - deleted
 				}
+
+				// return sizes instead of number of roots
+				deleted *= rhpv2.SectorSize
+				remaining *= rhpv2.SectorSize
 				return
 			})
 		})
@@ -324,12 +328,6 @@ func (w *worker) deleteContractRoots(t *rhpv2.Transport, rev *rhpv2.ContractRevi
 		return 0, nil
 	}
 
-	// record contract spending
-	var totalCost types.Currency
-	defer func() {
-		w.contractSpendingRecorder.Record(rev.ID(), rev.Revision.RevisionNumber, rev.Revision.Filesize, api.ContractSpending{Deletions: totalCost})
-	}()
-
 	// sort in descending order so that we can use 'range'
 	sort.Slice(indices, func(i, j int) bool {
 		return indices[i] > indices[j]
@@ -472,9 +470,11 @@ func (w *worker) deleteContractRoots(t *rhpv2.Transport, rev *rhpv2.ContractRevi
 			rev.Signatures[0].Signature = renterSig.Signature[:]
 			rev.Signatures[1].Signature = hostSig.Signature[:]
 
-			// update total cost
-			totalCost = totalCost.Add(cost)
+			// update deleted count
 			deleted += uint64(len(batch))
+
+			// record spending
+			w.contractSpendingRecorder.Record(rev.ID(), rev.Revision.RevisionNumber, rev.Revision.Filesize, api.ContractSpending{Deletions: cost})
 			return nil
 		}(); err != nil {
 			return
@@ -497,12 +497,6 @@ func (w *worker) fetchContractRoots(t *rhpv2.Transport, rev *rhpv2.ContractRevis
 	// derive the renter key
 	renterKey := w.deriveRenterKey(rev.HostKey())
 
-	// record contract spending
-	var totalCost types.Currency
-	defer func() {
-		w.contractSpendingRecorder.Record(rev.ID(), rev.Revision.RevisionNumber, rev.Revision.Filesize, api.ContractSpending{SectorRoots: totalCost})
-	}()
-
 	// download the full set of SectorRoots
 	numsectors := rev.NumSectors()
 	for offset := uint64(0); offset < numsectors; {
@@ -564,12 +558,12 @@ func (w *worker) fetchContractRoots(t *rhpv2.Transport, rev *rhpv2.ContractRevis
 			return nil, ErrInvalidMerkleProof
 		}
 
-		// update the total cost
-		totalCost = totalCost.Add(cost)
-
 		// append roots
 		roots = append(roots, rootsResp.SectorRoots...)
 		offset += n
+
+		// record spending
+		w.contractSpendingRecorder.Record(rev.ID(), rev.Revision.RevisionNumber, rev.Revision.Filesize, api.ContractSpending{SectorRoots: cost})
 	}
 	return
 }
diff --git a/worker/rhpv3.go b/worker/rhpv3.go
index aefd7c11c..efc8fa700 100644
--- a/worker/rhpv3.go
+++ b/worker/rhpv3.go
@@ -662,7 +662,7 @@ func (h *host) UploadSector(ctx context.Context, sector *[rhpv2.SectorSize]byte,
 
 	// record spending
 	h.contractSpendingRecorder.Record(rev.ParentID, rev.RevisionNumber, rev.Filesize, api.ContractSpending{Uploads: cost})
-	return root, err
+	return root, nil
 }
 
 // padBandwitdh pads the bandwidth to the next multiple of 1460 bytes.  1460