From 654aaa0522bec199054bc5ae930db0243d3ce6b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Viktor=20Levente=20T=C3=B3th?=
 <toth.viktor.levente@gmail.com>
Date: Tue, 26 Sep 2023 16:23:29 +0200
Subject: [PATCH] fix: set sample duration into status

---
 pkg/storageincentives/agent.go               | 8 ++++----
 pkg/storageincentives/redistributionstate.go | 3 ++-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/pkg/storageincentives/agent.go b/pkg/storageincentives/agent.go
index c6e746a5c45..0d151a66641 100644
--- a/pkg/storageincentives/agent.go
+++ b/pkg/storageincentives/agent.go
@@ -429,14 +429,17 @@ func (a *Agent) handleSample(ctx context.Context, round uint64) (bool, error) {
 		return false, nil
 	}
 
+	t := time.Now()
 	sample, err := a.makeSample(ctx, storageRadius)
 	if err != nil {
 		return false, err
 	}
+	dur := time.Since(t)
+	a.metrics.SampleDuration.Set(dur.Seconds())
 
 	a.logger.Info("produced sample", "hash", sample.ReserveSampleHash, "radius", sample.StorageRadius, "round", round)
 
-	a.state.SetSampleData(round, sample)
+	a.state.SetSampleData(round, sample, dur)
 
 	return true, nil
 }
@@ -452,13 +455,10 @@ func (a *Agent) makeSample(ctx context.Context, storageRadius uint8) (SampleData
 		return SampleData{}, err
 	}
 
-	t := time.Now()
 	rSample, err := a.store.ReserveSample(ctx, salt, storageRadius, uint64(timeLimiter), a.minBatchBalance())
 	if err != nil {
 		return SampleData{}, err
 	}
-	dur := time.Since(t)
-	a.metrics.SampleDuration.Set(dur.Seconds())
 
 	sampleHash, err := sampleHash(rSample.Items)
 	if err != nil {
diff --git a/pkg/storageincentives/redistributionstate.go b/pkg/storageincentives/redistributionstate.go
index 20f61c19cc0..b1f4b60dd03 100644
--- a/pkg/storageincentives/redistributionstate.go
+++ b/pkg/storageincentives/redistributionstate.go
@@ -244,13 +244,14 @@ func (r *RedistributionState) SampleData(round uint64) (SampleData, bool) {
 	return *rd.SampleData, true
 }
 
-func (r *RedistributionState) SetSampleData(round uint64, sd SampleData) {
+func (r *RedistributionState) SetSampleData(round uint64, sd SampleData, dur time.Duration) {
 	r.mtx.Lock()
 	defer r.mtx.Unlock()
 
 	rd := r.status.RoundData[round]
 	rd.SampleData = &sd
 	r.status.RoundData[round] = rd
+	r.status.SampleDuration = dur
 
 	r.save()
 }