Move agent main.go to use viper instead of yaml.v2 (#471)

* Move agent to use viper instead of yaml.v2

Author: Steven Karis

cmd/ocagent/main.go

@@ -19,7 +19,6 @@ package main
 import (
 	"context"
 	"fmt"
-	"io/ioutil"
 	"log"
 	"net"
 	"net/http"
@@ -80,19 +79,22 @@ func main() {
 }
 
 func runOCAgent() {
-	yamlBlob, err := ioutil.ReadFile(configYAMLFile)
+	viperCfg.SetConfigFile(configYAMLFile)
+	err := viperCfg.ReadInConfig()
 	if err != nil {
 		log.Fatalf("Cannot read the YAML file %v error: %v", configYAMLFile, err)
 	}
-	agentConfig, err := config.ParseOCAgentConfig(yamlBlob)
+
+	var agentConfig config.Config
+	err = viperCfg.Unmarshal(&agentConfig)
 	if err != nil {
-		log.Fatalf("Failed to parse own configuration %v error: %v", configYAMLFile, err)
+		log.Fatalf("Error unmarshalling yaml config file %v: %v", configYAMLFile, err)
 	}
 
 	// Ensure that we check and catch any logical errors with the
 	// configuration e.g. if an receiver shares the same address
 	// as an exporter which would cause a self DOS and waste resources.
-	if err := agentConfig.CheckLogicalConflicts(yamlBlob); err != nil {
+	if err := agentConfig.CheckLogicalConflicts(); err != nil {
 		log.Fatalf("Configuration logical error: %v", err)
 	}
 
@@ -110,12 +112,6 @@ func runOCAgent() {
 		log.Fatalf("Failed to start net/http/pprof: %v", err)
 	}
 
-	// TODO(skaris): move the rest of the configs to use viper
-	err = viperutils.LoadYAMLBytes(viperCfg, []byte(yamlBlob))
-	if err != nil {
-		log.Fatalf("Config: failed to create viper from YAML: %v", err)
-	}
-
 	traceExporters, metricsExporters, closeFns, err := config.ExportersFromViperConfig(logger, viperCfg)
 	if err != nil {
 		log.Fatalf("Config: failed to create exporters from YAML: %v", err)
@@ -125,7 +121,7 @@ func runOCAgent() {
 	commonMetricsSink := processor.NewMultiMetricsDataProcessor(metricsExporters)
 
 	// Add other receivers here as they are implemented
-	ocReceiverDoneFn, err := runOCReceiver(logger, agentConfig, commonSpanSink, commonMetricsSink)
+	ocReceiverDoneFn, err := runOCReceiver(logger, &agentConfig, commonSpanSink, commonMetricsSink)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -138,6 +134,7 @@ func runOCAgent() {
 		closeFns = append(closeFns, zCloseFn)
 	}
 
+	// TODO: Generalize the startup of these receivers when unifying them w/ collector
 	// If the Zipkin receiver is enabled, then run it
 	if agentConfig.ZipkinReceiverEnabled() {
 		zipkinReceiverAddr := agentConfig.ZipkinReceiverAddress()
@@ -167,7 +164,7 @@ func runOCAgent() {
 
 	// If the Prometheus receiver is enabled, then run it.
 	if agentConfig.PrometheusReceiverEnabled() {
-		promDoneFn, err := runPrometheusReceiver(agentConfig.PrometheusConfiguration(), commonMetricsSink)
+		promDoneFn, err := runPrometheusReceiver(viperCfg, commonMetricsSink)
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -228,10 +225,10 @@ func runOCReceiver(logger *zap.Logger, acfg *config.Config, tdp processor.TraceD
 		opencensusreceiver.WithCorsOrigins(corsOrigins))
 
 	if err != nil {
-		return nil, fmt.Errorf("Failed to create the OpenCensus receiver on address %q: error %v", addr, err)
+		return nil, fmt.Errorf("failed to create the OpenCensus receiver on address %q: error %v", addr, err)
 	}
 	if err := view.Register(observability.AllViews...); err != nil {
-		return nil, fmt.Errorf("Failed to register internal.AllViews: %v", err)
+		return nil, fmt.Errorf("failed to register internal.AllViews: %v", err)
 	}
 
 	// Temporarily disabling the grpc metrics since they do not provide good data at this moment,
@@ -245,19 +242,19 @@ func runOCReceiver(logger *zap.Logger, acfg *config.Config, tdp processor.TraceD
 	switch {
 	case acfg.CanRunOpenCensusTraceReceiver() && acfg.CanRunOpenCensusMetricsReceiver():
 		if err := ocr.Start(ctx, tdp, mdp); err != nil {
-			return nil, fmt.Errorf("Failed to start Trace and Metrics Receivers: %v", err)
+			return nil, fmt.Errorf("failed to start Trace and Metrics Receivers: %v", err)
 		}
 		log.Printf("Running OpenCensus Trace and Metrics receivers as a gRPC service at %q", addr)
 
 	case acfg.CanRunOpenCensusTraceReceiver():
 		if err := ocr.StartTraceReception(ctx, tdp); err != nil {
-			return nil, fmt.Errorf("Failed to start TraceReceiver: %v", err)
+			return nil, fmt.Errorf("failed to start TraceReceiver: %v", err)
 		}
 		log.Printf("Running OpenCensus Trace receiver as a gRPC service at %q", addr)
 
 	case acfg.CanRunOpenCensusMetricsReceiver():
 		if err := ocr.StartMetricsReception(ctx, mdp); err != nil {
-			return nil, fmt.Errorf("Failed to start MetricsReceiver: %v", err)
+			return nil, fmt.Errorf("failed to start MetricsReceiver: %v", err)
 		}
 		log.Printf("Running OpenCensus Metrics receiver as a gRPC service at %q", addr)
 	}
@@ -283,10 +280,10 @@ func runJaegerReceiver(collectorThriftPort, collectorHTTPPort int, next processo
 		// and not use their defaults of 5778, 6831, 6832
 	})
 	if err != nil {
-		return nil, fmt.Errorf("Failed to create new Jaeger receiver: %v", err)
+		return nil, fmt.Errorf("failed to create new Jaeger receiver: %v", err)
 	}
 	if err := jtr.StartTraceReception(context.Background(), next); err != nil {
-		return nil, fmt.Errorf("Failed to start Jaeger receiver: %v", err)
+		return nil, fmt.Errorf("failed to start Jaeger receiver: %v", err)
 	}
 	doneFn = func() error {
 		return jtr.StopTraceReception(context.Background())
@@ -298,11 +295,11 @@ func runJaegerReceiver(collectorThriftPort, collectorHTTPPort int, next processo
 func runZipkinReceiver(addr string, next processor.TraceDataProcessor) (doneFn func() error, err error) {
 	zi, err := zipkinreceiver.New(addr)
 	if err != nil {
-		return nil, fmt.Errorf("Failed to create the Zipkin receiver: %v", err)
+		return nil, fmt.Errorf("failed to create the Zipkin receiver: %v", err)
 	}
 
 	if err := zi.StartTraceReception(context.Background(), next); err != nil {
-		return nil, fmt.Errorf("Cannot start Zipkin receiver with address %q: %v", addr, err)
+		return nil, fmt.Errorf("cannot start Zipkin receiver with address %q: %v", addr, err)
 	}
 	doneFn = func() error {
 		return zi.StopTraceReception(context.Background())
@@ -314,11 +311,11 @@ func runZipkinReceiver(addr string, next processor.TraceDataProcessor) (doneFn f
 func runZipkinScribeReceiver(config *config.ScribeReceiverConfig, next processor.TraceDataProcessor) (doneFn func() error, err error) {
 	zs, err := scribe.NewReceiver(config.Address, config.Port, config.Category)
 	if err != nil {
-		return nil, fmt.Errorf("Failed to create the Zipkin Scribe receiver: %v", err)
+		return nil, fmt.Errorf("failed to create the Zipkin Scribe receiver: %v", err)
 	}
 
 	if err := zs.StartTraceReception(context.Background(), next); err != nil {
-		return nil, fmt.Errorf("Cannot start Zipkin Scribe receiver with %v: %v", config, err)
+		return nil, fmt.Errorf("cannot start Zipkin Scribe receiver with %v: %v", config, err)
 	}
 	doneFn = func() error {
 		return zs.StopTraceReception(context.Background())
@@ -327,8 +324,8 @@ func runZipkinScribeReceiver(config *config.ScribeReceiverConfig, next processor
 	return doneFn, nil
 }
 
-func runPrometheusReceiver(promConfig *prometheusreceiver.Configuration, next processor.MetricsDataProcessor) (doneFn func() error, err error) {
-	pmr, err := prometheusreceiver.New(promConfig)
+func runPrometheusReceiver(v *viper.Viper, next processor.MetricsDataProcessor) (doneFn func() error, err error) {
+	pmr, err := prometheusreceiver.New(v.Sub("receivers.prometheus"))
 	if err != nil {
 		return nil, err
 	}

internal/config/config.go

@@ -25,7 +25,6 @@ import (
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
-	yaml "gopkg.in/yaml.v2"
 
 	"github.com/census-instrumentation/opencensus-service/exporter/awsexporter"
 	"github.com/census-instrumentation/opencensus-service/exporter/datadogexporter"
@@ -87,9 +86,9 @@ var defaultScribeConfiguration = &ScribeReceiverConfig{
 // * ZPages
 // * Exporters
 type Config struct {
-	Receivers *Receivers    `yaml:"receivers"`
-	ZPages    *ZPagesConfig `yaml:"zpages"`
-	Exporters *Exporters    `yaml:"exporters"`
+	Receivers *Receivers    `mapstructure:"receivers"`
+	ZPages    *ZPagesConfig `mapstructure:"zpages"`
+	Exporters *Exporters    `mapstructure:"exporters"`
 }
 
 // Receivers denotes configurations for the various telemetry ingesters, such as:
@@ -98,10 +97,10 @@ type Config struct {
 // * Prometheus (metrics)
 // * Zipkin (traces)
 type Receivers struct {
-	OpenCensus *ReceiverConfig       `yaml:"opencensus"`
-	Zipkin     *ReceiverConfig       `yaml:"zipkin"`
-	Jaeger     *ReceiverConfig       `yaml:"jaeger"`
-	Scribe     *ScribeReceiverConfig `yaml:"zipkin-scribe"`
+	OpenCensus *ReceiverConfig       `mapstructure:"opencensus"`
+	Zipkin     *ReceiverConfig       `mapstructure:"zipkin"`
+	Jaeger     *ReceiverConfig       `mapstructure:"jaeger"`
+	Scribe     *ScribeReceiverConfig `mapstructure:"zipkin-scribe"`
 
 	// Prometheus contains the Prometheus configurations.
 	// Such as:
@@ -111,7 +110,7 @@ type Receivers struct {
 	//
 	//          static_configs:
 	//              - targets: ['localhost:9988']
-	Prometheus *prometheusreceiver.Configuration `yaml:"prometheus"`
+	Prometheus *prometheusreceiver.Configuration `mapstructure:"prometheus"`
 }
 
 // ReceiverConfig is the per-receiver configuration that identifies attributes
@@ -120,23 +119,23 @@ type Receivers struct {
 // * Various ports
 type ReceiverConfig struct {
 	// The address to which the OpenCensus receiver will be bound and run on.
-	Address             string `yaml:"address"`
-	CollectorHTTPPort   int    `yaml:"collector_http_port"`
-	CollectorThriftPort int    `yaml:"collector_thrift_port"`
+	Address             string `mapstructure:"address"`
+	CollectorHTTPPort   int    `mapstructure:"collector_http_port"`
+	CollectorThriftPort int    `mapstructure:"collector_thrift_port"`
 
 	// The allowed CORS origins for HTTP/JSON requests the grpc-gateway adapter
 	// for the OpenCensus receiver. See github.com/rs/cors
 	// An empty list means that CORS is not enabled at all. A wildcard (*) can be
 	// used to match any origin or one or more characters of an origin.
-	CorsAllowedOrigins []string `yaml:"cors_allowed_origins"`
+	CorsAllowedOrigins []string `mapstructure:"cors_allowed_origins"`
 
 	// DisableTracing disables trace receiving and is only applicable to trace receivers.
-	DisableTracing bool `yaml:"disable_tracing"`
+	DisableTracing bool `mapstructure:"disable_tracing"`
 	// DisableMetrics disables metrics receiving and is only applicable to metrics receivers.
-	DisableMetrics bool `yaml:"disable_metrics"`
+	DisableMetrics bool `mapstructure:"disable_metrics"`
 
 	// TLSCredentials is a (cert_file, key_file) configuration.
-	TLSCredentials *TLSCredentials `yaml:"tls_credentials"`
+	TLSCredentials *TLSCredentials `mapstructure:"tls_credentials"`
 }
 
 // ScribeReceiverConfig carries the settings for the Zipkin Scribe receiver.
@@ -147,23 +146,23 @@ type ScribeReceiverConfig struct {
 	// a listener for at most one of the host's IP addresses.
 	//
 	// The default value bind to all available interfaces on the local computer.
-	Address string `yaml:"address" mapstructure:"address"`
-	Port    uint16 `yaml:"port" mapstructure:"port"`
+	Address string `mapstructure:"address" mapstructure:"address"`
+	Port    uint16 `mapstructure:"port" mapstructure:"port"`
 	// Category is the string that will be used to identify the scribe log messages
 	// that contain Zipkin spans.
-	Category string `yaml:"category" mapstructure:"category"`
+	Category string `mapstructure:"category" mapstructure:"category"`
 }
 
 // Exporters denotes the configurations for the various backends
 // that this service exports observability signals to.
 type Exporters struct {
-	Zipkin *zipkinexporter.ZipkinConfig `yaml:"zipkin"`
+	Zipkin *zipkinexporter.ZipkinConfig `mapstructure:"zipkin"`
 }
 
 // ZPagesConfig denotes the configuration that zPages will be run with.
 type ZPagesConfig struct {
-	Disabled bool `yaml:"disabled"`
-	Port     int  `yaml:"port"`
+	Disabled bool `mapstructure:"disabled"`
+	Port     int  `mapstructure:"port"`
 }
 
 // OpenCensusReceiverAddress is a helper to safely retrieve the address
@@ -375,32 +374,17 @@ func (c *Config) OpenCensusReceiverTLSCredentialsServerOption() (opt opencensusr
 	return tlsCreds.ToOpenCensusReceiverServerOption()
 }
 
-// ParseOCAgentConfig unmarshals byte content in the YAML file format
-// to  retrieve the configuration that will be used to run the OpenCensus agent.
-func ParseOCAgentConfig(yamlBlob []byte) (*Config, error) {
-	var cfg Config
-	if err := yaml.Unmarshal(yamlBlob, &cfg); err != nil {
-		return nil, err
-	}
-	return &cfg, nil
-}
-
 // CheckLogicalConflicts serves to catch logical errors such as
 // if the Zipkin receiver port conflicts with that of the exporter,
 // lest we'll have a self DOS because spans will be exported "out" from
 // the exporter, yet be received from the receiver, then sent back out
 // and back in a never ending loop.
-func (c *Config) CheckLogicalConflicts(blob []byte) error {
-	var cfg Config
-	if err := yaml.Unmarshal(blob, &cfg); err != nil {
-		return err
-	}
-
-	if cfg.Exporters == nil || cfg.Exporters.Zipkin == nil || !c.ZipkinReceiverEnabled() {
+func (c *Config) CheckLogicalConflicts() error {
+	if c.Exporters == nil || c.Exporters.Zipkin == nil || !c.ZipkinReceiverEnabled() {
 		return nil
 	}
 
-	zc := cfg.Exporters.Zipkin
+	zc := c.Exporters.Zipkin
 
 	zExporterAddr := zc.EndpointURL()
 	zExporterURL, err := url.Parse(zExporterAddr)
@@ -485,7 +469,7 @@ func ExportersFromViperConfig(logger *zap.Logger, v *viper.Viper) ([]processor.T
 	for _, cfg := range parseFns {
 		tes, mes, tesDoneFns, err := cfg.fn(exportersViper)
 		if err != nil {
-			err = fmt.Errorf("Failed to create config for %q: %v", cfg.name, err)
+			err = fmt.Errorf("failed to create config for %q: %v", cfg.name, err)
 			return nil, nil, nil, err
 		}
 

internal/config/config_test.go

@@ -17,10 +17,11 @@ package config_test
 import (
 	"testing"
 
-	yaml "gopkg.in/yaml.v2"
+	"github.com/spf13/viper"
 
 	"github.com/census-instrumentation/opencensus-service/exporter/zipkinexporter"
 	"github.com/census-instrumentation/opencensus-service/internal/config"
+	"github.com/census-instrumentation/opencensus-service/internal/config/viperutils"
 )
 
 // Issue #233: Zipkin receiver and exporter loopback detection
@@ -39,11 +40,17 @@ exporters:
         endpoint: "http://localhost:9411/api/v2/spans"
 `)
 
-	cfg, err := config.ParseOCAgentConfig(regressionYAML)
+	v := viper.New()
+	err := viperutils.LoadYAMLBytes(v, regressionYAML)
 	if err != nil {
 		t.Fatalf("Unexpected YAML parse error: %v", err)
 	}
-	if err := cfg.CheckLogicalConflicts(regressionYAML); err != nil {
+	var cfg config.Config
+	err = v.Unmarshal(&cfg)
+	if err != nil {
+		t.Fatalf("Unexpected error unmarshaling viper: %s", err)
+	}
+	if err := cfg.CheckLogicalConflicts(); err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
 
@@ -53,10 +60,10 @@ exporters:
 
 	var ecfg struct {
 		Exporters *struct {
-			Zipkin *zipkinexporter.ZipkinConfig `yaml:"zipkin"`
-		} `yaml:"exporters"`
+			Zipkin *zipkinexporter.ZipkinConfig `mapstructure:"zipkin"`
+		} `mapstructure:"exporters"`
 	}
-	_ = yaml.Unmarshal(regressionYAML, &ecfg)
+	_ = v.Unmarshal(&ecfg)
 	if g, w := ecfg.Exporters.Zipkin.EndpointURL(), "http://localhost:9411/api/v2/spans"; g != w {
 		t.Errorf("Exporters.Zipkin.EndpointURL mismatch\nGot: %s\nWant:%s", g, w)
 	}
@@ -84,10 +91,15 @@ receivers:
     zipkin:
         address: "localhost:9410"`)
 
-	cfg, err := config.ParseOCAgentConfig(regressionYAML)
+	v := viper.New()
+	err := viperutils.LoadYAMLBytes(v, regressionYAML)
 	if err != nil {
 		t.Fatalf("Unexpected YAML parse error: %v", err)
 	}
+	err = v.Unmarshal(cfg)
+	if err != nil {
+		t.Fatalf("Unexpected error unmarshaling viper: %s", err)
+	}
 
 	if cfg.CanRunOpenCensusTraceReceiver() {
 		t.Fatal("yaml.CanRunOpenCensusTraceReceiver: Unexpected True for a nil Receiver.OpenCensus")

internal/config/tls_credentials.go

@@ -18,10 +18,10 @@ package config
 // that are used for starting a server.
 type TLSCredentials struct {
 	// CertFile is the file path containing the TLS certificate.
-	CertFile string `yaml:"cert_file"`
+	CertFile string `mapstructure:"cert_file"`
 
 	// KeyFile is the file path containing the TLS key.
-	KeyFile string `yaml:"key_file"`
+	KeyFile string `mapstructure:"key_file"`
 }
 
 // nonEmpty returns true if the TLSCredentials are non-nil and