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

@@ -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)
 }
 

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)

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.

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")

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
 }

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."

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", "")

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"`

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
+}

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")

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

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
+}

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 {

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)
+}

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)

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()