example/router.go

package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"strconv"
	"time"

	"gitlab.com/gitlab-org/labkit/correlation"
	"gitlab.com/gitlab-org/labkit/tracing"
	tracingcorrelation "gitlab.com/gitlab-org/labkit/tracing/correlation"
)

func main() {
	tracing.Initialize(tracing.WithServiceName("router"))

	tr := &http.Transport{
		MaxIdleConns:       10,
		IdleConnTimeout:    30 * time.Second,
		DisableCompression: true,
	}

	client := &http.Client{
		Transport: correlation.NewInstrumentedRoundTripper(tracing.NewRoundTripper(tr)),
	}

	// Listen and propagate traces
	http.Handle("/query",
		// Add the tracing middleware in
		correlation.InjectCorrelationID(
			tracing.Handler(
				tracingcorrelation.BaggageHandler(

					http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
						query := r.URL.Query()
						ttlString := query.Get("ttl")

						var ttl int

						var err error

						if ttlString == "" {
							ttl = 1
						} else {
							ttl, err = strconv.Atoi(ttlString)
							if err != nil {
								ttl = 1
							}
						}

						ttl--
						if ttl < 0 {
							fmt.Fprintf(w, "Hello")
							return
						}

						nextURL := fmt.Sprintf("http://localhost:8080/query?ttl=%d", ttl)

						req, err := http.NewRequest(http.MethodGet, nextURL, nil)
						if err != nil {
							w.WriteHeader(http.StatusInternalServerError)
							return
						}

						req = req.WithContext(r.Context())

						resp, err := client.Do(req)
						if err != nil {
							w.WriteHeader(http.StatusInternalServerError)
							return
						}
						defer resp.Body.Close()

						_, err = io.Copy(w, resp.Body)
						if err != nil {
							w.WriteHeader(http.StatusInternalServerError)
							return
						}
					})),
				// Use this route identifier with the tracing middleware
				tracing.WithRouteIdentifier("/query"),
			),
			correlation.WithPropagation()))

	log.Fatal(http.ListenAndServe(":8080", nil))
}