From 8f912b697e2e51a0b26c1f15898ef7700fb7b2ab Mon Sep 17 00:00:00 2001 From: Ashish Bhate Date: Fri, 3 Jun 2022 11:48:29 +0530 Subject: [PATCH] MM-43918: freemium limit for storage (#20225) Summary Adds an API end point to return storage usage Ticket Link https://mattermost.atlassian.net/browse/MM-43918 --- api4/usage.go | 23 +++++++++++- api4/usage_test.go | 14 ++++++++ app/app_iface.go | 5 ++- app/opentracing/opentracing_layer.go | 22 ++++++++++++ app/usage.go | 13 +++++-- i18n/en.json | 4 +++ model/client4.go | 15 +++++++- model/usage.go | 4 +++ store/localcachelayer/file_info_layer.go | 30 ++++++++++++++++ store/opentracinglayer/opentracinglayer.go | 18 ++++++++++ store/retrylayer/retrylayer.go | 21 +++++++++++ store/sqlstore/file_info_store.go | 17 +++++++++ store/store.go | 1 + store/storetest/file_info_store.go | 42 ++++++++++++++++++++++ store/storetest/mocks/FileInfoStore.go | 21 +++++++++++ store/timerlayer/timerlayer.go | 16 +++++++++ 16 files changed, 261 insertions(+), 5 deletions(-) diff --git a/api4/usage.go b/api4/usage.go index 9424331927..a1d86a5777 100644 --- a/api4/usage.go +++ b/api4/usage.go @@ -8,14 +8,16 @@ import ( "net/http" "github.com/mattermost/mattermost-server/v6/model" + "github.com/mattermost/mattermost-server/v6/utils" ) func (api *API) InitUsage() { // GET /api/v4/usage/posts api.BaseRoutes.Usage.Handle("/posts", api.APISessionRequired(getPostsUsage)).Methods("GET") + // GET /api/v4/usage/storage + api.BaseRoutes.Usage.Handle("/storage", api.APISessionRequired(getStorageUsage)).Methods("GET") // GET /api/v4/usage/teams api.BaseRoutes.Usage.Handle("/teams", api.APISessionRequired(getTeamsUsage)).Methods("GET") - // GET /api/v4/usage/integrations api.BaseRoutes.Usage.Handle("/integrations", api.APISessionRequired(getIntegrationsUsage)).Methods("GET") } @@ -36,6 +38,23 @@ func getPostsUsage(c *Context, w http.ResponseWriter, r *http.Request) { w.Write(json) } +func getStorageUsage(c *Context, w http.ResponseWriter, r *http.Request) { + usage, appErr := c.App.GetStorageUsage() + if appErr != nil { + c.Err = model.NewAppError("Api4.getStorageUsage", "app.usage.get_storage_usage.app_error", nil, appErr.Error(), http.StatusInternalServerError) + return + } + + usage = utils.RoundOffToZeroes(float64(usage)) + json, err := json.Marshal(&model.StorageUsage{Bytes: usage}) + if err != nil { + c.Err = model.NewAppError("Api4.getStorageUsage", "api.marshal_error", nil, err.Error(), http.StatusInternalServerError) + return + } + + w.Write(json) +} + func getTeamsUsage(c *Context, w http.ResponseWriter, r *http.Request) { teamsUsage, appErr := c.App.GetTeamsUsage() if appErr != nil { @@ -50,7 +69,9 @@ func getTeamsUsage(c *Context, w http.ResponseWriter, r *http.Request) { json, err := json.Marshal(teamsUsage) if err != nil { c.Err = model.NewAppError("Api4.getTeamsUsage", "api.marshal_error", nil, err.Error(), http.StatusInternalServerError) + return } + w.Write(json) } diff --git a/api4/usage_test.go b/api4/usage_test.go index b8fe12e168..84067db6dd 100644 --- a/api4/usage_test.go +++ b/api4/usage_test.go @@ -50,6 +50,20 @@ func TestGetPostsUsage(t *testing.T) { }) } +func TestGetStorageUsage(t *testing.T) { + t.Run("unauthenticated users cannot access", func(t *testing.T) { + th := Setup(t) + defer th.TearDown() + + th.Client.Logout() + + usage, r, err := th.Client.GetStorageUsage() + assert.Error(t, err) + assert.Nil(t, usage) + assert.Equal(t, http.StatusUnauthorized, r.StatusCode) + }) +} + func TestGetTeamsUsage(t *testing.T) { t.Run("unauthenticated users can not access", func(t *testing.T) { th := Setup(t) diff --git a/app/app_iface.go b/app/app_iface.go index 0589b33e7c..2af4770b6a 100644 --- a/app/app_iface.go +++ b/app/app_iface.go @@ -191,7 +191,8 @@ type AppIface interface { // To get the plugins environment when the plugins are disabled, manually acquire the plugins // lock instead. GetPluginsEnvironment() *plugin.Environment - // GetPostsUsage returns "rounded off" total posts count like returns 900 instead of 987 + // GetPostsUsage returns the total posts count rounded down to the most + // significant digit GetPostsUsage() (int64, *model.AppError) // GetProductNotices is called from the frontend to fetch the product notices that are relevant to the caller GetProductNotices(c *request.Context, userID, teamID string, client model.NoticeClientType, clientVersion string, locale string) (model.NoticeMessages, *model.AppError) @@ -204,6 +205,8 @@ type AppIface interface { // GetSessionLengthInMillis returns the session length, in milliseconds, // based on the type of session (Mobile, SSO, Web/LDAP). GetSessionLengthInMillis(session *model.Session) int64 + // GetStorageUsage returns the sum of files' sizes stored on this instance + GetStorageUsage() (int64, *model.AppError) // GetSuggestions returns suggestions for user input. GetSuggestions(c *request.Context, commandArgs *model.CommandArgs, commands []*model.Command, roleID string) []model.AutocompleteSuggestion // GetTeamGroupUsers returns the users who are associated to the team via GroupTeams and GroupMembers. diff --git a/app/opentracing/opentracing_layer.go b/app/opentracing/opentracing_layer.go index 474ff4b9ae..545d3261dd 100644 --- a/app/opentracing/opentracing_layer.go +++ b/app/opentracing/opentracing_layer.go @@ -9095,6 +9095,28 @@ func (a *OpenTracingAppLayer) GetStatusesByIds(userIDs []string) (map[string]int return resultVar0, resultVar1 } +func (a *OpenTracingAppLayer) GetStorageUsage() (int64, *model.AppError) { + origCtx := a.ctx + span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.GetStorageUsage") + + a.ctx = newCtx + a.app.Srv().Store.SetContext(newCtx) + defer func() { + a.app.Srv().Store.SetContext(origCtx) + a.ctx = origCtx + }() + + defer span.Finish() + resultVar0, resultVar1 := a.app.GetStorageUsage() + + if resultVar1 != nil { + span.LogFields(spanlog.Error(resultVar1)) + ext.Error.Set(span, true) + } + + return resultVar0, resultVar1 +} + func (a *OpenTracingAppLayer) GetSuggestions(c *request.Context, commandArgs *model.CommandArgs, commands []*model.Command, roleID string) []model.AutocompleteSuggestion { origCtx := a.ctx span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.GetSuggestions") diff --git a/app/usage.go b/app/usage.go index 599c94eb71..17cf74edad 100644 --- a/app/usage.go +++ b/app/usage.go @@ -45,7 +45,8 @@ func (ch *Channels) getIntegrationsUsage() (*model.IntegrationsUsage, *model.App return &model.IntegrationsUsage{Enabled: count}, nil } -// GetPostsUsage returns "rounded off" total posts count like returns 900 instead of 987 +// GetPostsUsage returns the total posts count rounded down to the most +// significant digit func (a *App) GetPostsUsage() (int64, *model.AppError) { count, err := a.Srv().Store.Post().AnalyticsPostCount(&model.PostCountOptions{ExcludeDeleted: true, UsersPostsOnly: true, AllowFromCache: true}) if err != nil { @@ -55,6 +56,15 @@ func (a *App) GetPostsUsage() (int64, *model.AppError) { return utils.RoundOffToZeroes(float64(count)), nil } +// GetStorageUsage returns the sum of files' sizes stored on this instance +func (a *App) GetStorageUsage() (int64, *model.AppError) { + usage, err := a.Srv().Store.FileInfo().GetStorageUsage(true, false) + if err != nil { + return 0, model.NewAppError("GetStorageUsage", "app.usage.get_storage_usage.app_error", nil, err.Error(), http.StatusInternalServerError) + } + return usage, nil +} + func (a *App) GetTeamsUsage() (*model.TeamsUsage, *model.AppError) { usage := &model.TeamsUsage{} includeDeleted := false @@ -79,6 +89,5 @@ func (a *App) GetTeamsUsage() (*model.TeamsUsage, *model.AppError) { } usage.CloudArchived = int64(cloudArchivedTeamCount) - return usage, nil } diff --git a/i18n/en.json b/i18n/en.json index 858af5229e..5f3a1d4a73 100644 --- a/i18n/en.json +++ b/i18n/en.json @@ -6379,6 +6379,10 @@ "id": "app.upload.upload_data.update.app_error", "translation": "Failed to update the upload session." }, + { + "id": "app.usage.get_storage_usage.app_error", + "translation": "Failed to get storage usage." + }, { "id": "app.user.analytics_daily_active_users.app_error", "translation": "Unable to get the active users during the requested period." diff --git a/model/client4.go b/model/client4.go index 27d7bdb62a..8b15c31d8c 100644 --- a/model/client4.go +++ b/model/client4.go @@ -8122,7 +8122,20 @@ func (c *Client4) GetPostsUsage() (*PostsUsage, *Response, error) { return usage, BuildResponse(r), err } -// GetTeamsUsage returns total usage of teams for the instance +// GetStorageUsage returns the file storage usage for the instance, +// rounded down the most signigicant digit +func (c *Client4) GetStorageUsage() (*StorageUsage, *Response, error) { + r, err := c.DoAPIGet(c.usageRoute()+"/storage", "") + if err != nil { + return nil, BuildResponse(r), err + } + defer closeBody(r) + + var usage *StorageUsage + err = json.NewDecoder(r.Body).Decode(&usage) + return usage, BuildResponse(r), err +} + // GetTeamsUsage returns total usage of teams for the instance func (c *Client4) GetTeamsUsage() (*TeamsUsage, *Response, error) { r, err := c.DoAPIGet(c.usageRoute()+"/teams", "") diff --git a/model/usage.go b/model/usage.go index 0a5fb07866..6ee0c4383f 100644 --- a/model/usage.go +++ b/model/usage.go @@ -7,6 +7,10 @@ type PostsUsage struct { Count int64 `json:"count"` } +type StorageUsage struct { + Bytes int64 `json:"bytes"` +} + type TeamsUsage struct { Active int64 `json:"active"` CloudArchived int64 `json:"cloud_archived"` diff --git a/store/localcachelayer/file_info_layer.go b/store/localcachelayer/file_info_layer.go index e077802df5..30c0e6f9e7 100644 --- a/store/localcachelayer/file_info_layer.go +++ b/store/localcachelayer/file_info_layer.go @@ -67,3 +67,33 @@ func (s LocalCacheFileInfoStore) InvalidateFileInfosForPostCache(postId string, s.rootStore.metrics.IncrementMemCacheInvalidationCounter("File Info Cache - Remove by PostId") } } + +func (s LocalCacheFileInfoStore) GetStorageUsage(allowFromCache, includeDeleted bool) (int64, error) { + storageUsageKey := "storage_usage" + if includeDeleted { + storageUsageKey += "_deleted" + } + + if !allowFromCache { + usage, err := s.FileInfoStore.GetStorageUsage(allowFromCache, includeDeleted) + if err != nil { + return 0, err + } + + s.rootStore.doStandardAddToCache(s.rootStore.fileInfoCache, storageUsageKey, usage) + return usage, nil + } + + var usage int64 + if err := s.rootStore.doStandardReadCache(s.rootStore.fileInfoCache, storageUsageKey, &usage); err == nil { + return usage, nil + } + + usage, err := s.FileInfoStore.GetStorageUsage(allowFromCache, includeDeleted) + if err != nil { + return 0, err + } + + s.rootStore.doStandardAddToCache(s.rootStore.fileInfoCache, storageUsageKey, usage) + return usage, nil +} diff --git a/store/opentracinglayer/opentracinglayer.go b/store/opentracinglayer/opentracinglayer.go index 0bd0248d68..32361f664f 100644 --- a/store/opentracinglayer/opentracinglayer.go +++ b/store/opentracinglayer/opentracinglayer.go @@ -3426,6 +3426,24 @@ func (s *OpenTracingLayerFileInfoStore) GetFromMaster(id string) (*model.FileInf return result, err } +func (s *OpenTracingLayerFileInfoStore) GetStorageUsage(allowFromCache bool, includeDeleted bool) (int64, error) { + origCtx := s.Root.Store.Context() + span, newCtx := tracing.StartSpanWithParentByContext(s.Root.Store.Context(), "FileInfoStore.GetStorageUsage") + s.Root.Store.SetContext(newCtx) + defer func() { + s.Root.Store.SetContext(origCtx) + }() + + defer span.Finish() + result, err := s.FileInfoStore.GetStorageUsage(allowFromCache, includeDeleted) + if err != nil { + span.LogFields(spanlog.Error(err)) + ext.Error.Set(span, true) + } + + return result, err +} + func (s *OpenTracingLayerFileInfoStore) GetWithOptions(page int, perPage int, opt *model.GetFileInfosOptions) ([]*model.FileInfo, error) { origCtx := s.Root.Store.Context() span, newCtx := tracing.StartSpanWithParentByContext(s.Root.Store.Context(), "FileInfoStore.GetWithOptions") diff --git a/store/retrylayer/retrylayer.go b/store/retrylayer/retrylayer.go index 23291b5945..6f816283e5 100644 --- a/store/retrylayer/retrylayer.go +++ b/store/retrylayer/retrylayer.go @@ -3841,6 +3841,27 @@ func (s *RetryLayerFileInfoStore) GetFromMaster(id string) (*model.FileInfo, err } +func (s *RetryLayerFileInfoStore) GetStorageUsage(allowFromCache bool, includeDeleted bool) (int64, error) { + + tries := 0 + for { + result, err := s.FileInfoStore.GetStorageUsage(allowFromCache, includeDeleted) + if err == nil { + return result, nil + } + if !isRepeatableError(err) { + return result, err + } + tries++ + if tries >= 3 { + err = errors.Wrap(err, "giving up after 3 consecutive repeatable transaction failures") + return result, err + } + timepkg.Sleep(100 * timepkg.Millisecond) + } + +} + func (s *RetryLayerFileInfoStore) GetWithOptions(page int, perPage int, opt *model.GetFileInfosOptions) ([]*model.FileInfo, error) { tries := 0 diff --git a/store/sqlstore/file_info_store.go b/store/sqlstore/file_info_store.go index ea20a24743..682db165e4 100644 --- a/store/sqlstore/file_info_store.go +++ b/store/sqlstore/file_info_store.go @@ -736,3 +736,20 @@ func (fs SqlFileInfoStore) GetFilesBatchForIndexing(startTime int64, startFileID return files, nil } + +func (fs SqlFileInfoStore) GetStorageUsage(allowFromCache, includeDeleted bool) (int64, error) { + query := fs.getQueryBuilder(). + Select("SUM(Size)"). + From("FileInfo") + + if !includeDeleted { + query = query.Where("DeleteAt = 0") + } + + var size int64 + err := fs.GetReplicaX().GetBuilder(&size, query) + if err != nil { + return int64(0), errors.Wrap(err, "failed to get storage usage") + } + return size, nil +} diff --git a/store/store.go b/store/store.go index 89ede5d848..752f6c9030 100644 --- a/store/store.go +++ b/store/store.go @@ -667,6 +667,7 @@ type FileInfoStore interface { CountAll() (int64, error) GetFilesBatchForIndexing(startTime int64, startFileID string, limit int) ([]*model.FileForIndexing, error) ClearCaches() + GetStorageUsage(allowFromCache, includeDeleted bool) (int64, error) } type UploadSessionStore interface { diff --git a/store/storetest/file_info_store.go b/store/storetest/file_info_store.go index 21c36d394c..467db6bc34 100644 --- a/store/storetest/file_info_store.go +++ b/store/storetest/file_info_store.go @@ -29,6 +29,7 @@ func TestFileInfoStore(t *testing.T, ss store.Store) { t.Run("FileInfoPermanentDeleteByUser", func(t *testing.T) { testFileInfoPermanentDeleteByUser(t, ss) }) t.Run("GetFilesBatchForIndexing", func(t *testing.T) { testFileInfoStoreGetFilesBatchForIndexing(t, ss) }) t.Run("CountAll", func(t *testing.T) { testFileInfoStoreCountAll(t, ss) }) + t.Run("GetStorageUsage", func(t *testing.T) { testFileInfoGetStorageUsage(t, ss) }) } func testFileInfoSaveGet(t *testing.T, ss store.Store) { @@ -725,3 +726,44 @@ func testFileInfoStoreCountAll(t *testing.T, ss store.Store) { require.NoError(t, err) require.Equal(t, int64(2), count) } + +func testFileInfoGetStorageUsage(t *testing.T, ss store.Store) { + _, err := ss.FileInfo().PermanentDeleteBatch(model.GetMillis(), 100000) + require.NoError(t, err) + f1, err := ss.FileInfo().Save(&model.FileInfo{ + PostId: model.NewId(), + CreatorId: model.NewId(), + Size: 10, + Path: "file1.txt", + }) + require.NoError(t, err) + + _, err = ss.FileInfo().Save(&model.FileInfo{ + PostId: model.NewId(), + CreatorId: model.NewId(), + Size: 10, + Path: "file2.txt", + }) + require.NoError(t, err) + _, err = ss.FileInfo().Save(&model.FileInfo{ + PostId: model.NewId(), + CreatorId: model.NewId(), + Size: 10, + Path: "file3.txt", + }) + require.NoError(t, err) + + usage, err := ss.FileInfo().GetStorageUsage(false, false) + require.NoError(t, err) + require.Equal(t, int64(30), usage) + + _, err = ss.FileInfo().DeleteForPost(f1.PostId) + require.NoError(t, err) + usage, err = ss.FileInfo().GetStorageUsage(false, false) + require.NoError(t, err) + require.Equal(t, int64(20), usage) + + usage, err = ss.FileInfo().GetStorageUsage(false, true) + require.NoError(t, err) + require.Equal(t, int64(30), usage) +} diff --git a/store/storetest/mocks/FileInfoStore.go b/store/storetest/mocks/FileInfoStore.go index 0c629bf215..97432a11f3 100644 --- a/store/storetest/mocks/FileInfoStore.go +++ b/store/storetest/mocks/FileInfoStore.go @@ -236,6 +236,27 @@ func (_m *FileInfoStore) GetFromMaster(id string) (*model.FileInfo, error) { return r0, r1 } +// GetStorageUsage provides a mock function with given fields: allowFromCache, includeDeleted +func (_m *FileInfoStore) GetStorageUsage(allowFromCache bool, includeDeleted bool) (int64, error) { + ret := _m.Called(allowFromCache, includeDeleted) + + var r0 int64 + if rf, ok := ret.Get(0).(func(bool, bool) int64); ok { + r0 = rf(allowFromCache, includeDeleted) + } else { + r0 = ret.Get(0).(int64) + } + + var r1 error + if rf, ok := ret.Get(1).(func(bool, bool) error); ok { + r1 = rf(allowFromCache, includeDeleted) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + // GetWithOptions provides a mock function with given fields: page, perPage, opt func (_m *FileInfoStore) GetWithOptions(page int, perPage int, opt *model.GetFileInfosOptions) ([]*model.FileInfo, error) { ret := _m.Called(page, perPage, opt) diff --git a/store/timerlayer/timerlayer.go b/store/timerlayer/timerlayer.go index ee0cf504bb..937aba6bb5 100644 --- a/store/timerlayer/timerlayer.go +++ b/store/timerlayer/timerlayer.go @@ -3132,6 +3132,22 @@ func (s *TimerLayerFileInfoStore) GetFromMaster(id string) (*model.FileInfo, err return result, err } +func (s *TimerLayerFileInfoStore) GetStorageUsage(allowFromCache bool, includeDeleted bool) (int64, error) { + start := timemodule.Now() + + result, err := s.FileInfoStore.GetStorageUsage(allowFromCache, includeDeleted) + + elapsed := float64(timemodule.Since(start)) / float64(timemodule.Second) + if s.Root.Metrics != nil { + success := "false" + if err == nil { + success = "true" + } + s.Root.Metrics.ObserveStoreMethodDuration("FileInfoStore.GetStorageUsage", success, elapsed) + } + return result, err +} + func (s *TimerLayerFileInfoStore) GetWithOptions(page int, perPage int, opt *model.GetFileInfosOptions) ([]*model.FileInfo, error) { start := timemodule.Now()