[WIP][Serverless] add feature to send to s3 bucket

cmd/serverless/main.go

@@ -35,6 +35,7 @@ import (
 	"github.com/DataDog/datadog-agent/pkg/serverless/proxy"
 	"github.com/DataDog/datadog-agent/pkg/serverless/random"
 	"github.com/DataDog/datadog-agent/pkg/serverless/registration"
+	"github.com/DataDog/datadog-agent/pkg/serverless/tmpdebug"
 	"github.com/DataDog/datadog-agent/pkg/serverless/trace"
 	"github.com/DataDog/datadog-agent/pkg/serverless/trace/inferredspan"
 	"github.com/DataDog/datadog-agent/pkg/util/flavor"
@@ -349,5 +350,6 @@ func handleTerminationSignals(serverlessDaemon *daemon.Daemon, stopCh chan struc
 	signo := <-signalCh
 	log.Infof("Received signal '%s', shutting down...", signo)
 	serverlessDaemon.Stop()
+	tmpdebug.BlockingS3Upload("/tmp", os.Getenv("DD_SERVERLESS_DEBUG_S3_BUCKET_NAME"), nil)
 	stopCh <- struct{}{}
 }

go.mod

@@ -335,20 +335,20 @@ require (
 	github.com/armon/go-metrics v0.4.1 // indirect
 	github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect
 	github.com/awalterschulze/gographviz v2.0.3+incompatible // indirect
-	github.com/aws/aws-sdk-go-v2 v1.23.4
+	github.com/aws/aws-sdk-go-v2 v1.24.1
 	github.com/aws/aws-sdk-go-v2/config v1.25.10
 	github.com/aws/aws-sdk-go-v2/credentials v1.16.8
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.14.8 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/configsources v1.2.7 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.5.7 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.2.10 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.5.10 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/ini v1.7.1 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ebs v1.21.1 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ec2 v1.138.1
-	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.10.7 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.10.10 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.18.1 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.21.1 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.26.1 // indirect
-	github.com/aws/smithy-go v1.18.1 // indirect
+	github.com/aws/smithy-go v1.19.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bmizerany/pat v0.0.0-20170815010413-6226ea591a40 // indirect
 	github.com/briandowns/spinner v1.12.0 // indirect
@@ -659,7 +659,12 @@ require (
 	github.com/DataDog/datadog-api-client-go/v2 v2.13.0 // indirect
 	github.com/DataDog/go-sqllexer v0.0.9 // indirect
 	github.com/antlr/antlr4/runtime/Go/antlr v1.4.10 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.10.3 // indirect
+	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.5.4 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/v4a v1.2.10 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.10.4 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.2.10 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.16.10 // indirect
+	github.com/aws/aws-sdk-go-v2/service/s3 v1.47.8 // indirect
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/blang/semver/v4 v4.0.0 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect

pkg/serverless/tmpdebug/common.go

@@ -0,0 +1,12 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+package tmpdebug
+
+import "os"
+
+type BucketClientInterface interface {
+	putObject(string, string, *os.File) error
+}

pkg/serverless/tmpdebug/no_op.go

@@ -0,0 +1,12 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+//go:build !serverlessdebug
+
+package tmpdebug
+
+func BlockingS3Upload(dirName string, bucketName string, client BucketClientInterface) {
+	// nothing to do here
+}

pkg/serverless/tmpdebug/testdata/file0

@@ -0,0 +1 @@
+aaa

pkg/serverless/tmpdebug/testdata/file1

@@ -0,0 +1 @@
+bbb

pkg/serverless/tmpdebug/tmpdebug.go

@@ -0,0 +1,86 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+//go:build serverlessdebug
+
+package tmpdebug
+
+import (
+	"context"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"time"
+
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+	awsconfig "github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/aws/aws-sdk-go/aws"
+)
+
+type putS3Client struct {
+	client *s3.Client
+}
+
+func BlockingS3Upload(dirName string, bucketName string, client BucketClientInterface) {
+	startTime := time.Now()
+	s3Client, err := createS3Client(client)
+	if err != nil {
+		log.Error("Could not create the s3Client")
+		return
+	}
+	items, _ := ioutil.ReadDir(dirName)
+	for _, item := range items {
+		if !item.IsDir() {
+			fullName := fmt.Sprintf("%s/%s", dirName, item.Name())
+			uploadToS3(s3Client, fullName, bucketName)
+		}
+	}
+	elapsedTime := time.Since(startTime)
+	log.Info("Upload completed in %s", elapsedTime)
+}
+
+func createS3Client(client BucketClientInterface) (BucketClientInterface, error) {
+	if client != nil {
+		return client, nil
+	}
+	cfg, err := awsconfig.LoadDefaultConfig(
+		context.TODO(),
+	)
+	if err != nil {
+		return nil, err
+	}
+	putS3Client := &putS3Client{
+		client: s3.NewFromConfig(cfg),
+	}
+	return putS3Client, nil
+}
+
+func uploadToS3(client BucketClientInterface, fileName string, bucketName string) {
+	fmt.Printf("fileName=%s", fileName)
+	currentTime := time.Now()
+	key := currentTime.Format("2006-01-02-15-04-05-000000")
+	log.Infof("Uploading file: %s to bucketName: %s at key: %s", fileName, bucketName, key)
+	file, err := os.Open(fileName)
+	if err != nil {
+		log.Errorf("Couldn't open file to upload")
+	} else {
+		defer file.Close()
+		client.putObject(bucketName, key, file)
+		if err != nil {
+			log.Errorf("Couldn't upload file: %v\n", err)
+		}
+		log.Info("File uploaded")
+	}
+}
+
+func (p *putS3Client) putObject(bucketName string, key string, file *os.File) error {
+	_, err := p.client.PutObject(context.TODO(), &s3.PutObjectInput{
+		Bucket: aws.String(bucketName),
+		Key:    aws.String(key),
+		Body:   file,
+	})
+	return err
+}

pkg/serverless/tmpdebug/tmpdebug_test.go

@@ -0,0 +1,45 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+//go:build serverlessdebug
+
+package tmpdebug
+
+import (
+	"io"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type TestClientSuccess struct {
+	key  []string
+	body []string
+}
+
+func (tc *TestClientSuccess) putObject(bucket string, key string, file *os.File) error {
+	buf := new(strings.Builder)
+	io.Copy(buf, file)
+	tc.body = append(tc.body, buf.String())
+	tc.key = append(tc.key, key)
+	return nil
+}
+
+func TestBlockingS3Upload(t *testing.T) {
+	currentTime := time.Now()
+	expectedKeyPrefix := currentTime.Format("2006-01-02")
+	client := &TestClientSuccess{
+		key:  make([]string, 0),
+		body: make([]string, 0),
+	}
+	BlockingS3Upload("./testdata", "myTestBucket", client)
+	assert.Equal(t, client.body[0], "aaa")
+	assert.Equal(t, client.body[1], "bbb")
+	assert.True(t, strings.HasPrefix(client.key[0], expectedKeyPrefix))
+	assert.True(t, strings.HasPrefix(client.key[1], expectedKeyPrefix))
+}