<commit_before>package myaws

import (
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/aws/aws-sdk-go/service/elbv2"
	"github.com/pkg/errors"
	funk "github.com/thoas/go-funk"
)

// WaitUntilECSContainerInstancesAreDrained is a helper function which waits until
// the ECS container instances are drained.
// Due to the current limitation of the implementation of `request.Waiter`,
// we need to wait it in two steps.
// 1. Wait until container instances are DRAINING state.
// 2. Wait until no running tasks on the container instances.
func (client *Client) WaitUntilECSContainerInstancesAreDrained(cluster string, containerInstances []*string) error {
	ctx := aws.BackgroundContext()

	input := &ecs.DescribeContainerInstancesInput{
		Cluster:            &cluster,
		ContainerInstances: containerInstances,
	}

	// make sure container instances are DRAINING state
	if err := client.waitUntilECSContainerInstancesStatusWithContext(ctx, input, "DRAINING"); err != nil {
		return err
	}

	// wait until no running tasks on the container instances
	if err := client.waitUntilECSContainerInstancesNoRunningTaskWithContext(ctx, input); err != nil {
		return err
	}

	return nil
}

func (client *Client) waitUntilECSContainerInstancesStatusWithContext(ctx aws.Context, input *ecs.DescribeContainerInstancesInput, status string, opts ...request.WaiterOption) error {
	w := request.Waiter{
		Name:        "WaitUntilECSContainerInstancesStatus",
		MaxAttempts: 20,
		Delay:       request.ConstantWaiterDelay(15 * time.Second),
		Acceptors: []request.WaiterAcceptor{
			{
				State:   request.SuccessWaiterState,
				Matcher: request.PathAllWaiterMatch, Argument: "ContainerInstances[].Status",
				Expected: status,
			},
		},
		Logger: client.config.Logger,
		NewRequest: func(opts []request.Option) (*request.Request, error) {
			var inCpy *ecs.DescribeContainerInstancesInput
			if input != nil {
				tmp := *input
				inCpy = &tmp
			}
			req, _ := client.ECS.DescribeContainerInstancesRequest(inCpy)
			req.SetContext(ctx)
			req.ApplyOptions(opts...)
			return req, nil
		},
	}
	w.ApplyOptions(opts...)

	return w.WaitWithContext(ctx)
}

func (client *Client) waitUntilECSContainerInstancesNoRunningTaskWithContext(ctx aws.Context, input *ecs.DescribeContainerInstancesInput, opts ...request.WaiterOption) error {
	w := request.Waiter{
		Name:        "WaitUntilECSContainerInstancesNoRunningTask",
		MaxAttempts: 20,
		Delay:       request.ConstantWaiterDelay(15 * time.Second),<commit_msg>ecs/node/drain: Extend wait timeout for drain<commit_after>package myaws

import (
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/aws/aws-sdk-go/service/elbv2"
	"github.com/pkg/errors"
	funk "github.com/thoas/go-funk"
)

// WaitUntilECSContainerInstancesAreDrained is a helper function which waits until
// the ECS container instances are drained.
// Due to the current limitation of the implementation of `request.Waiter`,
// we need to wait it in two steps.
// 1. Wait until container instances are DRAINING state.
// 2. Wait until no running tasks on the container instances.
func (client *Client) WaitUntilECSContainerInstancesAreDrained(cluster string, containerInstances []*string) error {
	ctx := aws.BackgroundContext()

	input := &ecs.DescribeContainerInstancesInput{
		Cluster:            &cluster,
		ContainerInstances: containerInstances,
	}

	// make sure container instances are DRAINING state
	if err := client.waitUntilECSContainerInstancesStatusWithContext(ctx, input, "DRAINING"); err != nil {
		return err
	}

	// wait until no running tasks on the container instances
	if err := client.waitUntilECSContainerInstancesNoRunningTaskWithContext(ctx, input); err != nil {
		return err
	}

	return nil
}

func (client *Client) waitUntilECSContainerInstancesStatusWithContext(ctx aws.Context, input *ecs.DescribeContainerInstancesInput, status string, opts ...request.WaiterOption) error {
	w := request.Waiter{
		Name:        "WaitUntilECSContainerInstancesStatus",
		MaxAttempts: 20,
		Delay:       request.ConstantWaiterDelay(15 * time.Second),
		Acceptors: []request.WaiterAcceptor{
			{
				State:   request.SuccessWaiterState,
				Matcher: request.PathAllWaiterMatch, Argument: "ContainerInstances[].Status",
				Expected: status,
			},
		},
		Logger: client.config.Logger,
		NewRequest: func(opts []request.Option) (*request.Request, error) {
			var inCpy *ecs.DescribeContainerInstancesInput
			if input != nil {
				tmp := *input
				inCpy = &tmp
			}
			req, _ := client.ECS.DescribeContainerInstancesRequest(inCpy)
			req.SetContext(ctx)
			req.ApplyOptions(opts...)
			return req, nil
		},
	}
	w.ApplyOptions(opts...)

	return w.WaitWithContext(ctx)
}

func (client *Client) waitUntilECSContainerInstancesNoRunningTaskWithContext(ctx aws.Context, input *ecs.DescribeContainerInstancesInput, opts ...request.WaiterOption) error {
	w := request.Waiter{
		Name:        "WaitUntilECSContainerInstancesNoRunningTask",
		MaxAttempts: 40,
		Delay:       request.ConstantWaiterDelay(15 * time.Second),
{'fraction_non_alphanumeric': 0.10783179012345678, 'fraction_numerical': 0.004243827160493827, 'mean_word_length': 4.367494824016563, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '875700', 'n_tokens_mistral': 1738, 'n_tokens_neox': 1619, 'n_words': 382}