Refactor app for multitenancy

- Add interfaces to allow for alternative implementations for Collector, ControlRouter
  and PipeRouter.
- Pass contexts on http handlers to these interfaces.  Although not used by the current
  (local, in-memory) implementations, the idea is this will be used to pass headers to
  implementations which support multitenancy (by, for instance, putting an authenticating
  reverse proxy in form of the app, and then inspecting the headers of the request for
  a used id).

app/api_report.go

@@ -2,12 +2,13 @@ package app
 
 import (
 	"net/http"
+
+	"golang.org/x/net/context"
 )
 
 // Raw report handler
-func makeRawReportHandler(rep Reporter) func(http.ResponseWriter, *http.Request) {
-	return func(w http.ResponseWriter, r *http.Request) {
-		// r.ParseForm()
-		respondWith(w, http.StatusOK, rep.Report())
+func makeRawReportHandler(rep Reporter) CtxHandlerFunc {
+	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
+		respondWith(w, http.StatusOK, rep.Report(ctx))
 	}
 }

app/api_topologies.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	"github.com/gorilla/mux"
+	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/render"
 	"github.com/weaveworks/scope/report"
@@ -192,9 +193,9 @@ func (r *registry) walk(f func(APITopologyDesc)) {
 }
 
 // makeTopologyList returns a handler that yields an APITopologyList.
-func (r *registry) makeTopologyList(rep Reporter) func(w http.ResponseWriter, r *http.Request) {
-	return func(w http.ResponseWriter, req *http.Request) {
-		topologies := r.renderTopologies(rep.Report(), req)
+func (r *registry) makeTopologyList(rep Reporter) CtxHandlerFunc {
+	return func(ctx context.Context, w http.ResponseWriter, req *http.Request) {
+		topologies := r.renderTopologies(rep.Report(ctx), req)
 		respondWith(w, http.StatusOK, topologies)
 	}
 }
@@ -252,27 +253,27 @@ func renderedForRequest(r *http.Request, topology APITopologyDesc) render.Render
 	return renderer
 }
 
-type reportRenderHandler func(Reporter, render.Renderer, http.ResponseWriter, *http.Request)
+type reportRenderHandler func(context.Context, Reporter, render.Renderer, http.ResponseWriter, *http.Request)
 
-func (r *registry) captureRenderer(rep Reporter, f reportRenderHandler) http.HandlerFunc {
-	return func(w http.ResponseWriter, req *http.Request) {
+func (r *registry) captureRenderer(rep Reporter, f reportRenderHandler) CtxHandlerFunc {
+	return func(ctx context.Context, w http.ResponseWriter, req *http.Request) {
 		topology, ok := r.get(mux.Vars(req)["topology"])
 		if !ok {
 			http.NotFound(w, req)
 			return
 		}
 		renderer := renderedForRequest(req, topology)
-		f(rep, renderer, w, req)
+		f(ctx, rep, renderer, w, req)
 	}
 }
 
-func (r *registry) captureRendererWithoutFilters(rep Reporter, topologyID string, f reportRenderHandler) http.HandlerFunc {
-	return func(w http.ResponseWriter, req *http.Request) {
+func (r *registry) captureRendererWithoutFilters(rep Reporter, topologyID string, f reportRenderHandler) CtxHandlerFunc {
+	return func(ctx context.Context, w http.ResponseWriter, req *http.Request) {
 		topology, ok := r.get(topologyID)
 		if !ok {
 			http.NotFound(w, req)
 			return
 		}
-		f(rep, topology.renderer, w, req)
+		f(ctx, rep, topology.renderer, w, req)
 	}
 }

app/api_topology.go

@@ -6,6 +6,7 @@ import (
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/gorilla/websocket"
+	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/common/xfer"
 	"github.com/weaveworks/scope/render"
@@ -28,14 +29,14 @@ type APINode struct {
 }
 
 // Full topology.
-func handleTopology(rep Reporter, renderer render.Renderer, w http.ResponseWriter, r *http.Request) {
+func handleTopology(ctx context.Context, rep Reporter, renderer render.Renderer, w http.ResponseWriter, r *http.Request) {
 	respondWith(w, http.StatusOK, APITopology{
-		Nodes: renderer.Render(rep.Report()).Prune(),
+		Nodes: renderer.Render(rep.Report(ctx)).Prune(),
 	})
 }
 
 // Websocket for the full topology. This route overlaps with the next.
-func handleWs(rep Reporter, renderer render.Renderer, w http.ResponseWriter, r *http.Request) {
+func handleWs(ctx context.Context, rep Reporter, renderer render.Renderer, w http.ResponseWriter, r *http.Request) {
 	if err := r.ParseForm(); err != nil {
 		respondWith(w, http.StatusInternalServerError, err.Error())
 		return
@@ -48,15 +49,15 @@ func handleWs(rep Reporter, renderer render.Renderer, w http.ResponseWriter, r *
 			return
 		}
 	}
-	handleWebsocket(w, r, rep, renderer, loop)
+	handleWebsocket(ctx, w, r, rep, renderer, loop)
 }
 
 // Individual nodes.
-func handleNode(nodeID string) func(Reporter, render.Renderer, http.ResponseWriter, *http.Request) {
-	return func(rep Reporter, renderer render.Renderer, w http.ResponseWriter, r *http.Request) {
+func handleNode(nodeID string) func(context.Context, Reporter, render.Renderer, http.ResponseWriter, *http.Request) {
+	return func(ctx context.Context, rep Reporter, renderer render.Renderer, w http.ResponseWriter, r *http.Request) {
 		var (
-			rpt      = rep.Report()
-			node, ok = renderer.Render(rep.Report())[nodeID]
+			rpt      = rep.Report(ctx)
+			node, ok = renderer.Render(rep.Report(ctx))[nodeID]
 		)
 		if !ok {
 			http.NotFound(w, r)
@@ -71,6 +72,7 @@ var upgrader = websocket.Upgrader{
 }
 
 func handleWebsocket(
+	ctx context.Context,
 	w http.ResponseWriter,
 	r *http.Request,
 	rep Reporter,
@@ -88,6 +90,7 @@ func handleWebsocket(
 	go func(c *websocket.Conn) {
 		for { // just discard everything the browser sends
 			if _, _, err := c.NextReader(); err != nil {
+				log.Println("err:", err)
 				close(quit)
 				break
 			}
@@ -99,15 +102,16 @@ func handleWebsocket(
 		tick         = time.Tick(loop)
 		wait         = make(chan struct{}, 1)
 	)
-	rep.WaitOn(wait)
-	defer rep.UnWait(wait)
+	rep.WaitOn(ctx, wait)
+	defer rep.UnWait(ctx, wait)
 
 	for {
-		newTopo := renderer.Render(rep.Report()).Prune()
+		newTopo := renderer.Render(rep.Report(ctx)).Prune()
 		diff := render.TopoDiff(previousTopo, newTopo)
 		previousTopo = newTopo
 
 		if err := conn.SetWriteDeadline(time.Now().Add(websocketTimeout)); err != nil {
+			log.Println("err:", err)
 			return
 		}
 

app/collector.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/spaolacci/murmur3"
+	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/report"
 )
@@ -13,15 +14,15 @@ import (
 // Reporter is something that can produce reports on demand. It's a convenient
 // interface for parts of the app, and several experimental components.
 type Reporter interface {
-	Report() report.Report
-	WaitOn(chan struct{})
-	UnWait(chan struct{})
+	Report(context.Context) report.Report
+	WaitOn(context.Context, chan struct{})
+	UnWait(context.Context, chan struct{})
 }
 
 // Adder is something that can accept reports. It's a convenient interface for
 // parts of the app, and several experimental components.
 type Adder interface {
-	Add(report.Report)
+	Add(context.Context, report.Report)
 }
 
 // A Collector is a Reporter and an Adder
@@ -45,13 +46,13 @@ type waitableCondition struct {
 	waiters map[chan struct{}]struct{}
 }
 
-func (wc *waitableCondition) WaitOn(waiter chan struct{}) {
+func (wc *waitableCondition) WaitOn(_ context.Context, waiter chan struct{}) {
 	wc.Lock()
 	wc.waiters[waiter] = struct{}{}
 	wc.Unlock()
 }
 
-func (wc *waitableCondition) UnWait(waiter chan struct{}) {
+func (wc *waitableCondition) UnWait(_ context.Context, waiter chan struct{}) {
 	wc.Lock()
 	delete(wc.waiters, waiter)
 	wc.Unlock()
@@ -82,7 +83,7 @@ func NewCollector(window time.Duration) Collector {
 var now = time.Now
 
 // Add adds a report to the collector's internal state. It implements Adder.
-func (c *collector) Add(rpt report.Report) {
+func (c *collector) Add(_ context.Context, rpt report.Report) {
 	c.mtx.Lock()
 	defer c.mtx.Unlock()
 	c.reports = append(c.reports, timestampReport{now(), rpt})
@@ -95,7 +96,7 @@ func (c *collector) Add(rpt report.Report) {
 
 // Report returns a merged report over all added reports. It implements
 // Reporter.
-func (c *collector) Report() report.Report {
+func (c *collector) Report(_ context.Context) report.Report {
 	c.mtx.Lock()
 	defer c.mtx.Unlock()
 

app/collector_test.go

@@ -4,6 +4,8 @@ import (
 	"testing"
 	"time"
 
+	"golang.org/x/net/context"
+
 	"github.com/weaveworks/scope/app"
 	"github.com/weaveworks/scope/report"
 	"github.com/weaveworks/scope/test"
@@ -11,6 +13,7 @@ import (
 )
 
 func TestCollector(t *testing.T) {
+	ctx := context.Background()
 	window := time.Millisecond
 	c := app.NewCollector(window)
 
@@ -20,32 +23,33 @@ func TestCollector(t *testing.T) {
 	r2 := report.MakeReport()
 	r2.Endpoint.AddNode("bar", report.MakeNode())
 
-	if want, have := report.MakeReport(), c.Report(); !reflect.DeepEqual(want, have) {
+	if want, have := report.MakeReport(), c.Report(ctx); !reflect.DeepEqual(want, have) {
 		t.Error(test.Diff(want, have))
 	}
 
-	c.Add(r1)
-	if want, have := r1, c.Report(); !reflect.DeepEqual(want, have) {
+	c.Add(ctx, r1)
+	if want, have := r1, c.Report(ctx); !reflect.DeepEqual(want, have) {
 		t.Error(test.Diff(want, have))
 	}
 
-	c.Add(r2)
+	c.Add(ctx, r2)
 
 	merged := report.MakeReport()
 	merged = merged.Merge(r1)
 	merged = merged.Merge(r2)
-	if want, have := merged, c.Report(); !reflect.DeepEqual(want, have) {
+	if want, have := merged, c.Report(ctx); !reflect.DeepEqual(want, have) {
 		t.Error(test.Diff(want, have))
 	}
 }
 
 func TestCollectorWait(t *testing.T) {
+	ctx := context.Background()
 	window := time.Millisecond
 	c := app.NewCollector(window)
 
 	waiter := make(chan struct{}, 1)
-	c.WaitOn(waiter)
-	defer c.UnWait(waiter)
+	c.WaitOn(ctx, waiter)
+	defer c.UnWait(ctx, waiter)
 	c.(interface {
 		Broadcast()
 	}).Broadcast()

app/control_router.go

@@ -0,0 +1,69 @@
+package app
+
+import (
+	"fmt"
+	"math/rand"
+	"sync"
+
+	"golang.org/x/net/context"
+
+	"github.com/weaveworks/scope/common/xfer"
+)
+
+// ControlRouter is a thing that can route control requests and responses
+// between the UI and a probe.
+type ControlRouter interface {
+	Handle(ctx context.Context, probeID string, req xfer.Request) (xfer.Response, error)
+	Register(ctx context.Context, probeID string, handler xfer.ControlHandlerFunc) (int64, error)
+	Deregister(ctx context.Context, probeID string, id int64) error
+}
+
+// NewLocalControlRouter creates a new ControlRouter that does everything
+// locally, in memory.
+func NewLocalControlRouter() ControlRouter {
+	return &localControlRouter{
+		probes: map[string]probe{},
+	}
+}
+
+type localControlRouter struct {
+	sync.Mutex
+	probes map[string]probe
+}
+
+type probe struct {
+	id      int64
+	handler xfer.ControlHandlerFunc
+}
+
+func (l *localControlRouter) Handle(_ context.Context, probeID string, req xfer.Request) (xfer.Response, error) {
+	l.Lock()
+	probe, ok := l.probes[probeID]
+	l.Unlock()
+	if !ok {
+		return xfer.Response{}, fmt.Errorf("Probe %s is not connected right now...", probeID)
+	}
+	return probe.handler(req), nil
+}
+
+func (l *localControlRouter) Register(_ context.Context, probeID string, handler xfer.ControlHandlerFunc) (int64, error) {
+	l.Lock()
+	defer l.Unlock()
+	id := rand.Int63()
+	l.probes[probeID] = probe{
+		id:      id,
+		handler: handler,
+	}
+	return id, nil
+}
+
+func (l *localControlRouter) Deregister(_ context.Context, probeID string, id int64) error {
+	l.Lock()
+	defer l.Unlock()
+	// NB probe might have reconnected in the mean time, need to ensure we do not
+	// delete new connection!  Also, it might have connected then deleted itself!
+	if l.probes[probeID].id == id {
+		delete(l.probes, probeID)
+	}
+	return nil
+}

app/controls.go

@@ -1,122 +1,84 @@
 package app
 
 import (
-	"math/rand"
 	"net/http"
 	"net/rpc"
-	"sync"
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/gorilla/mux"
+	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/common/xfer"
 )
 
 // RegisterControlRoutes registers the various control routes with a http mux.
-func RegisterControlRoutes(router *mux.Router) {
-	controlRouter := &controlRouter{
-		probes: map[string]controlHandler{},
-	}
-	router.Methods("GET").Path("/api/control/ws").HandlerFunc(controlRouter.handleProbeWS)
-	router.Methods("POST").MatcherFunc(URLMatcher("/api/control/{probeID}/{nodeID}/{control}")).HandlerFunc(controlRouter.handleControl)
-}
-
-type controlHandler struct {
-	id     int64
-	client *rpc.Client
-}
-
-type controlRouter struct {
-	sync.Mutex
-	probes map[string]controlHandler
-}
-
-func (ch *controlHandler) handle(req xfer.Request) xfer.Response {
-	var res xfer.Response
-	if err := ch.client.Call("control.Handle", req, &res); err != nil {
-		return xfer.ResponseError(err)
-	}
-	return res
-}
-
-func (cr *controlRouter) get(probeID string) (controlHandler, bool) {
-	cr.Lock()
-	defer cr.Unlock()
-	handler, ok := cr.probes[probeID]
-	return handler, ok
-}
-
-func (cr *controlRouter) set(probeID string, handler controlHandler) {
-	cr.Lock()
-	defer cr.Unlock()
-	cr.probes[probeID] = handler
-}
-
-func (cr *controlRouter) rm(probeID string, handler controlHandler) {
-	cr.Lock()
-	defer cr.Unlock()
-	// NB probe might have reconnected in the mean time, need to ensure we do not
-	// delete new connection!  Also, it might have connected then deleted itself!
-	if cr.probes[probeID].id == handler.id {
-		delete(cr.probes, probeID)
-	}
+func RegisterControlRoutes(router *mux.Router, cr ControlRouter) {
+	router.Methods("GET").Path("/api/control/ws").
+		HandlerFunc(requestContextDecorator(handleProbeWS(cr)))
+	router.Methods("POST").MatcherFunc(URLMatcher("/api/control/{probeID}/{nodeID}/{control}")).
+		HandlerFunc(requestContextDecorator(handleControl(cr)))
 }
 
 // handleControl routes control requests from the client to the appropriate
 // probe.  Its is blocking.
-func (cr *controlRouter) handleControl(w http.ResponseWriter, r *http.Request) {
-	var (
-		vars    = mux.Vars(r)
-		probeID = vars["probeID"]
-		nodeID  = vars["nodeID"]
-		control = vars["control"]
-	)
-	handler, ok := cr.get(probeID)
-	if !ok {
-		log.Errorf("Probe %s is not connected right now...", probeID)
-		http.NotFound(w, r)
-		return
+func handleControl(cr ControlRouter) CtxHandlerFunc {
+	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
+		var (
+			vars    = mux.Vars(r)
+			probeID = vars["probeID"]
+			nodeID  = vars["nodeID"]
+			control = vars["control"]
+		)
+		result, err := cr.Handle(ctx, probeID, xfer.Request{
+			AppID:   UniqueID,
+			NodeID:  nodeID,
+			Control: control,
+		})
+		if err != nil {
+			respondWith(w, http.StatusBadRequest, err.Error())
+			return
+		}
+		if result.Error != "" {
+			respondWith(w, http.StatusBadRequest, result.Error)
+			return
+		}
+		respondWith(w, http.StatusOK, result)
 	}
-
-	result := handler.handle(xfer.Request{
-		AppID:   UniqueID,
-		NodeID:  nodeID,
-		Control: control,
-	})
-	if result.Error != "" {
-		respondWith(w, http.StatusBadRequest, result.Error)
-		return
-	}
-	respondWith(w, http.StatusOK, result)
 }
 
 // handleProbeWS accepts websocket connections from the probe and registers
 // them in the control router, such that HandleControl calls can find them.
-func (cr *controlRouter) handleProbeWS(w http.ResponseWriter, r *http.Request) {
-	probeID := r.Header.Get(xfer.ScopeProbeIDHeader)
-	if probeID == "" {
-		respondWith(w, http.StatusBadRequest, xfer.ScopeProbeIDHeader)
-		return
+func handleProbeWS(cr ControlRouter) CtxHandlerFunc {
+	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
+		probeID := r.Header.Get(xfer.ScopeProbeIDHeader)
+		if probeID == "" {
+			respondWith(w, http.StatusBadRequest, xfer.ScopeProbeIDHeader)
+			return
+		}
+
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			log.Printf("Error upgrading control websocket: %v", err)
+			return
+		}
+		defer conn.Close()
+
+		codec := xfer.NewJSONWebsocketCodec(conn)
+		client := rpc.NewClientWithCodec(codec)
+		defer client.Close()
+
+		id, err := cr.Register(ctx, probeID, func(req xfer.Request) xfer.Response {
+			var res xfer.Response
+			if err := client.Call("control.Handle", req, &res); err != nil {
+				return xfer.ResponseError(err)
+			}
+			return res
+		})
+		if err != nil {
+			respondWith(w, http.StatusBadRequest, err.Error())
+			return
+		}
+		defer cr.Deregister(ctx, probeID, id)
+		codec.WaitForReadError()
 	}
-
-	conn, err := upgrader.Upgrade(w, r, nil)
-	if err != nil {
-		log.Errorf("Error upgrading to websocket: %v", err)
-		return
-	}
-	defer conn.Close()
-
-	codec := xfer.NewJSONWebsocketCodec(conn)
-	client := rpc.NewClientWithCodec(codec)
-	handler := controlHandler{
-		id:     rand.Int63(),
-		client: client,
-	}
-
-	cr.set(probeID, handler)
-
-	codec.WaitForReadError()
-
-	cr.rm(probeID, handler)
-	client.Close()
 }

app/controls_test.go

@@ -18,7 +18,7 @@ import (
 
 func TestControl(t *testing.T) {
 	router := mux.NewRouter()
-	app.RegisterControlRoutes(router)
+	app.RegisterControlRoutes(router, app.NewLocalControlRouter())
 	server := httptest.NewServer(router)
 	defer server.Close()
 

app/mock_reporter_test.go

@@ -3,12 +3,14 @@ package app_test
 import (
 	"github.com/weaveworks/scope/report"
 	"github.com/weaveworks/scope/test/fixture"
+
+	"golang.org/x/net/context"
 )
 
 // StaticReport is used as a fixture in tests. It emulates an xfer.Collector.
 type StaticReport struct{}
 
-func (s StaticReport) Report() report.Report { return fixture.Report }
-func (s StaticReport) Add(report.Report)     {}
-func (s StaticReport) WaitOn(chan struct{})  {}
-func (s StaticReport) UnWait(chan struct{})  {}
+func (s StaticReport) Report(context.Context) report.Report  { return fixture.Report }
+func (s StaticReport) Add(context.Context, report.Report)    {}
+func (s StaticReport) WaitOn(context.Context, chan struct{}) {}
+func (s StaticReport) UnWait(context.Context, chan struct{}) {}

app/pipe_router.go

@@ -0,0 +1,180 @@
+package app
+
+import (
+	"io"
+	"sync"
+	"time"
+
+	log "github.com/Sirupsen/logrus"
+	"golang.org/x/net/context"
+
+	"github.com/weaveworks/scope/common/mtime"
+	"github.com/weaveworks/scope/common/xfer"
+)
+
+const (
+	gcInterval  = 30 * time.Second // we check all the pipes every 30s
+	pipeTimeout = 1 * time.Minute  // pipes are closed when a client hasn't been connected for 1 minute
+	gcTimeout   = 10 * time.Minute // after another 10 minutes, tombstoned pipes are forgotten
+)
+
+// End is an enum for either end of the pipe.
+type End int
+
+// Valid values of type End
+const (
+	UIEnd = iota
+	ProbeEnd
+)
+
+// PipeRouter stores pipes and allows you to connect to either end of them.
+type PipeRouter interface {
+	Get(context.Context, string, End) (xfer.Pipe, io.ReadWriter, bool)
+	Release(context.Context, string, End)
+	Delete(context.Context, string)
+	Stop()
+}
+
+// PipeRouter connects incoming and outgoing pipes.
+type localPipeRouter struct {
+	sync.Mutex
+	wait  sync.WaitGroup
+	quit  chan struct{}
+	pipes map[string]*pipe
+}
+
+// for each end of the pipe, we keep a reference count & lastUsedTIme,
+// such that we can timeout pipes when either end is inactive.
+type pipe struct {
+	xfer.Pipe
+
+	tombstoneTime time.Time
+
+	ui, probe end
+}
+
+type end struct {
+	refCount     int
+	lastUsedTime time.Time
+}
+
+func (p *pipe) end(end End) (*end, io.ReadWriter) {
+	ui, probe := p.Ends()
+	if end == UIEnd {
+		return &p.ui, ui
+	}
+	return &p.probe, probe
+}
+
+// NewLocalPipeRouter returns a new local (in-memory) pipe router.
+func NewLocalPipeRouter() PipeRouter {
+	pipeRouter := &localPipeRouter{
+		quit:  make(chan struct{}),
+		pipes: map[string]*pipe{},
+	}
+	pipeRouter.wait.Add(1)
+	go pipeRouter.gcLoop()
+	return pipeRouter
+}
+
+func (pr *localPipeRouter) Get(_ context.Context, id string, e End) (xfer.Pipe, io.ReadWriter, bool) {
+	pr.Lock()
+	defer pr.Unlock()
+	p, ok := pr.pipes[id]
+	if !ok {
+		log.Infof("Creating pipe id %s", id)
+		p = &pipe{
+			ui:    end{lastUsedTime: mtime.Now()},
+			probe: end{lastUsedTime: mtime.Now()},
+			Pipe:  xfer.NewPipe(),
+		}
+		pr.pipes[id] = p
+	}
+	if p.Closed() {
+		return nil, nil, false
+	}
+	end, endIO := p.end(e)
+	end.refCount++
+	return p, endIO, true
+}
+
+func (pr *localPipeRouter) Release(_ context.Context, id string, e End) {
+	pr.Lock()
+	defer pr.Unlock()
+
+	p, ok := pr.pipes[id]
+	if !ok {
+		// uh oh
+		return
+	}
+
+	end, _ := p.end(e)
+	end.refCount--
+	if end.refCount > 0 {
+		return
+	}
+
+	if !p.Closed() {
+		end.lastUsedTime = mtime.Now()
+	}
+}
+
+func (pr *localPipeRouter) Delete(_ context.Context, id string) {
+	pr.Lock()
+	defer pr.Unlock()
+	p, ok := pr.pipes[id]
+	if !ok {
+		return
+	}
+	p.Close()
+	p.tombstoneTime = mtime.Now()
+}
+
+func (pr *localPipeRouter) Stop() {
+	close(pr.quit)
+	pr.wait.Wait()
+}
+
+func (pr *localPipeRouter) gcLoop() {
+	defer pr.wait.Done()
+	ticker := time.Tick(gcInterval)
+	for {
+		select {
+		case <-pr.quit:
+			return
+		case <-ticker:
+		}
+
+		pr.timeout()
+		pr.garbageCollect()
+	}
+}
+
+func (pr *localPipeRouter) timeout() {
+	pr.Lock()
+	defer pr.Unlock()
+	now := mtime.Now()
+	for id, pipe := range pr.pipes {
+		if pipe.Closed() || (pipe.ui.refCount > 0 && pipe.probe.refCount > 0) {
+			continue
+		}
+
+		if (pipe.ui.refCount == 0 && now.Sub(pipe.ui.lastUsedTime) >= pipeTimeout) ||
+			(pipe.probe.refCount == 0 && now.Sub(pipe.probe.lastUsedTime) >= pipeTimeout) {
+			log.Infof("Timing out pipe %s", id)
+			pipe.Close()
+			pipe.tombstoneTime = now
+		}
+	}
+}
+
+func (pr *localPipeRouter) garbageCollect() {
+	pr.Lock()
+	defer pr.Unlock()
+	now := mtime.Now()
+	for pipeID, pipe := range pr.pipes {
+		if pipe.Closed() && now.Sub(pipe.tombstoneTime) >= gcTimeout {
+			delete(pr.pipes, pipeID)
+		}
+	}
+}

app/pipes.go

@@ -1,199 +1,54 @@
 package app
 
 import (
-	"io"
 	"net/http"
-	"sync"
-	"time"
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/gorilla/mux"
-
-	"github.com/weaveworks/scope/common/mtime"
-	"github.com/weaveworks/scope/common/xfer"
+	"golang.org/x/net/context"
 )
 
-const (
-	gcInterval  = 30 * time.Second // we check all the pipes every 30s
-	pipeTimeout = 1 * time.Minute  // pipes are closed when a client hasn't been connected for 1 minute
-	gcTimeout   = 10 * time.Minute // after another 10 minutes, tombstoned pipes are forgotten
-)
-
-// PipeRouter connects incoming and outgoing pipes.
-type PipeRouter struct {
-	sync.Mutex
-	wait  sync.WaitGroup
-	quit  chan struct{}
-	pipes map[string]*pipe
-}
-
-// for each end of the pipe, we keep a reference count & lastUsedTIme,
-// such that we can timeout pipes when either end is inactive.
-type end struct {
-	refCount     int
-	lastUsedTime time.Time
-}
-
-type pipe struct {
-	ui, probe     end
-	tombstoneTime time.Time
-
-	xfer.Pipe
-}
-
 // RegisterPipeRoutes registers the pipe routes
-func RegisterPipeRoutes(router *mux.Router) *PipeRouter {
-	pipeRouter := &PipeRouter{
-		quit:  make(chan struct{}),
-		pipes: map[string]*pipe{},
-	}
-	pipeRouter.wait.Add(1)
-	go pipeRouter.gcLoop()
+func RegisterPipeRoutes(router *mux.Router, pr PipeRouter) {
 	router.Methods("GET").
 		Path("/api/pipe/{pipeID}").
-		HandlerFunc(pipeRouter.handleWs(func(p *pipe) (*end, io.ReadWriter) {
-		uiEnd, _ := p.Ends()
-		return &p.ui, uiEnd
-	}))
+		HandlerFunc(requestContextDecorator(handlePipeWs(pr, UIEnd)))
+
 	router.Methods("GET").
 		Path("/api/pipe/{pipeID}/probe").
-		HandlerFunc(pipeRouter.handleWs(func(p *pipe) (*end, io.ReadWriter) {
-		_, probeEnd := p.Ends()
-		return &p.probe, probeEnd
-	}))
+		HandlerFunc(requestContextDecorator(handlePipeWs(pr, ProbeEnd)))
+
 	router.Methods("DELETE", "POST").
 		Path("/api/pipe/{pipeID}").
-		HandlerFunc(pipeRouter.delete)
-	return pipeRouter
+		HandlerFunc(requestContextDecorator(deletePipe(pr)))
 }
 
-// Stop stops the pipeRouter
-func (pr *PipeRouter) Stop() {
-	close(pr.quit)
-	pr.wait.Wait()
-}
-
-func (pr *PipeRouter) gcLoop() {
-	defer pr.wait.Done()
-	ticker := time.Tick(gcInterval)
-	for {
-		select {
-		case <-pr.quit:
-			return
-		case <-ticker:
-		}
-
-		pr.timeout()
-		pr.garbageCollect()
-	}
-}
-
-func (pr *PipeRouter) timeout() {
-	pr.Lock()
-	defer pr.Unlock()
-	now := mtime.Now()
-	for id, pipe := range pr.pipes {
-		if pipe.Closed() || (pipe.ui.refCount > 0 && pipe.probe.refCount > 0) {
-			continue
-		}
-
-		if (pipe.ui.refCount == 0 && now.Sub(pipe.ui.lastUsedTime) >= pipeTimeout) ||
-			(pipe.probe.refCount == 0 && now.Sub(pipe.probe.lastUsedTime) >= pipeTimeout) {
-			log.Infof("Timing out pipe %s", id)
-			pipe.Close()
-			pipe.tombstoneTime = now
-		}
-	}
-}
-
-func (pr *PipeRouter) garbageCollect() {
-	pr.Lock()
-	defer pr.Unlock()
-	now := mtime.Now()
-	for pipeID, pipe := range pr.pipes {
-		if pipe.Closed() && now.Sub(pipe.tombstoneTime) >= gcTimeout {
-			delete(pr.pipes, pipeID)
-		}
-	}
-}
-
-func (pr *PipeRouter) getOrCreate(id string) (*pipe, bool) {
-	pr.Lock()
-	defer pr.Unlock()
-	p, ok := pr.pipes[id]
-	if !ok {
-		log.Infof("Creating pipe id %s", id)
-		p = &pipe{
-			ui:    end{lastUsedTime: mtime.Now()},
-			probe: end{lastUsedTime: mtime.Now()},
-			Pipe:  xfer.NewPipe(),
-		}
-		pr.pipes[id] = p
-	}
-	if p.Closed() {
-		return nil, false
-	}
-	return p, true
-}
-
-func (pr *PipeRouter) retain(id string, pipe *pipe, end *end) bool {
-	pr.Lock()
-	defer pr.Unlock()
-	if pipe.Closed() {
-		return false
-	}
-	end.refCount++
-	return true
-}
-
-func (pr *PipeRouter) release(id string, pipe *pipe, end *end) {
-	pr.Lock()
-	defer pr.Unlock()
-
-	end.refCount--
-	if end.refCount != 0 {
-		return
-	}
-
-	if !pipe.Closed() {
-		end.lastUsedTime = mtime.Now()
-	}
-}
-
-func (pr *PipeRouter) handleWs(endSelector func(*pipe) (*end, io.ReadWriter)) func(http.ResponseWriter, *http.Request) {
-	return func(w http.ResponseWriter, r *http.Request) {
-		pipeID := mux.Vars(r)["pipeID"]
-		pipe, ok := pr.getOrCreate(pipeID)
+func handlePipeWs(pr PipeRouter, end End) CtxHandlerFunc {
+	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
+		id := mux.Vars(r)["pipeID"]
+		pipe, endIO, ok := pr.Get(ctx, id, end)
 		if !ok {
 			http.NotFound(w, r)
 			return
 		}
-
-		endRef, endIO := endSelector(pipe)
-		if !pr.retain(pipeID, pipe, endRef) {
-			http.NotFound(w, r)
-			return
-		}
-		defer pr.release(pipeID, pipe, endRef)
+		defer pr.Release(ctx, id, end)
 
 		conn, err := upgrader.Upgrade(w, r, nil)
 		if err != nil {
-			log.Errorf("Error upgrading to websocket: %v", err)
+			log.Errorf("Error upgrading pipe %s (%d) websocket: %v", id, end, err)
 			return
 		}
 		defer conn.Close()
 
+		log.Infof("Pipe success %s (%d)", id, end)
 		pipe.CopyToWebsocket(endIO, conn)
 	}
 }
 
-func (pr *PipeRouter) delete(w http.ResponseWriter, r *http.Request) {
-	pipeID := mux.Vars(r)["pipeID"]
-	pipe, ok := pr.getOrCreate(pipeID)
-	if ok && pr.retain(pipeID, pipe, &pipe.ui) {
+func deletePipe(pr PipeRouter) CtxHandlerFunc {
+	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
+		pipeID := mux.Vars(r)["pipeID"]
 		log.Infof("Closing pipe %s", pipeID)
-		pipe.Close()
-		pipe.tombstoneTime = mtime.Now()
-		pr.release(pipeID, pipe, &pipe.ui)
+		pr.Delete(ctx, pipeID)
 	}
 }

app/pipes_internal_test.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/gorilla/mux"
 	"github.com/gorilla/websocket"
+	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/common/mtime"
 	"github.com/weaveworks/scope/common/xfer"
@@ -22,7 +23,8 @@ import (
 
 func TestPipeTimeout(t *testing.T) {
 	router := mux.NewRouter()
-	pr := RegisterPipeRoutes(router)
+	pr := NewLocalPipeRouter().(*localPipeRouter)
+	RegisterPipeRoutes(router, pr)
 	pr.Stop() // we don't want the loop running in the background
 
 	mtime.NowForce(time.Now())
@@ -30,7 +32,8 @@ func TestPipeTimeout(t *testing.T) {
 
 	// create a new pipe.
 	id := "foo"
-	pipe, ok := pr.getOrCreate(id)
+	ctx := context.Background()
+	pipe, _, ok := pr.Get(ctx, id, UIEnd)
 	if !ok {
 		t.Fatalf("not ok")
 	}
@@ -65,7 +68,8 @@ func (a adapter) PipeClose(_, pipeID string) error {
 
 func TestPipeClose(t *testing.T) {
 	router := mux.NewRouter()
-	pr := RegisterPipeRoutes(router)
+	pr := NewLocalPipeRouter()
+	RegisterPipeRoutes(router, pr)
 	defer pr.Stop()
 
 	server := httptest.NewServer(router)

app/router.go

@@ -10,6 +10,7 @@ import (
 	"github.com/PuerkitoBio/ghost/handlers"
 	"github.com/gorilla/mux"
 	"github.com/ugorji/go/codec"
+	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/common/hostname"
 	"github.com/weaveworks/scope/common/xfer"
@@ -24,6 +25,19 @@ var (
 	UniqueID = "0"
 )
 
+// RequestCtxKey is key used for request entry in context
+const RequestCtxKey = "request"
+
+// CtxHandlerFunc is a http.HandlerFunc, with added contexts
+type CtxHandlerFunc func(context.Context, http.ResponseWriter, *http.Request)
+
+func requestContextDecorator(f CtxHandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		ctx := context.WithValue(context.Background(), RequestCtxKey, r)
+		f(ctx, w, r)
+	}
+}
+
 // URLMatcher uses request.RequestURI (the raw, unparsed request) to attempt
 // to match pattern.  It does this as go's URL.Parse method is broken, and
 // mistakenly unescapes the Path before parsing it.  This breaks %2F (encoded
@@ -76,13 +90,13 @@ func gzipHandler(h http.HandlerFunc) http.HandlerFunc {
 // routes should be added to a router and passed to postRoutes.
 func TopologyHandler(c Reporter, preRoutes *mux.Router, postRoutes http.Handler) http.Handler {
 	get := preRoutes.Methods("GET").Subrouter()
-	get.HandleFunc("/api", gzipHandler(apiHandler))
-	get.HandleFunc("/api/topology", gzipHandler(topologyRegistry.makeTopologyList(c)))
+	get.HandleFunc("/api", gzipHandler(requestContextDecorator(apiHandler)))
+	get.HandleFunc("/api/topology", gzipHandler(requestContextDecorator(topologyRegistry.makeTopologyList(c))))
 	get.HandleFunc("/api/topology/{topology}",
-		gzipHandler(topologyRegistry.captureRenderer(c, handleTopology)))
+		gzipHandler(requestContextDecorator(topologyRegistry.captureRenderer(c, handleTopology))))
 	get.HandleFunc("/api/topology/{topology}/ws",
-		topologyRegistry.captureRenderer(c, handleWs)) // NB not gzip!
-	get.HandleFunc("/api/report", gzipHandler(makeRawReportHandler(c)))
+		requestContextDecorator(topologyRegistry.captureRenderer(c, handleWs))) // NB not gzip!
+	get.HandleFunc("/api/report", gzipHandler(requestContextDecorator(makeRawReportHandler(c))))
 
 	// We pull in the http.DefaultServeMux to get the pprof routes
 	preRoutes.PathPrefix("/debug/pprof").Handler(http.DefaultServeMux)
@@ -107,9 +121,9 @@ func TopologyHandler(c Reporter, preRoutes *mux.Router, postRoutes http.Handler)
 			return
 		}
 
-		handler := gzipHandler(topologyRegistry.captureRendererWithoutFilters(
+		handler := gzipHandler(requestContextDecorator(topologyRegistry.captureRendererWithoutFilters(
 			c, topologyID, handleNode(nodeID),
-		))
+		)))
 		handler.ServeHTTP(w, r)
 	})
 }
@@ -117,7 +131,7 @@ func TopologyHandler(c Reporter, preRoutes *mux.Router, postRoutes http.Handler)
 // RegisterReportPostHandler registers the handler for report submission
 func RegisterReportPostHandler(a Adder, router *mux.Router) {
 	post := router.Methods("POST").Subrouter()
-	post.HandleFunc("/api/report", func(w http.ResponseWriter, r *http.Request) {
+	post.HandleFunc("/api/report", requestContextDecorator(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
 		var (
 			rpt    report.Report
 			reader = r.Body
@@ -142,15 +156,15 @@ func RegisterReportPostHandler(a Adder, router *mux.Router) {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
-		a.Add(rpt)
+		a.Add(ctx, rpt)
 		if len(rpt.Pod.Nodes) > 0 {
 			topologyRegistry.enableKubernetesTopologies()
 		}
 		w.WriteHeader(http.StatusOK)
-	})
+	}))
 }
 
-func apiHandler(w http.ResponseWriter, r *http.Request) {
+func apiHandler(_ context.Context, w http.ResponseWriter, r *http.Request) {
 	respondWith(w, http.StatusOK, xfer.Details{
 		ID:       UniqueID,
 		Version:  Version,

app/router_test.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/gorilla/mux"
 	"github.com/ugorji/go/codec"
+	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/app"
 	"github.com/weaveworks/scope/test"
@@ -73,7 +74,8 @@ func TestReportPostHandler(t *testing.T) {
 			t.Fatalf("Error posting report: %d", resp.StatusCode)
 		}
 
-		if want, have := fixture.Report.Endpoint.Nodes, c.Report().Endpoint.Nodes; len(have) == 0 || len(want) != len(have) {
+		ctx := context.Background()
+		if want, have := fixture.Report.Endpoint.Nodes, c.Report(ctx).Endpoint.Nodes; len(have) == 0 || len(want) != len(have) {
 			t.Fatalf("Content-Type %s: %v", contentType, test.Diff(have, want))
 		}
 	}

prog/app.go

@@ -23,8 +23,8 @@ import (
 func router(c app.Collector) http.Handler {
 	router := mux.NewRouter()
 	app.RegisterReportPostHandler(c, router)
-	app.RegisterControlRoutes(router)
-	app.RegisterPipeRoutes(router)
+	app.RegisterControlRoutes(router, app.NewLocalControlRouter())
+	app.RegisterPipeRoutes(router, app.NewLocalPipeRouter())
 	return app.TopologyHandler(c, router, http.FileServer(FS(false)))
 }