From 8a255a0ee81d0ef0cf951f3c43f2b5ed946f67dd Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Mon, 3 Jun 2024 14:01:20 -0400
Subject: [PATCH] new bandwidth class interface to facilitate operating on
 global and class-based bandwidth limits in the same code, easily (#606)

---
 controller/limits/agent.go        | 10 ------
 controller/limits/limitClasses.go | 57 ++++++++++++++++++++++++-------
 controller/store/limitClass.go    | 43 +++++++++++++++++++++++
 3 files changed, 87 insertions(+), 23 deletions(-)

diff --git a/controller/limits/agent.go b/controller/limits/agent.go
index a7743391..ff0a5a56 100644
--- a/controller/limits/agent.go
+++ b/controller/limits/agent.go
@@ -98,16 +98,6 @@ func (a *Agent) CanCreateShare(acctId, envId int, reserved, uniqueName bool, _ s
 			return false, err
 		}
 
-		alc, err := a.str.FindLimitClassesForAccount(acctId, trx)
-		if err != nil {
-			logrus.Errorf("error finding limit classes for account with id '%d': %v", acctId, err)
-			return false, err
-		}
-		sortLimitClasses(alc)
-		if len(alc) > 0 {
-			logrus.Infof("selected limit class: %v", alc[0])
-		}
-
 		if a.cfg.Shares > Unlimited || (reserved && a.cfg.ReservedShares > Unlimited) || (reserved && uniqueName && a.cfg.UniqueNames > Unlimited) {
 			envs, err := a.str.FindEnvironmentsForAccount(acctId, trx)
 			if err != nil {
diff --git a/controller/limits/limitClasses.go b/controller/limits/limitClasses.go
index 95e0b15c..24f319dc 100644
--- a/controller/limits/limitClasses.go
+++ b/controller/limits/limitClasses.go
@@ -2,22 +2,53 @@ package limits
 
 import (
 	"github.com/openziti/zrok/controller/store"
-	"sort"
+	"github.com/openziti/zrok/sdk/golang/sdk"
 )
 
-func sortLimitClasses(lcs []*store.LimitClass) {
-	sort.Slice(lcs, func(i, j int) bool {
-		return modePoints(lcs[i]) > modePoints(lcs[j])
-	})
+type configBandwidthClass struct {
+	periodInMinutes int
+	bw              *Bandwidth
+	limitAction     store.LimitAction
 }
 
-func modePoints(lc *store.LimitClass) int {
-	points := 0
-	if lc.BackendMode != "" {
-		points += 1
+func newConfigBandwidthClasses(cfgClass *BandwidthPerPeriod) []store.BandwidthClass {
+	return []store.BandwidthClass{
+		&configBandwidthClass{
+			periodInMinutes: int(cfgClass.Period.Minutes()),
+			bw:              cfgClass.Warning,
+			limitAction:     store.WarningLimitAction,
+		},
 	}
-	if lc.ShareMode != "" {
-		points += 1
-	}
-	return points
+}
+
+func (bc *configBandwidthClass) IsGlobal() bool {
+	return true
+}
+
+func (bc *configBandwidthClass) GetShareMode() sdk.ShareMode {
+	return ""
+}
+
+func (bc *configBandwidthClass) GetBackendMode() sdk.BackendMode {
+	return ""
+}
+
+func (bc *configBandwidthClass) GetPeriodMinutes() int {
+	return bc.periodInMinutes
+}
+
+func (bc *configBandwidthClass) GetRxBytes() int64 {
+	return bc.bw.Rx
+}
+
+func (bc *configBandwidthClass) GetTxBytes() int64 {
+	return bc.bw.Tx
+}
+
+func (bc *configBandwidthClass) GetTotalBytes() int64 {
+	return bc.bw.Total
+}
+
+func (bc *configBandwidthClass) GetLimitAction() store.LimitAction {
+	return bc.limitAction
 }
diff --git a/controller/store/limitClass.go b/controller/store/limitClass.go
index 0290010d..9f9ae05e 100644
--- a/controller/store/limitClass.go
+++ b/controller/store/limitClass.go
@@ -7,6 +7,17 @@ import (
 	"github.com/pkg/errors"
 )
 
+type BandwidthClass interface {
+	IsGlobal() bool
+	GetShareMode() sdk.ShareMode
+	GetBackendMode() sdk.BackendMode
+	GetPeriodMinutes() int
+	GetRxBytes() int64
+	GetTxBytes() int64
+	GetTotalBytes() int64
+	GetLimitAction() LimitAction
+}
+
 type LimitClass struct {
 	Model
 	ShareMode      sdk.ShareMode
@@ -22,6 +33,38 @@ type LimitClass struct {
 	LimitAction    LimitAction
 }
 
+func (lc *LimitClass) IsGlobal() bool {
+	return false
+}
+
+func (lc *LimitClass) GetShareMode() sdk.ShareMode {
+	return lc.ShareMode
+}
+
+func (lc *LimitClass) GetBackendMode() sdk.BackendMode {
+	return lc.BackendMode
+}
+
+func (lc *LimitClass) GetPeriodMinutes() int {
+	return lc.PeriodMinutes
+}
+
+func (lc *LimitClass) GetRxBytes() int64 {
+	return lc.RxBytes
+}
+
+func (lc *LimitClass) GetTxBytes() int64 {
+	return lc.TxBytes
+}
+
+func (lc *LimitClass) GetTotalBytes() int64 {
+	return lc.TotalBytes
+}
+
+func (lc *LimitClass) GetLimitAction() LimitAction {
+	return lc.LimitAction
+}
+
 func (lc LimitClass) String() string {
 	out, err := json.MarshalIndent(&lc, "", "  ")
 	if err != nil {