From eee3a75b8b8a785e0ef0d1da75d0196c5aa240fa Mon Sep 17 00:00:00 2001
From: Michael Mandrus <41969079+mmandrus@users.noreply.github.com>
Date: Fri, 9 Aug 2024 10:38:08 -0400
Subject: [PATCH] CloudMigrations: Misc cleanup before codefreeze (#91725)

* fix retry logic

* slight adjustments

* fix disconnect and connect events

* Update pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go

Co-authored-by: Matheus Macabu <macabu@users.noreply.github.com>

---------

Co-authored-by: Matheus Macabu <macabu@users.noreply.github.com>
---
 .../cloudmigrationimpl/cloudmigration.go          |  2 +-
 .../cloudmigrationimpl/snapshot_mgmt.go           | 15 ++++++++++++---
 .../cloudmigrationimpl/xorm_store.go              |  6 +++++-
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration.go b/pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration.go
index 268c2dc7755..6a8fdfeb02e 100644
--- a/pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration.go
+++ b/pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration.go
@@ -391,7 +391,7 @@ func (s *Service) CreateSession(ctx context.Context, cmd cloudmigration.CloudMig
 		return nil, fmt.Errorf("error creating migration: %w", err)
 	}
 
-	s.report(ctx, cm, gmsclient.EventConnect, 0, nil)
+	s.report(ctx, &migration, gmsclient.EventConnect, 0, nil)
 
 	return &cloudmigration.CloudMigrationSessionResponse{
 		UID:     cm.UID,
diff --git a/pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go b/pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go
index 65bd07e564b..334a52eac96 100644
--- a/pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go
+++ b/pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go
@@ -355,10 +355,19 @@ func (s *Service) uploadUsingPresignedURL(ctx context.Context, uploadURL, key st
 }
 
 func (s *Service) updateSnapshotWithRetries(ctx context.Context, cmd cloudmigration.UpdateSnapshotCmd) (err error) {
+	maxRetries := 10
+	retries := 0
 	if err := retryer.Retry(func() (retryer.RetrySignal, error) {
-		err := s.store.UpdateSnapshot(ctx, cmd)
-		return retryer.FuncComplete, err
-	}, 10, time.Millisecond*100, time.Second*10); err != nil {
+		if err := s.store.UpdateSnapshot(ctx, cmd); err != nil {
+			s.log.Error("updating snapshot in retry loop", "error", err.Error())
+			retries++
+			if retries > maxRetries {
+				return retryer.FuncError, err
+			}
+			return retryer.FuncFailure, nil
+		}
+		return retryer.FuncComplete, nil
+	}, maxRetries, time.Millisecond*10, time.Second*5); err != nil {
 		s.log.Error("failed to update snapshot status", "snapshotUid", cmd.UID, "status", cmd.Status, "num_resources", len(cmd.Resources), "error", err.Error())
 		return fmt.Errorf("failed to update snapshot status: %w", err)
 	}
diff --git a/pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go b/pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go
index 0bc6daf9779..6929da5b830 100644
--- a/pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go
+++ b/pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go
@@ -164,10 +164,14 @@ func (ss *sqlStore) DeleteMigrationSessionByUID(ctx context.Context, uid string)
 		}
 		return nil
 	})
-
 	if err != nil {
 		return nil, nil, err
 	}
+
+	if err := ss.decryptToken(ctx, &c); err != nil {
+		return &c, snapshots, err
+	}
+
 	return &c, snapshots, nil
 }