Implement free storage in collStats, dbStats and aggregate `$co…

…llStats` (FerretDB#3594) Closes FerretDB#2342.
chilagrow · Oct 18, 2023 · e82aefb · e82aefb
1 parent f433e25
commit e82aefb
Show file tree

Hide file tree

Showing 16 changed files with 281 additions and 61 deletions.
diff --git a/integration/aggregate_stats_test.go b/integration/aggregate_stats_test.go
@@ -57,15 +57,13 @@ func TestAggregateCollStats(t *testing.T) {
 	assert.NotZero(t, must.NotFail(storageStats.Get("count")))
 	assert.NotZero(t, must.NotFail(storageStats.Get("avgObjSize")))
 	assert.NotZero(t, must.NotFail(storageStats.Get("storageSize")))
-	// TODO https://github.com/FerretDB/FerretDB/issues/2447
-	// assert.NotZero(t, must.NotFail(storageStats.Get("freeStorageSize")))
+	assert.Zero(t, must.NotFail(storageStats.Get("freeStorageSize")))
 	assert.Equal(t, false, must.NotFail(storageStats.Get("capped")))
 	assert.NotZero(t, must.NotFail(storageStats.Get("nindexes")))
 	assert.NotZero(t, must.NotFail(storageStats.Get("totalIndexSize")))
 	assert.NotZero(t, must.NotFail(storageStats.Get("totalSize")))
 	assert.NotZero(t, must.NotFail(storageStats.Get("indexSizes")))
-	// TODO https://github.com/FerretDB/FerretDB/issues/2447
-	// assert.Equal(t, int32(1), must.NotFail(storageStats.Get("scaleFactor")))
+	assert.Equal(t, int32(1), must.NotFail(storageStats.Get("scaleFactor")))
 }
 
 func TestAggregateCollStatsCommandErrors(t *testing.T) {

diff --git a/integration/commands_administration_compat_test.go b/integration/commands_administration_compat_test.go
@@ -221,3 +221,70 @@ func TestCommandsAdministrationCompatDBStatsWithScale(t *testing.T) {
 		})
 	}
 }
+
+func TestCommandsAdministrationCompatDBStatsFreeStorage(t *testing.T) {
+	t.Parallel()
+
+	s := setup.SetupCompatWithOpts(t, &setup.SetupCompatOpts{
+		Providers:                []shareddata.Provider{shareddata.DocumentsDocuments},
+		AddNonExistentCollection: true,
+	})
+
+	ctx, targetCollection, compatCollection := s.Ctx, s.TargetCollections[0], s.CompatCollections[0]
+
+	for name, tc := range map[string]struct { //nolint:vet // for readability
+		command bson.D // required, command to run
+		skip    string // optional, skip test with a specified reason
+	}{
+		"Unset": {
+			command: bson.D{{"dbStats", int32(1)}},
+		},
+		"Int32Zero": {
+			command: bson.D{{"dbStats", int32(1)}, {"freeStorage", int32(0)}},
+		},
+		"Int32One": {
+			command: bson.D{{"dbStats", int32(1)}, {"freeStorage", int32(1)}},
+		},
+		"Int32Negative": {
+			command: bson.D{{"dbStats", int32(1)}, {"freeStorage", int32(-1)}},
+		},
+		"True": {
+			command: bson.D{{"dbStats", int32(1)}, {"freeStorage", true}},
+		},
+		"False": {
+			command: bson.D{{"dbStats", int32(1)}, {"freeStorage", false}},
+		},
+		"Nil": {
+			command: bson.D{{"dbStats", int32(1)}, {"freeStorage", nil}},
+		},
+	} {
+		name, tc := name, tc
+
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			var targetRes bson.D
+			targetErr := targetCollection.Database().RunCommand(ctx, tc.command).Decode(&targetRes)
+
+			var compatRes bson.D
+			compatErr := compatCollection.Database().RunCommand(ctx, tc.command).Decode(&compatRes)
+
+			if targetErr != nil {
+				t.Logf("Target error: %v", targetErr)
+				t.Logf("Compat error: %v", compatErr)
+
+				// error messages are intentionally not compared
+				AssertMatchesCommandError(t, compatErr, targetErr)
+
+				return
+			}
+			require.NoError(t, compatErr, "compat error; target returned no error")
+
+			targetDoc := ConvertDocument(t, targetRes)
+			compatDoc := ConvertDocument(t, compatRes)
+
+			assert.Equal(t, compatDoc.Has("freeStorageSize"), targetDoc.Has("freeStorageSize"))
+			assert.Equal(t, compatDoc.Has("totalFreeStorageSize"), targetDoc.Has("totalFreeStorageSize"))
+		})
+	}
+}
diff --git a/integration/commands_administration_test.go b/integration/commands_administration_test.go
@@ -690,8 +690,7 @@ func TestCommandsAdministrationCollStatsEmpty(t *testing.T) {
 	assert.EqualValues(t, 0, must.NotFail(doc.Get("size")))
 	assert.EqualValues(t, 0, must.NotFail(doc.Get("count")))
 	assert.EqualValues(t, 0, must.NotFail(doc.Get("storageSize")))
-	// add assertion for freeStorageSize
-	// TODO https://github.com/FerretDB/FerretDB/issues/2447
+	assert.False(t, doc.Has("freeStorageSize"))
 	assert.EqualValues(t, 0, must.NotFail(doc.Get("nindexes")))
 	assert.EqualValues(t, 0, must.NotFail(doc.Get("totalIndexSize")))
 	assert.EqualValues(t, 0, must.NotFail(doc.Get("totalSize")))
@@ -727,8 +726,7 @@ func TestCommandsAdministrationCollStats(t *testing.T) {
 	assert.InDelta(t, 40_000, must.NotFail(doc.Get("size")), 39_900)
 	assert.InDelta(t, 2_400, must.NotFail(doc.Get("avgObjSize")), 2_370)
 	assert.InDelta(t, 40_000, must.NotFail(doc.Get("storageSize")), 39_900)
-	// add assertion for freeStorageSize
-	// TODO https://github.com/FerretDB/FerretDB/issues/2447
+	assert.EqualValues(t, 0, must.NotFail(doc.Get("freeStorageSize")))
 	assert.EqualValues(t, 1, must.NotFail(doc.Get("nindexes")))
 	assert.InDelta(t, 12_000, must.NotFail(doc.Get("totalIndexSize")), 11_000)
 	assert.InDelta(t, 32_000, must.NotFail(doc.Get("totalSize")), 30_000)
@@ -768,8 +766,7 @@ func TestCommandsAdministrationCollStatsWithScale(t *testing.T) {
 	assert.InDelta(t, 16, must.NotFail(doc.Get("size")), 16)
 	assert.InDelta(t, 2_400, must.NotFail(doc.Get("avgObjSize")), 2_370)
 	assert.InDelta(t, 24, must.NotFail(doc.Get("storageSize")), 24)
-	// add assertion for freeStorageSize
-	// TODO https://github.com/FerretDB/FerretDB/issues/2447
+	assert.Zero(t, must.NotFail(doc.Get("freeStorageSize")))
 	assert.EqualValues(t, 1, must.NotFail(doc.Get("nindexes")))
 	assert.InDelta(t, 8, must.NotFail(doc.Get("totalIndexSize")), 8)
 	assert.InDelta(t, 24, must.NotFail(doc.Get("totalSize")), 24)
@@ -1012,9 +1009,6 @@ func TestCommandsAdministrationDBStats(t *testing.T) {
 	freeStorageSize, _ := doc.Get("freeStorageSize")
 	assert.Nil(t, freeStorageSize)
 
-	indexFreeStorageSize, _ := doc.Get("indexFreeStorageSize")
-	assert.Nil(t, indexFreeStorageSize)
-
 	totalFreeStorageSize, _ := doc.Get("totalFreeStorageSize")
 	assert.Nil(t, totalFreeStorageSize)
 
@@ -1109,17 +1103,17 @@ func TestCommandsAdministrationDBStatsFreeStorage(t *testing.T) {
 
 	ctx, collection := setup.Setup(t, shareddata.DocumentsStrings)
 
-	var actual bson.D
+	var res bson.D
 	command := bson.D{{"dbStats", int32(1)}, {"freeStorage", int32(1)}}
-	err := collection.Database().RunCommand(ctx, command).Decode(&actual)
+	err := collection.Database().RunCommand(ctx, command).Decode(&res)
 	require.NoError(t, err)
 
-	doc := ConvertDocument(t, actual)
+	doc := ConvertDocument(t, res)
 
 	assert.Equal(t, float64(1), doc.Remove("scaleFactor"))
 	assert.Equal(t, float64(1), doc.Remove("ok"))
-	// assert freeStorageSize, indexFreeStorageSize and totalFreeStorageSize
-	// TODO https://github.com/FerretDB/FerretDB/issues/2447
+	assert.Zero(t, must.NotFail(doc.Get("freeStorageSize")))
+	assert.Zero(t, must.NotFail(doc.Get("totalFreeStorageSize")))
 }
 
 //nolint:paralleltest // we test a global server status

diff --git a/internal/backends/collection.go b/internal/backends/collection.go
@@ -240,11 +240,12 @@ type CollectionStatsParams struct {
 
 // CollectionStatsResult represents the results of Collection.Stats method.
 type CollectionStatsResult struct {
-	CountDocuments int64
-	SizeTotal      int64
-	SizeIndexes    int64
-	SizeCollection int64
-	IndexSizes     []IndexSize
+	CountDocuments  int64
+	SizeTotal       int64
+	SizeIndexes     int64
+	SizeCollection  int64
+	SizeFreeStorage int64
+	IndexSizes      []IndexSize
 }
 
 // IndexSize represents the name and the size of an index.

diff --git a/internal/backends/database.go b/internal/backends/database.go
@@ -207,6 +207,7 @@ type DatabaseStatsResult struct {
 	SizeTotal       int64
 	SizeIndexes     int64
 	SizeCollections int64
+	SizeFreeStorage int64
 }
 
 // Stats returns statistic estimations about the database.

diff --git a/internal/backends/postgresql/collection.go b/internal/backends/postgresql/collection.go
@@ -434,11 +434,12 @@ func (c *collection) Stats(ctx context.Context, params *backends.CollectionStats
 	}
 
 	return &backends.CollectionStatsResult{
-		CountDocuments: stats.countDocuments,
-		SizeTotal:      stats.sizeTables + stats.sizeIndexes,
-		SizeIndexes:    stats.sizeIndexes,
-		SizeCollection: stats.sizeTables,
-		IndexSizes:     indexSizes,
+		CountDocuments:  stats.countDocuments,
+		SizeTotal:       stats.sizeTables + stats.sizeIndexes,
+		SizeIndexes:     stats.sizeIndexes,
+		SizeCollection:  stats.sizeTables,
+		SizeFreeStorage: stats.sizeFreeStorage,
+		IndexSizes:      indexSizes,
 	}, nil
 }
 

diff --git a/internal/backends/postgresql/database.go b/internal/backends/postgresql/database.go
@@ -17,8 +17,6 @@ package postgresql
 import (
 	"context"
 
-	"github.com/AlekSi/pointer"
-
 	"github.com/FerretDB/FerretDB/internal/backends"
 	"github.com/FerretDB/FerretDB/internal/backends/postgresql/metadata"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
@@ -156,26 +154,23 @@ func (db *database) Stats(ctx context.Context, params *backends.DatabaseStatsPar
 	// See https://www.postgresql.org/docs/15/functions-admin.html#FUNCTIONS-ADMIN-DBOBJECT.
 	q := `
 		SELECT
-			SUM(pg_total_relation_size(quote_ident(schemaname) || '.' || quote_ident(tablename)))
+			COALESCE(SUM(pg_total_relation_size(quote_ident(schemaname) || '.' || quote_ident(tablename))), 0)
 		FROM pg_tables
 		WHERE schemaname = $1`
 	args := []any{db.name}
 	row := p.QueryRow(ctx, q, args...)
 
-	var schemaSize *int64
-	if err := row.Scan(&schemaSize); err != nil {
+	var sizeTotal int64
+	if err := row.Scan(&sizeTotal); err != nil {
 		return nil, lazyerrors.Error(err)
 	}
 
-	if schemaSize == nil {
-		schemaSize = pointer.ToInt64(0)
-	}
-
 	return &backends.DatabaseStatsResult{
 		CountDocuments:  stats.countDocuments,
-		SizeTotal:       *schemaSize,
+		SizeTotal:       sizeTotal,
 		SizeIndexes:     stats.sizeIndexes,
 		SizeCollections: stats.sizeTables,
+		SizeFreeStorage: stats.sizeFreeStorage,
 	}, nil
 }
 

diff --git a/internal/backends/postgresql/database_test.go b/internal/backends/postgresql/database_test.go
@@ -89,4 +89,29 @@ func TestDatabaseStats(t *testing.T) {
 		require.Equal(t, int64(1), res.CountDocuments)
 		require.NotZero(t, res.SizeIndexes)
 	})
+
+	t.Run("FreeStorageSize", func(t *testing.T) {
+		c, err := db.Collection(cNames[0])
+		require.NoError(t, err)
+
+		nInsert, deleteFromIndex, deleteToIndex := 50, 10, 40
+		ids := make([]any, nInsert)
+		toInsert := make([]*types.Document, nInsert)
+		for i := 0; i < nInsert; i++ {
+			ids[i] = types.NewObjectID()
+			toInsert[i] = must.NotFail(types.NewDocument("_id", ids[i], "v", "foo"))
+		}
+
+		_, err = c.InsertAll(ctx, &backends.InsertAllParams{Docs: toInsert})
+		require.NoError(t, err)
+
+		_, err = c.DeleteAll(ctx, &backends.DeleteAllParams{IDs: ids[deleteFromIndex:deleteToIndex]})
+		require.NoError(t, err)
+
+		res, err := db.Stats(ctx, new(backends.DatabaseStatsParams))
+		require.NoError(t, err)
+
+		t.Logf("freeStorage size: %d", res.SizeFreeStorage)
+		require.NotZero(t, res.SizeFreeStorage)
+	})
 }
diff --git a/internal/backends/postgresql/postgresql.go b/internal/backends/postgresql/postgresql.go
@@ -32,9 +32,10 @@ import (
 
 // stats represents information about statistics of tables and indexes.
 type stats struct {
-	countDocuments int64
-	sizeIndexes    int64
-	sizeTables     int64
+	countDocuments  int64
+	sizeIndexes     int64
+	sizeTables      int64
+	sizeFreeStorage int64
 }
 
 // collectionsStats returns statistics about tables and indexes for the given collections.
@@ -70,6 +71,8 @@ func collectionsStats(ctx context.Context, p *pgxpool.Pool, dbName string, list
 	// used, however it is not updated immediately after operation such as DELETE
 	// unless VACUUM is called, ANALYZE does not update pg_relation_size in this case.
 	//
+	// The free storage size is the size of free space map (fsm) of table relation.
+	//
 	// The smallest difference in size that `pg_relation_size` reports appears to be 8KB.
 	// Because of that inserting or deleting a single small object may not change the size.
 	//
@@ -82,6 +85,7 @@ func collectionsStats(ctx context.Context, p *pgxpool.Pool, dbName string, list
 		SELECT
 			COALESCE(SUM(c.reltuples), 0),
 			COALESCE(SUM(pg_relation_size(c.oid,'main')), 0),
+			COALESCE(SUM(pg_relation_size(c.oid,'fsm')), 0),
 			COALESCE(SUM(pg_indexes_size(c.oid)), 0)
 		FROM pg_tables AS t
 			LEFT JOIN pg_class AS c ON c.relname = t.tablename AND c.relnamespace = quote_ident(t.schemaname)::regnamespace
@@ -90,7 +94,7 @@ func collectionsStats(ctx context.Context, p *pgxpool.Pool, dbName string, list
 	)
 
 	row := p.QueryRow(ctx, q, args...)
-	if err := row.Scan(&s.countDocuments, &s.sizeTables, &s.sizeIndexes); err != nil {
+	if err = row.Scan(&s.countDocuments, &s.sizeTables, &s.sizeFreeStorage, &s.sizeIndexes); err != nil {
 		return nil, lazyerrors.Error(err)
 	}
 

diff --git a/internal/backends/sqlite/collection.go b/internal/backends/sqlite/collection.go
@@ -393,11 +393,12 @@ func (c *collection) Stats(ctx context.Context, params *backends.CollectionStats
 	}
 
 	return &backends.CollectionStatsResult{
-		CountDocuments: stats.countDocuments,
-		SizeTotal:      stats.sizeTables + stats.sizeIndexes,
-		SizeIndexes:    stats.sizeIndexes,
-		SizeCollection: stats.sizeTables,
-		IndexSizes:     indexSizes,
+		CountDocuments:  stats.countDocuments,
+		SizeTotal:       stats.sizeTables + stats.sizeIndexes,
+		SizeIndexes:     stats.sizeIndexes,
+		SizeCollection:  stats.sizeTables,
+		IndexSizes:      indexSizes,
+		SizeFreeStorage: stats.sizeFreeStorage,
 	}, nil
 }
 

diff --git a/internal/backends/sqlite/database.go b/internal/backends/sqlite/database.go
@@ -152,6 +152,7 @@ func (db *database) Stats(ctx context.Context, params *backends.DatabaseStatsPar
 		SizeTotal:       totalSize,
 		SizeIndexes:     stats.sizeIndexes,
 		SizeCollections: stats.sizeTables,
+		SizeFreeStorage: stats.sizeFreeStorage,
 	}, nil
 }