From c452f31eebb34830bfd9de2d6f080750d69a2940 Mon Sep 17 00:00:00 2001
From: Marcus Cobden <marcus@weave.works>
Date: Thu, 23 Aug 2018 13:02:09 +0100
Subject: [PATCH] Add middleware to detect trace context on HTTP

---
 prog/app.go | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/prog/app.go b/prog/app.go
index 91d0d09ff1..cce8ac6932 100644
--- a/prog/app.go
+++ b/prog/app.go
@@ -13,6 +13,8 @@ import (
 	"time"
 
 	"github.com/gorilla/mux"
+	"github.com/opentracing-contrib/go-stdlib/nethttp"
+	"github.com/opentracing/opentracing-go"
 	"github.com/prometheus/client_golang/prometheus"
 	log "github.com/sirupsen/logrus"
 	"github.com/tylerb/graceful"
@@ -70,11 +72,27 @@ func router(collector app.Collector, controlRouter app.ControlRouter, pipeRouter
 			uiHandler))
 	router.PathPrefix("/").Name("static").Handler(uiHandler)
 
-	instrument := middleware.Instrument{
+	instrumentDuration := middleware.Instrument{
 		RouteMatcher: router,
 		Duration:     requestDuration,
 	}
-	return instrument.Wrap(router)
+	instrumentTraces := middleware.Func(func(next http.Handler) http.Handler {
+		return nethttp.Middleware(opentracing.GlobalTracer(), next)
+	})
+
+	durationInstrumented := instrumentDuration.Wrap(router)
+	traceInstrumented := instrumentTraces.Wrap(durationInstrumented)
+	// Don't try and trace websocket requests because nethttp.Middleware
+	// doesn't support http.Hijack
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var next http.Handler
+		if middleware.IsWSHandshakeRequest(r) {
+			next = durationInstrumented
+		} else {
+			next = traceInstrumented
+		}
+		next.ServeHTTP(w, r)
+	})
 }
 
 func collectorFactory(userIDer multitenant.UserIDer, collectorURL, s3URL, natsHostname string,