chore: Add profilecli command to delete v1 blocks (#4320)

* chore: Add profilecli command to delete v1 blocks

As part of the v1 to v2 migration it is useful to clear out the bucket
of old v1 blocks. This command will help to dry run this step and then
also delete them.

* Address PR feedback

Author: simonswine

cmd/profilecli/main.go

@@ -102,6 +102,10 @@ func main() {
 	raftInfoCmd := raftCmd.Command("info", "Print info about a Raft node.")
 	raftInfoParams := addRaftInfoParams(raftInfoCmd)
 
+	v2MigrationCmd := adminCmd.Command("v2-migration", "Operation to aid the v1 to v2 storage migration.")
+	v2MigrationBucketCleanupCmd := v2MigrationCmd.Command("bucket-cleanup", "Clean up v1 artificats from data bucket.")
+	v2MigrationBucketCleanupParams := addV2MigrationBackupCleanupParam(v2MigrationBucketCleanupCmd)
+
 	// parse command line arguments
 	parsedCmd := kingpin.MustParse(app.Parse(os.Args[1:]))
 
@@ -181,6 +185,10 @@ func main() {
 		if err := raftInfo(ctx, raftInfoParams); err != nil {
 			os.Exit(checkError(err))
 		}
+	case v2MigrationBucketCleanupCmd.FullCommand():
+		if err := v2MigrationBucketCleanup(ctx, v2MigrationBucketCleanupParams); err != nil {
+			os.Exit(checkError(err))
+		}
 	default:
 		level.Error(logger).Log("msg", "unknown command", "cmd", parsedCmd)
 	}

cmd/profilecli/v2-migration.go

@@ -0,0 +1,176 @@
+package main
+
+import (
+	"container/list"
+	"context"
+	"flag"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/grafana/dskit/flagext"
+	"golang.org/x/sync/errgroup"
+
+	pyroscopecfg "github.com/grafana/pyroscope/pkg/cfg"
+	pyroscopeobj "github.com/grafana/pyroscope/pkg/objstore"
+	objstoreclient "github.com/grafana/pyroscope/pkg/objstore/client"
+)
+
+type v2MigrationBucketCleanupParams struct {
+	configFile      string
+	configExpandEnv bool
+	dryRun          string
+}
+
+func (p *v2MigrationBucketCleanupParams) isDryRun() bool {
+	return p.dryRun != "false"
+}
+
+type minimalConfig struct {
+	Bucket objstoreclient.Config `yaml:"storage"`
+}
+
+// Note: These are not the flags used, but we need to register them to get the defaults.
+func (c *minimalConfig) RegisterFlags(f *flag.FlagSet) {
+	c.Bucket.RegisterFlags(f)
+}
+
+func (c *minimalConfig) ApplyDynamicConfig() pyroscopecfg.Source {
+	return func(dst pyroscopecfg.Cloneable) error {
+		return nil
+	}
+}
+
+func (c *minimalConfig) Clone() flagext.Registerer {
+	return func(c minimalConfig) *minimalConfig {
+		return &c
+	}(*c)
+}
+
+func clientFromParams(ctx context.Context, params *v2MigrationBucketCleanupParams) (pyroscopeobj.Bucket, error) {
+	if params.configFile == "" {
+		return nil, fmt.Errorf("config file is required")
+	}
+	cfg := &minimalConfig{}
+	fs := flag.NewFlagSet("config-file-loader", flag.ContinueOnError)
+	if err := pyroscopecfg.Unmarshal(cfg,
+		pyroscopecfg.Defaults(fs),
+		pyroscopecfg.YAMLIgnoreUnknownFields(params.configFile, params.configExpandEnv),
+	); err != nil {
+		return nil, fmt.Errorf("failed parsing config: %w", err)
+	}
+
+	return objstoreclient.NewBucket(ctx, cfg.Bucket, "profilecli")
+}
+
+func addV2MigrationBackupCleanupParam(c commander) *v2MigrationBucketCleanupParams {
+	var (
+		params = &v2MigrationBucketCleanupParams{}
+	)
+	c.Flag("config.file", "The path to the pyroscope config").Default("/etc/pyroscope/config.yaml").StringVar(&params.configFile)
+	c.Flag("config.expand-env", "").Default("false").BoolVar(&params.configExpandEnv)
+	c.Flag("dry-run", "Dry run the operation.").Default("true").StringVar(&params.dryRun)
+	return params
+}
+
+func v2MigrationBucketCleanup(ctx context.Context, params *v2MigrationBucketCleanupParams) error {
+	client, err := clientFromParams(ctx, params)
+	if err != nil {
+		return fmt.Errorf("failed to create client: %w", err)
+	}
+
+	var pathsToDelete []string
+	// find prefix called "phlaredb/" on the second level
+	if err := client.Iter(ctx, "", func(name string) error {
+		if !strings.HasSuffix(name, "/") {
+			return nil
+		}
+		err := client.Iter(ctx, name, func(name string) error {
+			if strings.HasSuffix(name, "phlaredb/") {
+				pathsToDelete = append(pathsToDelete, name)
+			}
+			return nil
+		})
+		if err != nil {
+			return err
+		}
+
+		return nil
+	}); err != nil {
+		return fmt.Errorf("failed to list tenants: %w", err)
+	}
+
+	if len(pathsToDelete) == 0 {
+		fmt.Println("No paths to delete")
+		return nil
+	}
+
+	if params.isDryRun() {
+		fmt.Println("DRY-RUN: If ran with --dry-run=false, this will delete everything under:")
+	} else {
+		fmt.Println("This will delete everything under:")
+	}
+	for _, path := range pathsToDelete {
+		fmt.Println(" - ", path)
+	}
+
+	if params.isDryRun() {
+		fmt.Println("DRY-RUN: If ran with --dry-run=false, this will delete those object store keys:")
+		return recurse(ctx, client, func(key string) error {
+			fmt.Println(" - ", key)
+			return nil
+		}, pathsToDelete)
+	}
+
+	// We do actually delete here
+	fmt.Println("Last chance to cancel, waiting 3 seconds...")
+	<-time.After(3 * time.Second)
+
+	fmt.Println("Deleted object store keys:")
+	return recurse(ctx, client, func(key string) error {
+		if err := client.Delete(ctx, key); err != nil {
+			return fmt.Errorf("failed to delete %s: %w", key, err)
+		}
+		fmt.Println(" - ", key)
+		return nil
+	}, pathsToDelete)
+}
+
+const maxConcurrentActions = 16
+
+func recurse(ctx context.Context, b pyroscopeobj.Bucket, action func(key string) error, paths []string) error {
+	g, gctx := errgroup.WithContext(ctx)
+	g.SetLimit(maxConcurrentActions)
+
+	g.Go(func() error {
+		iters := list.New()
+		for _, path := range paths {
+			iters.PushBack(path)
+		}
+
+		for iters.Len() > 0 {
+			e := iters.Front()
+			path := e.Value.(string)
+
+			if err := b.Iter(gctx, path, func(path string) error {
+				if strings.HasSuffix(path, "/") {
+					iters.PushBack(path)
+					return nil
+				}
+
+				g.Go(func() error {
+					return action(path)
+				})
+
+				return nil
+			}); err != nil {
+				return fmt.Errorf("failed to iterate over %s: %w", path, err)
+			}
+			iters.Remove(e)
+		}
+
+		return nil
+	})
+
+	return g.Wait()
+}

pkg/cfg/cfg_test.go

@@ -18,7 +18,7 @@ server:
   timeout: 60h
 tls:
   key: YAML
-`))
+`), true)
 
 	fs := flag.NewFlagSet(t.Name(), flag.PanicOnError)
 	flagSource := dFlags(fs, []string{"-verbose", "-server.port=21"})
@@ -51,7 +51,7 @@ servers:
   timeoutz: 60h
 tls:
   keey: YAML
-`))
+`), true)
 
 	fs := flag.NewFlagSet(t.Name(), flag.PanicOnError)
 	flagSource := dFlags(fs, []string{"-verbose", "-server.port=21"})

pkg/cfg/files.go

@@ -41,6 +41,15 @@ func dJSON(y []byte) Source {
 // When expandEnvVars is true, variables in the supplied '.yaml\ file are expanded
 // using https://pkg.go.dev/github.com/drone/envsubst?tab=overview
 func YAML(f string, expandEnvVars bool) Source {
+	return yamlWithKnowFields(f, expandEnvVars, true)
+}
+
+// Like YAML but ignores fields that are not known
+func YAMLIgnoreUnknownFields(f string, expandEnvVars bool) Source {
+	return yamlWithKnowFields(f, expandEnvVars, false)
+}
+
+func yamlWithKnowFields(f string, expandEnvVars bool, knownFields bool) Source {
 	return func(dst Cloneable) error {
 		y, err := os.ReadFile(f)
 		if err != nil {
@@ -53,19 +62,19 @@ func YAML(f string, expandEnvVars bool) Source {
 			}
 			y = []byte(s)
 		}
-		err = dYAML(y)(dst)
+		err = dYAML(y, knownFields)(dst)
 		return errors.Wrap(err, f)
 	}
 }
 
 // dYAML returns a YAML source and allows dependency injection
-func dYAML(y []byte) Source {
+func dYAML(y []byte, knownFields bool) Source {
 	return func(dst Cloneable) error {
 		if len(y) == 0 {
 			return nil
 		}
 		dec := yaml.NewDecoder(bytes.NewReader(y))
-		dec.KnownFields(true)
+		dec.KnownFields(knownFields)
 		if err := dec.Decode(dst); err != nil {
 			return err
 		}

pkg/cfg/precedence_test.go

@@ -28,7 +28,7 @@ func TestYAMLOverDefaults(t *testing.T) {
 	fs := flag.NewFlagSet(t.Name(), flag.PanicOnError)
 	err := Unmarshal(&data,
 		Defaults(fs),
-		dYAML([]byte(y)),
+		dYAML([]byte(y), true),
 	)
 
 	require.NoError(t, err)
@@ -51,7 +51,7 @@ func TestFlagOverYAML(t *testing.T) {
 
 	err := Unmarshal(&data,
 		Defaults(fs),
-		dYAML([]byte(y)),
+		dYAML([]byte(y), true),
 		dFlags(fs, []string{"-verbose=false", "-tls.cert=CLI"}),
 	)