From db38802e83afc54ab6ed45d3f5ae43096270180d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enver=20Bi=C5=A1evac?= Date: Fri, 20 Dec 2024 10:28:47 +0000 Subject: [PATCH] feat: [code-2912]: usage service implementation (#3168) * feat: [code-2912]: rest and ui changes (#3178) * requested changes * rest and ui changes * requested changes * add midleware to raw and archieve endpoints * wire fixed * requested changes * Merge remote-tracking branch 'origin/main' into eb/code-2912 * added initial values from db * minor improvements * requested changes * remove check for bandwidth * wiring dep * improved test * limits check added * config added, minor improvements * usage service implementation --- .gitignore | 3 + .testapi/http-client.env.json | 5 + .testapi/space.http | 4 + app/api/controller/space/controller.go | 56 +++--- app/api/controller/space/usage.go | 55 ++++++ app/api/controller/space/wire.go | 4 +- app/api/handler/space/usage.go | 63 ++++++ app/api/openapi/space.go | 10 + app/router/api.go | 19 +- app/router/git.go | 6 +- app/router/wire.go | 5 +- app/services/usage/config.go | 51 +++++ app/services/usage/interface.go | 21 ++ app/services/usage/io.go | 123 ++++++++++++ app/services/usage/io_test.go | 103 ++++++++++ app/services/usage/middleware.go | 59 ++++++ app/services/usage/middleware_test.go | 87 +++++++++ app/services/usage/mocks.go | 105 ++++++++++ app/services/usage/queue.go | 59 ++++++ app/services/usage/usage.go | 255 +++++++++++++++++++++++++ app/services/usage/usage_test.go | 113 +++++++++++ app/services/usage/wire.go | 42 ++++ app/store/database.go | 9 + app/store/database/space.go | 22 +++ app/store/database/space_test.go | 23 +++ app/store/database/wire.go | 5 + cmd/gitness/wire.go | 2 + cmd/gitness/wire_gen.go | 7 +- go.mod | 4 +- go.sum | 4 +- types/config.go | 5 + 31 files changed, 1288 insertions(+), 41 deletions(-) create mode 100644 .testapi/http-client.env.json create mode 100644 .testapi/space.http create mode 100644 app/api/controller/space/usage.go create mode 100644 app/api/handler/space/usage.go create mode 100644 app/services/usage/config.go create mode 100644 app/services/usage/interface.go create mode 100644 app/services/usage/io.go create mode 100644 app/services/usage/io_test.go create mode 100644 app/services/usage/middleware.go create mode 100644 app/services/usage/middleware_test.go create mode 100644 app/services/usage/mocks.go create mode 100644 app/services/usage/queue.go create mode 100644 app/services/usage/usage.go create mode 100644 app/services/usage/usage_test.go create mode 100644 app/services/usage/wire.go diff --git a/.gitignore b/.gitignore index e2803bac1..cd103c130 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,6 @@ node_modules /distribution-spec /registry/distribution-spec /app/store/database/test.db + +# adding support for .http files +http-client.private.env.json \ No newline at end of file diff --git a/.testapi/http-client.env.json b/.testapi/http-client.env.json new file mode 100644 index 000000000..19370b800 --- /dev/null +++ b/.testapi/http-client.env.json @@ -0,0 +1,5 @@ +{ + "dev": { + "baseurl": "http://localhost:3000/api/v1" + } +} \ No newline at end of file diff --git a/.testapi/space.http b/.testapi/space.http new file mode 100644 index 000000000..96d546f39 --- /dev/null +++ b/.testapi/space.http @@ -0,0 +1,4 @@ +### Get metric for space + +GET {{baseurl}}/spaces/root/+/usage/metric +Authorization: {{token}} diff --git a/app/api/controller/space/controller.go b/app/api/controller/space/controller.go index 427790c84..fdf67016f 100644 --- a/app/api/controller/space/controller.go +++ b/app/api/controller/space/controller.go @@ -73,32 +73,33 @@ func (s SpaceOutput) MarshalJSON() ([]byte, error) { type Controller struct { nestedSpacesEnabled bool - tx dbtx.Transactor - urlProvider url.Provider - sseStreamer sse.Streamer - identifierCheck check.SpaceIdentifier - authorizer authz.Authorizer - spacePathStore store.SpacePathStore - pipelineStore store.PipelineStore - secretStore store.SecretStore - connectorStore store.ConnectorStore - templateStore store.TemplateStore - spaceStore store.SpaceStore - repoStore store.RepoStore - principalStore store.PrincipalStore - repoCtrl *repo.Controller - membershipStore store.MembershipStore - prListService *pullreq.ListService - importer *importer.Repository - exporter *exporter.Repository - resourceLimiter limiter.ResourceLimiter - publicAccess publicaccess.Service - auditService audit.Service - gitspaceSvc *gitspace.Service - labelSvc *label.Service - instrumentation instrument.Service - executionStore store.ExecutionStore - rulesSvc *rules.Service + tx dbtx.Transactor + urlProvider url.Provider + sseStreamer sse.Streamer + identifierCheck check.SpaceIdentifier + authorizer authz.Authorizer + spacePathStore store.SpacePathStore + pipelineStore store.PipelineStore + secretStore store.SecretStore + connectorStore store.ConnectorStore + templateStore store.TemplateStore + spaceStore store.SpaceStore + repoStore store.RepoStore + principalStore store.PrincipalStore + repoCtrl *repo.Controller + membershipStore store.MembershipStore + prListService *pullreq.ListService + importer *importer.Repository + exporter *exporter.Repository + resourceLimiter limiter.ResourceLimiter + publicAccess publicaccess.Service + auditService audit.Service + gitspaceSvc *gitspace.Service + labelSvc *label.Service + instrumentation instrument.Service + executionStore store.ExecutionStore + rulesSvc *rules.Service + usageMetricStore store.UsageMetricStore } func NewController(config *types.Config, tx dbtx.Transactor, urlProvider url.Provider, @@ -111,7 +112,7 @@ func NewController(config *types.Config, tx dbtx.Transactor, urlProvider url.Pro limiter limiter.ResourceLimiter, publicAccess publicaccess.Service, auditService audit.Service, gitspaceSvc *gitspace.Service, labelSvc *label.Service, instrumentation instrument.Service, executionStore store.ExecutionStore, - rulesSvc *rules.Service, + rulesSvc *rules.Service, usageMetricStore store.UsageMetricStore, ) *Controller { return &Controller{ nestedSpacesEnabled: config.NestedSpacesEnabled, @@ -141,6 +142,7 @@ func NewController(config *types.Config, tx dbtx.Transactor, urlProvider url.Pro instrumentation: instrumentation, executionStore: executionStore, rulesSvc: rulesSvc, + usageMetricStore: usageMetricStore, } } diff --git a/app/api/controller/space/usage.go b/app/api/controller/space/usage.go new file mode 100644 index 000000000..d559088c6 --- /dev/null +++ b/app/api/controller/space/usage.go @@ -0,0 +1,55 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package space + +import ( + "context" + "fmt" + + "github.com/harness/gitness/app/auth" + "github.com/harness/gitness/app/paths" + "github.com/harness/gitness/types" + "github.com/harness/gitness/types/enum" +) + +// GetUsageMetrics returns usage metrics for root space. +func (c *Controller) GetUsageMetrics( + ctx context.Context, + session *auth.Session, + spaceRef string, + startDate int64, + endDate int64, +) (*types.UsageMetric, error) { + rootSpaceRef, _, err := paths.DisectRoot(spaceRef) + if err != nil { + return nil, fmt.Errorf("could not find root space: %w", err) + } + space, err := c.getSpaceCheckAuth(ctx, session, rootSpaceRef, enum.PermissionSpaceView) + if err != nil { + return nil, fmt.Errorf("failed to acquire access to space: %w", err) + } + + metric, err := c.usageMetricStore.GetMetrics( + ctx, + space.ID, + startDate, + endDate, + ) + if err != nil { + return nil, fmt.Errorf("failed to retrieve usage metrics: %w", err) + } + + return metric, nil +} diff --git a/app/api/controller/space/wire.go b/app/api/controller/space/wire.go index c8ce21078..237536ac3 100644 --- a/app/api/controller/space/wire.go +++ b/app/api/controller/space/wire.go @@ -52,7 +52,7 @@ func ProvideController(config *types.Config, tx dbtx.Transactor, urlProvider url exporter *exporter.Repository, limiter limiter.ResourceLimiter, publicAccess publicaccess.Service, auditService audit.Service, gitspaceService *gitspace.Service, labelSvc *label.Service, instrumentation instrument.Service, executionStore store.ExecutionStore, - rulesSvc *rules.Service, + rulesSvc *rules.Service, usageMetricStore store.UsageMetricStore, ) *Controller { return NewController(config, tx, urlProvider, sseStreamer, identifierCheck, authorizer, @@ -63,6 +63,6 @@ func ProvideController(config *types.Config, tx dbtx.Transactor, urlProvider url exporter, limiter, publicAccess, auditService, gitspaceService, labelSvc, instrumentation, executionStore, - rulesSvc, + rulesSvc, usageMetricStore, ) } diff --git a/app/api/handler/space/usage.go b/app/api/handler/space/usage.go new file mode 100644 index 000000000..d912bdbc8 --- /dev/null +++ b/app/api/handler/space/usage.go @@ -0,0 +1,63 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package space + +import ( + "net/http" + "time" + + "github.com/harness/gitness/app/api/controller/space" + "github.com/harness/gitness/app/api/render" + "github.com/harness/gitness/app/api/request" +) + +func HandleUsageMetric(spaceCtrl *space.Controller) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + session, _ := request.AuthSessionFrom(ctx) + + spaceRef, err := request.GetSpaceRefFromPath(r) + if err != nil { + render.TranslatedUserError(ctx, w, err) + return + } + + now := time.Now() + start := now.Add(-30 * 24 * time.Hour).UnixMilli() + + startDate, ok, _ := request.QueryParamAsPositiveInt64(r, "start_date") + if !ok { + startDate = start + } + endDate, ok, _ := request.QueryParamAsPositiveInt64(r, "start_date") + if !ok { + endDate = now.UnixMilli() + } + + rule, err := spaceCtrl.GetUsageMetrics( + ctx, + session, + spaceRef, + startDate, + endDate, + ) + if err != nil { + render.TranslatedUserError(ctx, w, err) + return + } + + render.JSON(w, http.StatusOK, rule) + } +} diff --git a/app/api/openapi/space.go b/app/api/openapi/space.go index b9c1127e4..155243646 100644 --- a/app/api/openapi/space.go +++ b/app/api/openapi/space.go @@ -679,4 +679,14 @@ func spaceOperations(reflector *openapi3.Reflector) { _ = reflector.SetJSONResponse(&listPullReq, new(usererror.Error), http.StatusUnauthorized) _ = reflector.SetJSONResponse(&listPullReq, new(usererror.Error), http.StatusForbidden) _ = reflector.Spec.AddOperation(http.MethodGet, "/spaces/{repo_ref}/pullreq", listPullReq) + + opGetUsageMetrics := openapi3.Operation{} + opGetUsageMetrics.WithTags("space") + opGetUsageMetrics.WithMapOfAnything(map[string]interface{}{"operationId": "getSpaceUsageMetric"}) + _ = reflector.SetRequest(&opGetUsageMetrics, new(spaceRequest), http.MethodGet) + _ = reflector.SetJSONResponse(&opGetUsageMetrics, new(types.UsageMetric), http.StatusOK) + _ = reflector.SetJSONResponse(&opGetUsageMetrics, new(usererror.Error), http.StatusInternalServerError) + _ = reflector.SetJSONResponse(&opGetUsageMetrics, new(usererror.Error), http.StatusUnauthorized) + _ = reflector.SetJSONResponse(&opGetUsageMetrics, new(usererror.Error), http.StatusForbidden) + _ = reflector.Spec.AddOperation(http.MethodGet, "/spaces/{space_ref}/usage/metric", opGetUsageMetrics) } diff --git a/app/router/api.go b/app/router/api.go index fece99e28..385fdd154 100644 --- a/app/router/api.go +++ b/app/router/api.go @@ -85,6 +85,7 @@ import ( "github.com/harness/gitness/app/api/request" "github.com/harness/gitness/app/auth/authn" "github.com/harness/gitness/app/githook" + "github.com/harness/gitness/app/services/usage" "github.com/harness/gitness/audit" "github.com/harness/gitness/git" "github.com/harness/gitness/types" @@ -136,6 +137,7 @@ func NewAPIHandler( gitspaceCtrl *gitspace.Controller, aiagentCtrl *aiagent.Controller, capabilitiesCtrl *capabilities.Controller, + usageSender usage.Sender, ) http.Handler { // Use go-chi router for inner routing. r := chi.NewRouter() @@ -168,7 +170,7 @@ func NewAPIHandler( setupRoutesV1WithAuth(r, appCtx, config, repoCtrl, repoSettingsCtrl, executionCtrl, triggerCtrl, logCtrl, pipelineCtrl, connectorCtrl, templateCtrl, pluginCtrl, secretCtrl, spaceCtrl, pullreqCtrl, webhookCtrl, githookCtrl, git, saCtrl, userCtrl, principalCtrl, userGroupCtrl, checkCtrl, uploadCtrl, - searchCtrl, gitspaceCtrl, infraProviderCtrl, migrateCtrl, aiagentCtrl, capabilitiesCtrl) + searchCtrl, gitspaceCtrl, infraProviderCtrl, migrateCtrl, aiagentCtrl, capabilitiesCtrl, usageSender) }) }) @@ -220,11 +222,12 @@ func setupRoutesV1WithAuth(r chi.Router, migrateCtrl *migrate.Controller, aiagentCtrl *aiagent.Controller, capabilitiesCtrl *capabilities.Controller, + usageSender usage.Sender, ) { setupAccountWithAuth(r, userCtrl, config) setupSpaces(r, appCtx, spaceCtrl, userGroupCtrl, webhookCtrl, checkCtrl) setupRepos(r, repoCtrl, repoSettingsCtrl, pipelineCtrl, executionCtrl, triggerCtrl, - logCtrl, pullreqCtrl, webhookCtrl, checkCtrl, uploadCtrl) + logCtrl, pullreqCtrl, webhookCtrl, checkCtrl, uploadCtrl, usageSender) setupConnectors(r, connectorCtrl) setupTemplates(r, templateCtrl) setupSecrets(r, secretCtrl) @@ -296,6 +299,9 @@ func setupSpaces( SetupRulesSpace(r, spaceCtrl) r.Get("/checks/recent", handlercheck.HandleCheckListRecentSpace(checkCtrl)) + r.Route("/usage", func(r chi.Router) { + r.Get("/metric", handlerspace.HandleUsageMetric(spaceCtrl)) + }) }) }) } @@ -368,6 +374,7 @@ func setupRepos(r chi.Router, webhookCtrl *webhook.Controller, checkCtrl *check.Controller, uploadCtrl *upload.Controller, + usageSender usage.Sender, ) { r.Route("/repos", func(r chi.Router) { // Create takes path and parentId via body, not uri @@ -413,7 +420,9 @@ func setupRepos(r chi.Router, }) r.Route("/raw", func(r chi.Router) { - r.Get("/*", handlerrepo.HandleRaw(repoCtrl)) + r.With( + usage.Middleware(usageSender, false), + ).Get("/*", handlerrepo.HandleRaw(repoCtrl)) }) // commit operations @@ -464,7 +473,9 @@ func setupRepos(r chi.Router, r.Get("/codeowners/validate", handlerrepo.HandleCodeOwnersValidate(repoCtrl)) - r.Get(fmt.Sprintf("/archive/%s", request.PathParamArchiveGitRef), handlerrepo.HandleArchive(repoCtrl)) + r.With( + usage.Middleware(usageSender, false), + ).Get(fmt.Sprintf("/archive/%s", request.PathParamArchiveGitRef), handlerrepo.HandleArchive(repoCtrl)) SetupPullReq(r, pullreqCtrl) diff --git a/app/router/git.go b/app/router/git.go index d1d3a8b4a..5f38663f1 100644 --- a/app/router/git.go +++ b/app/router/git.go @@ -27,6 +27,7 @@ import ( "github.com/harness/gitness/app/api/middleware/logging" "github.com/harness/gitness/app/api/request" "github.com/harness/gitness/app/auth/authn" + "github.com/harness/gitness/app/services/usage" "github.com/harness/gitness/app/url" "github.com/harness/gitness/types" "github.com/harness/gitness/types/check" @@ -43,6 +44,7 @@ func NewGitHandler( urlProvider url.Provider, authenticator authn.Authenticator, repoCtrl *repo.Controller, + usageSender usage.Sender, ) http.Handler { // maxRepoDepth depends on config maxRepoDepth := check.MaxRepoPathDepth @@ -79,7 +81,9 @@ func NewGitHandler( r.Use(middlewareauthz.BlockSessionToken) // smart protocol - r.Post("/git-upload-pack", handlerrepo.HandleGitServicePack( + r.With( + usage.Middleware(usageSender, false), + ).Post("/git-upload-pack", handlerrepo.HandleGitServicePack( enum.GitServiceTypeUploadPack, repoCtrl, urlProvider)) r.Post("/git-receive-pack", handlerrepo.HandleGitServicePack( enum.GitServiceTypeReceivePack, repoCtrl, urlProvider)) diff --git a/app/router/wire.go b/app/router/wire.go index 83620c01b..681f12898 100644 --- a/app/router/wire.go +++ b/app/router/wire.go @@ -47,6 +47,7 @@ import ( "github.com/harness/gitness/app/api/controller/webhook" "github.com/harness/gitness/app/api/openapi" "github.com/harness/gitness/app/auth/authn" + "github.com/harness/gitness/app/services/usage" "github.com/harness/gitness/app/url" "github.com/harness/gitness/git" "github.com/harness/gitness/registry/app/api" @@ -112,6 +113,7 @@ func ProvideRouter( urlProvider url.Provider, openapi openapi.Service, registryRouter router.AppRouter, + usageSender usage.Sender, ) *Router { routers := make([]Interface, 4) @@ -121,6 +123,7 @@ func ProvideRouter( urlProvider, authenticator, repoCtrl, + usageSender, ) routers[0] = NewGitRouter(gitHandler, gitRoutingHost) routers[1] = router.NewRegistryRouter(registryRouter) @@ -130,7 +133,7 @@ func ProvideRouter( authenticator, repoCtrl, repoSettingsCtrl, executionCtrl, logCtrl, spaceCtrl, pipelineCtrl, secretCtrl, triggerCtrl, connectorCtrl, templateCtrl, pluginCtrl, pullreqCtrl, webhookCtrl, githookCtrl, git, saCtrl, userCtrl, principalCtrl, userGroupCtrl, checkCtrl, sysCtrl, blobCtrl, searchCtrl, - infraProviderCtrl, migrateCtrl, gitspaceCtrl, aiagentCtrl, capabilitiesCtrl) + infraProviderCtrl, migrateCtrl, gitspaceCtrl, aiagentCtrl, capabilitiesCtrl, usageSender) routers[2] = NewAPIRouter(apiHandler) webHandler := NewWebHandler(config, authenticator, openapi) diff --git a/app/services/usage/config.go b/app/services/usage/config.go new file mode 100644 index 000000000..9a14e9346 --- /dev/null +++ b/app/services/usage/config.go @@ -0,0 +1,51 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "github.com/harness/gitness/types" + + "github.com/alecthomas/units" +) + +type Config struct { + ChunkSize int64 + MaxWorkers int +} + +func NewConfig(global *types.Config) Config { + var err error + var n units.Base2Bytes + cfg := Config{ + MaxWorkers: global.UsageMetrics.MaxWorkers, + } + + if cfg.MaxWorkers == 0 { + cfg.MaxWorkers = 50 + } + + chunkSize := global.UsageMetrics.ChunkSize + if chunkSize == "" { + chunkSize = "10MiB" + } + + n, err = units.ParseBase2Bytes(chunkSize) + if err != nil { + panic(err) + } + cfg.ChunkSize = int64(n) + + return cfg +} diff --git a/app/services/usage/interface.go b/app/services/usage/interface.go new file mode 100644 index 000000000..0d8ac820f --- /dev/null +++ b/app/services/usage/interface.go @@ -0,0 +1,21 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import "context" + +type Sender interface { + Send(ctx context.Context, payload Metric) error +} diff --git a/app/services/usage/io.go b/app/services/usage/io.go new file mode 100644 index 000000000..0ac173a01 --- /dev/null +++ b/app/services/usage/io.go @@ -0,0 +1,123 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "context" + "io" + "net/http" +) + +type writeCounter struct { + ctx context.Context + w http.ResponseWriter + spaceRef string + intf Sender + isStorage bool +} + +func newWriter( + ctx context.Context, + w http.ResponseWriter, + spaceRef string, + intf Sender, + isStorage bool, +) *writeCounter { + return &writeCounter{ + ctx: ctx, + w: w, + spaceRef: spaceRef, + intf: intf, + isStorage: isStorage, + } +} + +func (c *writeCounter) Write(data []byte) (n int, err error) { + n, err = c.w.Write(data) + + m := Metric{ + SpaceRef: c.spaceRef, + Size: Size{ + Bandwidth: int64(n), + }, + } + if c.isStorage { + m.Storage = int64(n) + } + + sendErr := c.intf.Send(c.ctx, m) + if sendErr != nil { + return n, sendErr + } + + return n, err +} + +func (c *writeCounter) Header() http.Header { + return c.w.Header() +} + +func (c *writeCounter) WriteHeader(statusCode int) { + c.w.WriteHeader(statusCode) +} + +type readCounter struct { + ctx context.Context + r io.ReadCloser + spaceRef string + intf Sender + isStorage bool +} + +func newReader( + ctx context.Context, + r io.ReadCloser, + spaceRef string, + intf Sender, + isStorage bool, +) *readCounter { + return &readCounter{ + ctx: ctx, + r: r, + spaceRef: spaceRef, + intf: intf, + isStorage: isStorage, + } +} + +func (c *readCounter) Read(p []byte) (int, error) { + n, err := c.r.Read(p) + + m := Metric{ + SpaceRef: c.spaceRef, + Size: Size{ + Bandwidth: int64(n), + }, + } + if c.isStorage { + m.Storage = int64(n) + } + + sendErr := c.intf.Send(c.ctx, m) + if sendErr != nil { + return n, sendErr + } + + return n, err +} + +func (c *readCounter) Close() error { + return c.r.Close() +} diff --git a/app/services/usage/io_test.go b/app/services/usage/io_test.go new file mode 100644 index 000000000..50e5e1ef1 --- /dev/null +++ b/app/services/usage/io_test.go @@ -0,0 +1,103 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "bytes" + "context" + "io" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +func Test_writeCounter_Write(t *testing.T) { + size := 1 << 16 + var m Metric + mock := &mockInterface{ + SendFunc: func(_ context.Context, payload Metric) error { + m.Bandwidth += payload.Bandwidth + m.Storage += payload.Storage + return nil + }, + } + + // Create a buffer to hold the payload. + buffer := httptest.NewRecorder() + writer := newWriter( + context.Background(), + buffer, + spaceRef, + mock, + false, + ) + + expected := &bytes.Buffer{} + for i := 0; i < size; i += sampleLength { + if size-i < sampleLength { + // Write only the remaining characters to reach the exact size. + _, _ = writer.Write([]byte(sampleText[:size-i])) + expected.WriteString(sampleText[:size-i]) + break + } + _, _ = writer.Write([]byte(sampleText)) + expected.WriteString(sampleText) + } + + require.Equal(t, int64(size), m.Bandwidth, "expected %d, got %d", size, m.Bandwidth) + require.Equal(t, int64(0), m.Storage, "expected %d, got %d", size, m.Storage) + require.Equal(t, expected.Bytes(), buffer.Body.Bytes()) +} + +func Test_readCounter_Read(t *testing.T) { + size := 1 << 16 + var m Metric + mock := &mockInterface{ + SendFunc: func(_ context.Context, payload Metric) error { + m.Bandwidth += payload.Bandwidth + m.Storage += payload.Storage + return nil + }, + } + + buffer := &bytes.Buffer{} + reader := newReader( + context.Background(), + io.NopCloser(buffer), + spaceRef, + mock, + true, + ) + + for i := 0; i < size; i += sampleLength { + if size-i < sampleLength { + // Write only the remaining characters to reach the exact size. + buffer.WriteString(sampleText[:size-i]) + break + } + buffer.WriteString(sampleText) + } + + expected := buffer.Bytes() + got := &bytes.Buffer{} + + _, err := io.Copy(got, reader) + require.NoError(t, err) + + require.Equal(t, int64(size), m.Bandwidth, "expected %d, got %d", size, m.Bandwidth) + require.Equal(t, int64(size), m.Storage, "expected %d, got %d", size, m.Storage) + require.Equal(t, expected, got.Bytes()) +} diff --git a/app/services/usage/middleware.go b/app/services/usage/middleware.go new file mode 100644 index 000000000..a34718ab9 --- /dev/null +++ b/app/services/usage/middleware.go @@ -0,0 +1,59 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "net/http" + + "github.com/harness/gitness/app/api/request" + "github.com/harness/gitness/app/paths" + + "github.com/rs/zerolog/log" +) + +func Middleware(intf Sender, isStorage bool) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ref, err := request.GetRepoRefFromPath(r) + if err != nil { + log.Ctx(r.Context()).Warn().Err(err).Msg("unable to get space ref") + next.ServeHTTP(w, r) + return + } + rootSpace, _, err := paths.DisectRoot(ref) + if err != nil { + log.Ctx(r.Context()).Warn().Err(err).Msg("unable to get root space") + next.ServeHTTP(w, r) + return + } + writer := newWriter( + r.Context(), + w, + rootSpace, + intf, + isStorage, + ) + reader := newReader( + r.Context(), + r.Body, + rootSpace, + intf, + isStorage, + ) + r.Body = reader + next.ServeHTTP(writer, r) + }) + } +} diff --git a/app/services/usage/middleware_test.go b/app/services/usage/middleware_test.go new file mode 100644 index 000000000..1f96d68fb --- /dev/null +++ b/app/services/usage/middleware_test.go @@ -0,0 +1,87 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/harness/gitness/app/api/request" + + "github.com/go-chi/chi" + "github.com/stretchr/testify/require" +) + +func TestMiddleware(t *testing.T) { + var m Metric + mock := &mockInterface{ + SendFunc: func(_ context.Context, payload Metric) error { + m.Bandwidth += payload.Bandwidth + m.Storage += payload.Storage + return nil + }, + } + + r := chi.NewRouter() + r.Route(fmt.Sprintf("/testing/{%s}", request.PathParamRepoRef), func(r chi.Router) { + r.Use(Middleware(mock, false)) + r.Post("/", func(w http.ResponseWriter, r *http.Request) { + // read from body + _, _ = io.Copy(io.Discard, r.Body) + // write to response + _, _ = w.Write([]byte(sampleText)) + }) + }) + + ts := httptest.NewServer(r) + defer ts.Close() + + body := []byte(sampleText) + + _, _ = testRequest(t, ts, http.MethodPost, "/testing/"+spaceRef, bytes.NewReader(body)) + + // here we calculate upload/download so it is double size expected + require.Equal(t, int64(sampleLength*2), m.Bandwidth) +} + +func testRequest(t *testing.T, ts *httptest.Server, method, path string, body io.Reader) (*http.Response, string) { + t.Helper() + + req, err := http.NewRequest(method, ts.URL+path, body) + if err != nil { + t.Fatal(err) + return nil, "" + } + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + return nil, "" + } + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + return nil, "" + } + defer resp.Body.Close() + + return resp, string(respBody) +} diff --git a/app/services/usage/mocks.go b/app/services/usage/mocks.go new file mode 100644 index 000000000..dc34078a3 --- /dev/null +++ b/app/services/usage/mocks.go @@ -0,0 +1,105 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "context" + + "github.com/harness/gitness/types" +) + +const ( + sampleText = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789 " + sampleLength = len(sampleText) + spaceRef = "space1%2fspace2%2fspace3" +) + +type mockInterface struct { + SendFunc func( + ctx context.Context, + payload Metric, + ) error +} + +func (i *mockInterface) Send( + ctx context.Context, + payload Metric, +) error { + return i.SendFunc(ctx, payload) +} + +type SpaceStoreMock struct { + FindByRefFn func( + ctx context.Context, + spaceRef string, + ) (*types.Space, error) + FindByIDsFn func( + ctx context.Context, + ids ...int64, + ) ([]*types.Space, error) +} + +func (s *SpaceStoreMock) FindByRef( + ctx context.Context, + spaceRef string, +) (*types.Space, error) { + return s.FindByRefFn(ctx, spaceRef) +} + +func (s *SpaceStoreMock) FindByIDs( + ctx context.Context, + ids ...int64, +) ([]*types.Space, error) { + return s.FindByIDsFn(ctx, ids...) +} + +type MetricsMock struct { + UpsertOptimisticFn func(ctx context.Context, in *types.UsageMetric) error + GetMetricsFn func( + ctx context.Context, + rootSpaceID int64, + startDate int64, + endDate int64, + ) (*types.UsageMetric, error) + ListFn func( + ctx context.Context, + start int64, + end int64, + ) ([]types.UsageMetric, error) +} + +func (m *MetricsMock) GetMetrics( + ctx context.Context, + rootSpaceID int64, + startDate int64, + endDate int64, +) (*types.UsageMetric, error) { + return m.GetMetricsFn(ctx, rootSpaceID, startDate, endDate) +} + +func (m *MetricsMock) UpsertOptimistic( + ctx context.Context, + in *types.UsageMetric, +) error { + return m.UpsertOptimisticFn(ctx, in) +} + +func (m *MetricsMock) List( + ctx context.Context, + start int64, + end int64, +) ([]types.UsageMetric, error) { + return m.ListFn(ctx, start, end) +} diff --git a/app/services/usage/queue.go b/app/services/usage/queue.go new file mode 100644 index 000000000..5a4125e23 --- /dev/null +++ b/app/services/usage/queue.go @@ -0,0 +1,59 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import "context" + +type queue struct { + ch chan Metric +} + +func newQueue() *queue { + return &queue{ + ch: make(chan Metric, 1024), + } +} + +func (q *queue) Add(ctx context.Context, payload Metric) { + select { + case <-ctx.Done(): + return + case q.ch <- payload: + default: + // queue is full then wait in new go routine + // until one of consumer read from channel, + // we dont want to block caller goroutine + go func() { + q.ch <- payload + }() + } +} + +func (q *queue) Pop(ctx context.Context) (*Metric, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case payload := <-q.ch: + return &payload, nil + } +} + +func (q *queue) Close() { + close(q.ch) +} + +func (q *queue) Len() int { + return len(q.ch) +} diff --git a/app/services/usage/usage.go b/app/services/usage/usage.go new file mode 100644 index 000000000..b93f4b685 --- /dev/null +++ b/app/services/usage/usage.go @@ -0,0 +1,255 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "context" + "sync" + "time" + + "github.com/harness/gitness/types" + + "github.com/rs/zerolog/log" +) + +type Size struct { + Bandwidth int64 + Storage int64 +} + +type Metric struct { + SpaceRef string + Size +} + +type SpaceStore interface { + FindByRef(ctx context.Context, spaceRef string) (*types.Space, error) + FindByIDs(ctx context.Context, spaceIDs ...int64) ([]*types.Space, error) +} + +type Store interface { + UpsertOptimistic(ctx context.Context, in *types.UsageMetric) error + GetMetrics( + ctx context.Context, + rootSpaceID int64, + startDate int64, + endDate int64, + ) (*types.UsageMetric, error) + List( + ctx context.Context, + start int64, + end int64, + ) ([]types.UsageMetric, error) +} + +type LicenseFetcher interface { + Fetch(ctx context.Context, spaceID int64) (*Size, error) +} + +type Mediator struct { + queue *queue + + mux sync.RWMutex + chunks map[string]Size + spaces map[string]Size + workers []*worker + + spaceStore SpaceStore + usageMetricsStore Store + + wg sync.WaitGroup + + config Config +} + +func NewMediator( + ctx context.Context, + spaceStore SpaceStore, + usageMetricsStore Store, + config Config, +) *Mediator { + m := &Mediator{ + queue: newQueue(), + chunks: make(map[string]Size), + spaces: make(map[string]Size), + spaceStore: spaceStore, + usageMetricsStore: usageMetricsStore, + workers: make([]*worker, config.MaxWorkers), + config: config, + } + + m.initialize(ctx) + m.Start(ctx) + + return m +} + +func (m *Mediator) Start(ctx context.Context) { + for i := range m.workers { + w := newWorker(i, m.queue) + go w.start(ctx, m.process) + m.workers[i] = w + } +} + +func (m *Mediator) Stop() { + for i := range m.workers { + m.workers[i].stop() + } +} + +func (m *Mediator) Send(ctx context.Context, payload Metric) error { + m.wg.Add(1) + m.queue.Add(ctx, payload) + return nil +} + +func (m *Mediator) Wait() { + m.wg.Wait() +} + +func (m *Mediator) Size(space string) Size { + m.mux.RLock() + defer m.mux.RUnlock() + return m.spaces[space] +} + +// initialize will load when app is started all metrics for last 30 days. +func (m *Mediator) initialize(ctx context.Context) { + m.mux.Lock() + defer m.mux.Unlock() + + now := time.Now() + + metrics, err := m.usageMetricsStore.List(ctx, now.Add(-m.days30()).UnixMilli(), now.UnixMilli()) + if err != nil { + log.Ctx(ctx).Err(err).Msg("failed to list usage metrics") + return + } + + ids := make([]int64, len(metrics)) + values := make(map[int64]Size, len(metrics)) + for i, metric := range metrics { + ids[i] = metric.RootSpaceID + values[metric.RootSpaceID] = Size{ + Bandwidth: metric.Bandwidth, + Storage: metric.Storage, + } + } + + spaces, err := m.spaceStore.FindByIDs(ctx, ids...) + if err != nil { + log.Ctx(ctx).Err(err).Msg("failed to find spaces by id") + } + + for _, space := range spaces { + m.spaces[space.Identifier] = values[space.ID] + } +} + +func (m *Mediator) days30() time.Duration { + return time.Duration(30*24) * time.Hour +} + +func (m *Mediator) process(ctx context.Context, payload *Metric) { + defer m.wg.Done() + + m.mux.Lock() + defer m.mux.Unlock() + + size := m.chunks[payload.SpaceRef] + m.chunks[payload.SpaceRef] = Size{ + Bandwidth: size.Bandwidth + payload.Size.Bandwidth, + Storage: size.Storage + payload.Size.Storage, + } + + newSize := m.chunks[payload.SpaceRef] + + if newSize.Bandwidth < m.config.ChunkSize && newSize.Storage < m.config.ChunkSize { + return + } + + space, err := m.spaceStore.FindByRef(ctx, payload.SpaceRef) + if err != nil { + log.Ctx(ctx).Err(err).Msg("failed to find space") + return + } + + if err = m.usageMetricsStore.UpsertOptimistic(ctx, &types.UsageMetric{ + RootSpaceID: space.ID, + Bandwidth: newSize.Bandwidth, + Storage: newSize.Storage, + }); err != nil { + log.Ctx(ctx).Err(err).Msg("failed to upsert usage metrics") + } + + m.chunks[payload.SpaceRef] = Size{ + Bandwidth: 0, + Storage: 0, + } + + now := time.Now() + + metric, err := m.usageMetricsStore.GetMetrics(ctx, space.ID, now.Add(-m.days30()).UnixMilli(), now.UnixMilli()) + if err != nil { + log.Ctx(ctx).Err(err).Msg("failed to get usage metrics") + return + } + + m.spaces[space.Identifier] = Size{ + Bandwidth: metric.Bandwidth, + Storage: metric.Storage, + } +} + +type worker struct { + id int + queue *queue + stopCh chan struct{} +} + +func newWorker(id int, queue *queue) *worker { + return &worker{ + id: id, + queue: queue, + stopCh: make(chan struct{}), + } +} + +func (w *worker) start(ctx context.Context, fn func(context.Context, *Metric)) { + log.Ctx(ctx).Info().Int("worker", w.id).Msg("starting worker") + for { + select { + case <-ctx.Done(): + log.Ctx(ctx).Err(ctx.Err()).Msg("context canceled") + return + case <-w.stopCh: + log.Ctx(ctx).Warn().Int("worker", w.id).Msg("worker is stopped") + return + default: + payload, err := w.queue.Pop(ctx) + if err != nil { + log.Ctx(ctx).Err(err).Msg("failed to consume the queue") + return + } + fn(ctx, payload) + } + } +} + +func (w *worker) stop() { + defer close(w.stopCh) + w.stopCh <- struct{}{} +} diff --git a/app/services/usage/usage_test.go b/app/services/usage/usage_test.go new file mode 100644 index 000000000..7b31a2d77 --- /dev/null +++ b/app/services/usage/usage_test.go @@ -0,0 +1,113 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + + "github.com/harness/gitness/types" + + "github.com/stretchr/testify/require" +) + +func TestMediator_basic(t *testing.T) { + space := &types.Space{ + ID: 1, + Identifier: "space", + } + spaceMock := &SpaceStoreMock{ + FindByRefFn: func(context.Context, string) (*types.Space, error) { + return space, nil + }, + FindByIDsFn: func(context.Context, ...int64) ([]*types.Space, error) { + return []*types.Space{space}, nil + }, + } + initialBandwidth := int64(1024) + initialStorage := int64(1024) + bandwidth := atomic.Int64{} + storage := atomic.Int64{} + counter := atomic.Int64{} + usageMock := &MetricsMock{ + UpsertOptimisticFn: func(_ context.Context, in *types.UsageMetric) error { + if in.RootSpaceID != space.ID { + return fmt.Errorf("expected root space id to be %d, got %d", space.ID, in.RootSpaceID) + } + bandwidth.Add(in.Bandwidth) + storage.Add(in.Storage) + counter.Add(1) + return nil + }, + GetMetricsFn: func( + context.Context, + int64, // spaceID + int64, // startDate + int64, // endDate + ) (*types.UsageMetric, error) { + return &types.UsageMetric{ + Bandwidth: bandwidth.Load(), + Storage: storage.Load(), + }, nil + }, + ListFn: func(context.Context, int64, int64) ([]types.UsageMetric, error) { + bandwidth.Add(initialBandwidth) + storage.Add(initialStorage) + return []types.UsageMetric{ + { + RootSpaceID: space.ID, + Bandwidth: initialBandwidth, + Storage: initialStorage, + }, + }, nil + }, + } + + numRoutines := 10 + defaultSize := 512 + mediator := NewMediator( + context.Background(), + spaceMock, + usageMock, + Config{ + ChunkSize: 1024, + MaxWorkers: 10, + }, + ) + wg := sync.WaitGroup{} + for range numRoutines { + wg.Add(1) + go func() { + defer wg.Done() + _ = mediator.Send(context.Background(), Metric{ + SpaceRef: space.Identifier, + Size: Size{ + Bandwidth: int64(defaultSize), + Storage: int64(defaultSize), + }, + }) + }() + } + + wg.Wait() + mediator.Wait() + + require.Equal(t, int64(numRoutines*defaultSize/int(mediator.config.ChunkSize)), counter.Load()) + require.Equal(t, initialBandwidth+int64(numRoutines*defaultSize), bandwidth.Load()) + require.Equal(t, initialStorage+int64(numRoutines*defaultSize), storage.Load()) +} diff --git a/app/services/usage/wire.go b/app/services/usage/wire.go new file mode 100644 index 000000000..bf4c70a08 --- /dev/null +++ b/app/services/usage/wire.go @@ -0,0 +1,42 @@ +// Copyright 2023 Harness, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package usage + +import ( + "context" + + "github.com/harness/gitness/app/store" + "github.com/harness/gitness/types" + + "github.com/google/wire" +) + +var WireSet = wire.NewSet( + ProvideMediator, +) + +func ProvideMediator( + ctx context.Context, + config *types.Config, + spaceStore store.SpaceStore, + metricsStore store.UsageMetricStore, +) Sender { + return NewMediator( + ctx, + spaceStore, + metricsStore, + NewConfig(config), + ) +} diff --git a/app/store/database.go b/app/store/database.go index 73ffb0e3e..052b4af44 100644 --- a/app/store/database.go +++ b/app/store/database.go @@ -169,6 +169,9 @@ type ( // Find the space by id. Find(ctx context.Context, id int64) (*types.Space, error) + // FindByIDs finds all spaces with specified ids. + FindByIDs(ctx context.Context, ids ...int64) ([]*types.Space, error) + // FindByRef finds the space using the spaceRef as either the id or the space path. FindByRef(ctx context.Context, spaceRef string) (*types.Space, error) @@ -1276,11 +1279,17 @@ type ( UsageMetricStore interface { Upsert(ctx context.Context, in *types.UsageMetric) error + UpsertOptimistic(ctx context.Context, in *types.UsageMetric) error GetMetrics( ctx context.Context, rootSpaceID int64, startDate int64, endDate int64, ) (*types.UsageMetric, error) + List( + ctx context.Context, + start int64, + end int64, + ) ([]types.UsageMetric, error) } ) diff --git a/app/store/database/space.go b/app/store/database/space.go index 7f1e6b53a..429c104a2 100644 --- a/app/store/database/space.go +++ b/app/store/database/space.go @@ -93,6 +93,28 @@ func (s *SpaceStore) Find(ctx context.Context, id int64) (*types.Space, error) { return s.find(ctx, id, nil) } +// FindByIDs finds all spaces by ids. +func (s *SpaceStore) FindByIDs(ctx context.Context, ids ...int64) ([]*types.Space, error) { + stmt := database.Builder. + Select(spaceColumns). + From("spaces"). + Where(squirrel.Eq{"space_id": ids}) + + sql, args, err := stmt.ToSql() + if err != nil { + return nil, errors.Wrap(err, "Failed to convert query to sql") + } + + db := dbtx.GetAccessor(ctx, s.db) + + var dst []*space + if err = db.SelectContext(ctx, &dst, sql, args...); err != nil { + return nil, database.ProcessSQLErrorf(ctx, err, "Failed executing custom list query") + } + + return s.mapToSpaces(ctx, s.db, dst) +} + func (s *SpaceStore) find(ctx context.Context, id int64, deletedAt *int64) (*types.Space, error) { stmt := database.Builder. Select(spaceColumns). diff --git a/app/store/database/space_test.go b/app/store/database/space_test.go index bd5844f90..fd528d6c7 100644 --- a/app/store/database/space_test.go +++ b/app/store/database/space_test.go @@ -17,6 +17,8 @@ package database_test import ( "context" "testing" + + "github.com/stretchr/testify/require" ) func TestDatabase_GetRootSpace(t *testing.T) { @@ -41,3 +43,24 @@ func TestDatabase_GetRootSpace(t *testing.T) { } } } + +func TestSpaceStore_FindByIDs(t *testing.T) { + db, teardown := setupDB(t) + defer teardown() + + principalStore, spaceStore, spacePathStore, _ := setupStores(t, db) + + ctx := context.Background() + + createUser(ctx, t, principalStore) + + _ = createNestedSpaces(ctx, t, spaceStore, spacePathStore) + + spaces, err := spaceStore.FindByIDs(ctx, 4, 5, 6) + require.NoError(t, err) + + require.Len(t, spaces, 3) + require.Equal(t, int64(4), spaces[0].ID) + require.Equal(t, int64(5), spaces[1].ID) + require.Equal(t, int64(6), spaces[2].ID) +} diff --git a/app/store/database/wire.go b/app/store/database/wire.go index 051f407c7..22cbfd8ee 100644 --- a/app/store/database/wire.go +++ b/app/store/database/wire.go @@ -73,6 +73,7 @@ var WireSet = wire.NewSet( ProvidePullReqLabelStore, ProvideInfraProviderTemplateStore, ProvideInfraProvisionedStore, + ProvideUsageMetricStore, ) // migrator is helper function to set up the database by performing automated @@ -349,3 +350,7 @@ func ProvideInfraProviderTemplateStore(db *sqlx.DB) store.InfraProviderTemplateS func ProvideInfraProvisionedStore(db *sqlx.DB) store.InfraProvisionedStore { return NewInfraProvisionedStore(db) } + +func ProvideUsageMetricStore(db *sqlx.DB) store.UsageMetricStore { + return NewUsageMetricsStore(db) +} diff --git a/cmd/gitness/wire.go b/cmd/gitness/wire.go index f0edf78c7..225bb8b78 100644 --- a/cmd/gitness/wire.go +++ b/cmd/gitness/wire.go @@ -111,6 +111,7 @@ import ( "github.com/harness/gitness/app/services/settings" systemsvc "github.com/harness/gitness/app/services/system" "github.com/harness/gitness/app/services/trigger" + "github.com/harness/gitness/app/services/usage" usergroupservice "github.com/harness/gitness/app/services/usergroup" "github.com/harness/gitness/app/services/webhook" "github.com/harness/gitness/app/sse" @@ -283,6 +284,7 @@ func initSystem(ctx context.Context, config *types.Config) (*cliserver.System, e containerUser.WireSet, messagingservice.WireSet, runarg.WireSet, + usage.WireSet, ) return &cliserver.System{}, nil } diff --git a/cmd/gitness/wire_gen.go b/cmd/gitness/wire_gen.go index 7592b916f..7b7bb1414 100644 --- a/cmd/gitness/wire_gen.go +++ b/cmd/gitness/wire_gen.go @@ -102,6 +102,7 @@ import ( "github.com/harness/gitness/app/services/settings" system2 "github.com/harness/gitness/app/services/system" trigger2 "github.com/harness/gitness/app/services/trigger" + "github.com/harness/gitness/app/services/usage" "github.com/harness/gitness/app/services/usergroup" "github.com/harness/gitness/app/services/webhook" "github.com/harness/gitness/app/sse" @@ -338,7 +339,8 @@ func initSystem(ctx context.Context, config *types.Config) (*server.System, erro resolverFactory := secret.ProvideResolverFactory(passwordResolver) orchestratorOrchestrator := orchestrator.ProvideOrchestrator(scmSCM, platformConnector, infraProviderResourceStore, infraProvisioner, containerOrchestrator, eventsReporter, orchestratorConfig, ideFactory, resolverFactory) gitspaceService := gitspace.ProvideGitspace(transactor, gitspaceConfigStore, gitspaceInstanceStore, eventsReporter, gitspaceEventStore, spaceStore, infraproviderService, orchestratorOrchestrator, scmSCM, config) - spaceController := space.ProvideController(config, transactor, provider, streamer, spaceIdentifier, authorizer, spacePathStore, pipelineStore, secretStore, connectorStore, templateStore, spaceStore, repoStore, principalStore, repoController, membershipStore, listService, repository, exporterRepository, resourceLimiter, publicaccessService, auditService, gitspaceService, labelService, instrumentService, executionStore, rulesService) + usageMetricStore := database.ProvideUsageMetricStore(db) + spaceController := space.ProvideController(config, transactor, provider, streamer, spaceIdentifier, authorizer, spacePathStore, pipelineStore, secretStore, connectorStore, templateStore, spaceStore, repoStore, principalStore, repoController, membershipStore, listService, repository, exporterRepository, resourceLimiter, publicaccessService, auditService, gitspaceService, labelService, instrumentService, executionStore, rulesService, usageMetricStore) reporter3, err := events5.ProvideReporter(eventsSystem) if err != nil { return nil, err @@ -478,7 +480,8 @@ func initSystem(ctx context.Context, config *types.Config) (*server.System, erro cleanupPolicyRepository := database2.ProvideCleanupPolicyDao(db, transactor) apiHandler := router.APIHandlerProvider(registryRepository, upstreamProxyConfigRepository, tagRepository, manifestRepository, cleanupPolicyRepository, imageRepository, storageDriver, spaceStore, transactor, authenticator, provider, authorizer, auditService, spacePathStore) appRouter := router.AppRouterProvider(registryOCIHandler, apiHandler) - routerRouter := router2.ProvideRouter(ctx, config, authenticator, repoController, reposettingsController, executionController, logsController, spaceController, pipelineController, secretController, triggerController, connectorController, templateController, pluginController, pullreqController, webhookController, githookController, gitInterface, serviceaccountController, controller, principalController, usergroupController, checkController, systemController, uploadController, keywordsearchController, infraproviderController, gitspaceController, migrateController, aiagentController, capabilitiesController, provider, openapiService, appRouter) + sender := usage.ProvideMediator(ctx, config, spaceStore, usageMetricStore) + routerRouter := router2.ProvideRouter(ctx, config, authenticator, repoController, reposettingsController, executionController, logsController, spaceController, pipelineController, secretController, triggerController, connectorController, templateController, pluginController, pullreqController, webhookController, githookController, gitInterface, serviceaccountController, controller, principalController, usergroupController, checkController, systemController, uploadController, keywordsearchController, infraproviderController, gitspaceController, migrateController, aiagentController, capabilitiesController, provider, openapiService, appRouter, sender) serverServer := server2.ProvideServer(config, routerRouter) publickeyService := publickey.ProvidePublicKey(publicKeyStore, principalInfoCache) sshServer := ssh.ProvideServer(config, publickeyService, repoController) diff --git a/go.mod b/go.mod index 27eb16258..fc54e52fe 100644 --- a/go.mod +++ b/go.mod @@ -14,6 +14,7 @@ require ( github.com/distribution/reference v0.6.0 github.com/docker/docker v27.1.1+incompatible github.com/docker/go-connections v0.5.0 + github.com/docker/go-units v0.5.0 github.com/drone-runners/drone-runner-docker v1.8.4-0.20240815103043-c6c3a3e33ce3 github.com/drone/drone-go v1.7.1 github.com/drone/drone-yaml v1.2.3 @@ -106,7 +107,6 @@ require ( github.com/charmbracelet/lipgloss v0.12.1 // indirect github.com/charmbracelet/x/ansi v0.1.4 // indirect github.com/docker/distribution v2.8.2+incompatible // indirect - github.com/docker/go-units v0.5.0 // indirect github.com/drone/envsubst v1.0.3 // indirect github.com/fatih/semgroup v1.2.0 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect @@ -184,7 +184,7 @@ require ( cloud.google.com/go/profiler v0.3.1 github.com/Microsoft/go-winio v0.6.1 // indirect github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 // indirect - github.com/alecthomas/units v0.0.0-20231202071711-9a357b53e9c9 // indirect + github.com/alecthomas/units v0.0.0-20240927000941-0f3dac36c52b // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect diff --git a/go.sum b/go.sum index ca9c61ba8..e6075db4c 100644 --- a/go.sum +++ b/go.sum @@ -50,8 +50,8 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 h1:JYp7IbQjafo github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc= github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0= github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0= -github.com/alecthomas/units v0.0.0-20231202071711-9a357b53e9c9 h1:ez/4by2iGztzR4L0zgAOR8lTQK9VlyBVVd7G4omaOQs= -github.com/alecthomas/units v0.0.0-20231202071711-9a357b53e9c9/go.mod h1:OMCwj8VM1Kc9e19TLln2VL61YJF0x1XFtfdL4JdbSyE= +github.com/alecthomas/units v0.0.0-20240927000941-0f3dac36c52b h1:mimo19zliBX/vSQ6PWWSL9lK8qwHozUj03+zLoEB8O0= +github.com/alecthomas/units v0.0.0-20240927000941-0f3dac36c52b/go.mod h1:fvzegU4vN3H1qMT+8wDmzjAcDONcgo2/SZ/TyfdUOFs= github.com/andybalholm/brotli v1.0.5 h1:8uQZIdzKmjc/iuPu7O2ioW48L81FgatrcpfFmiq/cCs= github.com/andybalholm/brotli v1.0.5/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8= diff --git a/types/config.go b/types/config.go index d986c5fa1..72934ad72 100644 --- a/types/config.go +++ b/types/config.go @@ -502,4 +502,9 @@ type Config struct { Enable bool `envconfig:"GITNESS_INSTRUMENTATION_ENABLE" default:"false"` Cron string `envconfig:"GITNESS_INSTRUMENTATION_CRON" default:"0 0 * * *"` } + + UsageMetrics struct { + ChunkSize string `envconfig:"GITNESS_USAGE_METRICS_CHUNK_SIZE" default:"10MiB"` + MaxWorkers int `envconfig:"GITNESS_USAGE_METRICS_MAX_WORKERS" default:"50"` + } }