tracing.go

package tracing

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

	othttp "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	b3 "go.opentelemetry.io/contrib/propagators/b3"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/zipkin"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	"go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
)

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

	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() (*tracingConfig, 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 nil, err
	}

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

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

	return &config, 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
	}
	if s, err := os.Executable(); err == nil {
		return filepath.Base(s)
	}
	return "unknown_service"
}

func defaultResource(serviceName string) *resource.Resource {
	hostname, _ := os.Hostname()
	r, _ := resource.Merge(
		resource.Default(),
		resource.NewWithAttributes(
			semconv.SchemaURL,
			semconv.ServiceNameKey.String(serviceName),
			semconv.HostNameKey.String(hostname),
		),
	)
	return r
}

// 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 initTracing(serviceName string) {
	initOnce.Do(func() {
		// Kill switch from environment.
		if s := os.Getenv("TRACING_ENABLE"); s == "0" {
			return
		}

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

		ze, err := zipkin.New(config.ReportURL)
		if err != nil {
			log.Printf("error creating Zipkin exporter: %v", err)
			return
		}

		// The sampling policy only applies to incoming requests for
		// which tracing is not already enabled: in this case, we
		// always pass-through.
		var sampler trace.Sampler
		switch config.Sample {
		case "", "always":
			sampler = trace.AlwaysSample()
		case "never":
			sampler = 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
			}
			sampler = trace.TraceIDRatioBased(frac)
		}

		tp := trace.NewTracerProvider(
			trace.WithSampler(trace.ParentBased(sampler)),
			trace.WithBatcher(ze),
			trace.WithResource(defaultResource(serviceName)),
		)

		otel.SetTracerProvider(tp)

		otel.SetTextMapPropagator(
			propagation.NewCompositeTextMapPropagator(
				propagation.TraceContext{},
				propagation.Baggage{},
				b3.New()))

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

		Enabled = true
	})
}

// Init tracing support, if not using WrapHandler.
func Init() {
	initTracing(getServiceName())
}

// WrapTransport optionally wraps a http.RoundTripper with OpenCensus
// tracing functionality, if it is globally enabled.
//
// Must call Init() first.
func WrapTransport(t http.RoundTripper) http.RoundTripper {
	if Enabled {
		t = othttp.NewTransport(t)
	}
	return t
}

// WrapHandler wraps a http.Handler with OpenCensus tracing
// functionality, if globally enabled. Automatically calls Init().
func WrapHandler(h http.Handler, endpointAddr string) http.Handler {
	serviceName := getServiceName()
	initTracing(serviceName)
	if !Enabled {
		return h
	}

	// Format span names with the request URL path.
	return othttp.NewHandler(
		h, serviceName,
		othttp.WithSpanNameFormatter(func(op string, r *http.Request) string {
			return r.URL.Path
		}),
	)
}