refactor bandwidth throttling for replication target (#17980)

This refactor is to allow using the bandwidth throttling
for other purposes.

internal/bucket/bandwidth/monitor.go

@@ -25,27 +25,29 @@ import (
 	"golang.org/x/time/rate"
 )
 
-type throttle struct {
+type bucketThrottle struct {
 	*rate.Limiter
 	NodeBandwidthPerSec int64
 }
 
 // Monitor holds the state of the global bucket monitor
 type Monitor struct {
-	tlock                 sync.RWMutex // mutex for bucketThrottle
-	bucketThrottle        map[string]map[string]*throttle
-	mlock                 sync.RWMutex                             // mutex for activeBuckets map
-	activeBuckets         map[string]map[string]*bucketMeasurement // Buckets with objects in flight
-	bucketMovingAvgTicker *time.Ticker                             // Ticker for calculating moving averages
-	ctx                   context.Context                          // Context for generate
+	tlock sync.RWMutex // mutex for bucket throttling
+	mlock sync.RWMutex // mutex for bucket measurement
+
+	bucketsThrottle    map[BucketOptions]*bucketThrottle
+	bucketsMeasurement map[BucketOptions]*bucketMeasurement // Buckets with objects in flight
+
+	bucketMovingAvgTicker *time.Ticker    // Ticker for calculating moving averages
+	ctx                   context.Context // Context for generate
 	NodeCount             uint64
 }
 
 // NewMonitor returns a monitor with defaults.
 func NewMonitor(ctx context.Context, numNodes uint64) *Monitor {
 	m := &Monitor{
-		activeBuckets:         make(map[string]map[string]*bucketMeasurement),
-		bucketThrottle:        make(map[string]map[string]*throttle),
+		bucketsMeasurement:    make(map[BucketOptions]*bucketMeasurement),
+		bucketsThrottle:       make(map[BucketOptions]*bucketThrottle),
 		bucketMovingAvgTicker: time.NewTicker(2 * time.Second),
 		ctx:                   ctx,
 		NodeCount:             numNodes,
@@ -54,19 +56,16 @@ func NewMonitor(ctx context.Context, numNodes uint64) *Monitor {
 	return m
 }
 
-func (m *Monitor) updateMeasurement(bucket, arn string, bytes uint64) {
+func (m *Monitor) updateMeasurement(opts BucketOptions, bytes uint64) {
 	m.mlock.Lock()
 	defer m.mlock.Unlock()
-	tm, ok := m.activeBuckets[bucket]
+
+	tm, ok := m.bucketsMeasurement[opts]
 	if !ok {
-		tm = make(map[string]*bucketMeasurement)
+		tm = &bucketMeasurement{}
 	}
-	measurement, ok := tm[arn]
-	if !ok {
-		measurement = &bucketMeasurement{}
-	}
-	measurement.incrementBytes(bytes)
-	m.activeBuckets[bucket][arn] = measurement
+	tm.incrementBytes(bytes)
+	m.bucketsMeasurement[opts] = tm
 }
 
 // SelectionFunction for buckets
@@ -80,8 +79,8 @@ func SelectBuckets(buckets ...string) SelectionFunction {
 		}
 	}
 	return func(bucket string) bool {
-		for _, b := range buckets {
-			if b == "" || b == bucket {
+		for _, bkt := range buckets {
+			if bkt == bucket {
 				return true
 			}
 		}
@@ -97,7 +96,7 @@ type Details struct {
 
 // BucketBandwidthReport captures the details for all buckets.
 type BucketBandwidthReport struct {
-	BucketStats map[string]map[string]Details `json:"bucketStats,omitempty"`
+	BucketStats map[BucketOptions]Details `json:"bucketStats,omitempty"`
 }
 
 // GetReport gets the report for all bucket bandwidth details.
@@ -109,24 +108,18 @@ func (m *Monitor) GetReport(selectBucket SelectionFunction) *BucketBandwidthRepo
 
 func (m *Monitor) getReport(selectBucket SelectionFunction) *BucketBandwidthReport {
 	report := &BucketBandwidthReport{
-		BucketStats: make(map[string]map[string]Details),
+		BucketStats: make(map[BucketOptions]Details),
 	}
-	for bucket, bucketMeasurementMap := range m.activeBuckets {
-		if !selectBucket(bucket) {
+	for bucketOpts, bucketMeasurement := range m.bucketsMeasurement {
+		if !selectBucket(bucketOpts.Name) {
 			continue
 		}
 		m.tlock.RLock()
-		report.BucketStats[bucket] = make(map[string]Details)
-		if tgtThrottle, ok := m.bucketThrottle[bucket]; ok {
-			for arn, throttle := range tgtThrottle {
-				var currBw float64
-				if bucketMeasurement, ok := bucketMeasurementMap[arn]; ok {
-					currBw = bucketMeasurement.getExpMovingAvgBytesPerSecond()
-				}
-				report.BucketStats[bucket][arn] = Details{
-					LimitInBytesPerSecond:            throttle.NodeBandwidthPerSec * int64(m.NodeCount),
-					CurrentBandwidthInBytesPerSecond: currBw,
-				}
+		if tgtThrottle, ok := m.bucketsThrottle[bucketOpts]; ok {
+			currBw := bucketMeasurement.getExpMovingAvgBytesPerSecond()
+			report.BucketStats[bucketOpts] = Details{
+				LimitInBytesPerSecond:            tgtThrottle.NodeBandwidthPerSec * int64(m.NodeCount),
+				CurrentBandwidthInBytesPerSecond: currBw,
 			}
 		}
 		m.tlock.RUnlock()
@@ -149,92 +142,75 @@ func (m *Monitor) trackEWMA() {
 func (m *Monitor) updateMovingAvg() {
 	m.mlock.Lock()
 	defer m.mlock.Unlock()
-	for _, bucketMeasurement := range m.activeBuckets {
-		for _, measurement := range bucketMeasurement {
-			measurement.updateExponentialMovingAverage(time.Now())
-		}
+	for _, bucketMeasurement := range m.bucketsMeasurement {
+		bucketMeasurement.updateExponentialMovingAverage(time.Now())
 	}
 }
 
-func (m *Monitor) getBucketMeasurement(bucket, arn string, initTime time.Time) map[string]*bucketMeasurement {
-	bucketTracker, ok := m.activeBuckets[bucket]
-	if !ok {
-		bucketTracker = make(map[string]*bucketMeasurement)
-		bucketTracker[arn] = newBucketMeasurement(initTime)
-		m.activeBuckets[bucket] = bucketTracker
-	}
-	return bucketTracker
-}
-
-// track returns the measurement object for bucket
-func (m *Monitor) track(bucket, arn string) {
+func (m *Monitor) init(opts BucketOptions) {
 	m.mlock.Lock()
 	defer m.mlock.Unlock()
-	m.getBucketMeasurement(bucket, arn, time.Now())
+
+	_, ok := m.bucketsMeasurement[opts]
+	if !ok {
+		m.bucketsMeasurement[opts] = newBucketMeasurement(time.Now())
+	}
 }
 
 // DeleteBucket deletes monitoring the 'bucket'
 func (m *Monitor) DeleteBucket(bucket string) {
 	m.tlock.Lock()
-	delete(m.bucketThrottle, bucket)
+	for opts := range m.bucketsThrottle {
+		if opts.Name == bucket {
+			delete(m.bucketsThrottle, opts)
+		}
+	}
 	m.tlock.Unlock()
+
 	m.mlock.Lock()
-	delete(m.activeBuckets, bucket)
+	for opts := range m.bucketsMeasurement {
+		if opts.Name == bucket {
+			delete(m.bucketsMeasurement, opts)
+		}
+	}
 	m.mlock.Unlock()
 }
 
 // DeleteBucketThrottle deletes monitoring for a bucket's target
 func (m *Monitor) DeleteBucketThrottle(bucket, arn string) {
 	m.tlock.Lock()
-	if _, ok := m.bucketThrottle[bucket]; ok {
-		delete(m.bucketThrottle[bucket], arn)
-	}
+	delete(m.bucketsThrottle, BucketOptions{Name: bucket, ReplicationARN: arn})
 	m.tlock.Unlock()
 	m.mlock.Lock()
-	if _, ok := m.activeBuckets[bucket]; ok {
-		delete(m.activeBuckets[bucket], arn)
-	}
+	delete(m.bucketsMeasurement, BucketOptions{Name: bucket, ReplicationARN: arn})
 	m.mlock.Unlock()
 }
 
 // throttle returns currently configured throttle for this bucket
-func (m *Monitor) throttle(bucket, arn string) *throttle {
+func (m *Monitor) throttle(opts BucketOptions) *bucketThrottle {
 	m.tlock.RLock()
 	defer m.tlock.RUnlock()
-	return m.bucketThrottle[bucket][arn]
+	return m.bucketsThrottle[opts]
 }
 
 // SetBandwidthLimit sets the bandwidth limit for a bucket
 func (m *Monitor) SetBandwidthLimit(bucket, arn string, limit int64) {
 	m.tlock.Lock()
 	defer m.tlock.Unlock()
-	bw := limit / int64(m.NodeCount)
-	tgtMap, ok := m.bucketThrottle[bucket]
+	limitBytes := limit / int64(m.NodeCount)
+	throttle, ok := m.bucketsThrottle[BucketOptions{Name: bucket, ReplicationARN: arn}]
 	if !ok {
-		tgtMap = make(map[string]*throttle)
-		tgtMap[arn] = &throttle{
-			NodeBandwidthPerSec: bw,
-		}
+		throttle = &bucketThrottle{}
 	}
-	th, ok := tgtMap[arn]
-	if !ok {
-		th = &throttle{}
-	}
-	th.NodeBandwidthPerSec = bw
-	tgtMap[arn] = th
-	newlimit := rate.Every(time.Second / time.Duration(tgtMap[arn].NodeBandwidthPerSec))
-	tgtMap[arn].Limiter = rate.NewLimiter(newlimit, int(tgtMap[arn].NodeBandwidthPerSec))
-	m.bucketThrottle[bucket] = tgtMap
+	throttle.NodeBandwidthPerSec = limitBytes
+	throttle.Limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
+	m.bucketsThrottle[BucketOptions{Name: bucket, ReplicationARN: arn}] = throttle
 }
 
 // IsThrottled returns true if a bucket has bandwidth throttling enabled.
 func (m *Monitor) IsThrottled(bucket, arn string) bool {
 	m.tlock.RLock()
 	defer m.tlock.RUnlock()
-	th, ok := m.bucketThrottle[bucket]
-	if !ok {
-		return ok
-	}
-	_, ok = th[arn]
+	_, ok := m.bucketsThrottle[BucketOptions{Name: bucket, ReplicationARN: arn}]
 	return ok
 }