ext: ecsobserver Add Service Fetcher and mock

Author: pingleig

extension/observer/ecsobserver/fetcher.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"sort"
+	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/request"
@@ -33,6 +34,10 @@ const (
 	// Based on existing number from cloudwatch-agent
 	ec2CacheSize                   = 2000
 	describeContainerInstanceLimit = 100
+	describeServiceLimit           = 10
+	// NOTE: these constants are not defined in go sdk, there are three values for deployment status.
+	deploymentStatusActive  = "ACTIVE"
+	deploymentStatusPrimary = "PRIMARY"
 )
 
 // ecsClient includes API required by taskFetcher.
@@ -41,6 +46,8 @@ type ecsClient interface {
 	DescribeTasksWithContext(ctx context.Context, input *ecs.DescribeTasksInput, opts ...request.Option) (*ecs.DescribeTasksOutput, error)
 	DescribeTaskDefinitionWithContext(ctx context.Context, input *ecs.DescribeTaskDefinitionInput, opts ...request.Option) (*ecs.DescribeTaskDefinitionOutput, error)
 	DescribeContainerInstancesWithContext(ctx context.Context, input *ecs.DescribeContainerInstancesInput, opts ...request.Option) (*ecs.DescribeContainerInstancesOutput, error)
+	ListServicesWithContext(ctx context.Context, input *ecs.ListServicesInput, opts ...request.Option) (*ecs.ListServicesOutput, error)
+	DescribeServicesWithContext(ctx context.Context, input *ecs.DescribeServicesInput, opts ...request.Option) (*ecs.DescribeServicesOutput, error)
 }
 
 // ec2Client includes API required by TaskFetcher.
@@ -49,12 +56,13 @@ type ec2Client interface {
 }
 
 type taskFetcher struct {
-	logger       *zap.Logger
-	ecs          ecsClient
-	ec2          ec2Client
-	cluster      string
-	taskDefCache simplelru.LRUCache
-	ec2Cache     simplelru.LRUCache
+	logger            *zap.Logger
+	ecs               ecsClient
+	ec2               ec2Client
+	cluster           string
+	taskDefCache      simplelru.LRUCache
+	ec2Cache          simplelru.LRUCache
+	serviceNameFilter serviceNameFilter
 }
 
 type taskFetcherOptions struct {
@@ -85,6 +93,11 @@ func newTaskFetcher(opts taskFetcherOptions) (*taskFetcher, error) {
 		cluster:      opts.Cluster,
 		taskDefCache: taskDefCache,
 		ec2Cache:     ec2Cache,
+		// TODO: after the service matcher PR is merged, use actual service name filter here.
+		// For now, describe all the services
+		serviceNameFilter: func(name string) bool {
+			return true
+		},
 	}
 	// Return early if any clients are mocked, caller should overrides all the clients when mocking.
 	if fetcher.ecs != nil || fetcher.ec2 != nil {
@@ -105,9 +118,16 @@ func (f *taskFetcher) fetchAndDecorate(ctx context.Context) ([]*Task, error) {
 	}
 
 	// EC2
-	if err := f.attachContainerInstance(ctx, tasks); err != nil {
+	if err = f.attachContainerInstance(ctx, tasks); err != nil {
 		return nil, fmt.Errorf("attachContainerInstance failed: %w", err)
 	}
+
+	// Services
+	services, err := f.getAllServices(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("getAllServices failed: %w", err)
+	}
+	f.attachService(tasks, services)
 	return tasks, nil
 }
 
@@ -246,7 +266,7 @@ func (f *taskFetcher) describeContainerInstances(ctx context.Context, instanceLi
 		ContainerInstances: instanceList,
 	})
 	if err != nil {
-		return fmt.Errorf("ecs.DescribeContainerInstance faile: %w", err)
+		return fmt.Errorf("ecs.DescribeContainerInstance failed: %w", err)
 	}
 
 	// Create the index to map ec2 id back to container instance id.
@@ -282,6 +302,85 @@ func (f *taskFetcher) describeContainerInstances(ctx context.Context, instanceLi
 	return nil
 }
 
+// serviceNameFilter decides if we should get detail info for a service, i.e. make the describe API call.
+type serviceNameFilter func(name string) bool
+
+// getAllServices does not have cache like task definition or ec2 instances
+// because we need to get the deployment id to map service to task, which changes frequently.
+func (f *taskFetcher) getAllServices(ctx context.Context) ([]*ecs.Service, error) {
+	svc := f.ecs
+	cluster := aws.String(f.cluster)
+	// List and filter out services we need to desribe.
+	listReq := ecs.ListServicesInput{Cluster: cluster}
+	var servicesToDescribe []*string
+	for {
+		res, err := svc.ListServicesWithContext(ctx, &listReq)
+		if err != nil {
+			return nil, err
+		}
+		for _, arn := range res.ServiceArns {
+			segs := strings.Split(aws.StringValue(arn), "/")
+			name := segs[len(segs)-1]
+			if f.serviceNameFilter(name) {
+				servicesToDescribe = append(servicesToDescribe, arn)
+			}
+		}
+		if res.NextToken == nil {
+			break
+		}
+		listReq.NextToken = res.NextToken
+	}
+
+	// DescribeServices size limit is 10 so we need to do paging on client side.
+	var services []*ecs.Service
+	for i := 0; i < len(servicesToDescribe); i += describeServiceLimit {
+		end := minInt(i+describeServiceLimit, len(servicesToDescribe))
+		desc := &ecs.DescribeServicesInput{
+			Cluster:  cluster,
+			Services: servicesToDescribe[i:end],
+		}
+		res, err := svc.DescribeServicesWithContext(ctx, desc)
+		if err != nil {
+			return nil, fmt.Errorf("ecs.DescribeServices failed %w", err)
+		}
+		services = append(services, res.Services...)
+	}
+	return services, nil
+}
+
+// attachService map service to task using deployment id.
+// Each service can have multiple deployment and each task keep track of the deployment in task.StartedBy.
+func (f *taskFetcher) attachService(tasks []*Task, services []*ecs.Service) {
+	// Map deployment ID to service name
+	idToService := make(map[string]*ecs.Service)
+	for _, svc := range services {
+		for _, deployment := range svc.Deployments {
+			status := aws.StringValue(deployment.Status)
+			if status == deploymentStatusActive || status == deploymentStatusPrimary {
+				idToService[aws.StringValue(deployment.Id)] = svc
+				break
+			}
+		}
+	}
+
+	// Attach service to task
+	for _, t := range tasks {
+		// Task is created using RunTask i.e. not manged by a service.
+		if t.Task.StartedBy == nil {
+			continue
+		}
+		deploymentID := aws.StringValue(t.Task.StartedBy)
+		svc := idToService[deploymentID]
+		// Service not found happen a lot because we only fetch services defined in ServiceConfig.
+		// However, we fetch all the tasks, which could be started by other services no mentioned in config
+		// or started using RunTasks API directly.
+		if svc == nil {
+			continue
+		}
+		t.Service = svc
+	}
+}
+
 // Util Start
 
 func sortStringPointers(ps []*string) {

extension/observer/ecsobserver/fetcher_test.go

@@ -38,8 +38,7 @@ func TestFetcher_FetchAndDecorate(t *testing.T) {
 		ec2Override: c,
 	})
 	require.NoError(t, err)
-	// Create 1 task def and 10 tasks, 8 running on ec2, first 3 runs on fargate
-	// TODO: they will be changed to include service etc.
+	// Create 1 task def, 2 services and 10 tasks, 8 running on ec2, first 3 runs on fargate
 	nTasks := 11
 	nInstances := 2
 	nFargateInstances := 3
@@ -48,9 +47,11 @@ func TestFetcher_FetchAndDecorate(t *testing.T) {
 		ins := i % nInstances
 		if i < nFargateInstances {
 			task.LaunchType = aws.String(ecs.LaunchTypeFargate)
+			task.StartedBy = aws.String("deploy0")
 		} else {
 			task.LaunchType = aws.String(ecs.LaunchTypeEc2)
 			task.ContainerInstanceArn = aws.String(fmt.Sprintf("ci%d", ins))
+			task.StartedBy = aws.String("deploy1")
 		}
 		task.TaskDefinitionArn = aws.String("d0:1")
 	}))
@@ -59,11 +60,32 @@ func TestFetcher_FetchAndDecorate(t *testing.T) {
 		ci.Ec2InstanceId = aws.String(fmt.Sprintf("i-%d", i))
 	}))
 	c.SetEc2Instances(ecsmock.GenEc2Instances("i-", nInstances, nil))
+	// Service
+	c.SetServices(ecsmock.GenServices("s", 2, func(i int, s *ecs.Service) {
+		if i == 0 {
+			s.LaunchType = aws.String(ecs.LaunchTypeFargate)
+			s.Deployments = []*ecs.Deployment{
+				{
+					Status: aws.String("ACTIVE"),
+					Id:     aws.String("deploy0"),
+				},
+			}
+		} else {
+			s.LaunchType = aws.String(ecs.LaunchTypeEc2)
+			s.Deployments = []*ecs.Deployment{
+				{
+					Status: aws.String("ACTIVE"),
+					Id:     aws.String("deploy1"),
+				},
+			}
+		}
+	}))
 
 	ctx := context.Background()
 	tasks, err := f.fetchAndDecorate(ctx)
 	require.NoError(t, err)
 	assert.Equal(t, nTasks, len(tasks))
+	assert.Equal(t, "s0", aws.StringValue(tasks[0].Service.ServiceArn))
 }
 
 func TestFetcher_GetAllTasks(t *testing.T) {
@@ -220,3 +242,66 @@ func TestFetcher_AttachContainerInstance(t *testing.T) {
 		assert.Equal(t, "i-1", aws.StringValue(tasks[nFargateInstances].EC2.InstanceId))
 	})
 }
+
+func TestFetcher_GetAllServices(t *testing.T) {
+	c := ecsmock.NewCluster()
+	f, err := newTaskFetcher(taskFetcherOptions{
+		Logger:      zap.NewExample(),
+		Cluster:     "not used",
+		Region:      "not used",
+		ecsOverride: c,
+	})
+	require.NoError(t, err)
+	const nServices = 101
+	c.SetServices(ecsmock.GenServices("s", nServices, nil))
+	ctx := context.Background()
+	services, err := f.getAllServices(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, nServices, len(services))
+}
+
+func TestFetcher_AttachService(t *testing.T) {
+	c := ecsmock.NewCluster()
+	f, err := newTaskFetcher(taskFetcherOptions{
+		Logger:      zap.NewExample(),
+		Cluster:     "not used",
+		Region:      "not used",
+		ecsOverride: c,
+	})
+	require.NoError(t, err)
+	const nServices = 10
+	c.SetServices(ecsmock.GenServices("s", nServices, func(i int, s *ecs.Service) {
+		s.Deployments = []*ecs.Deployment{
+			{
+				Status: aws.String("ACTIVE"),
+				Id:     aws.String(fmt.Sprintf("deploy%d", i)),
+			},
+		}
+	}))
+	c.SetTaskDefinitions(ecsmock.GenTaskDefinitions("def", nServices, 1, nil))
+	const nTasks = 100
+	c.SetTasks(ecsmock.GenTasks("t", nTasks, func(i int, task *ecs.Task) {
+		// Last task is launched manually w/o service
+		if i == nTasks-1 {
+			return
+		}
+		deployID := i % nServices
+		task.TaskDefinitionArn = aws.String(fmt.Sprintf("def%d:1", deployID))
+		task.StartedBy = aws.String(fmt.Sprintf("deploy%d", deployID))
+
+	}))
+
+	ctx := context.Background()
+	rawTasks, err := f.getAllTasks(ctx)
+	require.NoError(t, err)
+	tasks, err := f.attachTaskDefinition(ctx, rawTasks)
+	require.NoError(t, err)
+	services, err := f.getAllServices(ctx)
+	require.NoError(t, err)
+	f.attachService(tasks, services)
+
+	// Just pick one
+	assert.Equal(t, "s0", aws.StringValue(tasks[0].Service.ServiceArn))
+	assert.NotNil(t, tasks[nTasks-2].Service)
+	assert.Nil(t, tasks[nTasks-1].Service)
+}

extension/observer/ecsobserver/internal/ecsmock/service.go

@@ -59,15 +59,18 @@ type Cluster struct {
 	containerInstanceList []*ecs.ContainerInstance
 	ec2Map                map[string]*ec2.Instance // key is instance id
 	ec2List               []*ec2.Instance
+	serviceMap            map[string]*ecs.Service
+	serviceList           []*ecs.Service
 	limit                 PageLimit
 	stats                 ClusterStats
 }
 
 // NewCluster creates a mock ECS cluster with default limits.
 func NewCluster() *Cluster {
 	return &Cluster{
-		taskMap: make(map[string]*ecs.Task),
-		limit:   DefaultPageLimit(),
+		// NOTE: we don't set the maps by design, they should be injected and API calls
+		// without setting up data should just panic.
+		limit: DefaultPageLimit(),
 	}
 }
 
@@ -196,6 +199,45 @@ func (c *Cluster) DescribeInstancesWithContext(_ context.Context, input *ec2.Des
 	}, nil
 }
 
+func (c *Cluster) ListServicesWithContext(_ context.Context, input *ecs.ListServicesInput, _ ...request.Option) (*ecs.ListServicesOutput, error) {
+	page, err := getPage(pageInput{
+		nextToken: input.NextToken,
+		size:      len(c.serviceList),
+		limit:     c.limit.ListServiceOutput,
+	})
+	if err != nil {
+		return nil, err
+	}
+	res := c.serviceList[page.start:page.end]
+	return &ecs.ListServicesOutput{
+		ServiceArns: getArns(res, func(i int) *string {
+			return res[i].ServiceArn
+		}),
+		NextToken: page.nextToken,
+	}, nil
+}
+
+func (c *Cluster) DescribeServicesWithContext(_ context.Context, input *ecs.DescribeServicesInput, _ ...request.Option) (*ecs.DescribeServicesOutput, error) {
+	var (
+		failures []*ecs.Failure
+		services []*ecs.Service
+	)
+	for i, serviceArn := range input.Services {
+		arn := aws.StringValue(serviceArn)
+		svc, ok := c.serviceMap[arn]
+		if !ok {
+			failures = append(failures, &ecs.Failure{
+				Arn:    serviceArn,
+				Detail: aws.String(fmt.Sprintf("service not found index %d arn %s total services %d", i, arn, len(c.serviceMap))),
+				Reason: aws.String("service not found"),
+			})
+			continue
+		}
+		services = append(services, svc)
+	}
+	return &ecs.DescribeServicesOutput{Failures: failures, Services: services}, nil
+}
+
 // API End
 
 // Hook Start
@@ -240,6 +282,16 @@ func (c *Cluster) SetEc2Instances(instances []*ec2.Instance) {
 	c.ec2List = instances
 }
 
+// SetServices updates the list and map.
+func (c *Cluster) SetServices(services []*ecs.Service) {
+	m := make(map[string]*ecs.Service, len(services))
+	for _, s := range services {
+		m[aws.StringValue(s.ServiceArn)] = s
+	}
+	c.serviceMap = m
+	c.serviceList = services
+}
+
 // Hook End
 
 // Generator Start
@@ -303,6 +355,20 @@ func GenEc2Instances(idPrefix string, count int, modifier func(i int, ins *ec2.I
 	return instances
 }
 
+func GenServices(arnPrefix string, count int, modifier func(i int, s *ecs.Service)) []*ecs.Service {
+	var services []*ecs.Service
+	for i := 0; i < count; i++ {
+		svc := &ecs.Service{
+			ServiceArn: aws.String(fmt.Sprintf("%s%d", arnPrefix, i)),
+		}
+		if modifier != nil {
+			modifier(i, svc)
+		}
+		services = append(services, svc)
+	}
+	return services
+}
+
 // Generator End
 
 // pagination Start