Merge pull request #7 from mashiike/feature/on-lambda-arn-make-so-lon…

…g-subject

When executing from a lambda function, the SNS subject is too long when using ARNs

go.mod

@@ -6,6 +6,12 @@ require (
 	github.com/aws/aws-lambda-go v1.28.0
 	github.com/aws/aws-sdk-go v1.42.12
 	github.com/pkg/errors v0.9.1
+	github.com/stretchr/testify v1.6.1
 )
 
-require github.com/jmespath/go-jmespath v0.4.0 // indirect
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/jmespath/go-jmespath v0.4.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776 // indirect
+)

go.sum

@@ -28,6 +28,7 @@ golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=

lambda.go

@@ -3,6 +3,9 @@ package tracer
 import (
 	"context"
 	"fmt"
+	"strings"
+
+	"github.com/aws/aws-sdk-go/aws/arn"
 )
 
 func (t *Tracer) LambdaHandlerFunc(opt *RunOption) func(ctx context.Context, event *ECSTaskEvent) error {
@@ -12,6 +15,31 @@ func (t *Tracer) LambdaHandlerFunc(opt *RunOption) func(ctx context.Context, eve
 		if lastStatus != "STOPPED" {
 			return nil
 		}
-		return t.Run(event.Detail.ClusterArn, event.Detail.TaskArn, opt)
+		cluster := extractClusterName(event.Detail.ClusterArn)
+		return t.Run(cluster, extractTaskID(cluster, event.Detail.TaskArn), opt)
+	}
+}
+
+func extractClusterName(clusterArn string) string {
+	parsed, err := arn.Parse(clusterArn)
+	if err != nil {
+		return clusterArn
+	}
+	prefix := "cluster/"
+	if parsed.Service == "ecs" && strings.HasPrefix(parsed.Resource, prefix) {
+		return strings.TrimPrefix(parsed.Resource, prefix)
+	}
+	return clusterArn
+}
+
+func extractTaskID(cluster, taskArn string) string {
+	parsed, err := arn.Parse(taskArn)
+	if err != nil {
+		return taskArn
+	}
+	prefix := "task/" + cluster + "/"
+	if parsed.Service == "ecs" && strings.HasPrefix(parsed.Resource, prefix) {
+		return strings.TrimPrefix(parsed.Resource, prefix)
 	}
+	return taskArn
 }

lambda_test.go

@@ -0,0 +1,56 @@
+package tracer
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestExtractClusterName(t *testing.T) {
+	cases := []struct {
+		input    string
+		expected string
+	}{
+		{
+			input:    "arn:aws:ecs:ap-northeast-1:012345678901:cluster/main",
+			expected: "main",
+		},
+		{
+			input:    "main",
+			expected: "main",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.input, func(t *testing.T) {
+			actual := extractClusterName(c.input)
+			require.Equal(t, c.expected, actual)
+		})
+	}
+}
+
+func TestExtractTaskID(t *testing.T) {
+	cases := []struct {
+		input    string
+		cluster  string
+		expected string
+	}{
+		{
+			input:    "0123456789abcdef0123456789abcdef",
+			cluster:  "main",
+			expected: "0123456789abcdef0123456789abcdef",
+		},
+		{
+			input:    "arn:aws:ecs:ap-northeast-1:012345678901:task/main/0123456789abcdef0123456789abcdef",
+			cluster:  "main",
+			expected: "0123456789abcdef0123456789abcdef",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.input, func(t *testing.T) {
+			actual := extractTaskID(c.cluster, c.input)
+			require.Equal(t, c.expected, actual)
+		})
+	}
+}

tracer.go

@@ -187,14 +187,24 @@ func subject(cluster, taskID string) string {
 	return s
 }
 
+const (
+	snsSubjectLimitLength = 100
+	ellipsisString        = "..."
+)
+
 func (t *Tracer) Publish(topicArn, cluster, taskID string) error {
 	msg := t.buf.String()
 	if len(msg) >= snsMaxPayloadSize {
 		msg = msg[:snsMaxPayloadSize]
 	}
+
+	s := subject(cluster, taskID)
+	if len(s) > snsSubjectLimitLength {
+		s = s[0:snsSubjectLimitLength-len(ellipsisString)] + ellipsisString
+	}
 	_, err := t.sns.PublishWithContext(t.ctx, &sns.PublishInput{
 		Message:  &msg,
-		Subject:  aws.String(subject(cluster, taskID)),
+		Subject:  aws.String(s),
 		TopicArn: &topicArn,
 	})
 	return err