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/util"
	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  = flag.String("path", "", "Path to metabase")
	force = flag.Bool("force", false, "Rewrite existing database")

	numContainers = flag.Uint("containers", 0, "Number of containers to be created")
	numObjects    = flag.Uint("objects", 0, "Number of objects per container")

	numAttributesPerObj = flag.Uint("attributes", 0, "Number of attributes pre object")

	// Need to populate indixes.
	numOwners     = flag.Uint("distinct-owners", 10, "Number of distinct owners to be used")
	numPayloads   = flag.Uint("distinct-payloads", 10, "Number of distinct payloads to be used")
	numAttributes = flag.Uint("distinct-attributes", 10, "Number of distinct attributes to be used")
)

func main() {
	flag.Parse()

	exitOnZero("must have payloads", *numPayloads)
	exitOnZero("must have attributes", *numAttributes)
	exitOnZero("must have owners", *numOwners)
	exitOnZero("path to metabase not specified", *path)

	if *numAttributesPerObj > *numAttributes {
		exit("object can't have more attributes than available")
	}

	info, err := os.Stat(*path)
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		exit("can't get path info")
	}

	// path exits
	if err == nil {
		if info.IsDir() {
			exit("path is a directory")
		}
		if !(*force) {
			exit("can't rewrite existing file, use '-force' flag")
		}

		err = os.Remove(*path)
		exitOnErr("can't remove existing file", err)
	}

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

	exitOnErr("can't open the metabase", db.Open(context.Background(), mode.ReadWrite))
	exitOnErr("can't init the metabase", db.Init())

	defer func() {
		exitOnErr("can't close the metabase", db.Close())
	}()

	exitOnErr("failed to populate database", populate(db))
}

func populate(db *meta.DB) error {
	payloads := util.GeneratePayloadPool(*numPayloads, 32)
	attributes := util.GenerateAttributePool(*numAttributes)
	owners := util.GenerateOwnerPool(*numOwners)

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

	g := &errgroup.Group{}

	var i uint
	for i = 0; i < *numContainers; i++ {
		cid := cidtest.ID()

		for _, objectType := range types {
			typ := objectType
			g.Go(func() error {
				return util.PopulateWithObjects(db, *numObjects, func() *objectSDK.Object {
					return util.GenerateObject(
						util.WithContainerID(cid),
						util.WithType(typ),
						util.WithPayloadFromPool(payloads),
						util.WithAttributesFromPool(attributes, *numAttributesPerObj),
						util.WithOwnerIDFromPool(owners),
					)
				})
			})
		}

		g.Go(func() error {
			return util.PopulateWithBigObjects(db, *numObjects, func() *objectSDK.Object {
				return util.GenerateObject(
					util.WithContainerID(cid),
					util.WithType(objectSDK.TypeRegular),
					util.WithAttributesFromPool(attributes, *numAttributesPerObj),
					util.WithOwnerIDFromPool(owners),
				)
			})
		})

		g.Go(func() error {
			return util.PopulateGraveyard(db, *numObjects, func() *objectSDK.Object {
				return util.GenerateObject(
					util.WithContainerID(cid),
					util.WithType(objectSDK.TypeRegular),
					util.WithAttributesFromPool(attributes, *numAttributesPerObj),
					util.WithOwnerIDFromPool(owners),
				)
			})
		})

		g.Go(func() error {
			return util.PopulateLocked(db, *numObjects, func() *objectSDK.Object {
				return util.GenerateObject(
					util.WithContainerID(cid),
					util.WithType(objectSDK.TypeRegular),
					util.WithAttributesFromPool(attributes, *numAttributesPerObj),
					util.WithOwnerIDFromPool(owners),
				)
			})
		})
	}

	return g.Wait()
}

func exit(msg string) {
	fmt.Fprintln(os.Stderr, msg)
	os.Exit(1)
}

func exitOnZero[T comparable](msg string, value T) {
	var zero T
	if value == zero {
		exit(msg)
	}
}

func exitOnErr(msg string, err error) {
	if err == nil {
		return
	}
	fmt.Fprintln(os.Stderr, fmt.Errorf("%s: %w", msg, err))
	os.Exit(1)
}