Implement DD_PROFILING_OUTPUT_DIR for dev/debug (DataDog#924)

This env variable can be used to quickly look at the pprof files
produced by the profiler locally. This is useful when developing new
profile types that are not yet understood by the backend. It might also
be useful when working with customers who are experiencing uploading
errors.

profiler/options.go

@@ -96,6 +96,7 @@ type config struct {
 	uploadTimeout time.Duration
 	mutexFraction int
 	blockRate     int
+	outputDir     string
 }
 
 func urlForSite(site string) (string, error) {
@@ -201,6 +202,10 @@ func defaultConfig() (*config, error) {
 	if v := os.Getenv("DD_PROFILING_URL"); v != "" {
 		WithURL(v)(&c)
 	}
+	// not for public use
+	if v := os.Getenv("DD_PROFILING_OUTPUT_DIR"); v != "" {
+		withOutputDir(v)(&c)
+	}
 	return &c, nil
 }
 
@@ -375,3 +380,12 @@ func WithUDS(socketPath string) Option {
 		},
 	})
 }
+
+// withOutputDir writes a copy of all uploaded profiles to the given
+// directory. This is intended for local development or debugging uploading
+// issues. The directory will keep growing, no cleanup is performed.
+func withOutputDir(dir string) Option {
+	return func(cfg *config) {
+		cfg.outputDir = dir
+	}
+}

profiler/options_test.go

@@ -6,6 +6,7 @@
 package profiler
 
 import (
+	"io/ioutil"
 	"net"
 	"os"
 	"path/filepath"
@@ -342,3 +343,35 @@ func TestAddProfileType(t *testing.T) {
 		assert.True(ok)
 	})
 }
+
+func TestWith_outputDir(t *testing.T) {
+	tmpDir, err := ioutil.TempDir("", "")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpDir)
+
+	// Use env to enable this like a user would.
+	os.Setenv("DD_PROFILING_OUTPUT_DIR", tmpDir)
+	defer os.Unsetenv("DD_PROFILING_OUTPUT_DIR")
+
+	p, err := unstartedProfiler()
+	require.NoError(t, err)
+	bat := batch{
+		end: time.Now(),
+		profiles: []*profile{
+			{name: "foo.pprof", data: []byte("foo")},
+			{name: "bar.pprof", data: []byte("bar")},
+		},
+	}
+	require.NoError(t, p.outputDir(bat))
+	files, err := filepath.Glob(filepath.Join(tmpDir, "*", "*.pprof"))
+	require.NoError(t, err)
+
+	fileData := map[string]string{}
+	for _, file := range files {
+		data, err := ioutil.ReadFile(file)
+		require.NoError(t, err)
+		fileData[filepath.Base(file)] = string(data)
+	}
+	want := map[string]string{"foo.pprof": "foo", "bar.pprof": "bar"}
+	require.Equal(t, want, fileData)
+}

profiler/profiler.go

@@ -8,7 +8,9 @@ package profiler
 import (
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"os"
+	"path/filepath"
 	"runtime"
 	"sync"
 	"time"
@@ -211,12 +213,37 @@ func (p *profiler) enqueueUpload(bat batch) {
 // send takes profiles from the output queue and uploads them.
 func (p *profiler) send() {
 	for bat := range p.out {
+		if err := p.outputDir(bat); err != nil {
+			log.Error("Failed to output profile to dir: %v", err)
+		}
 		if err := p.uploadFunc(bat); err != nil {
 			log.Error("Failed to upload profile: %v", err)
 		}
 	}
 }
 
+func (p *profiler) outputDir(bat batch) error {
+	if p.cfg.outputDir == "" {
+		return nil
+	}
+	// Basic ISO 8601 Format in UTC as the name for the directories.
+	dir := bat.end.UTC().Format("20060102T150405Z")
+	dirPath := filepath.Join(p.cfg.outputDir, dir)
+	// 0755 is what mkdir does, should be reasonable for the use cases here.
+	if err := os.MkdirAll(dirPath, 0755); err != nil {
+		return err
+	}
+
+	for _, prof := range bat.profiles {
+		filePath := filepath.Join(dirPath, prof.name)
+		// 0644 is what touch does, should be reasonable for the use cases here.
+		if err := ioutil.WriteFile(filePath, prof.data, 0644); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // stop stops the profiler.
 func (p *profiler) stop() {
 	p.stopOnce.Do(func() {