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

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