From 7bd42bc569dfd8073a951cb8974b129d90551b41 Mon Sep 17 00:00:00 2001
From: schulze <thefeli73@hotmail.com>
Date: Tue, 11 Apr 2023 12:56:07 +0200
Subject: [PATCH] Wait for image ready (migration working

---
 mtdaws/mtd.go   |  8 ++++++++
 mtdaws/utils.go | 49 +++++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/mtdaws/mtd.go b/mtdaws/mtd.go
index cbd1d4c..ab60d50 100644
--- a/mtdaws/mtd.go
+++ b/mtdaws/mtd.go
@@ -2,6 +2,7 @@ package mtdaws
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/thefeli73/polemos/state"
@@ -27,6 +28,13 @@ func AWSMoveInstance(config state.Config) (state.Config) {
 	}
 	fmt.Println("Created image: ", imageName)
 
+	err = waitForImageReady(svc, imageName, 5*time.Minute)
+	if err != nil {
+		fmt.Println("Error waiting for image to be ready:", err)
+		return config
+	}
+	fmt.Println("Image is ready:", imageName)
+
 	newInstanceID, err := launchInstance(svc, realInstance, imageName)
 	if err != nil {
 		fmt.Println("Error launching instance:", err)
diff --git a/mtdaws/utils.go b/mtdaws/utils.go
index 2f8b440..429ffb3 100644
--- a/mtdaws/utils.go
+++ b/mtdaws/utils.go
@@ -2,6 +2,7 @@ package mtdaws
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"strings"
@@ -129,21 +130,49 @@ func createImage(svc *ec2.Client, instanceID string) (string, error) {
 	return aws.ToString(output.ImageId), nil
 }
 
-// launchInstance launches a instance specified by id with parameters
-func launchInstance(svc *ec2.Client, instance *types.Instance, imageID string) (string, error) {
-	securityGroupIds := make([]string, len(instance.SecurityGroups))
-	for i, sg := range instance.SecurityGroups {
+// waitForImageReady polls every second to see if the image is ready
+func waitForImageReady(svc *ec2.Client, imageID string, timeout time.Duration) error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return errors.New("timed out waiting for image to be ready")
+		case <-time.After(1 * time.Second):
+			input := &ec2.DescribeImagesInput{
+				ImageIds: []string{imageID},
+			}
+			output, err := svc.DescribeImages(ctx, input)
+			if err != nil {
+				return err
+			}
+
+			if len(output.Images) > 0 && output.Images[0].State == types.ImageStateAvailable {
+				return nil
+			}
+		}
+	}
+}
+
+// launchInstance launches a instance based on an oldInstance and AMI (duplicating the instance)
+func launchInstance(svc *ec2.Client, oldInstance *types.Instance, imageID string) (string, error) {
+	securityGroupIds := make([]string, len(oldInstance.SecurityGroups))
+	for i, sg := range oldInstance.SecurityGroups {
 		securityGroupIds[i] = aws.ToString(sg.GroupId)
 	}
-
+	// TODO: select random zone that is not the current one.
+	availabilityZone := "us-east-1d"
 	input := &ec2.RunInstancesInput{
 		ImageId:         aws.String(imageID),
-		InstanceType:    instance.InstanceType,
+		InstanceType:    oldInstance.InstanceType,
 		MinCount:        aws.Int32(1),
 		MaxCount:        aws.Int32(1),
-		KeyName:         instance.KeyName,
-		SubnetId:        instance.SubnetId,
+		KeyName:         oldInstance.KeyName,
 		SecurityGroupIds: securityGroupIds,
+		Placement: &types.Placement{
+			AvailabilityZone: aws.String(availabilityZone),
+		},
 	}
 
 	output, err := svc.RunInstances(context.TODO(), input)
@@ -151,6 +180,8 @@ func launchInstance(svc *ec2.Client, instance *types.Instance, imageID string) (
 		return "", err
 	}
 
+	// TODO: save/index config for the new instance
+
 	return aws.ToString(output.Instances[0].InstanceId), nil
 }
 
@@ -161,6 +192,8 @@ func terminateInstance(svc *ec2.Client, instanceID string) error {
 	}
 
 	_, err := svc.TerminateInstances(context.TODO(), input)
+
+	// TODO: remove config for old instance
 	return err
 }