frostfs-node/scripts/populate-metabase/main.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"os"

	meta "git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/metabase"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/shard/mode"
	"git.frostfs.info/TrueCloudLab/frostfs-node/scripts/populate-metabase/internal"
	cidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id/test"
	objectSDK "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	"golang.org/x/sync/errgroup"
)

var (
	path  string
	force bool
	jobs  uint

	numContainers,
	numObjects,
	numAttributesPerObj,
	numOwners,
	numPayloads,
	numAttributes uint
)

func main() {
	flag.StringVar(&path, "path", "", "Path to metabase")
	flag.BoolVar(&force, "force", false, "Rewrite existing database")
	flag.UintVar(&jobs, "j", 10000, "Number of jobs to run")

	flag.UintVar(&numContainers, "containers", 0, "Number of containers to be created")
	flag.UintVar(&numObjects, "objects", 0, "Number of objects per container")
	flag.UintVar(&numAttributesPerObj, "attributes", 0, "Number of attributes per object")

	flag.UintVar(&numOwners, "distinct-owners", 10, "Number of distinct owners to be used")
	flag.UintVar(&numPayloads, "distinct-payloads", 10, "Number of distinct payloads to be used")
	flag.UintVar(&numAttributes, "distinct-attributes", 10, "Number of distinct attributes to be used")

	flag.Parse()

	exitIf(numPayloads == 0, "must have payloads\n")
	exitIf(numAttributes == 0, "must have attributes\n")
	exitIf(numOwners == 0, "must have owners\n")
	exitIf(len(path) == 0, "path to metabase not specified\n")
	exitIf(
		numAttributesPerObj > numAttributes,
		"object can't have more attributes than available\n",
	)

	info, err := os.Stat(path)
	exitIf(
		err != nil && !errors.Is(err, os.ErrNotExist),
		"couldn't get path info: %s\n", err,
	)

	// Path exits.
	if err == nil {
		exitIf(info.IsDir(), "path is a directory\n")
		exitIf(!force, "couldn't rewrite existing file, use '-force' flag\n")

		err = os.Remove(path)
		exitIf(err != nil, "couldn't remove existing file: %s\n", err)
	}

	err = populate()
	exitIf(err != nil, "couldn't populate the metabase: %s\n", err)
}

func getObjectFactory(opts ...internal.ObjectOption) func() *objectSDK.Object {
	return func() *objectSDK.Object {
		return internal.GenerateObject(opts...)
	}
}

func populate() (err error) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	db := meta.New([]meta.Option{
		meta.WithPath(path),
		meta.WithPermissions(0o600),
		meta.WithEpochState(internal.EpochState{}),
	}...)

	if err = db.Open(ctx, mode.ReadWrite); err != nil {
		return fmt.Errorf("couldn't open the metabase: %w", err)
	}
	defer func() {
		if errOnClose := db.Close(ctx); errOnClose != nil {
			err = errors.Join(
				err,
				fmt.Errorf("couldn't close the metabase: %w", db.Close(ctx)),
			)
		}
	}()

	if err = db.Init(ctx); err != nil {
		return fmt.Errorf("couldn't init the metabase: %w", err)
	}

	payloads := internal.GeneratePayloadPool(numPayloads, 32)
	attributes := internal.GenerateAttributePool(numAttributes)
	owners := internal.GenerateOwnerPool(numOwners)

	types := []objectSDK.Type{
		objectSDK.TypeRegular,
		objectSDK.TypeLock,
		objectSDK.TypeTombstone,
	}

	eg, ctx := errgroup.WithContext(ctx)
	eg.SetLimit(int(jobs))

	for range numContainers {
		cid := cidtest.ID()

		for _, typ := range types {
			internal.PopulateWithObjects(ctx, db, eg, numObjects, getObjectFactory(
				internal.WithContainerID(cid),
				internal.WithType(typ),
				internal.WithPayloadFromPool(payloads),
				internal.WithOwnerIDFromPool(owners),
				internal.WithAttributesFromPool(attributes, numAttributesPerObj),
			))
		}
		internal.PopulateWithBigObjects(ctx, db, eg, numObjects, getObjectFactory(
			internal.WithContainerID(cid),
			internal.WithType(objectSDK.TypeRegular),
			internal.WithAttributesFromPool(attributes, numAttributesPerObj),
			internal.WithOwnerIDFromPool(owners),
		))
		internal.PopulateGraveyard(ctx, db, eg, int(jobs), numObjects, getObjectFactory(
			internal.WithContainerID(cid),
			internal.WithType(objectSDK.TypeRegular),
			internal.WithAttributesFromPool(attributes, numAttributesPerObj),
			internal.WithOwnerIDFromPool(owners),
		))
		internal.PopulateLocked(ctx, db, eg, int(jobs), numObjects, getObjectFactory(
			internal.WithContainerID(cid),
			internal.WithType(objectSDK.TypeRegular),
			internal.WithAttributesFromPool(attributes, numAttributesPerObj),
			internal.WithOwnerIDFromPool(owners),
		))
	}

	return eg.Wait()
}

func exitIf(cond bool, format string, args ...any) {
	if cond {
		fmt.Fprintf(os.Stderr, format, args...)
		os.Exit(1)
	}
}