From 3fadc99629348be5df668b3cfbcbfca46f55bba1 Mon Sep 17 00:00:00 2001
From: Tiago Silva <tiago.silva@goteleport.com>
Date: Tue, 24 Sep 2024 09:44:23 -0700
Subject: [PATCH] retry migration every 12h (max)

---
 lib/backend/firestore/firestorebk.go | 25 ++++++++++++++++++++-----
 lib/backend/firestore/migration.go   |  9 ++++++++-
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/lib/backend/firestore/firestorebk.go b/lib/backend/firestore/firestorebk.go
index 33e2079de7488..d3f660f5afc0f 100644
--- a/lib/backend/firestore/firestorebk.go
+++ b/lib/backend/firestore/firestorebk.go
@@ -46,6 +46,8 @@ import (
 	"github.com/gravitational/teleport/api/utils/retryutils"
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/defaults"
+	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/interval"
 )
 
 func init() {
@@ -415,8 +417,24 @@ func New(ctx context.Context, params backend.Params, options Options) (*Backend,
 	}
 
 	// Migrate incorrect key types to the correct type.
-	// Start the migration after a delay to allow the backend to start up and won't be affected by the migration.
-	_ = b.clock.AfterFunc(5*time.Minute, b.migrateIncorrectKeyTypes)
+	// TODO(tigrato|rosstimothy): DELETE in 19.0.0
+	go func() {
+		migrationInterval := interval.New(interval.Config{
+			Duration:      time.Hour * 12,
+			FirstDuration: utils.FullJitter(time.Minute * 5),
+			Jitter:        retryutils.NewSeventhJitter(),
+			Clock:         b.clock,
+		})
+		defer migrationInterval.Stop()
+		for {
+			select {
+			case <-migrationInterval.Next():
+				b.migrateIncorrectKeyTypes()
+			case <-b.clientContext.Done():
+				return
+			}
+		}
+	}()
 
 	l.Info("Backend created.")
 	return b, nil
@@ -645,7 +663,6 @@ func (b *Backend) CompareAndSwap(ctx context.Context, expected backend.Item, rep
 
 		return nil
 	}, firestore.MaxAttempts(maxTxnAttempts))
-
 	if err != nil {
 		if status.Code(err) == codes.Aborted {
 			// RunTransaction does not officially document what error is returned if MaxAttempts is exceeded,
@@ -710,7 +727,6 @@ func (b *Backend) ConditionalDelete(ctx context.Context, key backend.Key, rev st
 
 		return nil
 	}, firestore.MaxAttempts(maxTxnAttempts))
-
 	if err != nil {
 		if status.Code(err) == codes.Aborted {
 			// RunTransaction does not officially document what error is returned if MaxAttempts is exceeded,
@@ -777,7 +793,6 @@ func (b *Backend) ConditionalUpdate(ctx context.Context, item backend.Item) (*ba
 
 		return nil
 	}, firestore.MaxAttempts(maxTxnAttempts))
-
 	if err != nil {
 		if status.Code(err) == codes.Aborted {
 			// RunTransaction does not officially document what error is returned if MaxAttempts is exceeded,
diff --git a/lib/backend/firestore/migration.go b/lib/backend/firestore/migration.go
index 2e782b73a76cd..34dbedcc461cb 100644
--- a/lib/backend/firestore/migration.go
+++ b/lib/backend/firestore/migration.go
@@ -103,7 +103,14 @@ func migrateKeyType[T any](ctx context.Context, b *Backend, newKey func([]byte)
 	// handle the migration in batches of 300 documents per second
 	t := time.NewTimer(time.Second)
 	defer t.Stop()
-	for range t.C {
+	for {
+
+		select {
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		case <-t.C:
+		}
+
 		docs, err := b.svc.Collection(b.CollectionName).
 			// passing the key type here forces the client to map the key to the underlying type
 			// and return all the keys in that share the same underlying type.