[prometheus-reporter] Add ability to infer listen port based on env var

prometheus/config.go

@@ -38,6 +38,11 @@ type Configuration struct {
 	// handler on the default HTTP serve mux without listening.
 	ListenAddress string `yaml:"listenAddress"`
 
+	// DynamicListenAddress if specified will be used instead of just registering the
+	// handler on the default HTTP serve mux without listening.
+	// Note: if DynamicListenAddress is specified, ListenAddress is ignored.
+	DynamicListenAddress *ListenAddressConfiguration `yaml:"dynamicListenAddress"`
+
 	// TimerType is the default Prometheus type to use for Tally timers.
 	TimerType string `yaml:"timerType"`
 
@@ -129,7 +134,12 @@ func (c Configuration) NewReporter(
 		path = handlerPath
 	}
 
-	if addr := strings.TrimSpace(c.ListenAddress); addr == "" {
+	addr, resolved, err := c.resolveListenAddress()
+	if err != nil {
+		return nil, err
+	}
+
+	if !resolved {
 		http.Handle(path, reporter.HTTPHandler())
 	} else {
 		mux := http.NewServeMux()
@@ -143,3 +153,22 @@ func (c Configuration) NewReporter(
 
 	return reporter, nil
 }
+
+func (c Configuration) resolveListenAddress() (addr string, resolved bool, err error) {
+	// first try DynamicListenAddress
+	if c.DynamicListenAddress != nil {
+		addr, err := c.DynamicListenAddress.Resolve()
+		if err != nil {
+			return "", false, err
+		}
+		return addr, true, nil
+	}
+
+	// next, try ListenAddress
+	addr = strings.TrimSpace(c.ListenAddress)
+	if addr == "" {
+		return "", false, nil
+	}
+
+	return addr, true, nil
+}

prometheus/config_test.go

@@ -0,0 +1,121 @@
+// Copyright (c) 2018 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package prometheus
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestOldStyleListenAddressConfig(t *testing.T) {
+	go func() {
+		if err := http.ListenAndServe(":8080", nil); err != nil {
+			panic(err)
+		}
+	}()
+
+	config := Configuration{}
+	_, err := config.NewReporter(ConfigurationOptions{})
+	require.NoError(t, err)
+
+	start := time.Now()
+	testTimeout := time.Second
+	timeCheckInterval := 50 * time.Millisecond
+	getSuccess := false
+	for time.Since(start) < testTimeout {
+		resp, err := http.Get("http://0.0.0.0:8080/metrics") // default test address
+		if err == nil && resp.StatusCode == 200 {
+			getSuccess = true
+			break
+		}
+		time.Sleep(timeCheckInterval)
+	}
+
+	require.True(t, getSuccess)
+}
+
+func TestNewStyleListenAddressConfig(t *testing.T) {
+	port := 12321
+	config := Configuration{
+		DynamicListenAddress: &ListenAddressConfiguration{
+			Hostname: "0.0.0.0",
+			Port: &PortConfiguration{
+				PortType: ConfigResolver,
+				Value:    &port,
+			},
+		},
+	}
+	_, err := config.NewReporter(ConfigurationOptions{})
+	require.NoError(t, err)
+
+	start := time.Now()
+	testTimeout := time.Second
+	timeCheckInterval := 50 * time.Millisecond
+	getSuccess := false
+	for time.Since(start) < testTimeout {
+		resp, err := http.Get("http://0.0.0.0:12321/metrics")
+		if err == nil && resp.StatusCode == 200 {
+			getSuccess = true
+			break
+		}
+		time.Sleep(timeCheckInterval)
+	}
+
+	require.True(t, getSuccess)
+}
+
+func TestNewStyleEnvVarBasedListenAddressConfig(t *testing.T) {
+	port := 13331
+	envVarName := "SOMETHINGLONGANDABSURD"
+	require.NoError(t, os.Setenv(envVarName, fmt.Sprintf("%d", port)))
+
+	config := Configuration{
+		DynamicListenAddress: &ListenAddressConfiguration{
+			Hostname: "0.0.0.0",
+			Port: &PortConfiguration{
+				PortType:   EnvironmentResolver,
+				EnvVarName: &envVarName,
+			},
+		},
+	}
+	_, err := config.NewReporter(ConfigurationOptions{})
+	require.NoError(t, err)
+
+	start := time.Now()
+	testTimeout := time.Second
+	timeCheckInterval := 50 * time.Millisecond
+	getSuccess := false
+	for time.Since(start) < testTimeout {
+		resp, err := http.Get("http://0.0.0.0:13331/metrics")
+		if err == nil && resp.StatusCode == 200 {
+			getSuccess = true
+			break
+		}
+		time.Sleep(timeCheckInterval)
+	}
+
+	require.True(t, getSuccess)
+}

prometheus/listen_config.go

@@ -0,0 +1,98 @@
+// Copyright (c) 2018 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package prometheus
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"strconv"
+)
+
+// ListenAddressResolver is a type of port resolver
+type ListenAddressResolver string
+
+const (
+	// ConfigResolver resolves port using a value provided in config
+	ConfigResolver ListenAddressResolver = "config"
+	// EnvironmentResolver resolves port using an environment variable
+	// of which the name is provided in config
+	EnvironmentResolver ListenAddressResolver = "environment"
+)
+
+// ListenAddressConfiguration is the configuration for resolving a listen address.
+type ListenAddressConfiguration struct {
+	// Hostname is the hostname to use
+	Hostname string `yaml:"hostname" validate:"nonzero"`
+
+	// port specifies the port to listen on
+	Port *PortConfiguration `yaml:"port" validate:"nonzero"`
+}
+
+// PortConfiguration specifies the port to use
+type PortConfiguration struct {
+	// PortType is the port type for the port
+	PortType ListenAddressResolver `yaml:"portType" validate:"nonzero"`
+
+	// Value is the config specified port if using config port type.
+	Value *int `yaml:"value"`
+
+	// EnvVarName is the environment specified port if using environment port type.
+	EnvVarName *string `yaml:"envVarName"`
+}
+
+// Resolve returns the resolved listen address given the configuration.
+func (c ListenAddressConfiguration) Resolve() (string, error) {
+	if c.Port == nil {
+		return "", errors.New("missing port config")
+	}
+	p := c.Port
+
+	var port int
+	switch p.PortType {
+	case ConfigResolver:
+		if p.Value == nil {
+			err := fmt.Errorf("missing port type using: resolver=%s",
+				string(p.PortType))
+			return "", err
+		}
+		port = *p.Value
+	case EnvironmentResolver:
+		if p.EnvVarName == nil {
+			err := fmt.Errorf("missing port env var name using: resolver=%s",
+				string(p.PortType))
+			return "", err
+		}
+		portStr := os.Getenv(*p.EnvVarName)
+		var err error
+		port, err = strconv.Atoi(portStr)
+		if err != nil {
+			err := fmt.Errorf("invalid port env var value using: resolver=%s, name=%s",
+				string(p.PortType), *p.EnvVarName)
+			return "", err
+		}
+	default:
+		return "", fmt.Errorf("unknown port type: resolver=%s",
+			string(p.PortType))
+	}
+
+	return fmt.Sprintf("%s:%d", c.Hostname, port), nil
+}