backend: add migration tool to migrate between any two backends

lib/backend/migration/migration.go

@@ -0,0 +1,155 @@
+package migration
+
+import (
+	"context"
+	"sync/atomic"
+	"time"
+
+	"github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/api/utils/retryutils"
+	"github.com/gravitational/teleport/lib/backend"
+	"github.com/gravitational/trace"
+)
+
+// Migration manages a migration between two [backend.Backend] interfaces.
+type Migration struct {
+	src      backend.Backend
+	dst      backend.Backend
+	parallel int
+	total    int
+	migrated atomic.Int64
+	log      logrus.FieldLogger
+}
+
+// MigrationConfig configures a [Migration] with a source and destination backend.
+// All items from the source are copied to the destination. All Teleport Auth
+// Service instances should be stopped when running a migration to avoid data
+// inconsistencies.
+type MigrationConfig struct {
+	// Source is the backend [backend.Config] items are migrated from.
+	Source backend.Config `yaml:"src"`
+	// Destination is the [backend.Config] items are migrated to.
+	Destination backend.Config `yaml:"dst"`
+	// Parallel is the number of items that will be migraated in parallel.
+	Parallel int `yaml:"parallel"`
+	// Log logs the progress of a [Migration]
+	Log logrus.FieldLogger
+}
+
+// New returns a [Migration] based on the provided [MigrationConfig].
+func New(ctx context.Context, config MigrationConfig) (*Migration, error) {
+	src, err := backend.New(ctx, config.Source.Type, config.Source.Params)
+	if err != nil {
+		return nil, trace.Wrap(err, "failed to create source backend")
+	}
+	dst, err := backend.New(ctx, config.Destination.Type, config.Destination.Params)
+	if err != nil {
+		return nil, trace.Wrap(err, "failed to create destination backend")
+	}
+	migration := &Migration{
+		src:      src,
+		dst:      dst,
+		parallel: config.Parallel,
+		log:      config.Log,
+	}
+	if migration.parallel == 0 {
+		migration.parallel = 1
+	}
+	if migration.log == nil {
+		migration.log = logrus.WithField(teleport.ComponentKey, "migration")
+	}
+	return migration, nil
+}
+
+// Run runs a [Migration] until complete.
+func (m *Migration) Run(ctx context.Context) error {
+	var all []backend.Item
+	start := backend.Key("")
+	err := retry(ctx, 3, func() error {
+		result, err := m.src.GetRange(ctx, start, backend.RangeEnd(start), 0)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+		all = result.Items
+		return nil
+	})
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	m.total = len(all)
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	group, ctx := errgroup.WithContext(ctx)
+	group.SetLimit(m.parallel)
+
+	logProgress := func() {
+		m.log.Info("Migrated %d/%d", m.migrated.Load(), m.total)
+	}
+	defer logProgress()
+	go func() {
+		ticker := time.NewTicker(time.Second)
+		defer ticker.Stop()
+		select {
+		case <-ticker.C:
+			logProgress()
+		case <-ctx.Done():
+			return
+		}
+	}()
+
+	for _, item := range all {
+		item := item
+		group.Go(func() error {
+			if err := retry(ctx, 3, func() error {
+				if _, err := m.dst.Put(ctx, item); err != nil {
+					return trace.Wrap(err)
+				}
+				return nil
+			}); err != nil {
+				return trace.Wrap(err)
+			}
+			m.migrated.Add(1)
+			return nil
+		})
+		if err := ctx.Err(); err != nil {
+			break
+		}
+	}
+
+	if err := group.Wait(); err != nil {
+		return trace.Wrap(err)
+	}
+	return nil
+}
+
+func retry(ctx context.Context, attempts int, fn func() error) error {
+	retry, err := retryutils.NewRetryV2(retryutils.RetryV2Config{
+		Driver: retryutils.NewExponentialDriver(time.Millisecond * 100),
+		Max:    time.Second * 2,
+	})
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	if attempts <= 0 {
+		return trace.Errorf("retry attempts must be > 0")
+	}
+
+	for i := 0; i < attempts; i++ {
+		err = fn()
+		if err == nil {
+			return nil
+		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-retry.After():
+			retry.Inc()
+		}
+	}
+	return trace.Wrap(err)
+}

lib/backend/migration/migration_test.go

@@ -0,0 +1,56 @@
+package migration
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/require"
+
+	"github.com/gravitational/teleport/lib/backend"
+	"github.com/gravitational/teleport/lib/backend/memory"
+)
+
+func TestMigration(t *testing.T) {
+	ctx := context.Background()
+	src, err := memory.New(memory.Config{})
+	require.NoError(t, err)
+
+	dst, err := memory.New(memory.Config{})
+	require.NoError(t, err)
+
+	itemCount := 1111
+	items := make([]backend.Item, itemCount)
+
+	for i := 0; i < itemCount; i++ {
+		item := backend.Item{
+			Key:   backend.Key(fmt.Sprintf("key-%05d", i)),
+			Value: []byte(fmt.Sprintf("value-%d", i)),
+		}
+		_, err := src.Put(ctx, item)
+		require.NoError(t, err)
+		items[i] = item
+	}
+
+	migration := Migration{
+		src:      src,
+		dst:      dst,
+		parallel: 10,
+		log:      logrus.New(),
+	}
+
+	err = migration.Run(ctx)
+	require.NoError(t, err)
+
+	start := backend.Key("")
+	result, err := dst.GetRange(ctx, start, backend.RangeEnd(start), 0)
+	require.NoError(t, err)
+
+	diff := cmp.Diff(items, result.Items, cmpopts.IgnoreFields(backend.Item{}, "Revision", "ID"))
+	require.Empty(t, diff)
+	require.Equal(t, itemCount, migration.total)
+	require.Equal(t, itemCount, int(migration.migrated.Load()))
+}

tool/teleport/common/migrate.go

@@ -0,0 +1,31 @@
+package common
+
+import (
+	"context"
+	"os"
+
+	"gopkg.in/yaml.v3"
+
+	"github.com/gravitational/teleport/lib/backend/migration"
+	"github.com/gravitational/trace"
+)
+
+func runMigration(ctx context.Context, path string) error {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	config := migration.MigrationConfig{}
+	if err := yaml.Unmarshal(data, &config); err != nil {
+		return trace.Wrap(err)
+	}
+
+	migration, err := migration.New(ctx, config)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	if err := migration.Run(ctx); err != nil {
+		return trace.Wrap(err)
+	}
+	return nil
+}

tool/teleport/common/teleport.go

@@ -107,8 +107,13 @@ func Run(options Options) (app *kingpin.Application, executedCommand string, con
 	remoteForward := app.Command(teleport.RemoteForwardSubCommand, "Used internally by Teleport to re-exec itself to remote port forward.").Hidden()
 	checkHomeDir := app.Command(teleport.CheckHomeDirSubCommand, "Used internally by Teleport to re-exec itself to check access to a directory.").Hidden()
 	park := app.Command(teleport.ParkSubCommand, "Used internally by Teleport to re-exec itself to do nothing.").Hidden()
+	migrate := app.Command("migrate", "Migrate data between two backends.").Hidden()
 	app.HelpFlag.Short('h')
 
+	migrate.Flag("config", "Path to migration config file.").
+		Short('c').
+		StringVar(&ccf.ConfigFile)
+
 	// define start flags:
 	start.Flag("debug", "Enable verbose logging to stderr").
 		Short('d').
@@ -655,6 +660,8 @@ func Run(options Options) (app *kingpin.Application, executedCommand string, con
 		err = onGetLogLevel(ccf.ConfigFile)
 	case collectProfilesCmd.FullCommand():
 		err = onCollectProfiles(ccf.ConfigFile, ccf.Profiles, ccf.ProfileSeconds)
+	case migrate.FullCommand():
+		err = runMigration(context.Background(), ccf.ConfigFile)
 	}
 	if err != nil {
 		utils.FatalError(err)