tracing.go

package tracing

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"sync"

	openzipkin "github.com/openzipkin/zipkin-go"
	zipkinHTTP "github.com/openzipkin/zipkin-go/reporter/http"
	"go.opencensus.io/exporter/zipkin"
	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/trace"
)

var (
	// Enabled reports whether tracing is globally enabled or not.
	Enabled bool

	// The active tracing configuration, if Enabled is true.
	config tracingConfig

	initOnce sync.Once
)

const globalTracingConfigPath = "/etc/tracing/client.conf"

type tracingConfig struct {
	ReportURL string `json:"report_url"`
	Sample    string `json:"sample"`
}

// Read the global tracing configuration file. Its location is
// hardcoded, but it can be overriden using the TRACING_CONFIG
// environment variable.
func readTracingConfig() error {
	// Read and decode configuration.
	cfgPath := globalTracingConfigPath
	if s := os.Getenv("TRACING_CONFIG"); s != "" {
		cfgPath = s
	}
	data, err := ioutil.ReadFile(cfgPath)
	if err != nil {
		return err
	}

	if err := json.Unmarshal(data, &config); err != nil {
		log.Printf("warning: error in tracing configuration: %v, tracing disabled", err)
		return err
	}

	if config.ReportURL == "" {
		log.Printf("warning: tracing configuration contains no report_url, tracing disabled")
		return errors.New("no report_url")
	}

	return nil
}

// Compute the service name for Zipkin: this is usually the program
// name (without path), but it can be overriden by the TRACING_SERVICE
// environment variable.
func getServiceName() string {
	if s := os.Getenv("TRACING_SERVICE"); s != "" {
		return s
	}
	return filepath.Base(os.Args[0])
}

// Initialize tracing. Tracing will be enabled if the system-wide
// tracing configuration file is present and valid. Explicitly set
// TRACING_ENABLE=0 in the environment to disable tracing.
//
// We need to check the configuration as soon as possible, because
// it's likely that client transports are created before HTTP servers,
// and we need to wrap them with opencensus at creation time.
func init() {
	// Kill switch from environment.
	if s := os.Getenv("TRACING_ENABLE"); s == "0" {
		return
	}

	if err := readTracingConfig(); err != nil {
		return
	}

	Enabled = true
}

func initTracing(endpointAddr string) {
	initOnce.Do(func() {
		localEndpoint, err := openzipkin.NewEndpoint(getServiceName(), endpointAddr)
		if err != nil {
			log.Printf("warning: error creating tracing endpoint: %v, tracing disabled", err)
			return
		}

		reporter := zipkinHTTP.NewReporter(config.ReportURL)
		ze := zipkin.NewExporter(reporter, localEndpoint)
		trace.RegisterExporter(ze)

		var tc trace.Config
		switch config.Sample {
		case "", "always":
			tc.DefaultSampler = trace.AlwaysSample()
		case "never":
			tc.DefaultSampler = trace.NeverSample()
		default:
			frac, err := strconv.ParseFloat(config.Sample, 64)
			if err != nil {
				log.Printf("warning: error in tracing configuration: sample: %v, tracing disabled", err)
				return
			}
			tc.DefaultSampler = trace.ProbabilitySampler(frac)
		}
		trace.ApplyConfig(tc)

		log.Printf("tracing enabled (report_url %s)", config.ReportURL)

		Enabled = true
	})
}

// WrapTransport optionally wraps a http.RoundTripper with OpenCensus
// tracing functionality, if it is globally enabled.
func WrapTransport(t http.RoundTripper) http.RoundTripper {
	if Enabled {
		t = &ochttp.Transport{Base: t}
	}
	return t
}

// WrapHandler wraps a http.Handler with OpenCensus tracing
// functionality, if globally enabled.
func WrapHandler(h http.Handler, endpointAddr string) http.Handler {
	if Enabled {
		initTracing(endpointAddr)
		h = &ochttp.Handler{Handler: h}
	}
	return h
}