From ea5670b4aefd92c9d27f818b3c824e20c184e2b4 Mon Sep 17 00:00:00 2001 From: Michael Quigley Date: Thu, 30 May 2024 14:27:39 -0400 Subject: [PATCH] new bandwidth_limit_journal table (#606) --- controller/limits/agent.go | 32 +------- controller/store/accountLimitJournal_test.go | 79 ------------------- .../025_v0_4_31_bandwidth_limit_journal.sql | 19 +++++ .../025_v0_4_31_bandwidth_limit_journal.sql | 16 ++++ 4 files changed, 37 insertions(+), 109 deletions(-) delete mode 100644 controller/store/accountLimitJournal_test.go create mode 100644 controller/store/sql/postgresql/025_v0_4_31_bandwidth_limit_journal.sql create mode 100644 controller/store/sql/sqlite3/025_v0_4_31_bandwidth_limit_journal.sql diff --git a/controller/limits/agent.go b/controller/limits/agent.go index 6fddc103..123da3bb 100644 --- a/controller/limits/agent.go +++ b/controller/limits/agent.go @@ -1,14 +1,12 @@ package limits import ( - "fmt" "github.com/jmoiron/sqlx" "github.com/openziti/zrok/controller/emailUi" "github.com/openziti/zrok/controller/metrics" "github.com/openziti/zrok/controller/store" "github.com/openziti/zrok/controller/zrokEdgeSdk" "github.com/openziti/zrok/sdk/golang/sdk" - "github.com/openziti/zrok/util" "github.com/pkg/errors" "github.com/sirupsen/logrus" "reflect" @@ -83,7 +81,7 @@ func (a *Agent) CanCreateEnvironment(acctId int, trx *sqlx.Tx) (bool, error) { return true, nil } -func (a *Agent) CanCreateShare(acctId, envId int, reserved, uniqueName bool, shareMode sdk.ShareMode, backendMode sdk.BackendMode, trx *sqlx.Tx) (bool, error) { +func (a *Agent) CanCreateShare(acctId, envId int, reserved, uniqueName bool, _ sdk.ShareMode, _ sdk.BackendMode, trx *sqlx.Tx) (bool, error) { if a.cfg.Enforcing { if err := a.str.LimitCheckLock(acctId, trx); err != nil { return false, err @@ -416,7 +414,7 @@ func (a *Agent) relax() error { func (a *Agent) checkAccountLimit(acctId int64) (enforce, warning bool, rxBytes, txBytes int64, err error) { period := 24 * time.Hour limit := DefaultBandwidthPerPeriod() - if a.cfg.Bandwidth != nil && a.cfg.Bandwidth != nil { + if a.cfg.Bandwidth != nil { limit = a.cfg.Bandwidth } if limit.Period > 0 { @@ -454,29 +452,3 @@ func (a *Agent) checkLimit(cfg *BandwidthPerPeriod, rx, tx int64) (enforce, warn return false, false } - -func describeLimit(cfg *BandwidthPerPeriod, rx, tx int64) string { - out := "" - - if cfg.Limit.Rx != Unlimited && rx > cfg.Limit.Rx { - out += fmt.Sprintf("['%v' over rx limit '%v']", util.BytesToSize(rx), util.BytesToSize(cfg.Limit.Rx)) - } - if cfg.Limit.Tx != Unlimited && tx > cfg.Limit.Tx { - out += fmt.Sprintf("['%v' over tx limit '%v']", util.BytesToSize(tx), util.BytesToSize(cfg.Limit.Tx)) - } - if cfg.Limit.Total != Unlimited && rx+tx > cfg.Limit.Total { - out += fmt.Sprintf("['%v' over total limit '%v']", util.BytesToSize(rx+tx), util.BytesToSize(cfg.Limit.Total)) - } - - if cfg.Warning.Rx != Unlimited && rx > cfg.Warning.Rx { - out += fmt.Sprintf("['%v' over rx warning '%v']", util.BytesToSize(rx), util.BytesToSize(cfg.Warning.Rx)) - } - if cfg.Warning.Tx != Unlimited && tx > cfg.Warning.Tx { - out += fmt.Sprintf("['%v' over tx warning '%v']", util.BytesToSize(tx), util.BytesToSize(cfg.Warning.Tx)) - } - if cfg.Warning.Total != Unlimited && rx+tx > cfg.Warning.Total { - out += fmt.Sprintf("['%v' over total warning '%v']", util.BytesToSize(rx+tx), util.BytesToSize(cfg.Warning.Total)) - } - - return out -} diff --git a/controller/store/accountLimitJournal_test.go b/controller/store/accountLimitJournal_test.go deleted file mode 100644 index add540e2..00000000 --- a/controller/store/accountLimitJournal_test.go +++ /dev/null @@ -1,79 +0,0 @@ -package store - -import ( - "github.com/stretchr/testify/assert" - "testing" -) - -func TestAccountLimitJournal(t *testing.T) { - str, err := Open(&Config{Path: ":memory:", Type: "sqlite3"}) - assert.Nil(t, err) - assert.NotNil(t, str) - - trx, err := str.Begin() - assert.Nil(t, err) - assert.NotNil(t, trx) - - aljEmpty, err := str.IsAccountLimitJournalEmpty(1, trx) - assert.Nil(t, err) - assert.True(t, aljEmpty) - - acctId, err := str.CreateAccount(&Account{Email: "nobody@nowehere.com", Salt: "salt", Password: "password", Token: "token", Limitless: false, Deleted: false}, trx) - assert.Nil(t, err) - - _, err = str.CreateAccountLimitJournal(&AccountLimitJournal{AccountId: acctId, RxBytes: 1024, TxBytes: 2048, Action: WarningLimitAction}, trx) - assert.Nil(t, err) - - aljEmpty, err = str.IsAccountLimitJournalEmpty(acctId, trx) - assert.Nil(t, err) - assert.False(t, aljEmpty) - - latestAlj, err := str.FindLatestAccountLimitJournal(acctId, trx) - assert.Nil(t, err) - assert.NotNil(t, latestAlj) - assert.Equal(t, int64(1024), latestAlj.RxBytes) - assert.Equal(t, int64(2048), latestAlj.TxBytes) - assert.Equal(t, WarningLimitAction, latestAlj.Action) - - _, err = str.CreateAccountLimitJournal(&AccountLimitJournal{AccountId: acctId, RxBytes: 2048, TxBytes: 4096, Action: LimitLimitAction}, trx) - assert.Nil(t, err) - - latestAlj, err = str.FindLatestAccountLimitJournal(acctId, trx) - assert.Nil(t, err) - assert.NotNil(t, latestAlj) - assert.Equal(t, int64(2048), latestAlj.RxBytes) - assert.Equal(t, int64(4096), latestAlj.TxBytes) - assert.Equal(t, LimitLimitAction, latestAlj.Action) -} - -func TestFindAllLatestAccountLimitJournal(t *testing.T) { - str, err := Open(&Config{Path: ":memory:", Type: "sqlite3"}) - assert.Nil(t, err) - assert.NotNil(t, str) - - trx, err := str.Begin() - assert.Nil(t, err) - assert.NotNil(t, trx) - - acctId1, err := str.CreateAccount(&Account{Email: "nobody@nowehere.com", Salt: "salt1", Password: "password1", Token: "token1", Limitless: false, Deleted: false}, trx) - assert.Nil(t, err) - - _, err = str.CreateAccountLimitJournal(&AccountLimitJournal{AccountId: acctId1, RxBytes: 2048, TxBytes: 4096, Action: WarningLimitAction}, trx) - assert.Nil(t, err) - _, err = str.CreateAccountLimitJournal(&AccountLimitJournal{AccountId: acctId1, RxBytes: 2048, TxBytes: 4096, Action: ClearLimitAction}, trx) - assert.Nil(t, err) - aljId13, err := str.CreateAccountLimitJournal(&AccountLimitJournal{AccountId: acctId1, RxBytes: 2048, TxBytes: 4096, Action: LimitLimitAction}, trx) - assert.Nil(t, err) - - acctId2, err := str.CreateAccount(&Account{Email: "someone@somewhere.com", Salt: "salt2", Password: "password2", Token: "token2", Limitless: false, Deleted: false}, trx) - assert.Nil(t, err) - - aljId21, err := str.CreateAccountLimitJournal(&AccountLimitJournal{AccountId: acctId2, RxBytes: 2048, TxBytes: 4096, Action: WarningLimitAction}, trx) - assert.Nil(t, err) - - aljs, err := str.FindAllLatestAccountLimitJournal(trx) - assert.Nil(t, err) - assert.Equal(t, 2, len(aljs)) - assert.Equal(t, aljId13, aljs[0].Id) - assert.Equal(t, aljId21, aljs[1].Id) -} diff --git a/controller/store/sql/postgresql/025_v0_4_31_bandwidth_limit_journal.sql b/controller/store/sql/postgresql/025_v0_4_31_bandwidth_limit_journal.sql new file mode 100644 index 00000000..c234d9a4 --- /dev/null +++ b/controller/store/sql/postgresql/025_v0_4_31_bandwidth_limit_journal.sql @@ -0,0 +1,19 @@ +-- +migrate Up + +drop table account_limit_journal; +drop table environment_limit_journal; +drop table share_limit_journal; + +drop type limit_action_type; +create type limit_action_type as enum ('warning', 'limit'); + +create table bandwidth_limit_journal ( + id serial primary key, + account_id integer references accounts (id) not null, + limit_class integer references limit_classes, + action limit_action_type not null, + rx_bytes bigint not null, + tx_bytes bigint not null, + created_at timestamptz not null default(current_timestamp), + updated_at timestamptz not null default(current_timestamp) +); \ No newline at end of file diff --git a/controller/store/sql/sqlite3/025_v0_4_31_bandwidth_limit_journal.sql b/controller/store/sql/sqlite3/025_v0_4_31_bandwidth_limit_journal.sql new file mode 100644 index 00000000..6d9bd0fe --- /dev/null +++ b/controller/store/sql/sqlite3/025_v0_4_31_bandwidth_limit_journal.sql @@ -0,0 +1,16 @@ +-- +migrate Up + +drop table account_limit_journal; +drop table environment_limit_journal; +drop table share_limit_journal; + +create table bandwidth_limit_journal ( + id serial primary key, + account_id integer references accounts (id) not null, + limit_class integer references limit_classes, + action string not null, + rx_bytes bigint not null, + tx_bytes bigint not null, + created_at timestamptz not null default(current_timestamp), + updated_at timestamptz not null default(current_timestamp) +); \ No newline at end of file