frostfs-node/scripts/populate-metabase/main.go
Aleksey Savchuk b9043433a0 [#1223] scripts: Add script to populate metabase
Signed-off-by: Aleksey Savchuk <a.savchuk@yadro.com>
2024-09-05 08:03:52 +00:00

159 lines
4.6 KiB
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(); 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)
}
}