From 1ecf559b713bc27704e1fba45be715a95d1a7805 Mon Sep 17 00:00:00 2001
From: liranp <liran.polak@gmail.com>
Date: Tue, 28 Apr 2020 19:54:57 +0300
Subject: [PATCH] feat(spot/ocean): add support for grace period

---
 pkg/model/spotinstmodel/instance_group.go  | 11 +++++++++++
 upup/pkg/fi/cloudup/spotinsttasks/ocean.go | 22 ++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/pkg/model/spotinstmodel/instance_group.go b/pkg/model/spotinstmodel/instance_group.go
index 183a51296f..b1eec12b54 100644
--- a/pkg/model/spotinstmodel/instance_group.go
+++ b/pkg/model/spotinstmodel/instance_group.go
@@ -53,6 +53,11 @@ const (
 	// be enabled.
 	InstanceGroupLabelFallbackToOnDemand = "spotinst.io/fallback-to-ondemand"
 
+	// InstanceGroupLabelGracePeriod is the metadata label used on the
+	// instance group to specify a period of time, in seconds, that Ocean
+	// should wait before applying instance health checks.
+	InstanceGroupLabelGracePeriod = "spotinst.io/grace-period"
+
 	// InstanceGroupLabelHealthCheckType is the metadata label used on the
 	// instance group to specify the type of the health check that should be used.
 	InstanceGroupLabelHealthCheckType = "spotinst.io/health-check-type"
@@ -355,6 +360,12 @@ func (b *InstanceGroupModelBuilder) buildOcean(c *fi.ModelBuilderContext, igs ..
 				return err
 			}
 
+		case InstanceGroupLabelGracePeriod:
+			ocean.GracePeriod, err = parseInt(v)
+			if err != nil {
+				return err
+			}
+
 		case InstanceGroupLabelOceanInstanceTypesWhitelist:
 			ocean.InstanceTypesWhitelist, err = parseStringSlice(v)
 			if err != nil {
diff --git a/upup/pkg/fi/cloudup/spotinsttasks/ocean.go b/upup/pkg/fi/cloudup/spotinsttasks/ocean.go
index b4dc508fd0..e1395d575c 100644
--- a/upup/pkg/fi/cloudup/spotinsttasks/ocean.go
+++ b/upup/pkg/fi/cloudup/spotinsttasks/ocean.go
@@ -47,6 +47,7 @@ type Ocean struct {
 	SpotPercentage           *float64
 	UtilizeReservedInstances *bool
 	FallbackToOnDemand       *bool
+	GracePeriod              *int64
 	InstanceTypesWhitelist   []string
 	InstanceTypesBlacklist   []string
 	Tags                     map[string]string
@@ -145,6 +146,10 @@ func (o *Ocean) Find(c *fi.Context) (*Ocean, error) {
 			actual.SpotPercentage = strategy.SpotPercentage
 			actual.FallbackToOnDemand = strategy.FallbackToOnDemand
 			actual.UtilizeReservedInstances = strategy.UtilizeReservedInstances
+
+			if strategy.GracePeriod != nil {
+				actual.GracePeriod = fi.Int64(int64(fi.IntValue(strategy.GracePeriod)))
+			}
 		}
 	}
 
@@ -366,6 +371,10 @@ func (_ *Ocean) create(cloud awsup.AWSCloud, a, e, changes *Ocean) error {
 		ocean.Strategy.SetSpotPercentage(e.SpotPercentage)
 		ocean.Strategy.SetFallbackToOnDemand(e.FallbackToOnDemand)
 		ocean.Strategy.SetUtilizeReservedInstances(e.UtilizeReservedInstances)
+
+		if e.GracePeriod != nil {
+			ocean.Strategy.SetGracePeriod(fi.Int(int(*e.GracePeriod)))
+		}
 	}
 
 	// Compute.
@@ -610,6 +619,17 @@ func (_ *Ocean) update(cloud awsup.AWSCloud, a, e, changes *Ocean) error {
 			changes.UtilizeReservedInstances = nil
 			changed = true
 		}
+
+		// Grace period.
+		if changes.GracePeriod != nil {
+			if ocean.Strategy == nil {
+				ocean.Strategy = new(aws.Strategy)
+			}
+
+			ocean.Strategy.SetGracePeriod(fi.Int(int(*e.GracePeriod)))
+			changes.GracePeriod = nil
+			changed = true
+		}
 	}
 
 	// Compute.
@@ -977,6 +997,7 @@ type terraformOceanStrategy struct {
 	SpotPercentage           *float64 `json:"spot_percentage,omitempty" cty:"spot_percentage"`
 	FallbackToOnDemand       *bool    `json:"fallback_to_ondemand,omitempty" cty:"fallback_to_ondemand"`
 	UtilizeReservedInstances *bool    `json:"utilize_reserved_instances,omitempty" cty:"utilize_reserved_instances"`
+	GracePeriod              *int64   `json:"grace_period,omitempty" cty:"grace_period"`
 }
 
 type terraformOceanLaunchSpec struct {
@@ -1012,6 +1033,7 @@ func (_ *Ocean) RenderTerraform(t *terraform.TerraformTarget, a, e, changes *Oce
 			SpotPercentage:           e.SpotPercentage,
 			FallbackToOnDemand:       e.FallbackToOnDemand,
 			UtilizeReservedInstances: e.UtilizeReservedInstances,
+			GracePeriod:              e.GracePeriod,
 		},
 		terraformOceanLaunchSpec: &terraformOceanLaunchSpec{},
 	}