Merge pull request #421 from SiaFoundation/pj/contract-spending

Contract Spending
SiaFoundation · Jun 16, 2023 · d50d6a2 · d50d6a2
2 parents cf3e113 + ae50fee
commit d50d6a2
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 22 deletions.
diff --git a/internal/testing/cluster_test.go b/internal/testing/cluster_test.go
@@ -524,6 +524,12 @@ func TestUploadDownloadSpending(t *testing.T) {
 
 		nFunded := 0
 		for _, c := range cms {
+			if !c.Spending.Uploads.IsZero() {
+				t.Fatal("upload spending should be zero")
+			}
+			if !c.Spending.Downloads.IsZero() {
+				t.Fatal("download spending should be zero")
+			}
 			if !c.Spending.FundAccount.IsZero() {
 				nFunded++
 				if c.RevisionNumber == 0 {
@@ -654,8 +660,8 @@ func TestUploadDownloadSpending(t *testing.T) {
 		}
 
 		for _, c := range cms {
-			if !c.Spending.Uploads.IsZero() {
-				t.Fatal("upload spending should be zero")
+			if c.Spending.Uploads.IsZero() {
+				t.Fatal("upload spending shouldn't be zero")
 			}
 			if !c.Spending.Downloads.IsZero() {
 				t.Fatal("download spending should be zero")

diff --git a/worker/rhpv3.go b/worker/rhpv3.go
@@ -545,27 +545,27 @@ func (h *host) priceTable(ctx context.Context, rev *types.FileContractRevision)
 	return pt.HostPriceTable, nil
 }
 
-func (r *host) DownloadSector(ctx context.Context, w io.Writer, root types.Hash256, offset, length uint64) (err error) {
-	pt, err := r.priceTable(ctx, nil)
+func (h *host) DownloadSector(ctx context.Context, w io.Writer, root types.Hash256, offset, length uint64) (err error) {
+	pt, err := h.priceTable(ctx, nil)
 	if err != nil {
 		return err
 	}
 	// return errBalanceInsufficient if balance insufficient
 	defer func() {
 		if isBalanceInsufficient(err) {
-			err = fmt.Errorf("%w %v, err: %v", errBalanceInsufficient, r.HostKey(), err)
+			err = fmt.Errorf("%w %v, err: %v", errBalanceInsufficient, h.HostKey(), err)
 		}
 	}()
 
-	return r.acc.WithWithdrawal(ctx, func() (amount types.Currency, err error) {
-		err = r.transportPool.withTransportV3(ctx, r.HostKey(), r.siamuxAddr, func(t *transportV3) error {
+	return h.acc.WithWithdrawal(ctx, func() (amount types.Currency, err error) {
+		err = h.transportPool.withTransportV3(ctx, h.HostKey(), h.siamuxAddr, func(t *transportV3) error {
 			cost, err := readSectorCost(pt)
 			if err != nil {
 				return err
 			}
 
 			var refund types.Currency
-			payment := rhpv3.PayByEphemeralAccount(r.acc.id, cost, pt.HostBlockHeight+defaultWithdrawalExpiryBlocks, r.accountKey)
+			payment := rhpv3.PayByEphemeralAccount(h.acc.id, cost, pt.HostBlockHeight+defaultWithdrawalExpiryBlocks, h.accountKey)
 			cost, refund, err = RPCReadSector(ctx, t, w, pt, &payment, offset, length, root, true)
 			amount = cost.Sub(refund)
 			return err
@@ -575,9 +575,9 @@ func (r *host) DownloadSector(ctx context.Context, w io.Writer, root types.Hash2
 }
 
 // UploadSector uploads a sector to the host.
-func (r *host) UploadSector(ctx context.Context, sector *[rhpv2.SectorSize]byte, rev types.FileContractRevision) (root types.Hash256, err error) {
+func (h *host) UploadSector(ctx context.Context, sector *[rhpv2.SectorSize]byte, rev types.FileContractRevision) (root types.Hash256, err error) {
 	// fetch price table
-	pt, err := r.priceTable(ctx, nil)
+	pt, err := h.priceTable(ctx, nil)
 	if err != nil {
 		return types.Hash256{}, err
 	}
@@ -593,15 +593,22 @@ func (r *host) UploadSector(ctx context.Context, sector *[rhpv2.SectorSize]byte,
 	if rev.RevisionNumber == math.MaxUint64 {
 		return types.Hash256{}, errors.New("revision number has reached max")
 	}
-	payment, ok := rhpv3.PayByContract(&rev, expectedCost, r.acc.id, r.renterKey)
+	payment, ok := rhpv3.PayByContract(&rev, expectedCost, h.acc.id, h.renterKey)
 	if !ok {
 		return types.Hash256{}, errors.New("failed to create payment")
 	}
 
-	err = r.transportPool.withTransportV3(ctx, r.HostKey(), r.siamuxAddr, func(t *transportV3) error {
-		root, _, err = RPCAppendSector(ctx, t, r.renterKey, pt, rev, &payment, sector)
+	var cost types.Currency
+	err = h.transportPool.withTransportV3(ctx, h.HostKey(), h.siamuxAddr, func(t *transportV3) error {
+		root, cost, err = RPCAppendSector(ctx, t, h.renterKey, pt, rev, &payment, sector)
 		return err
 	})
+	if err != nil {
+		return types.Hash256{}, err
+	}
+
+	// record spending
+	h.contractSpendingRecorder.Record(rev.ParentID, rev.RevisionNumber, api.ContractSpending{Uploads: cost})
 	return root, err
 }
 

diff --git a/worker/spending.go b/worker/spending.go
@@ -41,12 +41,6 @@ func recordContractSpending(ctx context.Context, rev *types.FileContractRevision
 	}
 }
 
-// WithContractSpendingRecorder returns a context with the
-// ContractSpendingRecorder attached.
-func WithContractSpendingRecorder(ctx context.Context, sr ContractSpendingRecorder) context.Context {
-	return context.WithValue(ctx, keyContractSpendingRecorder, sr)
-}
-
 func (w *worker) initContractSpendingRecorder() {
 	if w.contractSpendingRecorder != nil {
 		panic("contractSpendingRecorder already initialized") // developer error

diff --git a/worker/worker.go b/worker/worker.go
@@ -820,9 +820,6 @@ func (w *worker) slabMigrateHandler(jc jape.Context) {
 	// attach gouging checker to the context
 	ctx = WithGougingChecker(ctx, w.bus, up.GougingParams)
 
-	// attach contract spending recorder to the context.
-	ctx = WithContractSpendingRecorder(ctx, w.contractSpendingRecorder)
-
 	// fetch all contracts
 	dlContracts, err := w.bus.Contracts(ctx)
 	if jc.Check("couldn't fetch contracts from bus", err) != nil {