From 4b925ca55acc75d51098a7e02ea1e3abeb9bab76 Mon Sep 17 00:00:00 2001
From: Sanaz Taheri <35961250+staheri14@users.noreply.github.com>
Date: Thu, 26 Oct 2023 11:22:48 -0700
Subject: [PATCH] feat: Adds sample rate getter in Monitor & computes
 send/receive rates using configured sample rate (#1116)

This PR replaces the hardcoded sample rate in traffic rate tests with
the configured rate for each connection. This ensures the tests remain
relevant if the sampling rate changes in the future.
Closes https://github.com/celestiaorg/celestia-core/issues/1114
---
 libs/flowrate/flowrate.go   |  7 +++++++
 p2p/conn/connection_test.go | 31 +++++++++++++++++--------------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/libs/flowrate/flowrate.go b/libs/flowrate/flowrate.go
index 23252f4e33..9b063a40eb 100644
--- a/libs/flowrate/flowrate.go
+++ b/libs/flowrate/flowrate.go
@@ -64,6 +64,13 @@ func New(sampleRate, windowSize time.Duration) *Monitor {
 	}
 }
 
+// GetSampleRate returns the current sampling rate.
+func (m *Monitor) GetSampleRate() time.Duration {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.sRate
+}
+
 // Update records the transfer of n bytes and returns n. It should be called
 // after each Read/Write operation, even if n is 0.
 func (m *Monitor) Update(n int) int {
diff --git a/p2p/conn/connection_test.go b/p2p/conn/connection_test.go
index 8de1c7748d..f628b8aec9 100644
--- a/p2p/conn/connection_test.go
+++ b/p2p/conn/connection_test.go
@@ -161,13 +161,12 @@ func TestMConnectionSendRate(t *testing.T) {
 // maxSendRate returns the maximum send rate in bytes per second based on the MConnection's SendRate and other configs. It is used to calculate the highest expected value for the peak send rate.
 // The returned value is slightly higher than the configured SendRate.
 func (c *MConnection) maxSendRate() int64 {
-	// the sample rate is set when creating the MConnection and setting up its send monitor i.e., `c.sendMonitor`
-	// it defaults to 100ms which is what we use here
-	sampleRate := 100 * time.Millisecond
-	sendRate := round(float64(c.config.SendRate) * sampleRate.Seconds())
-	batchSizeBytes := int64(numBatchPacketMsgs * c._maxPacketMsgSize)
-	effectiveRatePerSample := int64(math.Ceil(float64(sendRate)/float64(batchSizeBytes))) * batchSizeBytes
-	effectiveSendRate := 10 * effectiveRatePerSample
+	sampleRate := c.sendMonitor.GetSampleRate().Seconds()
+	numberOfSamplePerSecond := 1 / sampleRate
+	sendRate := float64(round(float64(c.config.SendRate) * sampleRate))
+	batchSizeBytes := float64(numBatchPacketMsgs * c._maxPacketMsgSize)
+	effectiveRatePerSample := math.Ceil(sendRate/batchSizeBytes) * batchSizeBytes
+	effectiveSendRate := round(numberOfSamplePerSecond * effectiveRatePerSample)
 
 	return effectiveSendRate
 }
@@ -229,6 +228,11 @@ func TestMConnectionReceiveRate(t *testing.T) {
 	maxRecvRate := clientConn.maxRecvRate()
 
 	assert.True(t, peakRecvRate <= maxRecvRate, fmt.Sprintf("peakRecvRate %d > maxRecvRate %d", peakRecvRate, maxRecvRate))
+
+	peakSendRate := clientConn.sendMonitor.Status().PeakRate
+	maxSendRate := clientConn.maxSendRate()
+
+	assert.True(t, peakSendRate <= maxSendRate, fmt.Sprintf("peakSendRate %d > maxSendRate %d", peakSendRate, maxSendRate))
 }
 
 // maxRecvRate returns the maximum receive rate in bytes per second based on
@@ -236,13 +240,12 @@ func TestMConnectionReceiveRate(t *testing.T) {
 // It is used to calculate the highest expected value for the peak receive rate.
 // Note that the returned value is slightly higher than the configured RecvRate.
 func (c *MConnection) maxRecvRate() int64 {
-	// the sample rate is set when creating the MConnection and setting up its receive monitor i.e., `c.recvMonitor`
-	// it defaults to 100ms which is what we use here
-	sampleRate := 100 * time.Millisecond
-	recvRate := round(float64(c.config.RecvRate) * sampleRate.Seconds())
-	batchSizeBytes := int64(c._maxPacketMsgSize)
-	effectiveRecvRatePerSample := int64(math.Ceil(float64(recvRate)/float64(batchSizeBytes))) * batchSizeBytes
-	effectiveRecvRate := 10 * effectiveRecvRatePerSample
+	sampleRate := c.recvMonitor.GetSampleRate().Seconds()
+	numberOfSamplePerSeccond := 1 / sampleRate
+	recvRate := float64(round(float64(c.config.RecvRate) * sampleRate))
+	batchSizeBytes := float64(c._maxPacketMsgSize)
+	effectiveRecvRatePerSample := math.Ceil(recvRate/batchSizeBytes) * batchSizeBytes
+	effectiveRecvRate := round(numberOfSamplePerSeccond * effectiveRecvRatePerSample)
 
 	return effectiveRecvRate
 }