From 221bf8a270fa76375646e9477b5d8797d343e8f8 Mon Sep 17 00:00:00 2001
From: Ryan <ryanford@poluglottos.com>
Date: Wed, 3 May 2023 10:08:21 +0200
Subject: [PATCH] fix: exporting method instead of exporting field on
 middleware

---
 share/p2p/middleware.go      | 15 ++++++++++-----
 share/p2p/shrexeds/server.go |  6 +-----
 share/p2p/shrexnd/server.go  |  6 +-----
 3 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/share/p2p/middleware.go b/share/p2p/middleware.go
index 651263a756..df0a690af7 100644
--- a/share/p2p/middleware.go
+++ b/share/p2p/middleware.go
@@ -10,12 +10,12 @@ import (
 var log = logging.Logger("shrex/middleware")
 
 type Middleware struct {
-	// NumRateLimited is the number of requests that were rate limited. It is reset to 0 every time
-	// it is read and observed into metrics.
-	NumRateLimited atomic.Int64
-
+	// concurrencyLimit is the maximum number of requests that can be processed at once.
 	concurrencyLimit int64
+	// parallelRequests is the number of requests currently being processed.
 	parallelRequests atomic.Int64
+	// numRateLimited is the number of requests that were rate limited.
+	numRateLimited atomic.Int64
 }
 
 func NewMiddleware(concurrencyLimit int) *Middleware {
@@ -24,13 +24,18 @@ func NewMiddleware(concurrencyLimit int) *Middleware {
 	}
 }
 
+// DrainCounter returns the current value of the rate limit counter and resets it to 0.
+func (m *Middleware) DrainCounter() int64 {
+	return m.numRateLimited.Swap(0)
+}
+
 func (m *Middleware) RateLimitHandler(handler network.StreamHandler) network.StreamHandler {
 	return func(stream network.Stream) {
 		current := m.parallelRequests.Add(1)
 		defer m.parallelRequests.Add(-1)
 
 		if current > m.concurrencyLimit {
-			m.NumRateLimited.Add(1)
+			m.numRateLimited.Add(1)
 			log.Debug("concurrency limit reached")
 			err := stream.Close()
 			if err != nil {
diff --git a/share/p2p/shrexeds/server.go b/share/p2p/shrexeds/server.go
index 58d79a7361..6b2f2dcafb 100644
--- a/share/p2p/shrexeds/server.go
+++ b/share/p2p/shrexeds/server.go
@@ -63,11 +63,7 @@ func (s *Server) Stop(context.Context) error {
 }
 
 func (s *Server) observeRateLimitedRequests(ctx context.Context) {
-	var numRateLimited int64
-	if s.metrics != nil {
-		numRateLimited = s.middleware.NumRateLimited.Swap(0)
-	}
-
+	numRateLimited := s.middleware.DrainCounter()
 	if numRateLimited > 0 {
 		s.metrics.ObserveRequests(ctx, numRateLimited, p2p.StatusRateLimited)
 	}
diff --git a/share/p2p/shrexnd/server.go b/share/p2p/shrexnd/server.go
index 43576600ee..899ccd32f9 100644
--- a/share/p2p/shrexnd/server.go
+++ b/share/p2p/shrexnd/server.go
@@ -75,11 +75,7 @@ func (srv *Server) Stop(context.Context) error {
 }
 
 func (srv *Server) observeRateLimitedRequests(ctx context.Context) {
-	var numRateLimited int64
-	if srv.metrics != nil {
-		numRateLimited = srv.middleware.NumRateLimited.Swap(0)
-	}
-
+	numRateLimited := srv.middleware.DrainCounter()
 	if numRateLimited > 0 {
 		srv.metrics.ObserveRequests(ctx, numRateLimited, p2p.StatusRateLimited)
 	}