From 7dfe19c1cb233ed50bd3a49bb6d79c9285283c6c Mon Sep 17 00:00:00 2001 From: Alex Guerrieri Date: Fri, 13 Dec 2024 12:18:45 +0800 Subject: [PATCH] One value for usage --- client.go | 21 ++- go.work.sum | 9 +- handler.go | 87 ++++++--- handler_test.go | 38 ++-- internal/usage/usage.go | 34 ++-- proto/proto.go | 24 +-- proto/proto_test.go | 152 +++++++-------- proto/quotacontrol.gen.go | 388 ++++++++++++++++++++++++++++++++++---- proto/quotacontrol.gen.ts | 115 ++++++++++- proto/quotacontrol.ridl | 21 ++- tests/mock/mem.go | 32 ++-- 11 files changed, 687 insertions(+), 234 deletions(-) diff --git a/client.go b/client.go index 6661bcd..f235f98 100644 --- a/client.go +++ b/client.go @@ -255,7 +255,6 @@ func (c *Client) SpendQuota(ctx context.Context, quota *proto.AccessQuota, cost if err != nil { // limit exceeded if errors.Is(err, proto.ErrQuotaExceeded) { - c.usage.AddKeyUsage(accessKey, now, proto.AccessUsage{LimitedCompute: cost}) return false, total, proto.ErrQuotaExceeded } // ping the server to prepare usage @@ -281,20 +280,24 @@ func (c *Client) SpendQuota(ctx context.Context, quota *proto.AccessQuota, cost } - usage, event := cfg.GetSpendResult(cost, total) - if quota.AccessKey.AccessKey == "" { - c.usage.AddProjectUsage(quota.AccessKey.ProjectID, now, usage) - } else { - c.usage.AddKeyUsage(accessKey, now, usage) - } - if usage.LimitedCompute != 0 { - return false, total, proto.ErrQuotaExceeded + ok, usage, event := cfg.GetSpendResult(cost, total) + if usage != 0 { + if quota.AccessKey.AccessKey == "" { + c.usage.AddProjectUsage(quota.AccessKey.ProjectID, now, usage) + } else { + c.usage.AddKeyUsage(accessKey, now, usage) + } } + if event != nil { if _, err := c.quotaClient.NotifyEvent(ctx, quota.AccessKey.ProjectID, *event); err != nil { logger.Error("notify event failed", slog.Any("error", err)) } } + + if !ok { + return false, total, proto.ErrQuotaExceeded + } return true, total, nil } logger.Error("operation timed out") diff --git a/go.work.sum b/go.work.sum index 478acbb..618c54f 100644 --- a/go.work.sum +++ b/go.work.sum @@ -83,11 +83,8 @@ github.com/supranational/blst v0.3.12/go.mod h1:jZJtfjgudtNl4en1tzwPIV3KjUnQUvG3 github.com/supranational/blst v0.3.13/go.mod h1:jZJtfjgudtNl4en1tzwPIV3KjUnQUvG3/j+w+fVonLw= github.com/tyler-smith/go-bip39 v1.1.0/go.mod h1:gUYDtqQw1JS3ZJ8UWVcGTGqqr6YIN3CWg+kkNaLt55U= github.com/webrpc/gen-golang v0.15.0/go.mod h1:qy1qEWMlTvrRzjSuQLy+176RqNaX1ymUULDtlo7Dapo= -github.com/webrpc/gen-golang v0.17.0/go.mod h1:qy1qEWMlTvrRzjSuQLy+176RqNaX1ymUULDtlo7Dapo= github.com/webrpc/gen-typescript v0.14.1/go.mod h1:xQzYnVaSMfcygDXA5SuW8eYyCLHBHkj15wCF7gcJF5Y= -github.com/webrpc/gen-typescript v0.16.1/go.mod h1:xQzYnVaSMfcygDXA5SuW8eYyCLHBHkj15wCF7gcJF5Y= github.com/webrpc/webrpc v0.20.3/go.mod h1:nplt2iyIXMsMCGyQx1NgEXwQOifWqKNLXcFnQnG9dBU= -github.com/webrpc/webrpc v0.22.0/go.mod h1:eeABnLz9BC4F9GGw6UKebVPkzkFYLrZRlcOvh6o8n10= go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo= go.opentelemetry.io/contrib/detectors/gcp v1.29.0/go.mod h1:GW2aWZNwR2ZxDLdv8OyC2G8zkRoQBuURgV7RPQgcPoU= go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0/go.mod h1:B9yO6b04uB80CzjedvewuqDhxJxi11s7/GtiGa8bAjI= @@ -99,6 +96,7 @@ go.opentelemetry.io/otel/sdk/metric v1.29.0/go.mod h1:6zZLdCl2fkauYoZIOn/soQIDSW go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8= golang.org/x/crypto v0.26.0/go.mod h1:GY7jblb9wI+FOo5y8/S2oY4zWP07AkOJ4+jxCqdqn54= +golang.org/x/crypto v0.27.0/go.mod h1:1Xngt8kV6Dvbssa53Ziq6Eqn0HqbZi5Z6R0ZpwQzt70= golang.org/x/exp v0.0.0-20240222234643-814bf88cf225/go.mod h1:CxmFvTBINI24O/j8iY7H1xHzx2i4OsyguNBmN/uPtqc= golang.org/x/exp v0.0.0-20240808152545-0cdaa3abc0fa/go.mod h1:akd2r19cwCdwSwWeIdzYQGa/EZZyqcOdwWiwj5L5eKQ= golang.org/x/exp v0.0.0-20240909161429-701f63a606c0/go.mod h1:2TbTHSBQa924w8M6Xs1QcRcFwyucIwBGpK1p2f1YFFY= @@ -106,18 +104,23 @@ golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac= golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM= golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg= +golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ= +golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.23.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE= golang.org/x/telemetry v0.0.0-20240521205824-bda55230c457/go.mod h1:pRgIJT+bRLFKnoM1ldnzKoxTIn14Yxz928LQRYYgIN0= golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY= golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0= golang.org/x/term v0.23.0/go.mod h1:DgV24QBUrK6jhZXl+20l6UWznPlwAHm1Q1mGHtydmSk= golang.org/x/term v0.24.0/go.mod h1:lOBK/LVxemqiMij05LGJ0tzNr8xlmwBRJ81PX6wVLH8= +golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM= golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI= golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= golang.org/x/text v0.18.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= +golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= golang.org/x/time v0.6.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.11.0/go.mod h1:anzJrxPjNtfgiYQYirP2CPGzGLxrH2u2QBhn6Bf3qY8= golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d h1:vU5i/LfpvrRCpgM/VPfJLg5KjxD3E+hfT1SH+d9zLwg= diff --git a/handler.go b/handler.go index 3d9a7e6..75c3432 100644 --- a/handler.go +++ b/handler.go @@ -27,9 +27,9 @@ type AccessKeyStore interface { } type UsageStore interface { - GetAccessKeyUsage(ctx context.Context, projectID uint64, accessKey string, service *proto.Service, min, max time.Time) (proto.AccessUsage, error) - GetAccountUsage(ctx context.Context, projectID uint64, service *proto.Service, min, max time.Time) (proto.AccessUsage, error) - UpdateAccessUsage(ctx context.Context, projectID uint64, accessKey string, service proto.Service, time time.Time, usage proto.AccessUsage) error + GetUsageAccessKey(ctx context.Context, projectID uint64, accessKey string, service *proto.Service, min, max time.Time) (int64, error) + GetUsageProject(ctx context.Context, projectID uint64, service *proto.Service, min, max time.Time) (int64, error) + UpdateUsage(ctx context.Context, projectID uint64, accessKey string, service proto.Service, time time.Time, usage int64) error } type CycleStore interface { @@ -103,59 +103,86 @@ func (h handler) GetTimeRange(ctx context.Context, projectID uint64, from, to *t return *from, from.Add(duration), nil } -func (h handler) GetAccountUsage(ctx context.Context, projectID uint64, service *proto.Service, from, to *time.Time) (*proto.AccessUsage, error) { +func (h handler) GetUsageProject(ctx context.Context, projectID uint64, service *proto.Service, from, to *time.Time) (int64, error) { min, max, err := h.GetTimeRange(ctx, projectID, from, to) if err != nil { - return nil, err + return 0, fmt.Errorf("get time range: %w", err) + } + + usage, err := h.store.UsageStore.GetUsageProject(ctx, projectID, service, min, max) + if err != nil { + return 0, fmt.Errorf("get project usage: %w", err) } + return usage, nil +} - usage, err := h.store.UsageStore.GetAccountUsage(ctx, projectID, service, min, max) +// deprecated: use GetUsageProject instead +func (h handler) GetAccountUsage(ctx context.Context, projectID uint64, service *proto.Service, from, to *time.Time) (*proto.AccessUsage, error) { + usage, err := h.GetUsageProject(ctx, projectID, service, from, to) if err != nil { return nil, err } - return &usage, nil + return &proto.AccessUsage{ValidCompute: usage}, nil } -func (h handler) GetAsyncUsage(ctx context.Context, projectID uint64, service *proto.Service, from, to *time.Time) (*proto.AccessUsage, error) { +func (h handler) GetUsageAsync(ctx context.Context, projectID uint64, service *proto.Service, from, to *time.Time) (int64, error) { min, max, err := h.GetTimeRange(ctx, projectID, from, to) if err != nil { - return nil, err + return 0, fmt.Errorf("get time range: %w", err) } - usage, err := h.store.UsageStore.GetAccessKeyUsage(ctx, projectID, "", service, min, max) + usage, err := h.store.UsageStore.GetUsageAccessKey(ctx, projectID, "", service, min, max) + if err != nil { + return 0, fmt.Errorf("get async usage: %w", err) + } + return usage, nil +} + +// deprecated: use GetUsageAsync instead +func (h handler) GetAsyncUsage(ctx context.Context, projectID uint64, service *proto.Service, from, to *time.Time) (*proto.AccessUsage, error) { + usage, err := h.GetUsageAsync(ctx, projectID, service, from, to) if err != nil { return nil, err } - return &usage, nil + return &proto.AccessUsage{ValidCompute: usage}, nil } -func (h handler) GetAccessKeyUsage(ctx context.Context, accessKey string, service *proto.Service, from, to *time.Time) (*proto.AccessUsage, error) { +func (h handler) GetUsageAccessKey(ctx context.Context, accessKey string, service *proto.Service, from, to *time.Time) (int64, error) { projectID, err := proto.GetProjectID(accessKey) if err != nil { - return nil, err + return 0, fmt.Errorf("get project id: %w", err) } min, max, err := h.GetTimeRange(ctx, projectID, from, to) if err != nil { - return nil, err + return 0, fmt.Errorf("get time range: %w", err) + } + + usage, err := h.store.UsageStore.GetUsageAccessKey(ctx, projectID, accessKey, service, min, max) + if err != nil { + return 0, fmt.Errorf("get access key usage: %w", err) } + return usage, nil +} - usage, err := h.store.UsageStore.GetAccessKeyUsage(ctx, projectID, accessKey, service, min, max) +// deprecated: use GetUsageAccessKey instead +func (h handler) GetAccessKeyUsage(ctx context.Context, accessKey string, service *proto.Service, from, to *time.Time) (*proto.AccessUsage, error) { + usage, err := h.GetUsageAccessKey(ctx, accessKey, service, from, to) if err != nil { return nil, err } - return &usage, nil + return &proto.AccessUsage{ValidCompute: usage}, nil } func (h handler) PrepareUsage(ctx context.Context, projectID uint64, cycle *proto.Cycle, now time.Time) (bool, error) { min, max := cycle.GetStart(now), cycle.GetEnd(now) - usage, err := h.GetAccountUsage(ctx, projectID, nil, &min, &max) + usage, err := h.store.GetUsageProject(ctx, projectID, nil, min, max) if err != nil { return false, err } key := getQuotaKey(projectID, cycle, now) - if err := h.cache.UsageCache.SetUsage(ctx, key, usage.GetTotalUsage()); err != nil { + if err := h.cache.UsageCache.SetUsage(ctx, key, usage); err != nil { return false, err } return true, nil @@ -230,11 +257,20 @@ func (h handler) NotifyEvent(ctx context.Context, projectID uint64, eventType pr return true, nil } +// deprecated: use StoreUsageProject instead func (h handler) UpdateProjectUsage(ctx context.Context, service proto.Service, now time.Time, usage map[uint64]*proto.AccessUsage) (map[uint64]bool, error) { + input := make(map[uint64]int64, len(usage)) + for projectID, accessUsage := range usage { + input[projectID] = accessUsage.GetTotalUsage() + } + return h.StoreUsageProject(ctx, service, now, input) +} + +func (h handler) StoreUsageProject(ctx context.Context, service proto.Service, now time.Time, usage map[uint64]int64) (map[uint64]bool, error) { var errs []error m := make(map[uint64]bool, len(usage)) - for projectID, accessUsage := range usage { - err := h.store.UsageStore.UpdateAccessUsage(ctx, projectID, "", service, now, *accessUsage) + for projectID, u := range usage { + err := h.store.UsageStore.UpdateUsage(ctx, projectID, "", service, now, u) if err != nil { errs = append(errs, fmt.Errorf("%d: %w", projectID, err)) } @@ -246,7 +282,16 @@ func (h handler) UpdateProjectUsage(ctx context.Context, service proto.Service, return m, nil } +// deprecated: use UpdateKeyUsage instead func (h handler) UpdateKeyUsage(ctx context.Context, service proto.Service, now time.Time, usage map[string]*proto.AccessUsage) (map[string]bool, error) { + input := make(map[string]int64, len(usage)) + for key, accessUsage := range usage { + input[key] = accessUsage.GetTotalUsage() + } + return h.StoreUsageAccessKey(ctx, service, now, input) +} + +func (h handler) StoreUsageAccessKey(ctx context.Context, service proto.Service, now time.Time, usage map[string]int64) (map[string]bool, error) { var errs []error m := make(map[string]bool, len(usage)) for key, u := range usage { @@ -255,7 +300,7 @@ func (h handler) UpdateKeyUsage(ctx context.Context, service proto.Service, now errs = append(errs, fmt.Errorf("%s: %w", key, err)) continue } - if err = h.store.UsageStore.UpdateAccessUsage(ctx, projectID, key, service, now, *u); err != nil { + if err = h.store.UsageStore.UpdateUsage(ctx, projectID, key, service, now, u); err != nil { errs = append(errs, fmt.Errorf("%s: %w", key, err)) } m[key] = err == nil diff --git a/handler_test.go b/handler_test.go index d89f801..5619851 100644 --- a/handler_test.go +++ b/handler_test.go @@ -89,7 +89,7 @@ func TestMiddlewareUseAccessKey(t *testing.T) { r.Handle("/*", &counter) - expectedUsage := proto.AccessUsage{} + expectedUsage := int64(0) t.Run("WithAccessKey", func(t *testing.T) { go client.Run(context.Background()) @@ -106,7 +106,7 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.Equal(t, strconv.FormatInt(limit.FreeMax-i, 10), headers.Get(middleware.HeaderQuotaRemaining)) assert.Equal(t, "", headers.Get(middleware.HeaderQuotaOverage)) assert.Empty(t, server.GetEvents(ProjectID), i) - expectedUsage.Add(proto.AccessUsage{ValidCompute: _credits}) + expectedUsage += _credits } // Go over free CU @@ -117,7 +117,7 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.Equal(t, "0", headers.Get(middleware.HeaderQuotaRemaining)) assert.Equal(t, "", headers.Get(middleware.HeaderQuotaOverage)) assert.Contains(t, server.GetEvents(ProjectID), proto.EventType_FreeMax) - expectedUsage.Add(proto.AccessUsage{ValidCompute: _credits}) + expectedUsage += _credits // Get close to soft quota for i := limit.FreeWarn + _credits; i < limit.OverWarn; i += _credits { @@ -128,7 +128,7 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.Equal(t, "0", headers.Get(middleware.HeaderQuotaRemaining)) assert.Equal(t, strconv.FormatInt(i-limit.FreeWarn, 10), headers.Get(middleware.HeaderQuotaOverage)) assert.Len(t, server.GetEvents(ProjectID), 1) - expectedUsage.Add(proto.AccessUsage{OverCompute: _credits}) + expectedUsage += _credits } // Go over soft quota @@ -139,7 +139,7 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.Equal(t, "0", headers.Get(middleware.HeaderQuotaRemaining)) assert.Equal(t, strconv.FormatInt(limit.OverWarn-limit.FreeWarn, 10), headers.Get(middleware.HeaderQuotaOverage)) assert.Contains(t, server.GetEvents(ProjectID), proto.EventType_OverWarn) - expectedUsage.Add(proto.AccessUsage{OverCompute: _credits}) + expectedUsage += _credits // Get close to hard quota for i := limit.OverWarn + _credits; i < limit.OverMax; i += _credits { @@ -150,7 +150,7 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.Equal(t, "0", headers.Get(middleware.HeaderQuotaRemaining)) assert.Equal(t, strconv.FormatInt(i-limit.FreeWarn, 10), headers.Get(middleware.HeaderQuotaOverage)) assert.Len(t, server.GetEvents(ProjectID), 2) - expectedUsage.Add(proto.AccessUsage{OverCompute: _credits}) + expectedUsage += _credits } // Go over hard quota @@ -161,7 +161,7 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.Equal(t, "0", headers.Get(middleware.HeaderQuotaRemaining)) assert.Equal(t, strconv.FormatInt(limit.OverMax-limit.FreeWarn, 10), headers.Get(middleware.HeaderQuotaOverage)) assert.Contains(t, server.GetEvents(ProjectID), proto.EventType_OverMax) - expectedUsage.Add(proto.AccessUsage{OverCompute: _credits}) + expectedUsage += _credits // Denied for i := 0; i < 10; i++ { @@ -171,15 +171,13 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.Equal(t, strconv.FormatInt(limit.FreeMax, 10), headers.Get(middleware.HeaderQuotaLimit)) assert.Equal(t, "0", headers.Get(middleware.HeaderQuotaRemaining)) assert.Equal(t, strconv.FormatInt(limit.OverMax-limit.FreeWarn, 10), headers.Get(middleware.HeaderQuotaOverage)) - expectedUsage.Add(proto.AccessUsage{LimitedCompute: _credits}) } // check the usage client.Stop(context.Background()) - usage, err := server.Store.GetAccountUsage(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) + usage, err := server.Store.GetUsageProject(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) assert.NoError(t, err) - assert.Equal(t, int64(expectedUsage.GetTotalUsage()), _credits*counter.GetValue()) - assert.Equal(t, &expectedUsage, &usage) + assert.Equal(t, expectedUsage, usage) }) t.Run("ChangeLimits", func(t *testing.T) { @@ -204,10 +202,10 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.Equal(t, "0", headers.Get(middleware.HeaderQuotaLimit)) client.Stop(context.Background()) - usage, err := server.Store.GetAccountUsage(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) + usage, err := server.Store.GetUsageProject(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) assert.NoError(t, err) - expectedUsage.Add(proto.AccessUsage{ValidCompute: 0, OverCompute: _credits, LimitedCompute: 0}) - assert.Equal(t, int64(expectedUsage.GetTotalUsage()), _credits*counter.GetValue()) + expectedUsage += _credits + assert.Equal(t, expectedUsage, _credits*counter.GetValue()) assert.Equal(t, &expectedUsage, &usage) }) @@ -230,9 +228,9 @@ func TestMiddlewareUseAccessKey(t *testing.T) { } client.Stop(context.Background()) - usage, err := server.Store.GetAccountUsage(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) + usage, err := server.Store.GetUsageProject(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) assert.NoError(t, err) - assert.Equal(t, int64(expectedUsage.GetTotalUsage()), _credits*counter.GetValue()) + assert.Equal(t, expectedUsage, _credits*counter.GetValue()) assert.Equal(t, &expectedUsage, &usage) }) @@ -270,9 +268,9 @@ func TestMiddlewareUseAccessKey(t *testing.T) { server.ErrPrepareUsage = nil client.Stop(context.Background()) - usage, err := server.Store.GetAccountUsage(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) + usage, err := server.Store.GetUsageProject(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) assert.NoError(t, err) - assert.Equal(t, int64(expectedUsage.GetTotalUsage()), _credits*counter.GetValue()) + assert.Equal(t, expectedUsage, _credits*counter.GetValue()) assert.Equal(t, &expectedUsage, &usage) }) @@ -290,9 +288,9 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.NoError(t, err) client.Stop(context.Background()) - usage, err := server.Store.GetAccountUsage(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) + usage, err := server.Store.GetUsageProject(ctx, ProjectID, &Service, now.Add(-time.Hour), now.Add(time.Hour)) assert.NoError(t, err) - assert.Equal(t, int64(expectedUsage.GetTotalUsage()), _credits*counter.GetValue()) + assert.Equal(t, expectedUsage, _credits*counter.GetValue()) assert.Equal(t, &expectedUsage, &usage) }) } diff --git a/internal/usage/usage.go b/internal/usage/usage.go index 908daec..27792df 100644 --- a/internal/usage/usage.go +++ b/internal/usage/usage.go @@ -11,20 +11,20 @@ import ( // UsageUpdater is an interface that allows to update the usage of a service type UsageUpdater interface { - UpdateKeyUsage(ctx context.Context, service proto.Service, now time.Time, usage map[string]*proto.AccessUsage) (map[string]bool, error) - UpdateProjectUsage(ctx context.Context, service proto.Service, now time.Time, usage map[uint64]*proto.AccessUsage) (map[uint64]bool, error) + StoreUsageAccessKey(ctx context.Context, service proto.Service, now time.Time, usage map[string]int64) (map[string]bool, error) + StoreUsageProject(ctx context.Context, service proto.Service, now time.Time, usage map[uint64]int64) (map[uint64]bool, error) } func NewRecord() Record { return Record{ - ByProjectID: make(map[uint64]*proto.AccessUsage), - ByAccessKey: make(map[string]*proto.AccessUsage), + ByProjectID: make(map[uint64]int64), + ByAccessKey: make(map[string]int64), } } type Record struct { - ByProjectID map[uint64]*proto.AccessUsage - ByAccessKey map[string]*proto.AccessUsage + ByProjectID map[uint64]int64 + ByAccessKey map[string]int64 } func NewTracker() *Tracker { @@ -44,28 +44,22 @@ type Tracker struct { } // AddUsage adds the usage of a access key. -func (u *Tracker) AddKeyUsage(accessKey string, now time.Time, usage proto.AccessUsage) { +func (u *Tracker) AddKeyUsage(accessKey string, now time.Time, usage int64) { u.dataMutex.Lock() if _, ok := u.usage[now]; !ok { u.usage[now] = NewRecord() } - if _, ok := u.usage[now].ByAccessKey[accessKey]; !ok { - u.usage[now].ByAccessKey[accessKey] = &proto.AccessUsage{} - } - u.usage[now].ByAccessKey[accessKey].Add(usage) + u.usage[now].ByAccessKey[accessKey] += usage u.dataMutex.Unlock() } // AddUsage adds the usage of a access key. -func (u *Tracker) AddProjectUsage(projectID uint64, now time.Time, usage proto.AccessUsage) { +func (u *Tracker) AddProjectUsage(projectID uint64, now time.Time, usage int64) { u.dataMutex.Lock() if _, ok := u.usage[now]; !ok { u.usage[now] = NewRecord() } - if _, ok := u.usage[now].ByProjectID[projectID]; !ok { - u.usage[now].ByProjectID[projectID] = &proto.AccessUsage{} - } - u.usage[now].ByProjectID[projectID].Add(usage) + u.usage[now].ByProjectID[projectID] += usage u.dataMutex.Unlock() } @@ -84,7 +78,7 @@ func (u *Tracker) SyncUsage(ctx context.Context, updater UsageUpdater, service p defer u.syncMutex.Unlock() var errList []error for now, usages := range u.GetUpdates() { - keyResult, err := updater.UpdateKeyUsage(ctx, service, now, usages.ByAccessKey) + keyResult, err := updater.StoreUsageAccessKey(ctx, service, now, usages.ByAccessKey) if err != nil { errList = append(errList, err) } @@ -93,10 +87,10 @@ func (u *Tracker) SyncUsage(ctx context.Context, updater UsageUpdater, service p if v { continue } - u.AddKeyUsage(accessKey, now, *usages.ByAccessKey[accessKey]) + u.AddKeyUsage(accessKey, now, usages.ByAccessKey[accessKey]) } - projectResult, err := updater.UpdateProjectUsage(ctx, service, now, usages.ByProjectID) + projectResult, err := updater.StoreUsageProject(ctx, service, now, usages.ByProjectID) if err != nil { errList = append(errList, err) } @@ -105,7 +99,7 @@ func (u *Tracker) SyncUsage(ctx context.Context, updater UsageUpdater, service p if v { continue } - u.AddProjectUsage(projectID, now, *usages.ByProjectID[projectID]) + u.AddProjectUsage(projectID, now, usages.ByProjectID[projectID]) } } diff --git a/proto/proto.go b/proto/proto.go index 83a32e6..69f98eb 100644 --- a/proto/proto.go +++ b/proto/proto.go @@ -84,36 +84,38 @@ func getOverThreshold(v, total, threshold int64) (int64, bool) { return max(0, total-threshold), true } -func (l *Limit) GetSpendResult(v, total int64) (AccessUsage, *EventType) { +func (l *Limit) GetSpendResult(v, total int64) (bool, int64, *EventType) { // valid usage if total < l.FreeMax { // threshold of included alert if _, ok := getOverThreshold(v, total, l.FreeWarn); ok { - return AccessUsage{ValidCompute: v}, Ptr(EventType_FreeWarn) + return true, v, Ptr(EventType_FreeWarn) } // normal valid usage - return AccessUsage{ValidCompute: v}, nil + return true, v, nil } // overage usage if total < l.OverMax { // threshold of included limit - if over, ok := getOverThreshold(v, total, l.FreeMax); ok { - return AccessUsage{ValidCompute: v - over, OverCompute: over}, Ptr(EventType_FreeMax) + if _, ok := getOverThreshold(v, total, l.FreeMax); ok { + return true, v, Ptr(EventType_FreeMax) } // threshold of overage alert - if _, ok := getOverThreshold(v, total, l.OverWarn); ok { - return AccessUsage{OverCompute: v}, Ptr(EventType_OverWarn) + if l.OverWarn != l.OverMax { + if _, ok := getOverThreshold(v, total, l.OverWarn); ok { + return true, v, Ptr(EventType_OverWarn) + } + // normal overage usage + return true, v, nil } - // normal overage usage - return AccessUsage{OverCompute: v}, nil } // limited usage if over, ok := getOverThreshold(v, total, l.OverMax); ok { - return AccessUsage{LimitedCompute: over, OverCompute: v - over}, Ptr(EventType_OverMax) + return over == 0, v - over, Ptr(EventType_OverMax) } - return AccessUsage{LimitedCompute: v}, nil + return false, 0, nil } func (q *AccessQuota) IsActive() bool { diff --git a/proto/proto_test.go b/proto/proto_test.go index c0741f7..5d3015e 100644 --- a/proto/proto_test.go +++ b/proto/proto_test.go @@ -6,7 +6,6 @@ import ( "github.com/0xsequence/quotacontrol/proto" "github.com/goware/validation" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" ) func TestAccessKeyValidateOrigin(t *testing.T) { @@ -75,121 +74,86 @@ func TestGetSpendResult(t *testing.T) { Name string Total int64 - Usage proto.AccessUsage + Ok bool + Spent int64 Event *proto.EventType } for _, tc := range []TestCase{ // Include Alert { - Name: "Within_IncludedAlert", - Total: limit.FreeWarn - 1, - Usage: proto.AccessUsage{ValidCompute: _CU}, - Event: nil, + Name: "Within_IncludedAlert", Total: limit.FreeWarn - 1, + Ok: true, Spent: _CU, Event: nil, }, { - Name: "Within_IncludedAlert_Exact", - Total: limit.FreeWarn, - Usage: proto.AccessUsage{ValidCompute: _CU}, - Event: proto.Ptr(proto.EventType_FreeWarn), + Name: "Within_IncludedAlert_Exact", Total: limit.FreeWarn, + Ok: true, Spent: _CU, Event: proto.Ptr(proto.EventType_FreeWarn), }, { - Name: "Above_IncludedAlert", - Total: limit.FreeWarn + 1, - Usage: proto.AccessUsage{ValidCompute: _CU}, - Event: proto.Ptr(proto.EventType_FreeWarn), + Name: "Above_IncludedAlert", Total: limit.FreeWarn + 1, + Ok: true, Spent: _CU, Event: proto.Ptr(proto.EventType_FreeWarn), }, { - Name: "Above_IncludedAlert_Exact", - Total: limit.FreeWarn + _CU, - Usage: proto.AccessUsage{ValidCompute: _CU}, - Event: nil, + Name: "Above_IncludedAlert_Exact", Total: limit.FreeWarn + _CU, + Ok: true, Spent: _CU, Event: nil, }, // Include Limit { - Name: "Within_IncludedLimit", - Total: limit.FreeWarn - 1, - Usage: proto.AccessUsage{ValidCompute: _CU}, - Event: nil, + Name: "Within_IncludedLimit", Total: limit.FreeWarn - 1, + Ok: true, Spent: _CU, Event: nil, }, { - Name: "Within_IncludedLimit_Exact", - Total: limit.FreeMax, - Usage: proto.AccessUsage{ValidCompute: _CU}, - Event: proto.Ptr(proto.EventType_FreeMax), + Name: "Within_IncludedLimit_Exact", Total: limit.FreeMax, + Ok: true, Spent: _CU, Event: proto.Ptr(proto.EventType_FreeMax), }, { - Name: "Above_IncludedLimit", - Total: limit.FreeMax + 1, - Usage: proto.AccessUsage{ValidCompute: _CU - 1, OverCompute: 1}, - Event: proto.Ptr(proto.EventType_FreeMax), + Name: "Above_IncludedLimit", Total: limit.FreeMax + 1, + Ok: true, Spent: _CU, Event: proto.Ptr(proto.EventType_FreeMax), }, { - Name: "Above_IncludedLimit_Exact", - Total: limit.FreeMax + _CU, - Usage: proto.AccessUsage{OverCompute: _CU}, - Event: nil, + Name: "Above_IncludedLimit_Exact", Total: limit.FreeMax + _CU, + Ok: true, Spent: _CU, Event: nil, }, // Overage Alert { - Name: "Within_OverageAlert", - Total: limit.OverWarn - 1, - Usage: proto.AccessUsage{OverCompute: _CU}, - Event: nil, + Name: "Within_OverageAlert", Total: limit.OverWarn - 1, + Ok: true, Spent: _CU, Event: nil, }, { - Name: "Within_OverageAlert_Exact", - Total: limit.OverWarn, - Usage: proto.AccessUsage{OverCompute: _CU}, - Event: proto.Ptr(proto.EventType_OverWarn), + Name: "Within_OverageAlert_Exact", Total: limit.OverWarn, + Ok: true, Spent: _CU, Event: proto.Ptr(proto.EventType_OverWarn), }, { - Name: "Above_OverageAlert", - Total: limit.OverWarn + 2, - Usage: proto.AccessUsage{OverCompute: _CU}, - Event: proto.Ptr(proto.EventType_OverWarn), + Name: "Above_OverageAlert", Total: limit.OverWarn + 2, + Ok: true, Spent: _CU, Event: proto.Ptr(proto.EventType_OverWarn), }, { - Name: "Above_OverageAlert_Exact", - Total: limit.OverWarn + _CU, - Usage: proto.AccessUsage{OverCompute: _CU}, - Event: nil, + Name: "Above_OverageAlert_Exact", Total: limit.OverWarn + _CU, + Ok: true, Spent: _CU, Event: nil, }, // Overage Limit { - Name: "Within_OverageLimit", - Total: limit.OverMax - 1, - Usage: proto.AccessUsage{OverCompute: _CU}, - Event: nil, + Name: "Within_OverageLimit", Total: limit.OverMax - 1, + Ok: true, Spent: _CU, Event: nil, }, { - Name: "Above_OverageLimit_Exact", - Total: limit.OverMax, - Usage: proto.AccessUsage{OverCompute: _CU}, - Event: proto.Ptr(proto.EventType_OverMax), + Name: "Above_OverageLimit_Exact", Total: limit.OverMax, + Ok: true, Spent: _CU, Event: proto.Ptr(proto.EventType_OverMax), }, { - Name: "Above_OverageLimit", - Total: limit.OverMax + 2, - Usage: proto.AccessUsage{OverCompute: _CU - 2, LimitedCompute: 2}, - Event: proto.Ptr(proto.EventType_OverMax), + Name: "Above_OverageLimit", Total: limit.OverMax + 2, + Ok: false, Spent: 3, Event: proto.Ptr(proto.EventType_OverMax), }, { - Name: "Above_OverageLimit_More", - Total: limit.OverMax + _CU, - Usage: proto.AccessUsage{LimitedCompute: _CU}, - Event: nil, + Name: "Above_OverageLimit_More", Total: limit.OverMax + _CU, + Ok: false, Spent: 0, Event: nil, }, } { t.Run(tc.Name, func(t *testing.T) { - u, evt := limit.GetSpendResult(_CU, tc.Total) - assert.Equal(t, tc.Usage, u) - if tc.Event == nil { - assert.Nil(t, evt) - return - } - require.NotNil(t, evt) - assert.Equal(t, tc.Event.String(), evt.String()) + ok, u, evt := limit.GetSpendResult(_CU, tc.Total) + assert.Equal(t, tc.Ok, ok) + assert.Equal(t, tc.Spent, u) + assert.Equal(t, tc.Event, evt) }) } @@ -198,24 +162,40 @@ func TestGetSpendResult(t *testing.T) { t.Run("NoFreeWarn", func(t *testing.T) { // it works for freeMax and 0 limit.FreeWarn = limit.FreeMax - u, evt := limit.GetSpendResult(1, limit.FreeMax) - assert.Equal(t, proto.AccessUsage{ValidCompute: 1}, u) - assert.Equal(t, proto.EventType_FreeMax.String(), evt.String()) + ok, u, evt := limit.GetSpendResult(1, limit.FreeMax) + assert.True(t, ok) + assert.Equal(t, int64(1), u) + assert.Equal(t, proto.EventType_FreeMax, *evt) limit.FreeWarn = 0 - u, evt = limit.GetSpendResult(1, limit.FreeMax) - assert.Equal(t, proto.AccessUsage{ValidCompute: 1}, u) - assert.Equal(t, proto.EventType_FreeMax.String(), evt.String()) + ok, u, evt = limit.GetSpendResult(1, limit.FreeMax) + assert.Equal(t, int64(1), u) + assert.Equal(t, proto.EventType_FreeMax, *evt) }) t.Run("NoOverWarn", func(t *testing.T) { // it works for overMax and 0 limit.OverWarn = limit.OverMax - u, evt := limit.GetSpendResult(1, limit.OverMax) - assert.Equal(t, proto.AccessUsage{OverCompute: 1}, u) - assert.Equal(t, proto.EventType_OverMax.String(), evt.String()) + + ok, u, evt := limit.GetSpendResult(1, limit.OverMax) + assert.True(t, ok) + assert.Equal(t, int64(1), u) + assert.Equal(t, proto.EventType_OverMax, *evt) + + ok, u, evt = limit.GetSpendResult(2, limit.OverMax+1) + assert.False(t, ok) + assert.Equal(t, int64(1), u) + assert.Equal(t, proto.EventType_OverMax, *evt) + limit.OverWarn = 0 - u, evt = limit.GetSpendResult(1, limit.OverMax) - assert.Equal(t, proto.AccessUsage{OverCompute: 1}, u) - assert.Equal(t, proto.EventType_OverMax.String(), evt.String()) + + ok, u, evt = limit.GetSpendResult(1, limit.OverMax) + assert.True(t, ok) + assert.Equal(t, int64(1), u) + assert.Equal(t, proto.EventType_OverMax, *evt) + + ok, u, evt = limit.GetSpendResult(2, limit.OverMax+1) + assert.False(t, ok) + assert.Equal(t, int64(1), u) + assert.Equal(t, proto.EventType_OverMax, *evt) }) }) } diff --git a/proto/quotacontrol.gen.go b/proto/quotacontrol.gen.go index 0278887..ddfcddd 100644 --- a/proto/quotacontrol.gen.go +++ b/proto/quotacontrol.gen.go @@ -1,4 +1,4 @@ -// quota-control v0.19.2 4b5f19d012ae4365adc25b547b4d39ac58a974e8 +// quota-control v0.19.2 db40d579f0863c38fafda257ec579b72c8751b69 // -- // Code generated by webrpc-gen@v0.22.0 with golang@v0.17.0 generator. DO NOT EDIT. // @@ -37,7 +37,7 @@ func WebRPCSchemaVersion() string { // Schema hash generated from your RIDL schema func WebRPCSchemaHash() string { - return "4b5f19d012ae4365adc25b547b4d39ac58a974e8" + return "db40d579f0863c38fafda257ec579b72c8751b69" } type WebrpcGenVersions struct { @@ -386,6 +386,21 @@ var methods = map[string]method{ Service: "QuotaControl", Annotations: map[string]string{}, }, + "/rpc/QuotaControl/GetUsageProject": { + Name: "GetUsageProject", + Service: "QuotaControl", + Annotations: map[string]string{}, + }, + "/rpc/QuotaControl/GetUsageAccessKey": { + Name: "GetUsageAccessKey", + Service: "QuotaControl", + Annotations: map[string]string{}, + }, + "/rpc/QuotaControl/GetUsageAsync": { + Name: "GetUsageAsync", + Service: "QuotaControl", + Annotations: map[string]string{}, + }, "/rpc/QuotaControl/PrepareUsage": { Name: "PrepareUsage", Service: "QuotaControl", @@ -411,8 +426,13 @@ var methods = map[string]method{ Service: "QuotaControl", Annotations: map[string]string{}, }, - "/rpc/QuotaControl/UpdateUsage": { - Name: "UpdateUsage", + "/rpc/QuotaControl/StoreUsageProject": { + Name: "StoreUsageProject", + Service: "QuotaControl", + Annotations: map[string]string{}, + }, + "/rpc/QuotaControl/StoreUsageAccessKey": { + Name: "StoreUsageAccessKey", Service: "QuotaControl", Annotations: map[string]string{}, }, @@ -449,12 +469,16 @@ var WebRPCServices = map[string][]string{ "GetAccountUsage", "GetAccessKeyUsage", "GetAsyncUsage", + "GetUsageProject", + "GetUsageAccessKey", + "GetUsageAsync", "PrepareUsage", "ClearUsage", "NotifyEvent", "UpdateProjectUsage", "UpdateKeyUsage", - "UpdateUsage", + "StoreUsageProject", + "StoreUsageAccessKey", "GetUserPermission", }, } @@ -475,20 +499,28 @@ type QuotaControl interface { UpdateDefaultAccessKey(ctx context.Context, projectID uint64, accessKey string) (bool, error) ListAccessKeys(ctx context.Context, projectId uint64, active *bool, service *Service) ([]*AccessKey, error) DisableAccessKey(ctx context.Context, accessKey string) (bool, error) - // Usage + // Quota GetProjectQuota(ctx context.Context, projectId uint64, now time.Time) (*AccessQuota, error) GetAccessQuota(ctx context.Context, accessKey string, now time.Time) (*AccessQuota, error) ClearAccessQuotaCache(ctx context.Context, projectID uint64) (bool, error) + // deprecated: this is now GetUsageProject, and it will be removed in the future GetAccountUsage(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (*AccessUsage, error) + // deprecated: this is now GetUsageAccessKey, and it will be removed in the future GetAccessKeyUsage(ctx context.Context, accessKey string, service *Service, from *time.Time, to *time.Time) (*AccessUsage, error) + // deprecated: this is now GetUsageAsync, and it will be removed in the future GetAsyncUsage(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (*AccessUsage, error) + GetUsageProject(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (int64, error) + GetUsageAccessKey(ctx context.Context, accessKey string, service *Service, from *time.Time, to *time.Time) (int64, error) + GetUsageAsync(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (int64, error) PrepareUsage(ctx context.Context, projectID uint64, cycle *Cycle, now time.Time) (bool, error) ClearUsage(ctx context.Context, projectID uint64, now time.Time) (bool, error) NotifyEvent(ctx context.Context, projectID uint64, eventType EventType) (bool, error) + // deprecated: this is now StoreUsageProject, and it will be removed in the future UpdateProjectUsage(ctx context.Context, service Service, now time.Time, usage map[uint64]*AccessUsage) (map[uint64]bool, error) + // deprecated: this is now StoreUsageAccessKey, and it will be removed in the future UpdateKeyUsage(ctx context.Context, service Service, now time.Time, usage map[string]*AccessUsage) (map[string]bool, error) - // DEPRECATED: this is now UpdateKeyUsage, and it will be removed in the future - UpdateUsage(ctx context.Context, service Service, now time.Time, usage map[string]*AccessUsage) (map[string]bool, error) + StoreUsageProject(ctx context.Context, service Service, now time.Time, usage map[uint64]int64) (map[uint64]bool, error) + StoreUsageAccessKey(ctx context.Context, service Service, now time.Time, usage map[string]int64) (map[string]bool, error) // User permissions for a projectId GetUserPermission(ctx context.Context, projectId uint64, userId string) (UserPermission, *ResourceAccess, error) } @@ -509,20 +541,28 @@ type QuotaControlClient interface { UpdateDefaultAccessKey(ctx context.Context, projectID uint64, accessKey string) (bool, error) ListAccessKeys(ctx context.Context, projectId uint64, active *bool, service *Service) ([]*AccessKey, error) DisableAccessKey(ctx context.Context, accessKey string) (bool, error) - // Usage + // Quota GetProjectQuota(ctx context.Context, projectId uint64, now time.Time) (*AccessQuota, error) GetAccessQuota(ctx context.Context, accessKey string, now time.Time) (*AccessQuota, error) ClearAccessQuotaCache(ctx context.Context, projectID uint64) (bool, error) + // deprecated: this is now GetUsageProject, and it will be removed in the future GetAccountUsage(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (*AccessUsage, error) + // deprecated: this is now GetUsageAccessKey, and it will be removed in the future GetAccessKeyUsage(ctx context.Context, accessKey string, service *Service, from *time.Time, to *time.Time) (*AccessUsage, error) + // deprecated: this is now GetUsageAsync, and it will be removed in the future GetAsyncUsage(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (*AccessUsage, error) + GetUsageProject(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (int64, error) + GetUsageAccessKey(ctx context.Context, accessKey string, service *Service, from *time.Time, to *time.Time) (int64, error) + GetUsageAsync(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (int64, error) PrepareUsage(ctx context.Context, projectID uint64, cycle *Cycle, now time.Time) (bool, error) ClearUsage(ctx context.Context, projectID uint64, now time.Time) (bool, error) NotifyEvent(ctx context.Context, projectID uint64, eventType EventType) (bool, error) + // deprecated: this is now StoreUsageProject, and it will be removed in the future UpdateProjectUsage(ctx context.Context, service Service, now time.Time, usage map[uint64]*AccessUsage) (map[uint64]bool, error) + // deprecated: this is now StoreUsageAccessKey, and it will be removed in the future UpdateKeyUsage(ctx context.Context, service Service, now time.Time, usage map[string]*AccessUsage) (map[string]bool, error) - // DEPRECATED: this is now UpdateKeyUsage, and it will be removed in the future - UpdateUsage(ctx context.Context, service Service, now time.Time, usage map[string]*AccessUsage) (map[string]bool, error) + StoreUsageProject(ctx context.Context, service Service, now time.Time, usage map[uint64]int64) (map[uint64]bool, error) + StoreUsageAccessKey(ctx context.Context, service Service, now time.Time, usage map[string]int64) (map[string]bool, error) // User permissions for a projectId GetUserPermission(ctx context.Context, projectId uint64, userId string) (UserPermission, *ResourceAccess, error) } @@ -597,6 +637,12 @@ func (s *quotaControlServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { handler = s.serveGetAccessKeyUsageJSON case "/rpc/QuotaControl/GetAsyncUsage": handler = s.serveGetAsyncUsageJSON + case "/rpc/QuotaControl/GetUsageProject": + handler = s.serveGetUsageProjectJSON + case "/rpc/QuotaControl/GetUsageAccessKey": + handler = s.serveGetUsageAccessKeyJSON + case "/rpc/QuotaControl/GetUsageAsync": + handler = s.serveGetUsageAsyncJSON case "/rpc/QuotaControl/PrepareUsage": handler = s.servePrepareUsageJSON case "/rpc/QuotaControl/ClearUsage": @@ -607,8 +653,10 @@ func (s *quotaControlServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { handler = s.serveUpdateProjectUsageJSON case "/rpc/QuotaControl/UpdateKeyUsage": handler = s.serveUpdateKeyUsageJSON - case "/rpc/QuotaControl/UpdateUsage": - handler = s.serveUpdateUsageJSON + case "/rpc/QuotaControl/StoreUsageProject": + handler = s.serveStoreUsageProjectJSON + case "/rpc/QuotaControl/StoreUsageAccessKey": + handler = s.serveStoreUsageAccessKeyJSON case "/rpc/QuotaControl/GetUserPermission": handler = s.serveGetUserPermissionJSON default: @@ -1315,6 +1363,144 @@ func (s *quotaControlServer) serveGetAsyncUsageJSON(ctx context.Context, w http. w.Write(respBody) } +func (s *quotaControlServer) serveGetUsageProjectJSON(ctx context.Context, w http.ResponseWriter, r *http.Request) { + ctx = context.WithValue(ctx, MethodNameCtxKey, "GetUsageProject") + + reqBody, err := io.ReadAll(r.Body) + if err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to read request data: %w", err)) + return + } + defer r.Body.Close() + + reqPayload := struct { + Arg0 uint64 `json:"projectID"` + Arg1 *Service `json:"service"` + Arg2 *time.Time `json:"from"` + Arg3 *time.Time `json:"to"` + }{} + if err := json.Unmarshal(reqBody, &reqPayload); err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to unmarshal request data: %w", err)) + return + } + + // Call service method implementation. + ret0, err := s.QuotaControl.GetUsageProject(ctx, reqPayload.Arg0, reqPayload.Arg1, reqPayload.Arg2, reqPayload.Arg3) + if err != nil { + rpcErr, ok := err.(WebRPCError) + if !ok { + rpcErr = ErrWebrpcEndpoint.WithCause(err) + } + s.sendErrorJSON(w, r, rpcErr) + return + } + + respPayload := struct { + Ret0 int64 `json:"usage"` + }{ret0} + respBody, err := json.Marshal(respPayload) + if err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadResponse.WithCausef("failed to marshal json response: %w", err)) + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(respBody) +} + +func (s *quotaControlServer) serveGetUsageAccessKeyJSON(ctx context.Context, w http.ResponseWriter, r *http.Request) { + ctx = context.WithValue(ctx, MethodNameCtxKey, "GetUsageAccessKey") + + reqBody, err := io.ReadAll(r.Body) + if err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to read request data: %w", err)) + return + } + defer r.Body.Close() + + reqPayload := struct { + Arg0 string `json:"accessKey"` + Arg1 *Service `json:"service"` + Arg2 *time.Time `json:"from"` + Arg3 *time.Time `json:"to"` + }{} + if err := json.Unmarshal(reqBody, &reqPayload); err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to unmarshal request data: %w", err)) + return + } + + // Call service method implementation. + ret0, err := s.QuotaControl.GetUsageAccessKey(ctx, reqPayload.Arg0, reqPayload.Arg1, reqPayload.Arg2, reqPayload.Arg3) + if err != nil { + rpcErr, ok := err.(WebRPCError) + if !ok { + rpcErr = ErrWebrpcEndpoint.WithCause(err) + } + s.sendErrorJSON(w, r, rpcErr) + return + } + + respPayload := struct { + Ret0 int64 `json:"usage"` + }{ret0} + respBody, err := json.Marshal(respPayload) + if err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadResponse.WithCausef("failed to marshal json response: %w", err)) + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(respBody) +} + +func (s *quotaControlServer) serveGetUsageAsyncJSON(ctx context.Context, w http.ResponseWriter, r *http.Request) { + ctx = context.WithValue(ctx, MethodNameCtxKey, "GetUsageAsync") + + reqBody, err := io.ReadAll(r.Body) + if err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to read request data: %w", err)) + return + } + defer r.Body.Close() + + reqPayload := struct { + Arg0 uint64 `json:"projectID"` + Arg1 *Service `json:"service"` + Arg2 *time.Time `json:"from"` + Arg3 *time.Time `json:"to"` + }{} + if err := json.Unmarshal(reqBody, &reqPayload); err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to unmarshal request data: %w", err)) + return + } + + // Call service method implementation. + ret0, err := s.QuotaControl.GetUsageAsync(ctx, reqPayload.Arg0, reqPayload.Arg1, reqPayload.Arg2, reqPayload.Arg3) + if err != nil { + rpcErr, ok := err.(WebRPCError) + if !ok { + rpcErr = ErrWebrpcEndpoint.WithCause(err) + } + s.sendErrorJSON(w, r, rpcErr) + return + } + + respPayload := struct { + Ret0 int64 `json:"usage"` + }{ret0} + respBody, err := json.Marshal(respPayload) + if err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadResponse.WithCausef("failed to marshal json response: %w", err)) + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(respBody) +} + func (s *quotaControlServer) servePrepareUsageJSON(ctx context.Context, w http.ResponseWriter, r *http.Request) { ctx = context.WithValue(ctx, MethodNameCtxKey, "PrepareUsage") @@ -1538,8 +1724,8 @@ func (s *quotaControlServer) serveUpdateKeyUsageJSON(ctx context.Context, w http w.Write(respBody) } -func (s *quotaControlServer) serveUpdateUsageJSON(ctx context.Context, w http.ResponseWriter, r *http.Request) { - ctx = context.WithValue(ctx, MethodNameCtxKey, "UpdateUsage") +func (s *quotaControlServer) serveStoreUsageProjectJSON(ctx context.Context, w http.ResponseWriter, r *http.Request) { + ctx = context.WithValue(ctx, MethodNameCtxKey, "StoreUsageProject") reqBody, err := io.ReadAll(r.Body) if err != nil { @@ -1549,9 +1735,54 @@ func (s *quotaControlServer) serveUpdateUsageJSON(ctx context.Context, w http.Re defer r.Body.Close() reqPayload := struct { - Arg0 Service `json:"service"` - Arg1 time.Time `json:"now"` - Arg2 map[string]*AccessUsage `json:"usage"` + Arg0 Service `json:"service"` + Arg1 time.Time `json:"now"` + Arg2 map[uint64]int64 `json:"usage"` + }{} + if err := json.Unmarshal(reqBody, &reqPayload); err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to unmarshal request data: %w", err)) + return + } + + // Call service method implementation. + ret0, err := s.QuotaControl.StoreUsageProject(ctx, reqPayload.Arg0, reqPayload.Arg1, reqPayload.Arg2) + if err != nil { + rpcErr, ok := err.(WebRPCError) + if !ok { + rpcErr = ErrWebrpcEndpoint.WithCause(err) + } + s.sendErrorJSON(w, r, rpcErr) + return + } + + respPayload := struct { + Ret0 map[uint64]bool `json:"ok"` + }{ret0} + respBody, err := json.Marshal(respPayload) + if err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadResponse.WithCausef("failed to marshal json response: %w", err)) + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(respBody) +} + +func (s *quotaControlServer) serveStoreUsageAccessKeyJSON(ctx context.Context, w http.ResponseWriter, r *http.Request) { + ctx = context.WithValue(ctx, MethodNameCtxKey, "StoreUsageAccessKey") + + reqBody, err := io.ReadAll(r.Body) + if err != nil { + s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to read request data: %w", err)) + return + } + defer r.Body.Close() + + reqPayload := struct { + Arg0 Service `json:"service"` + Arg1 time.Time `json:"now"` + Arg2 map[string]int64 `json:"usage"` }{} if err := json.Unmarshal(reqBody, &reqPayload); err != nil { s.sendErrorJSON(w, r, ErrWebrpcBadRequest.WithCausef("failed to unmarshal request data: %w", err)) @@ -1559,7 +1790,7 @@ func (s *quotaControlServer) serveUpdateUsageJSON(ctx context.Context, w http.Re } // Call service method implementation. - ret0, err := s.QuotaControl.UpdateUsage(ctx, reqPayload.Arg0, reqPayload.Arg1, reqPayload.Arg2) + ret0, err := s.QuotaControl.StoreUsageAccessKey(ctx, reqPayload.Arg0, reqPayload.Arg1, reqPayload.Arg2) if err != nil { rpcErr, ok := err.(WebRPCError) if !ok { @@ -1661,12 +1892,12 @@ const QuotaControlPathPrefix = "/rpc/QuotaControl/" type quotaControlClient struct { client HTTPClient - urls [22]string + urls [26]string } func NewQuotaControlClient(addr string, client HTTPClient) QuotaControlClient { prefix := urlBase(addr) + QuotaControlPathPrefix - urls := [22]string{ + urls := [26]string{ prefix + "GetProjectStatus", prefix + "GetAccessKey", prefix + "GetDefaultAccessKey", @@ -1682,12 +1913,16 @@ func NewQuotaControlClient(addr string, client HTTPClient) QuotaControlClient { prefix + "GetAccountUsage", prefix + "GetAccessKeyUsage", prefix + "GetAsyncUsage", + prefix + "GetUsageProject", + prefix + "GetUsageAccessKey", + prefix + "GetUsageAsync", prefix + "PrepareUsage", prefix + "ClearUsage", prefix + "NotifyEvent", prefix + "UpdateProjectUsage", prefix + "UpdateKeyUsage", - prefix + "UpdateUsage", + prefix + "StoreUsageProject", + prefix + "StoreUsageAccessKey", prefix + "GetUserPermission", } return "aControlClient{ @@ -2001,6 +2236,72 @@ func (c *quotaControlClient) GetAsyncUsage(ctx context.Context, projectID uint64 return out.Ret0, err } +func (c *quotaControlClient) GetUsageProject(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (int64, error) { + in := struct { + Arg0 uint64 `json:"projectID"` + Arg1 *Service `json:"service"` + Arg2 *time.Time `json:"from"` + Arg3 *time.Time `json:"to"` + }{projectID, service, from, to} + out := struct { + Ret0 int64 `json:"usage"` + }{} + + resp, err := doHTTPRequest(ctx, c.client, c.urls[15], in, &out) + if resp != nil { + cerr := resp.Body.Close() + if err == nil && cerr != nil { + err = ErrWebrpcRequestFailed.WithCausef("failed to close response body: %w", cerr) + } + } + + return out.Ret0, err +} + +func (c *quotaControlClient) GetUsageAccessKey(ctx context.Context, accessKey string, service *Service, from *time.Time, to *time.Time) (int64, error) { + in := struct { + Arg0 string `json:"accessKey"` + Arg1 *Service `json:"service"` + Arg2 *time.Time `json:"from"` + Arg3 *time.Time `json:"to"` + }{accessKey, service, from, to} + out := struct { + Ret0 int64 `json:"usage"` + }{} + + resp, err := doHTTPRequest(ctx, c.client, c.urls[16], in, &out) + if resp != nil { + cerr := resp.Body.Close() + if err == nil && cerr != nil { + err = ErrWebrpcRequestFailed.WithCausef("failed to close response body: %w", cerr) + } + } + + return out.Ret0, err +} + +func (c *quotaControlClient) GetUsageAsync(ctx context.Context, projectID uint64, service *Service, from *time.Time, to *time.Time) (int64, error) { + in := struct { + Arg0 uint64 `json:"projectID"` + Arg1 *Service `json:"service"` + Arg2 *time.Time `json:"from"` + Arg3 *time.Time `json:"to"` + }{projectID, service, from, to} + out := struct { + Ret0 int64 `json:"usage"` + }{} + + resp, err := doHTTPRequest(ctx, c.client, c.urls[17], in, &out) + if resp != nil { + cerr := resp.Body.Close() + if err == nil && cerr != nil { + err = ErrWebrpcRequestFailed.WithCausef("failed to close response body: %w", cerr) + } + } + + return out.Ret0, err +} + func (c *quotaControlClient) PrepareUsage(ctx context.Context, projectID uint64, cycle *Cycle, now time.Time) (bool, error) { in := struct { Arg0 uint64 `json:"projectID"` @@ -2011,7 +2312,7 @@ func (c *quotaControlClient) PrepareUsage(ctx context.Context, projectID uint64, Ret0 bool `json:"ok"` }{} - resp, err := doHTTPRequest(ctx, c.client, c.urls[15], in, &out) + resp, err := doHTTPRequest(ctx, c.client, c.urls[18], in, &out) if resp != nil { cerr := resp.Body.Close() if err == nil && cerr != nil { @@ -2031,7 +2332,7 @@ func (c *quotaControlClient) ClearUsage(ctx context.Context, projectID uint64, n Ret0 bool `json:"ok"` }{} - resp, err := doHTTPRequest(ctx, c.client, c.urls[16], in, &out) + resp, err := doHTTPRequest(ctx, c.client, c.urls[19], in, &out) if resp != nil { cerr := resp.Body.Close() if err == nil && cerr != nil { @@ -2051,7 +2352,7 @@ func (c *quotaControlClient) NotifyEvent(ctx context.Context, projectID uint64, Ret0 bool `json:"ok"` }{} - resp, err := doHTTPRequest(ctx, c.client, c.urls[17], in, &out) + resp, err := doHTTPRequest(ctx, c.client, c.urls[20], in, &out) if resp != nil { cerr := resp.Body.Close() if err == nil && cerr != nil { @@ -2072,7 +2373,7 @@ func (c *quotaControlClient) UpdateProjectUsage(ctx context.Context, service Ser Ret0 map[uint64]bool `json:"ok"` }{} - resp, err := doHTTPRequest(ctx, c.client, c.urls[18], in, &out) + resp, err := doHTTPRequest(ctx, c.client, c.urls[21], in, &out) if resp != nil { cerr := resp.Body.Close() if err == nil && cerr != nil { @@ -2093,7 +2394,7 @@ func (c *quotaControlClient) UpdateKeyUsage(ctx context.Context, service Service Ret0 map[string]bool `json:"ok"` }{} - resp, err := doHTTPRequest(ctx, c.client, c.urls[19], in, &out) + resp, err := doHTTPRequest(ctx, c.client, c.urls[22], in, &out) if resp != nil { cerr := resp.Body.Close() if err == nil && cerr != nil { @@ -2104,17 +2405,38 @@ func (c *quotaControlClient) UpdateKeyUsage(ctx context.Context, service Service return out.Ret0, err } -func (c *quotaControlClient) UpdateUsage(ctx context.Context, service Service, now time.Time, usage map[string]*AccessUsage) (map[string]bool, error) { +func (c *quotaControlClient) StoreUsageProject(ctx context.Context, service Service, now time.Time, usage map[uint64]int64) (map[uint64]bool, error) { in := struct { - Arg0 Service `json:"service"` - Arg1 time.Time `json:"now"` - Arg2 map[string]*AccessUsage `json:"usage"` + Arg0 Service `json:"service"` + Arg1 time.Time `json:"now"` + Arg2 map[uint64]int64 `json:"usage"` + }{service, now, usage} + out := struct { + Ret0 map[uint64]bool `json:"ok"` + }{} + + resp, err := doHTTPRequest(ctx, c.client, c.urls[23], in, &out) + if resp != nil { + cerr := resp.Body.Close() + if err == nil && cerr != nil { + err = ErrWebrpcRequestFailed.WithCausef("failed to close response body: %w", cerr) + } + } + + return out.Ret0, err +} + +func (c *quotaControlClient) StoreUsageAccessKey(ctx context.Context, service Service, now time.Time, usage map[string]int64) (map[string]bool, error) { + in := struct { + Arg0 Service `json:"service"` + Arg1 time.Time `json:"now"` + Arg2 map[string]int64 `json:"usage"` }{service, now, usage} out := struct { Ret0 map[string]bool `json:"ok"` }{} - resp, err := doHTTPRequest(ctx, c.client, c.urls[20], in, &out) + resp, err := doHTTPRequest(ctx, c.client, c.urls[24], in, &out) if resp != nil { cerr := resp.Body.Close() if err == nil && cerr != nil { @@ -2135,7 +2457,7 @@ func (c *quotaControlClient) GetUserPermission(ctx context.Context, projectId ui Ret1 *ResourceAccess `json:"resourceAccess"` }{} - resp, err := doHTTPRequest(ctx, c.client, c.urls[21], in, &out) + resp, err := doHTTPRequest(ctx, c.client, c.urls[25], in, &out) if resp != nil { cerr := resp.Body.Close() if err == nil && cerr != nil { diff --git a/proto/quotacontrol.gen.ts b/proto/quotacontrol.gen.ts index cac0cce..f99105c 100644 --- a/proto/quotacontrol.gen.ts +++ b/proto/quotacontrol.gen.ts @@ -1,5 +1,5 @@ /* eslint-disable */ -// quota-control v0.19.2 4b5f19d012ae4365adc25b547b4d39ac58a974e8 +// quota-control v0.19.2 db40d579f0863c38fafda257ec579b72c8751b69 // -- // Code generated by webrpc-gen@v0.22.0 with typescript@v0.16.1 generator. DO NOT EDIT. // @@ -16,7 +16,7 @@ export const WebRPCVersion = "v1" export const WebRPCSchemaVersion = "v0.19.2" // Schema hash generated from your RIDL schema -export const WebRPCSchemaHash = "4b5f19d012ae4365adc25b547b4d39ac58a974e8" +export const WebRPCSchemaHash = "db40d579f0863c38fafda257ec579b72c8751b69" type WebrpcGenVersions = { webrpcGenVersion: string; @@ -171,12 +171,16 @@ export interface QuotaControl { getAccountUsage(args: GetAccountUsageArgs, headers?: object, signal?: AbortSignal): Promise getAccessKeyUsage(args: GetAccessKeyUsageArgs, headers?: object, signal?: AbortSignal): Promise getAsyncUsage(args: GetAsyncUsageArgs, headers?: object, signal?: AbortSignal): Promise + getUsageProject(args: GetUsageProjectArgs, headers?: object, signal?: AbortSignal): Promise + getUsageAccessKey(args: GetUsageAccessKeyArgs, headers?: object, signal?: AbortSignal): Promise + getUsageAsync(args: GetUsageAsyncArgs, headers?: object, signal?: AbortSignal): Promise prepareUsage(args: PrepareUsageArgs, headers?: object, signal?: AbortSignal): Promise clearUsage(args: ClearUsageArgs, headers?: object, signal?: AbortSignal): Promise notifyEvent(args: NotifyEventArgs, headers?: object, signal?: AbortSignal): Promise updateProjectUsage(args: UpdateProjectUsageArgs, headers?: object, signal?: AbortSignal): Promise updateKeyUsage(args: UpdateKeyUsageArgs, headers?: object, signal?: AbortSignal): Promise - updateUsage(args: UpdateUsageArgs, headers?: object, signal?: AbortSignal): Promise + storeUsageProject(args: StoreUsageProjectArgs, headers?: object, signal?: AbortSignal): Promise + storeUsageAccessKey(args: StoreUsageAccessKeyArgs, headers?: object, signal?: AbortSignal): Promise getUserPermission(args: GetUserPermissionArgs, headers?: object, signal?: AbortSignal): Promise } @@ -305,6 +309,36 @@ export interface GetAsyncUsageArgs { export interface GetAsyncUsageReturn { usage: AccessUsage } +export interface GetUsageProjectArgs { + projectID: number + service?: Service + from?: string + to?: string +} + +export interface GetUsageProjectReturn { + usage: number +} +export interface GetUsageAccessKeyArgs { + accessKey: string + service?: Service + from?: string + to?: string +} + +export interface GetUsageAccessKeyReturn { + usage: number +} +export interface GetUsageAsyncArgs { + projectID: number + service?: Service + from?: string + to?: string +} + +export interface GetUsageAsyncReturn { + usage: number +} export interface PrepareUsageArgs { projectID: number cycle: Cycle @@ -348,13 +382,22 @@ export interface UpdateKeyUsageArgs { export interface UpdateKeyUsageReturn { ok: {[key: string]: boolean} } -export interface UpdateUsageArgs { +export interface StoreUsageProjectArgs { service: Service now: string - usage: {[key: string]: AccessUsage} + usage: {[key: number]: number} +} + +export interface StoreUsageProjectReturn { + ok: {[key: number]: boolean} +} +export interface StoreUsageAccessKeyArgs { + service: Service + now: string + usage: {[key: string]: number} } -export interface UpdateUsageReturn { +export interface StoreUsageAccessKeyReturn { ok: {[key: string]: boolean} } export interface GetUserPermissionArgs { @@ -596,6 +639,48 @@ export class QuotaControl implements QuotaControl { }) } + getUsageProject = (args: GetUsageProjectArgs, headers?: object, signal?: AbortSignal): Promise => { + return this.fetch( + this.url('GetUsageProject'), + createHTTPRequest(args, headers, signal)).then((res) => { + return buildResponse(res).then(_data => { + return { + usage: (_data.usage), + } + }) + }, (error) => { + throw WebrpcRequestFailedError.new({ cause: `fetch(): ${error.message || ''}` }) + }) + } + + getUsageAccessKey = (args: GetUsageAccessKeyArgs, headers?: object, signal?: AbortSignal): Promise => { + return this.fetch( + this.url('GetUsageAccessKey'), + createHTTPRequest(args, headers, signal)).then((res) => { + return buildResponse(res).then(_data => { + return { + usage: (_data.usage), + } + }) + }, (error) => { + throw WebrpcRequestFailedError.new({ cause: `fetch(): ${error.message || ''}` }) + }) + } + + getUsageAsync = (args: GetUsageAsyncArgs, headers?: object, signal?: AbortSignal): Promise => { + return this.fetch( + this.url('GetUsageAsync'), + createHTTPRequest(args, headers, signal)).then((res) => { + return buildResponse(res).then(_data => { + return { + usage: (_data.usage), + } + }) + }, (error) => { + throw WebrpcRequestFailedError.new({ cause: `fetch(): ${error.message || ''}` }) + }) + } + prepareUsage = (args: PrepareUsageArgs, headers?: object, signal?: AbortSignal): Promise => { return this.fetch( this.url('PrepareUsage'), @@ -666,9 +751,23 @@ export class QuotaControl implements QuotaControl { }) } - updateUsage = (args: UpdateUsageArgs, headers?: object, signal?: AbortSignal): Promise => { + storeUsageProject = (args: StoreUsageProjectArgs, headers?: object, signal?: AbortSignal): Promise => { + return this.fetch( + this.url('StoreUsageProject'), + createHTTPRequest(args, headers, signal)).then((res) => { + return buildResponse(res).then(_data => { + return { + ok: <{[key: number]: boolean}>(_data.ok), + } + }) + }, (error) => { + throw WebrpcRequestFailedError.new({ cause: `fetch(): ${error.message || ''}` }) + }) + } + + storeUsageAccessKey = (args: StoreUsageAccessKeyArgs, headers?: object, signal?: AbortSignal): Promise => { return this.fetch( - this.url('UpdateUsage'), + this.url('StoreUsageAccessKey'), createHTTPRequest(args, headers, signal)).then((res) => { return buildResponse(res).then(_data => { return { diff --git a/proto/quotacontrol.ridl b/proto/quotacontrol.ridl index 57708f8..dd52708 100644 --- a/proto/quotacontrol.ridl +++ b/proto/quotacontrol.ridl @@ -112,20 +112,35 @@ service QuotaControl - ListAccessKeys(projectId: uint64, active?: bool, service?: Service) => (accessKeys: []AccessKey) - DisableAccessKey(accessKey: string) => (ok: bool) - # Usage + # Quota - GetProjectQuota(projectId: uint64, now: timestamp) => (accessQuota: AccessQuota) - GetAccessQuota(accessKey: string, now: timestamp) => (accessQuota: AccessQuota) - ClearAccessQuotaCache(projectID: uint64) => (ok: bool) + + # Usage + + # deprecated: this is now GetUsageProject, and it will be removed in the future - GetAccountUsage(projectID: uint64, service?: Service, from?: timestamp, to?: timestamp) => (usage: AccessUsage) + # deprecated: this is now GetUsageAccessKey, and it will be removed in the future - GetAccessKeyUsage(accessKey: string, service?: Service, from?: timestamp, to?: timestamp) => (usage: AccessUsage) + # deprecated: this is now GetUsageAsync, and it will be removed in the future - GetAsyncUsage(projectID: uint64, service?: Service, from?: timestamp, to?: timestamp) => (usage: AccessUsage) + + - GetUsageProject(projectID: uint64, service?: Service, from?: timestamp, to?: timestamp) => (usage: int64) + - GetUsageAccessKey(accessKey: string, service?: Service, from?: timestamp, to?: timestamp) => (usage: int64) + - GetUsageAsync(projectID: uint64, service?: Service, from?: timestamp, to?: timestamp) => (usage: int64) + - PrepareUsage(projectID: uint64, cycle: Cycle, now: timestamp) => (ok: bool) - ClearUsage(projectID: uint64, now: timestamp) => (ok: bool) - NotifyEvent(projectID: uint64, eventType: EventType) => (ok: bool) + + # deprecated: this is now StoreUsageProject, and it will be removed in the future - UpdateProjectUsage(service: Service, now: timestamp, usage: map) => (ok: map) + # deprecated: this is now StoreUsageAccessKey, and it will be removed in the future - UpdateKeyUsage(service: Service, now: timestamp, usage: map) => (ok: map) - # DEPRECATED: this is now UpdateKeyUsage, and it will be removed in the future - - UpdateUsage(service: Service, now: timestamp, usage: map) => (ok: map) + + - StoreUsageProject(service: Service, now: timestamp, usage: map) => (ok: map) + - StoreUsageAccessKey(service: Service, now: timestamp, usage: map) => (ok: map) # User permissions for a projectId - GetUserPermission(projectId: uint64, userId: string) => (permission: UserPermission, resourceAccess: ResourceAccess) diff --git a/tests/mock/mem.go b/tests/mock/mem.go index fc83f0e..ffcd9dc 100644 --- a/tests/mock/mem.go +++ b/tests/mock/mem.go @@ -121,52 +121,44 @@ func (m *MemoryStore) ListAccessKeys(ctx context.Context, projectID uint64, acti return accessKeys, nil } -func (m *MemoryStore) GetAccountUsage(ctx context.Context, projectID uint64, service *proto.Service, min, max time.Time) (proto.AccessUsage, error) { +func (m *MemoryStore) GetUsageProject(ctx context.Context, projectID uint64, service *proto.Service, min, max time.Time) (int64, error) { m.Lock() defer m.Unlock() - usage := proto.AccessUsage{} - if m.usage.ByProjectID[projectID] != nil { - usage.Add(*m.usage.ByProjectID[projectID]) - } + + usage, _ := m.usage.ByProjectID[projectID] + for _, v := range m.accessKeys { if v.ProjectID == projectID { u, ok := m.usage.ByAccessKey[v.AccessKey] if !ok { continue } - usage.Add(*u) + usage += u } } return usage, nil } -func (m *MemoryStore) GetAccessKeyUsage(ctx context.Context, projectID uint64, accessKey string, service *proto.Service, min, max time.Time) (proto.AccessUsage, error) { +func (m *MemoryStore) GetUsageAccessKey(ctx context.Context, projectID uint64, accessKey string, service *proto.Service, min, max time.Time) (int64, error) { m.Lock() defer m.Unlock() if _, ok := m.accessKeys[accessKey]; !ok { - return proto.AccessUsage{}, proto.ErrAccessKeyNotFound + return 0, proto.ErrAccessKeyNotFound } - usage, ok := m.usage.ByAccessKey[accessKey] - if !ok { - return proto.AccessUsage{}, nil - } - return *usage, nil + usage, _ := m.usage.ByAccessKey[accessKey] + return usage, nil } -func (m *MemoryStore) UpdateAccessUsage(ctx context.Context, projectID uint64, accessKey string, service proto.Service, time time.Time, usage proto.AccessUsage) error { +func (m *MemoryStore) UpdateUsage(ctx context.Context, projectID uint64, accessKey string, service proto.Service, time time.Time, usage int64) error { m.Lock() defer m.Unlock() - if _, ok := m.usage.ByAccessKey[accessKey]; !ok { - m.usage.ByAccessKey[accessKey] = &usage - return nil - } - m.usage.ByAccessKey[accessKey].Add(usage) + m.usage.ByAccessKey[accessKey] += usage return nil } func (m *MemoryStore) ResetUsage(ctx context.Context, accessKey string) error { m.Lock() - m.usage.ByAccessKey[accessKey] = &proto.AccessUsage{} + m.usage.ByAccessKey[accessKey] = 0 m.Unlock() return nil }