(WIP) feature: Drain nodes prior to termination when using RollingUpgrade strategy

api/v1alpha1/instancegroup_types.go

@@ -106,12 +106,19 @@ var (
 		LaunchTemplate,
 	}
 
-	DefaultRollingUpdateStrategy = &RollingUpdateStrategy{
+	DefaultRollingUpgradeTimeoutSeconds int64 = 900
+	DefaultRollingUpgradeForce                = true
+	DefaultRollingUpdateStrategy              = &RollingUpdateStrategy{
 		MaxUnavailable: &intstr.IntOrString{
 			Type:   intstr.Int,
 			IntVal: 1,
 		},
+		DrainOptions: RollingUpgradeDrainOptions{
+			TimeoutSeconds: &DefaultRollingUpgradeTimeoutSeconds,
+			Force:          &DefaultRollingUpgradeForce,
+		},
 	}
+
 	DefaultCRDStrategyMaxRetries = 3
 
 	AllowedContainerRuntimes            = []ContainerRuntime{ContainerDRuntime, DockerRuntime}
@@ -160,7 +167,12 @@ type AwsUpgradeStrategy struct {
 }
 
 type RollingUpdateStrategy struct {
-	MaxUnavailable *intstr.IntOrString `json:"maxUnavailable,omitempty"`
+	MaxUnavailable *intstr.IntOrString        `json:"maxUnavailable,omitempty"`
+	DrainOptions   RollingUpgradeDrainOptions `json:"drainOptions,omitempty"`
+}
+
+func (s *RollingUpdateStrategy) GetDrainOptions() RollingUpgradeDrainOptions {
+	return s.DrainOptions
 }
 
 func (s *RollingUpdateStrategy) GetMaxUnavailable() *intstr.IntOrString {
@@ -171,6 +183,27 @@ func (s *RollingUpdateStrategy) SetMaxUnavailable(value *intstr.IntOrString) {
 	s.MaxUnavailable = value
 }
 
+type RollingUpgradeDrainOptions struct {
+	TimeoutSeconds *int64 `json:"timeoutSeconds,omitempty"`
+	Force          *bool  `json:"force,omitempty"`
+}
+
+func (d *RollingUpgradeDrainOptions) GetTimeoutSeconds() int64 {
+	return *d.TimeoutSeconds
+}
+
+func (d *RollingUpgradeDrainOptions) SetTimeoutSeconds(n *int64) {
+	d.TimeoutSeconds = n
+}
+
+func (d *RollingUpgradeDrainOptions) GetForce() bool {
+	return *d.Force
+}
+
+func (d *RollingUpgradeDrainOptions) SetForce(condition *bool) {
+	d.Force = condition
+}
+
 type CRDUpdateStrategy struct {
 	Spec                string `json:"spec,omitempty"`
 	CRDName             string `json:"crdName,omitempty"`
@@ -760,12 +793,14 @@ func (ig *InstanceGroup) Validate() error {
 		}
 	}
 
-	if strings.EqualFold(s.AwsUpgradeStrategy.Type, RollingUpdateStrategyName) && s.AwsUpgradeStrategy.RollingUpdateType == nil {
-		s.AwsUpgradeStrategy.RollingUpdateType = DefaultRollingUpdateStrategy
+	if strings.EqualFold(s.AwsUpgradeStrategy.Type, RollingUpdateStrategyName) {
+		if err := s.AwsUpgradeStrategy.RollingUpdateType.Validate(); err != nil {
+			return err
+		}
 	}
-
 	return nil
 }
+
 func (c *EKSConfiguration) GetRoleName() string {
 	return c.ExistingRoleName
 }
@@ -950,6 +985,26 @@ func (s *AwsUpgradeStrategy) SetCRDType(crd *CRDUpdateStrategy) {
 	s.CRDType = crd
 }
 
+func (r *RollingUpdateStrategy) Validate() error {
+	if r == nil {
+		r = DefaultRollingUpdateStrategy
+	}
+
+	if r.MaxUnavailable == nil {
+		r.MaxUnavailable = DefaultRollingUpdateStrategy.MaxUnavailable
+	}
+
+	if r.DrainOptions.TimeoutSeconds == nil {
+		r.DrainOptions.TimeoutSeconds = DefaultRollingUpdateStrategy.DrainOptions.TimeoutSeconds
+	}
+
+	if r.DrainOptions.Force == nil {
+		r.DrainOptions.Force = DefaultRollingUpdateStrategy.DrainOptions.Force
+	}
+
+	return nil
+}
+
 func (c *CRDUpdateStrategy) Validate() error {
 	if c.GetSpec() == "" {
 		return errors.New("spec is empty")

config/crd/bases/instancemgr.keikoproj.io_instancegroups.yaml

@@ -403,6 +403,14 @@ spec:
                     type: object
                   rollingUpdate:
                     properties:
+                      drainOptions:
+                        properties:
+                          force:
+                            type: boolean
+                          timeoutSeconds:
+                            format: int64
+                            type: integer
+                        type: object
                       maxUnavailable:
                         anyOf:
                         - type: integer

controllers/instancegroup_controller.go

@@ -51,7 +51,9 @@ type InstanceGroupReconciler struct {
 	Auth                   *InstanceGroupAuthenticator
 	ConfigMap              *corev1.ConfigMap
 	Namespaces             map[string]corev1.Namespace
-	NamespacesLock         *sync.RWMutex
+	NamespacesLock         *sync.Mutex
+	DrainGroupMapper       *sync.Map
+	DrainErrorMapper       *sync.Map
 	ConfigRetention        int
 	Metrics                *common.MetricsCollector
 }
@@ -132,6 +134,7 @@ func (r *InstanceGroupReconciler) Reconcile(ctxt context.Context, req ctrl.Reque
 
 	// set/unset finalizer
 	r.SetFinalizer(instanceGroup)
+	namespacedName := instanceGroup.NamespacedName()
 
 	input := provisioners.ProvisionerInput{
 		AwsWorker:       r.Auth.Aws,
@@ -149,6 +152,13 @@ func (r *InstanceGroupReconciler) Reconcile(ctxt context.Context, req ctrl.Reque
 	)
 	status.SetConfigHash(configHash)
 
+	drainGroup, _ := r.DrainGroupMapper.LoadOrStore(namespacedName, &sync.WaitGroup{})
+	drainErrs, _ := r.DrainErrorMapper.LoadOrStore(namespacedName, make(chan error))
+	input.DrainManager = kubeprovider.DrainManager{
+		DrainErrors: drainErrs.(chan error),
+		DrainGroup:  drainGroup.(*sync.WaitGroup),
+	}
+
 	if !reflect.DeepEqual(*r.ConfigMap, corev1.ConfigMap{}) {
 		// Configmap exist - apply defaults/boundaries if namespace is not excluded
 		namespace := instanceGroup.GetNamespace()

controllers/providers/aws/predicates.go

@@ -16,6 +16,8 @@ limitations under the License.
 package aws
 
 import (
+	"strings"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/autoscaling"
 	"github.com/aws/aws-sdk-go/service/eks"
@@ -127,6 +129,21 @@ func IsUsingMixedInstances(group *autoscaling.Group) bool {
 	return false
 }
 
+func IsDesiredInService(scalingGroup *autoscaling.Group) bool {
+	var (
+		desired        = aws.Int64Value(scalingGroup.DesiredCapacity)
+		inServiceCount int64
+	)
+
+	for _, instance := range scalingGroup.Instances {
+		lifecycle := aws.StringValue(instance.LifecycleState)
+		if strings.EqualFold(lifecycle, autoscaling.LifecycleStateInService) {
+			inServiceCount++
+		}
+	}
+
+	return desired == inServiceCount
+
 func IsUsingWarmPool(group *autoscaling.Group) bool {
 	if group.WarmPoolConfiguration != nil {
 		return true