From 9738c9e7809be62e54f04731798d5f431721bd33 Mon Sep 17 00:00:00 2001
From: Ryuzo Yamamoto <ryuzo.yamamoto@gmail.com>
Date: Wed, 7 Nov 2018 23:48:51 +0900
Subject: [PATCH] config: enable setting "MaxTPS" from environment variable
 (#519)

---
 config/merge_env.go      | 11 +++++++++++
 config/merge_env_test.go | 22 ++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/config/merge_env.go b/config/merge_env.go
index fd49d42fa..adcb8ec28 100644
--- a/config/merge_env.go
+++ b/config/merge_env.go
@@ -26,6 +26,7 @@ const (
 	envLogLevel        = "DD_LOG_LEVEL"             // logging level
 	envAnalyzedSpans   = "DD_APM_ANALYZED_SPANS"    // spans to analyze for transactions
 	envConnectionLimit = "DD_CONNECTION_LIMIT"      // (deprecated) limit of unique connections
+	envMaxTPS          = "DD_MAX_TPS"               // maximum limit to the total number of traces per second to sample (MaxTPS)
 )
 
 // loadEnv applies overrides from environment variables to the trace agent configuration
@@ -122,6 +123,16 @@ func (c *AgentConfig) loadEnv() {
 			}
 		}
 	}
+
+	if v := os.Getenv(envMaxTPS); v != "" {
+		maxTPS, err := strconv.ParseFloat(v, 64)
+		if err != nil {
+			log.Errorf("Failed to parse %s: it should be a float number", envMaxTPS)
+		} else {
+			c.MaxTPS = maxTPS
+		}
+	}
+
 }
 
 func parseNameAndRate(token string) (string, float64, error) {
diff --git a/config/merge_env_test.go b/config/merge_env_test.go
index c67de00b7..7f4cc28ce 100644
--- a/config/merge_env_test.go
+++ b/config/merge_env_test.go
@@ -1,6 +1,7 @@
 package config
 
 import (
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -45,3 +46,24 @@ func TestAnalyzedSpansEnvConfigParsing(t *testing.T) {
 		assert.NotNil(err)
 	})
 }
+
+func TestLoadEnvMaxTPS(t *testing.T) {
+	assert := assert.New(t)
+
+	t.Run("default", func(t *testing.T) {
+		ac := New()
+		ac.loadEnv()
+		assert.EqualValues(10.0, ac.MaxTPS)
+	})
+
+	t.Run("env", func(t *testing.T) {
+		if err := os.Setenv("DD_MAX_TPS", "123.4"); err != nil {
+			t.Fatal(err)
+		}
+		defer os.Unsetenv("DD_MAX_TPS")
+
+		ac := New()
+		ac.loadEnv()
+		assert.EqualValues(123.4, ac.MaxTPS)
+	})
+}