[forwarder] Allow configuration of multiple endpoints

cmd/agent/app/runloop.go

@@ -76,11 +76,9 @@ func StartAgent() {
 	api.StartServer()
 
 	// setup the forwarder
-	// for now we handle only one key and one domain
-	keysPerDomain := map[string][]string{
-		config.Datadog.GetString("dd_url"): {
-			config.Datadog.GetString("api_key"),
-		},
+	keysPerDomain, err := config.GetMultipleEndpoints()
+	if err != nil {
+		log.Error("Misconfiguration of agent endpoints: ", err)
 	}
 	common.Forwarder = forwarder.NewDefaultForwarder(keysPerDomain)
 	log.Debugf("Starting forwarder")

cmd/dogstatsd/main.go

@@ -89,11 +89,10 @@ func start(cmd *cobra.Command, args []string) error {
 		return nil
 	}
 
-	// for now we handle only one key and one domain
-	keysPerDomain := map[string][]string{
-		config.Datadog.GetString("dd_url"): {
-			config.Datadog.GetString("api_key"),
-		},
+	// setup the forwarder
+	keysPerDomain, err := config.GetMultipleEndpoints()
+	if err != nil {
+		log.Error("Misconfiguration of agent endpoints: ", err)
 	}
 	f := forwarder.NewDefaultForwarder(keysPerDomain)
 	f.Start()

pkg/config/config.go

@@ -1,8 +1,10 @@
 package config
 
 import (
+	"strings"
 	"time"
 
+	log "github.com/cihub/seelog"
 	"github.com/spf13/viper"
 )
 
@@ -58,3 +60,56 @@ func init() {
 	Datadog.BindEnv("dogstatsd_socket")
 	Datadog.BindEnv("dogstatsd_non_local_traffic")
 }
+
+// GetMultipleEndpoints returns the api keys per domain specified in the main agent config
+func GetMultipleEndpoints() (map[string][]string, error) {
+	return getMultipleEndpoints(Datadog)
+}
+
+// getMultipleEndpoints implements the logic to extract the api keys per domain from an agent config
+func getMultipleEndpoints(config *viper.Viper) (map[string][]string, error) {
+	keysPerDomain := map[string][]string{
+		config.GetString("dd_url"): {
+			config.GetString("api_key"),
+		},
+	}
+
+	var additionalEndpoints map[string][]string
+	err := config.UnmarshalKey("additional_endpoints", &additionalEndpoints)
+	if err != nil {
+		return keysPerDomain, err
+	}
+
+	// merge additional endpoints into keysPerDomain
+	for domain, apiKeys := range additionalEndpoints {
+		if _, ok := keysPerDomain[domain]; ok {
+			for _, apiKey := range apiKeys {
+				keysPerDomain[domain] = append(keysPerDomain[domain], apiKey)
+			}
+		} else {
+			keysPerDomain[domain] = apiKeys
+		}
+	}
+
+	// dedupe api keys and remove domains with no api keys (or empty ones)
+	for domain, apiKeys := range keysPerDomain {
+		dedupedAPIKeys := make([]string, 0, len(apiKeys))
+		seen := make(map[string]bool)
+		for _, apiKey := range apiKeys {
+			trimmedAPIKey := strings.TrimSpace(apiKey)
+			if _, ok := seen[trimmedAPIKey]; !ok && trimmedAPIKey != "" {
+				seen[trimmedAPIKey] = true
+				dedupedAPIKeys = append(dedupedAPIKeys, trimmedAPIKey)
+			}
+		}
+
+		if len(dedupedAPIKeys) > 0 {
+			keysPerDomain[domain] = dedupedAPIKeys
+		} else {
+			log.Infof("No API key provided for domain \"%s\", removing domain from endpoints", domain)
+			delete(keysPerDomain, domain)
+		}
+	}
+
+	return keysPerDomain, nil
+}

pkg/config/config_test.go

@@ -1,11 +1,138 @@
 package config
 
 import (
+	"bytes"
 	"testing"
 
+	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestDefaults(t *testing.T) {
 	assert.Equal(t, Datadog.GetString("dd_url"), "http://localhost:17123")
 }
+
+func setupViperConf(yamlConfig string) *viper.Viper {
+	conf := viper.New()
+	conf.SetConfigType("yaml")
+	conf.ReadConfig(bytes.NewBuffer([]byte(yamlConfig)))
+	return conf
+}
+
+func TestGetMultipleEndpoints(t *testing.T) {
+	datadogYaml := `
+dd_url: "https://app.datadoghq.com"
+api_key: fakeapikey
+
+additional_endpoints:
+  "https://app.datadoghq.com":
+  - fakeapikey2
+  - fakeapikey3
+  "https://foo.datadoghq.com":
+  - someapikey
+`
+
+	testConfig := setupViperConf(datadogYaml)
+
+	multipleEndpoints, err := getMultipleEndpoints(testConfig)
+
+	expectedMultipleEndpoints := map[string][]string{
+		"https://app.datadoghq.com": {
+			"fakeapikey",
+			"fakeapikey2",
+			"fakeapikey3",
+		},
+		"https://foo.datadoghq.com": {
+			"someapikey",
+		},
+	}
+
+	assert.Nil(t, err)
+	assert.EqualValues(t, expectedMultipleEndpoints, multipleEndpoints)
+}
+
+func TestGetMultipleEndpointsWithNoAdditionalEndpoints(t *testing.T) {
+	datadogYaml := `
+dd_url: "https://app.datadoghq.com"
+api_key: fakeapikey
+`
+
+	testConfig := setupViperConf(datadogYaml)
+
+	multipleEndpoints, err := getMultipleEndpoints(testConfig)
+
+	expectedMultipleEndpoints := map[string][]string{
+		"https://app.datadoghq.com": {
+			"fakeapikey",
+		},
+	}
+
+	assert.Nil(t, err)
+	assert.EqualValues(t, expectedMultipleEndpoints, multipleEndpoints)
+}
+
+func TestGetMultipleEndpointseIgnoresDomainWithoutApiKey(t *testing.T) {
+	datadogYaml := `
+dd_url: "https://app.datadoghq.com"
+api_key: fakeapikey
+
+additional_endpoints:
+  "https://app.datadoghq.com":
+  - fakeapikey2
+  "https://foo.datadoghq.com":
+  - someapikey
+  "https://bar.datadoghq.com":
+  - ""
+`
+
+	testConfig := setupViperConf(datadogYaml)
+
+	multipleEndpoints, err := getMultipleEndpoints(testConfig)
+
+	expectedMultipleEndpoints := map[string][]string{
+		"https://app.datadoghq.com": {
+			"fakeapikey",
+			"fakeapikey2",
+		},
+		"https://foo.datadoghq.com": {
+			"someapikey",
+		},
+	}
+
+	assert.Nil(t, err)
+	assert.EqualValues(t, expectedMultipleEndpoints, multipleEndpoints)
+}
+
+func TestGetMultipleEndpointsApiKeyDeduping(t *testing.T) {
+	datadogYaml := `
+dd_url: "https://app.datadoghq.com"
+api_key: fakeapikey
+
+additional_endpoints:
+  "https://app.datadoghq.com":
+  - fakeapikey2
+  - fakeapikey
+  "https://foo.datadoghq.com":
+  - someapikey
+  - someotherapikey
+  - someapikey
+`
+
+	testConfig := setupViperConf(datadogYaml)
+
+	multipleEndpoints, err := getMultipleEndpoints(testConfig)
+
+	expectedMultipleEndpoints := map[string][]string{
+		"https://app.datadoghq.com": {
+			"fakeapikey",
+			"fakeapikey2",
+		},
+		"https://foo.datadoghq.com": {
+			"someapikey",
+			"someotherapikey",
+		},
+	}
+
+	assert.Nil(t, err)
+	assert.EqualValues(t, expectedMultipleEndpoints, multipleEndpoints)
+}

pkg/forwarder/forwarder.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"sort"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -171,6 +172,14 @@ func (f *DefaultForwarder) Start() error {
 	go f.handleFailedTransactions()
 	f.internalState = Started
 	log.Infof("DefaultForwarder started (%v workers)", f.NumberOfWorkers)
+
+	// log endpoints configuration
+	endpointLogs := make([]string, 0, len(f.KeysPerDomains))
+	for domain, apiKeys := range f.KeysPerDomains {
+		endpointLogs = append(endpointLogs, fmt.Sprintf("\"%s\" (%v api key(s))", domain, len(apiKeys)))
+	}
+	log.Infof("DefaultForwarder sending to %v endpoint(s): %s", len(endpointLogs), strings.Join(endpointLogs, " ; "))
+
 	return nil
 }