From 37ef9039672c9f6aa475f00f76f8c5d18a10d23e Mon Sep 17 00:00:00 2001
From: Kyle Nusbaum <kyle.nusbaum@datadoghq.com>
Date: Fri, 15 May 2020 12:30:05 -0500
Subject: [PATCH] profiler: add WithSite option and DD_SITE environment
 variable. (#644)

WithSite and DD_SITE will determine the apiURL that profiles are sent to.
---
 profiler/options.go      | 27 ++++++++++++++++++++++++++-
 profiler/options_test.go | 21 +++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/profiler/options.go b/profiler/options.go
index 23a63de172..2e2321b2e6 100644
--- a/profiler/options.go
+++ b/profiler/options.go
@@ -7,11 +7,14 @@ package profiler
 
 import (
 	"fmt"
+	"net/url"
 	"os"
 	"path/filepath"
 	"strings"
 	"time"
 
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+
 	"github.com/DataDog/datadog-go/statsd"
 )
 
@@ -52,6 +55,12 @@ type config struct {
 	blockRate     int
 }
 
+func urlForSite(site string) (string, error) {
+	u := fmt.Sprintf("https://intake.profile.%s/v1/input", site)
+	_, err := url.Parse(u)
+	return u, err
+}
+
 func (c *config) addProfileType(t ProfileType) {
 	if c.types == nil {
 		c.types = make(map[ProfileType]struct{})
@@ -61,8 +70,8 @@ func (c *config) addProfileType(t ProfileType) {
 
 func defaultConfig() *config {
 	c := config{
-		apiURL:        defaultAPIURL,
 		env:           defaultEnv,
+		apiURL:        defaultAPIURL,
 		service:       filepath.Base(os.Args[0]),
 		statsd:        &statsd.NoOpClient{},
 		period:        DefaultPeriod,
@@ -75,6 +84,9 @@ func defaultConfig() *config {
 		c.addProfileType(t)
 	}
 
+	if v := os.Getenv("DD_SITE"); v != "" {
+		WithSite(v)(&c)
+	}
 	if v := os.Getenv("DD_ENV"); v != "" {
 		WithEnv(v)(&c)
 	}
@@ -174,3 +186,16 @@ func WithStatsd(client StatsdClient) Option {
 		cfg.statsd = client
 	}
 }
+
+// WithSite specifies the datadog site (datadoghq.com, datadoghq.eu, etc.)
+// which profiles will be sent to.
+func WithSite(site string) Option {
+	return func(cfg *config) {
+		u, err := urlForSite(site)
+		if err != nil {
+			log.Error("profiler: invalid site provided, using %s (%s)", defaultAPIURL, err)
+			return
+		}
+		cfg.apiURL = u
+	}
+}
diff --git a/profiler/options_test.go b/profiler/options_test.go
index 4eef7f758d..14741dbc5c 100644
--- a/profiler/options_test.go
+++ b/profiler/options_test.go
@@ -62,6 +62,20 @@ func TestOptions(t *testing.T) {
 		assert.Equal(t, "serviceName", cfg.service)
 	})
 
+	t.Run("WithSite", func(t *testing.T) {
+		var cfg config
+		WithSite("datadog.eu")(&cfg)
+		assert.Equal(t, "https://intake.profile.datadog.eu/v1/input", cfg.apiURL)
+	})
+
+	t.Run("WithSite/override", func(t *testing.T) {
+		os.Setenv("DD_SITE", "wrong.site")
+		defer os.Unsetenv("DD_SITE")
+		cfg := defaultConfig()
+		WithSite("datadog.eu")(cfg)
+		assert.Equal(t, "https://intake.profile.datadog.eu/v1/input", cfg.apiURL)
+	})
+
 	t.Run("WithEnv", func(t *testing.T) {
 		var cfg config
 		WithEnv("envName")(&cfg)
@@ -112,6 +126,13 @@ func TestOptions(t *testing.T) {
 }
 
 func TestEnvVars(t *testing.T) {
+	t.Run("DD_SITE", func(t *testing.T) {
+		os.Setenv("DD_SITE", "datadog.eu")
+		defer os.Unsetenv("DD_SITE")
+		cfg := defaultConfig()
+		assert.Equal(t, "https://intake.profile.datadog.eu/v1/input", cfg.apiURL)
+	})
+
 	t.Run("DD_ENV", func(t *testing.T) {
 		os.Setenv("DD_ENV", "someEnv")
 		defer os.Unsetenv("DD_ENV")