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

package engine

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

	"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"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/internal/testutil"
	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"
	cidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id/test"
	objectSDK "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id/test"
	"git.frostfs.info/TrueCloudLab/hrw"
	"github.com/panjf2000/ants/v2"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest"
)

type epochState struct{}

func (s epochState) CurrentEpoch() uint64 {
	return 0
}

func BenchmarkExists(b *testing.B) {
	b.Run("2 shards", func(b *testing.B) {
		benchmarkExists(b, 2)
	})
	b.Run("4 shards", func(b *testing.B) {
		benchmarkExists(b, 4)
	})
	b.Run("8 shards", func(b *testing.B) {
		benchmarkExists(b, 8)
	})
}

func benchmarkExists(b *testing.B, shardNum int) {
	shards := make([]*shard.Shard, shardNum)
	for i := 0; i < shardNum; i++ {
		shards[i] = testNewShard(b, i)
	}

	e := testNewEngine(b).setInitializedShards(b, shards...).engine
	b.Cleanup(func() {
		_ = e.Close()
		_ = os.RemoveAll(b.Name())
	})

	addr := oidtest.Address()
	for i := 0; i < 100; i++ {
		obj := testutil.GenerateObjectWithCID(cidtest.ID())
		err := Put(context.Background(), e, obj)
		if err != nil {
			b.Fatal(err)
		}
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ok, err := e.exists(context.Background(), addr)
		if err != nil || ok {
			b.Fatalf("%t %v", ok, err)
		}
	}
}

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

func testNewEngine(t testing.TB, opts ...Option) *testEngineWrapper {
	engine := New(WithLogger(&logger.Logger{Logger: zaptest.NewLogger(t)}))
	for _, opt := range opts {
		opt(engine.cfg)
	}
	return &testEngineWrapper{
		engine: engine,
	}
}

func (te *testEngineWrapper) setInitializedShards(t testing.TB, shards ...*shard.Shard) *testEngineWrapper {
	for _, s := range shards {
		pool, err := ants.NewPool(10, ants.WithNonblocking(true))
		require.NoError(t, err)

		te.engine.shards[s.ID().String()] = hashedShard{
			shardWrapper: shardWrapper{
				errorCount: new(atomic.Uint32),
				Shard:      s,
			},
			hash: hrw.StringHash(s.ID().String()),
		}
		te.engine.shardPools[s.ID().String()] = pool
		te.shardIDs = append(te.shardIDs, s.ID())
	}
	return te
}

func (te *testEngineWrapper) setShardsNum(t testing.TB, num int) *testEngineWrapper {
	shards := make([]*shard.Shard, 0, num)

	for i := 0; i < num; i++ {
		shards = append(shards, testNewShard(t, i))
	}

	return te.setInitializedShards(t, shards...)
}

func (te *testEngineWrapper) setShardsNumOpts(t testing.TB, num int, shardOpts func(id int) []shard.Option) *testEngineWrapper {
	for i := 0; i < num; i++ {
		opts := shardOpts(i)
		id, err := te.engine.AddShard(opts...)
		require.NoError(t, err)
		te.shardIDs = append(te.shardIDs, id)
	}
	return te
}

func (te *testEngineWrapper) setShardsNumAdditionalOpts(t testing.TB, num int, shardOpts func(id int) []shard.Option) *testEngineWrapper {
	for i := 0; i < num; i++ {
		defaultOpts := testDefaultShardOptions(t, i)
		opts := append(defaultOpts, shardOpts(i)...)
		id, err := te.engine.AddShard(opts...)
		require.NoError(t, err)
		te.shardIDs = append(te.shardIDs, id)
	}
	return te
}

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

func newTestStorages(root string, smallSize uint64) ([]blobstor.SubStorage, *teststore.TestStore, *teststore.TestStore) {
	smallFileStorage := teststore.New(
		teststore.WithSubstorage(blobovniczatree.NewBlobovniczaTree(
			blobovniczatree.WithRootPath(filepath.Join(root, "blobovnicza")),
			blobovniczatree.WithBlobovniczaShallowDepth(1),
			blobovniczatree.WithBlobovniczaShallowWidth(1),
			blobovniczatree.WithPermissions(0700)),
		))
	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
}

func testNewShard(t testing.TB, id int) *shard.Shard {
	sid, err := generateShardID()
	require.NoError(t, err)

	shardOpts := append([]shard.Option{shard.WithID(sid)}, testDefaultShardOptions(t, id)...)
	s := shard.New(shardOpts...)

	require.NoError(t, s.Open())
	require.NoError(t, s.Init(context.Background()))

	return s
}

func testDefaultShardOptions(t testing.TB, id int) []shard.Option {
	return []shard.Option{
		shard.WithLogger(&logger.Logger{Logger: zap.L()}),
		shard.WithBlobStorOptions(
			blobstor.WithStorages(
				newStorages(t.TempDir(), 1<<20))),
		shard.WithPiloramaOptions(pilorama.WithPath(filepath.Join(t.TempDir(), "pilorama"))),
		shard.WithMetaBaseOptions(
			meta.WithPath(filepath.Join(t.TempDir(), "metabase")),
			meta.WithPermissions(0700),
			meta.WithEpochState(epochState{}),
		)}
}