[#1223] scripts: Add script to populate metabase

Signed-off-by: Aleksey Savchuk <a.savchuk@yadro.com>
2024-07-15 14:07:32 +03:00 · 2024-07-15 14:07:32 +03:00 · b9043433a0
commit b9043433a0
parent a4fb7f085b
3 changed files with 554 additions and 0 deletions
--- a/scripts/populate-metabase/main.go
+++ b/scripts/populate-metabase/main.go
@ -0,0 +1,159 @@
+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(); errOnClose != nil {
+			err = errors.Join(
+				err,
+				fmt.Errorf("couldn't close the metabase: %w", db.Close()),
+			)
+		}
+	}()
+
+	if err = db.Init(); 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 i := uint(0); i < numContainers; i++ {
+		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)
+	}
+}