frostfs-node/pkg/local_object_storage/engine/engine_test.go

package engine

import (
	"context"
	"path/filepath"
	"sync/atomic"
	"testing"

	"git.frostfs.info/TrueCloudLab/frostfs-node/internal/qos"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/blobovniczatree"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/fstree"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/teststore"
	meta "git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/metabase"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/pilorama"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/shard"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/util/logger/test"
	objectSDK "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	"github.com/stretchr/testify/require"
)

type epochState struct {
	currEpoch uint64
}

func (s epochState) CurrentEpoch() uint64 {
	return s.currEpoch
}

type testEngineWrapper struct {
	engine   *StorageEngine
	shards   []*shard.Shard
	shardIDs []*shard.ID
}

func testNewEngine(t testing.TB, opts ...Option) *testEngineWrapper {
	opts = append(testGetDefaultEngineOptions(t), opts...)
	return &testEngineWrapper{engine: New(opts...)}
}

func (te *testEngineWrapper) setShardsNum(t testing.TB, num int) *testEngineWrapper {
	return te.setShardsNumOpts(t, num, func(_ int) []shard.Option {
		return testGetDefaultShardOptions(t)
	})
}

func (te *testEngineWrapper) setShardsNumOpts(
	t testing.TB, num int, shardOpts func(id int) []shard.Option,
) *testEngineWrapper {
	te.shards = make([]*shard.Shard, num)
	te.shardIDs = make([]*shard.ID, num)
	for i := range num {
		shard, err := te.engine.createShard(context.Background(), shardOpts(i))
		require.NoError(t, err)
		require.NoError(t, te.engine.addShard(shard))
		te.shards[i] = shard
		te.shardIDs[i] = shard.ID()
	}
	require.Len(t, te.engine.shards, num)
	return te
}

func (te *testEngineWrapper) setShardsNumAdditionalOpts(
	t testing.TB, num int, shardOpts func(id int) []shard.Option,
) *testEngineWrapper {
	return te.setShardsNumOpts(t, num, func(id int) []shard.Option {
		return append(testGetDefaultShardOptions(t), shardOpts(id)...)
	})
}

// prepare calls Open and Init on the created engine.
func (te *testEngineWrapper) prepare(t testing.TB) *testEngineWrapper {
	require.NoError(t, te.engine.Open(context.Background()))
	require.NoError(t, te.engine.Init(context.Background()))
	return te
}

func testGetDefaultEngineOptions(t testing.TB) []Option {
	return []Option{
		WithLogger(test.NewLogger(t)),
	}
}

func testGetDefaultShardOptions(t testing.TB) []shard.Option {
	return []shard.Option{
		shard.WithLogger(test.NewLogger(t)),
		shard.WithBlobStorOptions(
			blobstor.WithStorages(
				newStorages(t, t.TempDir(), 1<<20)),
			blobstor.WithLogger(test.NewLogger(t)),
		),
		shard.WithPiloramaOptions(pilorama.WithPath(filepath.Join(t.TempDir(), "pilorama"))),
		shard.WithMetaBaseOptions(testGetDefaultMetabaseOptions(t)...),
		shard.WithLimiter(&testQoSLimiter{t: t}),
	}
}

func testGetDefaultMetabaseOptions(t testing.TB) []meta.Option {
	return []meta.Option{
		meta.WithPath(filepath.Join(t.TempDir(), "metabase")),
		meta.WithPermissions(0o700),
		meta.WithEpochState(epochState{}),
		meta.WithLogger(test.NewLogger(t)),
	}
}

func newStorages(t testing.TB, root string, smallSize uint64) []blobstor.SubStorage {
	return []blobstor.SubStorage{
		{
			Storage: blobovniczatree.NewBlobovniczaTree(
				context.Background(),
				blobovniczatree.WithRootPath(filepath.Join(root, "blobovnicza")),
				blobovniczatree.WithBlobovniczaShallowDepth(1),
				blobovniczatree.WithBlobovniczaShallowWidth(1),
				blobovniczatree.WithPermissions(0o700),
				blobovniczatree.WithLogger(test.NewLogger(t))),
			Policy: func(_ *objectSDK.Object, data []byte) bool {
				return uint64(len(data)) < smallSize
			},
		},
		{
			Storage: fstree.New(
				fstree.WithPath(root),
				fstree.WithDepth(1),
				fstree.WithLogger(test.NewLogger(t))),
		},
	}
}

func newTestStorages(root string, smallSize uint64) ([]blobstor.SubStorage, *teststore.TestStore, *teststore.TestStore) {
	smallFileStorage := teststore.New(
		teststore.WithSubstorage(blobovniczatree.NewBlobovniczaTree(
			context.Background(),
			blobovniczatree.WithRootPath(filepath.Join(root, "blobovnicza")),
			blobovniczatree.WithBlobovniczaShallowDepth(1),
			blobovniczatree.WithBlobovniczaShallowWidth(1),
			blobovniczatree.WithPermissions(0o700)),
		))
	largeFileStorage := teststore.New(
		teststore.WithSubstorage(fstree.New(
			fstree.WithPath(root),
			fstree.WithDepth(1)),
		))
	return []blobstor.SubStorage{
		{
			Storage: smallFileStorage,
			Policy: func(_ *objectSDK.Object, data []byte) bool {
				return uint64(len(data)) < smallSize
			},
		},
		{
			Storage: largeFileStorage,
		},
	}, smallFileStorage, largeFileStorage
}

var _ qos.Limiter = (*testQoSLimiter)(nil)

type testQoSLimiter struct {
	t     testing.TB
	read  atomic.Int64
	write atomic.Int64
}

func (t *testQoSLimiter) SetMetrics(qos.Metrics) {}

func (t *testQoSLimiter) Close() {
	require.Equal(t.t, int64(0), t.read.Load(), "read requests count after limiter close must be 0")
	require.Equal(t.t, int64(0), t.write.Load(), "write requests count after limiter close must be 0")
}

func (t *testQoSLimiter) ReadRequest(context.Context) (qos.ReleaseFunc, error) {
	t.read.Add(1)
	return func() { t.read.Add(-1) }, nil
}

func (t *testQoSLimiter) WriteRequest(context.Context) (qos.ReleaseFunc, error) {
	t.write.Add(1)
	return func() { t.write.Add(-1) }, nil
}

func (t *testQoSLimiter) SetParentID(string) {}