Review feedback

app/multitenant/consul_client.go

@@ -0,0 +1,201 @@
+package multitenant
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"time"
+
+	log "github.com/Sirupsen/logrus"
+	consul "github.com/hashicorp/consul/api"
+
+	"github.com/weaveworks/scope/common/mtime"
+)
+
+const (
+	longPollDuration = 10 * time.Second
+)
+
+// ConsulClient is a high-level client for Consul, that exposes operations
+// such as CAS and Watch which take callbacks.  It also deals with serialisation.
+type ConsulClient interface {
+	Get(key string, out interface{}) error
+	CAS(key string, out interface{}, f CASCallback) error
+	Watch(key string, deadline time.Time, out interface{}, f func(interface{}) (bool, error)) error
+	WatchPrefix(prefix string, out interface{}, f func(string, interface{}) bool)
+}
+
+// CASCallback is the type of the callback to CAS.  If err is nil, out must be non-nil.
+type CASCallback func(in interface{}) (out interface{}, retry bool, err error)
+
+// NewConsulClient returns a new ConsulClient
+func NewConsulClient(addr string) (ConsulClient, error) {
+	client, err := consul.NewClient(&consul.Config{
+		Address: addr,
+		Scheme:  "http",
+	})
+	if err != nil {
+		return nil, err
+	}
+	return (*consulClient)(client), nil
+}
+
+var (
+	queryOptions = &consul.QueryOptions{
+		RequireConsistent: true,
+	}
+	writeOptions = &consul.WriteOptions{}
+
+	// ErrNotFound is returned by ConsulClient.Get
+	ErrNotFound = fmt.Errorf("Not found")
+)
+
+type consulClient consul.Client
+
+// Get and deserialise a JSON value from consul.
+func (c *consulClient) Get(key string, out interface{}) error {
+	kv := (*consul.Client)(c).KV()
+	kvp, _, err := kv.Get(key, queryOptions)
+	if err != nil {
+		return err
+	}
+	if kvp == nil {
+		return ErrNotFound
+	}
+	if err := json.NewDecoder(bytes.NewReader(kvp.Value)).Decode(out); err != nil {
+		return err
+	}
+	return nil
+}
+
+// CAS atomically modify a value in a callback.
+// If value doesn't exist you'll get nil as a argument to your callback.
+func (c *consulClient) CAS(key string, out interface{}, f CASCallback) error {
+	var (
+		index        = uint64(0)
+		kv           = (*consul.Client)(c).KV()
+		retries      = 10
+		retry        = true
+		intermediate interface{}
+	)
+	for i := 0; i < retries; i++ {
+		kvp, _, err := kv.Get(key, queryOptions)
+		if err != nil {
+			log.Errorf("Error getting %s: %v", key, err)
+			continue
+		}
+		if kvp != nil {
+			if err := json.NewDecoder(bytes.NewReader(kvp.Value)).Decode(out); err != nil {
+				log.Errorf("Error deserialising %s: %v", key, err)
+				continue
+			}
+			index = kvp.ModifyIndex // if key doesn't exist, index will be 0
+			intermediate = out
+		}
+
+		intermediate, retry, err = f(intermediate)
+		if err != nil {
+			log.Errorf("Error CASing %s: %v", key, err)
+			if !retry {
+				return err
+			}
+			continue
+		}
+
+		if intermediate == nil {
+			panic("Callback must instantiate value!")
+		}
+
+		value := bytes.Buffer{}
+		if err := json.NewEncoder(&value).Encode(intermediate); err != nil {
+			log.Errorf("Error serialising value for %s: %v", key, err)
+			continue
+		}
+		ok, _, err := kv.CAS(&consul.KVPair{
+			Key:         key,
+			Value:       value.Bytes(),
+			ModifyIndex: index,
+		}, writeOptions)
+		if err != nil {
+			log.Errorf("Error CASing %s: %v", key, err)
+			continue
+		}
+		if !ok {
+			log.Errorf("Error CASing %s, trying again %d", key, index)
+			continue
+		}
+		return nil
+	}
+	return fmt.Errorf("Failed to CAS %s", key)
+}
+
+// Watch a given key value and trigger a callback when it changes.
+// if callback returns false or error, exit (with the error).
+func (c *consulClient) Watch(key string, deadline time.Time, out interface{}, f func(interface{}) (bool, error)) error {
+	var (
+		index = uint64(0)
+		kv    = (*consul.Client)(c).KV()
+	)
+	for deadline.After(mtime.Now()) {
+		// Do a (blocking) long poll waiting for the entry to get updated. index here
+		// is really a version number; this call will wait for the key to be updated
+		// past said version.  As we always start from version 0, we're guaranteed
+		// not to miss any updates - in fact we will always call the callback with
+		// the current value of the key immediately.
+		kvp, meta, err := kv.Get(key, &consul.QueryOptions{
+			RequireConsistent: true,
+			WaitIndex:         index,
+			WaitTime:          longPollDuration,
+		})
+		if err != nil {
+			return fmt.Errorf("Error getting %s: %v", key, err)
+		}
+		if kvp == nil {
+			return ErrNotFound
+		}
+
+		if err := json.NewDecoder(bytes.NewReader(kvp.Value)).Decode(out); err != nil {
+			return err
+		}
+		if ok, err := f(out); !ok {
+			return nil
+		} else if err != nil {
+			return err
+		}
+
+		index = meta.LastIndex
+	}
+	return fmt.Errorf("Timed out waiting on %s", key)
+}
+
+func (c *consulClient) WatchPrefix(prefix string, out interface{}, f func(string, interface{}) bool) {
+	var (
+		index = uint64(0)
+		kv    = (*consul.Client)(c).KV()
+	)
+	for {
+		kvps, meta, err := kv.List(prefix, &consul.QueryOptions{
+			RequireConsistent: true,
+			WaitIndex:         index,
+			WaitTime:          longPollDuration,
+		})
+		if err != nil {
+			log.Errorf("Error getting path %s: %v", prefix, err)
+			continue
+		}
+		if index == meta.LastIndex {
+			continue
+		}
+		index = meta.LastIndex
+
+		for _, kvp := range kvps {
+			if err := json.NewDecoder(bytes.NewReader(kvp.Value)).Decode(out); err != nil {
+				log.Errorf("Error deserialising %s: %v", kvp.Key, err)
+				continue
+			}
+			if !f(kvp.Key, out) {
+				return
+			}
+		}
+	}
+}

app/multitenant/consul_pipe_router.go

@@ -1,8 +1,6 @@
 package multitenant
 
 import (
-	"bytes"
-	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
@@ -13,7 +11,6 @@ import (
 	log "github.com/Sirupsen/logrus"
 	"github.com/gorilla/mux"
 	"github.com/gorilla/websocket"
-	consul "github.com/hashicorp/consul/api"
 	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/app"
@@ -23,68 +20,44 @@ import (
 )
 
 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
-	longPollDuration = 10 * time.Second
+	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
 
 	privateAPIPort = 4444
 )
 
 var (
-	queryOptions = &consul.QueryOptions{
-		RequireConsistent: true,
-	}
-	writeOptions = &consul.WriteOptions{}
-	wsDialer     = &websocket.Dialer{}
+	wsDialer = &websocket.Dialer{}
 )
 
 // TODO deal with garbage collection
 type consulPipe struct {
 	CreatedAt, DeletedAt time.Time
-	UIEnd, ProbeEnd      string // Addrs where each end is connected
+	UIAddr, ProbeAddr    string // Addrs where each end is connected
 	UIRef, ProbeRef      int    // Ref counts
 }
 
-func (c *consulPipe) toBytes() ([]byte, error) {
-	buf := bytes.Buffer{}
-	if err := json.NewEncoder(&buf).Encode(c); err != nil {
-		return nil, err
-	}
-	return buf.Bytes(), nil
-}
-
-func (c *consulPipe) fromBytes(bs []byte) error {
-	return json.NewDecoder(bytes.NewReader(bs)).Decode(c)
-}
-
-func (c *consulPipe) setEnd(e app.End, addr string) {
+func (c *consulPipe) setAddrFor(e app.End, addr string) {
 	if e == app.UIEnd {
-		c.UIEnd = addr
+		c.UIAddr = addr
 	} else {
-		c.ProbeEnd = addr
+		c.ProbeAddr = addr
 	}
 }
 
-func (c *consulPipe) end(e app.End) string {
+func (c *consulPipe) addrFor(e app.End) string {
 	if e == app.UIEnd {
-		return c.UIEnd
+		return c.UIAddr
 	}
-	return c.ProbeEnd
-}
-
-func (c *consulPipe) otherEnd(e app.End) string {
-	if e == app.UIEnd {
-		return c.ProbeEnd
-	}
-	return c.UIEnd
+	return c.ProbeAddr
 }
 
 func (c *consulPipe) eitherEndFor(addr string) bool {
-	return c.end(app.UIEnd) == addr || c.end(app.ProbeEnd) == addr
+	return c.addrFor(app.UIEnd) == addr || c.addrFor(app.ProbeEnd) == addr
 }
 
-func (c *consulPipe) incr(e app.End) int {
+func (c *consulPipe) acquire(e app.End) int {
 	if e == app.UIEnd {
 		c.UIRef++
 		return c.UIRef
@@ -93,7 +66,7 @@ func (c *consulPipe) incr(e app.End) int {
 	return c.ProbeRef
 }
 
-func (c *consulPipe) decr(e app.End) int {
+func (c *consulPipe) release(e app.End) int {
 	if e == app.UIEnd {
 		c.UIRef--
 		return c.UIRef
@@ -105,12 +78,12 @@ func (c *consulPipe) decr(e app.End) int {
 type consulPipeRouter struct {
 	prefix    string
 	advertise string // Address of this pipe router to advertise in consul
-	client    *consul.Client
+	client    ConsulClient
 	userIDer  UserIDer
 
-	pipes     map[string]xfer.Pipe // Active pipes
-	bridges   map[string]*bridgeConnection
-	actorChan chan func()
+	activePipes map[string]xfer.Pipe
+	bridges     map[string]*bridgeConnection
+	actorChan   chan func()
 
 	// Used by Stop()
 	quit chan struct{}
@@ -118,28 +91,21 @@ type consulPipeRouter struct {
 }
 
 // NewConsulPipeRouter returns a new consul based router
-func NewConsulPipeRouter(addr, prefix, inf string, userIDer UserIDer) (app.PipeRouter, error) {
+func NewConsulPipeRouter(client ConsulClient, prefix, inf string, userIDer UserIDer) (app.PipeRouter, error) {
 	advertise, err := network.GetFirstAddressOf(inf)
 	if err != nil {
 		return nil, err
 	}
-	client, err := consul.NewClient(&consul.Config{
-		Address: addr,
-		Scheme:  "http",
-	})
-	if err != nil {
-		return nil, err
-	}
 	pipeRouter := &consulPipeRouter{
 		prefix:    prefix,
 		advertise: advertise,
 		client:    client,
 		userIDer:  userIDer,
 
-		pipes:     map[string]xfer.Pipe{},
-		bridges:   map[string]*bridgeConnection{},
-		actorChan: make(chan func()),
-		quit:      make(chan struct{}),
+		activePipes: map[string]xfer.Pipe{},
+		bridges:     map[string]*bridgeConnection{},
+		actorChan:   make(chan func()),
+		quit:        make(chan struct{}),
 	}
 	pipeRouter.wait.Add(2)
 	go pipeRouter.watchAll()
@@ -171,55 +137,28 @@ func (pr *consulPipeRouter) actor() {
 // and all the methods are implemented as CAS's on consul, to
 // trigger an event in this loop.
 func (pr *consulPipeRouter) watchAll() {
-	var (
-		index = uint64(0)
-		kv    = pr.client.KV()
-	)
-	for {
+	defer pr.wait.Done()
+	pr.client.WatchPrefix(pr.prefix, &consulPipe{}, func(key string, value interface{}) bool {
 		select {
 		case <-pr.quit:
-			pr.wait.Done()
-			return
+			return false
 		default:
 		}
 
-		kvps, meta, err := kv.List(pr.prefix, &consul.QueryOptions{
-			RequireConsistent: true,
-			WaitIndex:         index,
-			WaitTime:          longPollDuration,
-		})
-		if err != nil {
-			log.Errorf("Error getting path %s: %v", pr.prefix, err)
-			continue
-		}
-		if index == meta.LastIndex {
-			continue
-		}
-		index = meta.LastIndex
-
-		for _, kvp := range kvps {
-			//log.Infof("Got background update to %s (%d)", kvp.Key, index)
-
-			cp := consulPipe{}
-			if err := cp.fromBytes(kvp.Value); err != nil {
-				log.Errorf("Error deserialising pipe %s: %s", kvp.Key, err)
-				continue
-			}
-
-			pr.actorChan <- func() { pr.handlePipeUpdate(kvp.Key, cp) }
-		}
-	}
+		pr.actorChan <- func() { pr.handlePipeUpdate(key, value.(*consulPipe)) }
+		return true
+	})
 }
 
-func (pr *consulPipeRouter) handlePipeUpdate(key string, cp consulPipe) {
+func (pr *consulPipeRouter) handlePipeUpdate(key string, cp *consulPipe) {
 	log.Infof("Got update to pipe %s", key)
 
 	// 1. If this pipe is closed, or we're not one of the ends, we
 	//    should ensure our local pipe (and bridge) is closed.
 	if !cp.DeletedAt.IsZero() || !cp.eitherEndFor(pr.advertise) {
 		log.Infof("Pipe %s not in use on this node.", key)
-		pipe, ok := pr.pipes[key]
-		delete(pr.pipes, key)
+		pipe, ok := pr.activePipes[key]
+		delete(pr.activePipes, key)
 		if ok {
 			pipe.Close()
 		}
@@ -237,23 +176,23 @@ func (pr *consulPipeRouter) handlePipeUpdate(key string, cp consulPipe) {
 	}
 
 	// 2. If this pipe if for us, we should have a pipe for it.
-	pipe, ok := pr.pipes[key]
+	pipe, ok := pr.activePipes[key]
 	if !ok {
 		pipe = xfer.NewPipe()
-		pr.pipes[key] = pipe
+		pr.activePipes[key] = pipe
 	}
 
 	// 3. Ensure there is a bridging connection for this pipe.
 	//    Semantics are the owner of the UIEnd connects to the owner of the ProbeEnd
 	shouldBridge := cp.DeletedAt.IsZero() &&
-		cp.end(app.UIEnd) != cp.end(app.ProbeEnd) &&
-		cp.end(app.UIEnd) == pr.advertise &&
-		cp.end(app.ProbeEnd) != ""
+		cp.addrFor(app.UIEnd) != cp.addrFor(app.ProbeEnd) &&
+		cp.addrFor(app.UIEnd) == pr.advertise &&
+		cp.addrFor(app.ProbeEnd) != ""
 	bridge, ok := pr.bridges[key]
 
 	// If we shouldn't be bridging but are, or we should be bridging but are pointing
 	// at the wrong place, stop the current bridge.
-	if (!shouldBridge && ok) || (shouldBridge && ok && bridge.addr != cp.end(app.ProbeEnd)) {
+	if (!shouldBridge && ok) || (shouldBridge && ok && bridge.addr != cp.addrFor(app.ProbeEnd)) {
 		log.Infof("Stopping bridge connection for %s", key)
 		delete(pr.bridges, key)
 		bridge.stop()
@@ -263,7 +202,7 @@ func (pr *consulPipeRouter) handlePipeUpdate(key string, cp consulPipe) {
 	// If we should be bridging and are not, start a new bridge
 	if shouldBridge && !ok {
 		log.Infof("Starting bridge connection for %s", key)
-		bridge = newBridgeConnection(key, cp.end(app.ProbeEnd), pipe)
+		bridge = newBridgeConnection(key, cp.addrFor(app.ProbeEnd), pipe)
 		pr.bridges[key] = bridge
 	}
 }
@@ -278,7 +217,7 @@ func (pr *consulPipeRouter) privateAPI() {
 				pc  = make(chan xfer.Pipe)
 			)
 			pr.actorChan <- func() {
-				pc <- pr.pipes[key]
+				pc <- pr.activePipes[key]
 			}
 			pipe := <-pc
 			if pipe == nil {
@@ -304,131 +243,20 @@ func (pr *consulPipeRouter) privateAPI() {
 	log.Infof("Private API terminated: %v", http.ListenAndServe(addr, router))
 }
 
-// Atomically modify a pipe in a callback.
-// If pipe doesn't exist you'll get nil in callback.
-func (pr *consulPipeRouter) get(key string) (*consulPipe, error) {
-	var (
-		kv   = pr.client.KV()
-		pipe consulPipe
-	)
-	kvp, _, err := kv.Get(key, queryOptions)
-	if err != nil {
-		return nil, err
-	}
-	if kvp == nil {
-		return nil, nil
-	}
-	if err := pipe.fromBytes(kvp.Value); err != nil {
-		return nil, err
-	}
-	return &pipe, nil
-}
-
-// Atomically modify a pipe in a callback.
-// If pipe doesn't exist you'll get nil in callback.
-func (pr *consulPipeRouter) cas(key string, f func(*consulPipe) (*consulPipe, bool, error)) (*consulPipe, error) {
-	var (
-		index   = uint64(0)
-		kv      = pr.client.KV()
-		pipe    *consulPipe
-		retries = 10
-		retry   = true
-	)
-	for i := 0; i < retries; i++ {
-		kvp, _, err := kv.Get(key, queryOptions)
-		if err != nil {
-			log.Errorf("Error getting %s: %v", key, err)
-			continue
-		}
-		if kvp != nil {
-			pipe = &consulPipe{}
-			if err := pipe.fromBytes(kvp.Value); err != nil {
-				log.Errorf("Error deserialising pipe %s: %v", key, err)
-				continue
-			}
-			index = kvp.ModifyIndex // if it doesn't exist, it will be 0
-		}
-
-		if pipe, retry, err = f(pipe); err != nil {
-			log.Errorf("Error CASing pipe %s: %v", key, err)
-			if !retry {
-				return nil, err
-			}
-			continue
-		}
-
-		if pipe == nil {
-			panic("Callback must instantiate pipe!")
-		}
-
-		value, err := pipe.toBytes()
-		if err != nil {
-			log.Errorf("Error serialising pipe %s: %v", key, err)
-			continue
-		}
-		ok, _, err := kv.CAS(&consul.KVPair{
-			Key:         key,
-			Value:       value,
-			ModifyIndex: index,
-		}, writeOptions)
-		if err != nil {
-			log.Errorf("Error CASing pipe %s: %v", key, err)
-			continue
-		}
-		if !ok {
-			log.Errorf("Error CASing pipe %s, trying again %d", key, index)
-			continue
-		}
-		return pipe, nil
-	}
-	return nil, fmt.Errorf("Failed to aquire pipe")
-}
-
-// Watch a given pipe and trigger a callback when it changes.
-// if callback returns false or error, exit (with the error).
-func (pr *consulPipeRouter) watch(key string, deadline time.Time, f func(*consulPipe) (bool, error)) (*consulPipe, error) {
-	var (
-		index = uint64(0)
-		kv    = pr.client.KV()
-	)
-	for deadline.After(mtime.Now()) {
-		// Poll waiting for the entry to get updated
-		kvp, meta, err := kv.Get(key, &consul.QueryOptions{
-			RequireConsistent: true,
-			WaitIndex:         index,
-			WaitTime:          longPollDuration,
-		})
-		if err != nil {
-			return nil, fmt.Errorf("Error getting %s: %v", key, err)
-		}
-		if kvp == nil {
-			return nil, fmt.Errorf("Pipe %s unexpectedly deleted!", key)
-		}
-
-		pipe := &consulPipe{}
-		pipe.fromBytes(kvp.Value)
-		if ok, err := f(pipe); !ok {
-			return pipe, nil
-		} else if err != nil {
-			return pipe, err
-		}
-
-		index = meta.LastIndex
-	}
-	return nil, fmt.Errorf("Timed out waiting on %s", key)
-}
-
 func (pr *consulPipeRouter) Exists(ctx context.Context, id string) (bool, error) {
 	userID, err := pr.userIDer(ctx)
 	if err != nil {
 		return false, err
 	}
 	key := fmt.Sprintf("%s%s-%s", pr.prefix, userID, id)
-	consulPipe, err := pr.get(key)
-	if err != nil {
+	consulPipe := consulPipe{}
+	err = pr.client.Get(key, &consulPipe)
+	if err == ErrNotFound {
+		return false, nil
+	} else if err != nil {
 		return false, err
 	}
-	return consulPipe == nil || consulPipe.DeletedAt.IsZero(), nil
+	return consulPipe.DeletedAt.IsZero(), nil
 }
 
 func (pr *consulPipeRouter) Get(ctx context.Context, id string, e app.End) (xfer.Pipe, io.ReadWriter, error) {
@@ -441,22 +269,25 @@ func (pr *consulPipeRouter) Get(ctx context.Context, id string, e app.End) (xfer
 
 	// Try to ensure the given end of the given pipe
 	// is 'owned' by this pipe service replica in consul.
-	_, err = pr.cas(key, func(p *consulPipe) (*consulPipe, bool, error) {
-		if p == nil {
-			p = &consulPipe{
+	err = pr.client.CAS(key, &consulPipe{}, func(in interface{}) (interface{}, bool, error) {
+		var pipe *consulPipe
+		if in == nil {
+			pipe = &consulPipe{
 				CreatedAt: mtime.Now(),
 			}
+		} else {
+			pipe = in.(*consulPipe)
 		}
-		if !p.DeletedAt.IsZero() {
+		if !pipe.DeletedAt.IsZero() {
 			return nil, false, fmt.Errorf("Pipe %s has been deleted", key)
 		}
-		end := p.end(e)
+		end := pipe.addrFor(e)
 		if end != "" && end != pr.advertise {
 			return nil, true, fmt.Errorf("Error: Pipe %s has existing connection to %s", key, end)
 		}
-		p.setEnd(e, pr.advertise)
-		p.incr(e)
-		return p, false, nil
+		pipe.setAddrFor(e, pr.advertise)
+		pipe.acquire(e)
+		return pipe, false, nil
 	})
 	if err != nil {
 		return nil, nil, err
@@ -465,10 +296,10 @@ func (pr *consulPipeRouter) Get(ctx context.Context, id string, e app.End) (xfer
 	// next see if we already have a active pipe
 	pc := make(chan xfer.Pipe)
 	pr.actorChan <- func() {
-		pipe, ok := pr.pipes[key]
+		pipe, ok := pr.activePipes[key]
 		if !ok {
 			pipe = xfer.NewPipe()
-			pr.pipes[key] = pipe
+			pr.activePipes[key] = pipe
 		}
 		pc <- pipe
 	}
@@ -490,20 +321,20 @@ func (pr *consulPipeRouter) Release(ctx context.Context, id string, e app.End) e
 	log.Infof("Release %s:%s", key, e)
 
 	// atomically clear my end of the pipe in consul
-	_, err = pr.cas(key, func(p *consulPipe) (*consulPipe, bool, error) {
-		if p == nil {
+	return pr.client.CAS(key, &consulPipe{}, func(in interface{}) (interface{}, bool, error) {
+		if in == nil {
 			return nil, false, fmt.Errorf("Pipe %s not found", id)
 		}
-		if p.end(e) != pr.advertise {
+		p := in.(*consulPipe)
+		if p.addrFor(e) != pr.advertise {
 			return nil, false, fmt.Errorf("Pipe %s not owned by us!", id)
 		}
-		refs := p.decr(e)
+		refs := p.release(e)
 		if refs == 0 {
-			p.setEnd(e, "")
+			p.setAddrFor(e, "")
 		}
 		return p, true, nil
 	})
-	return err
 }
 
 func (pr *consulPipeRouter) Delete(ctx context.Context, id string) error {
@@ -514,16 +345,19 @@ func (pr *consulPipeRouter) Delete(ctx context.Context, id string) error {
 	key := fmt.Sprintf("%s%s-%s", pr.prefix, userID, id)
 	log.Infof("Delete %s", key)
 
-	_, err = pr.cas(key, func(p *consulPipe) (*consulPipe, bool, error) {
-		if p == nil {
+	return pr.client.CAS(key, &consulPipe{}, func(in interface{}) (interface{}, bool, error) {
+		if in == nil {
 			return nil, false, fmt.Errorf("Pipe %s not found", id)
 		}
+		p := in.(*consulPipe)
 		p.DeletedAt = mtime.Now()
 		return p, false, nil
 	})
-	return err
 }
 
+// A bridgeConnection represents a connection between two pipe router replicas.
+// They are created & destroyed in response to events from consul, which in turn
+// are triggered when UIs or Probes connect to various pipe routers.
 type bridgeConnection struct {
 	key  string
 	addr string // address to connect to

app/multitenant/sqs_control_router.go

@@ -30,9 +30,9 @@ var (
 // probe queue.  When probe recieves a request, handles it and posts the
 // response back to the response queue.
 type sqsControlRouter struct {
-	service  *sqs.SQS
-	queueURL *string
-	userIDer UserIDer
+	service          *sqs.SQS
+	responseQueueURL *string
+	userIDer         UserIDer
 
 	mtx          sync.Mutex
 	responses    map[string]chan xfer.Response
@@ -57,10 +57,10 @@ func NewSQSControlRouter(url, region string, creds *credentials.Credentials, use
 			WithEndpoint(url).
 			WithRegion(region).
 			WithCredentials(creds))),
-		queueURL:     nil,
-		userIDer:     userIDer,
-		responses:    map[string]chan xfer.Response{},
-		probeWorkers: map[int64]*probeWorker{},
+		responseQueueURL: nil,
+		userIDer:         userIDer,
+		responses:        map[string]chan xfer.Response{},
+		probeWorkers:     map[int64]*probeWorker{},
 	}
 	go result.loop()
 	return result
@@ -70,26 +70,20 @@ func (cr *sqsControlRouter) Stop() error {
 	return nil
 }
 
-func (cr *sqsControlRouter) setQueueURL(url *string) {
+func (cr *sqsControlRouter) setResponseQueueURL(url *string) {
 	cr.mtx.Lock()
 	defer cr.mtx.Unlock()
-	cr.queueURL = url
+	cr.responseQueueURL = url
 }
 
-func (cr *sqsControlRouter) getQueueURL() *string {
+func (cr *sqsControlRouter) getResponseQueueURL() *string {
 	cr.mtx.Lock()
 	defer cr.mtx.Unlock()
-	return cr.queueURL
+	return cr.responseQueueURL
 }
 
 func (cr *sqsControlRouter) getOrCreateQueue(name string) (*string, error) {
-	getQueueURLRes, err := cr.service.GetQueueUrl(&sqs.GetQueueUrlInput{
-		QueueName: aws.String(name),
-	})
-	if err == nil {
-		return getQueueURLRes.QueueUrl, nil
-	}
-
+	// CreateQueue creates a queue or if it already exists, returns url of said queue
 	createQueueRes, err := cr.service.CreateQueue(&sqs.CreateQueueInput{
 		QueueName: aws.String(name),
 	})
@@ -100,35 +94,39 @@ func (cr *sqsControlRouter) getOrCreateQueue(name string) (*string, error) {
 }
 
 func (cr *sqsControlRouter) loop() {
+	var (
+		responseQueueURL *string
+		err              error
+	)
 	for {
 		// This app has a random id and uses this as a return path for all responses from probes.
 		name := fmt.Sprintf("control-app-%d", rand.Int63())
-		queueURL, err := cr.getOrCreateQueue(name)
+		responseQueueURL, err = cr.getOrCreateQueue(name)
 		if err != nil {
 			log.Errorf("Failed to create queue: %v", err)
 			time.Sleep(1 * time.Second)
 			continue
 		}
-		cr.setQueueURL(queueURL)
+		cr.setResponseQueueURL(responseQueueURL)
 		break
 	}
 
 	for {
 		res, err := cr.service.ReceiveMessage(&sqs.ReceiveMessageInput{
-			QueueUrl:        cr.queueURL,
+			QueueUrl:        responseQueueURL,
 			WaitTimeSeconds: longPollTime,
 		})
 		if err != nil {
-			log.Errorf("Error recieving message from %s: %v", *cr.queueURL, err)
+			log.Errorf("Error receiving message from %s: %v", *responseQueueURL, err)
 			continue
 		}
 		if len(res.Messages) == 0 {
 			continue
 		}
-		cr.handleResponses(res)
-		if err := cr.deleteMessages(cr.queueURL, res.Messages); err != nil {
-			log.Errorf("Error deleting message from %s: %v", *cr.queueURL, err)
+		if err := cr.deleteMessages(responseQueueURL, res.Messages); err != nil {
+			log.Errorf("Error deleting message from %s: %v", *responseQueueURL, err)
 		}
+		cr.handleResponses(res)
 	}
 }
 
@@ -148,21 +146,17 @@ func (cr *sqsControlRouter) deleteMessages(queueURL *string, messages []*sqs.Mes
 }
 
 func (cr *sqsControlRouter) handleResponses(res *sqs.ReceiveMessageOutput) {
-	sqsResponses := []sqsResponseMessage{}
+	cr.mtx.Lock()
+	defer cr.mtx.Unlock()
+
 	for _, message := range res.Messages {
 		var sqsResponse sqsResponseMessage
 		if err := json.NewDecoder(bytes.NewBufferString(*message.Body)).Decode(&sqsResponse); err != nil {
 			log.Errorf("Error decoding message: %v", err)
 			continue
 		}
-		sqsResponses = append(sqsResponses, sqsResponse)
-	}
 
-	for _, sqsResponse := range sqsResponses {
-		cr.mtx.Lock()
 		waiter, ok := cr.responses[sqsResponse.ID]
-		cr.mtx.Unlock()
-
 		if !ok {
 			log.Errorf("Dropping response %s - no one waiting for it!", sqsResponse.ID)
 			continue
@@ -192,10 +186,11 @@ func (cr *sqsControlRouter) Handle(ctx context.Context, probeID string, req xfer
 	}
 
 	// Get the queue url for the local (control app) queue, and for the probe.
-	queueURL := cr.getQueueURL()
-	if queueURL == nil {
+	responseQueueURL := cr.getResponseQueueURL()
+	if responseQueueURL == nil {
 		return xfer.Response{}, fmt.Errorf("No SQS queue yet!")
 	}
+
 	probeQueueName := fmt.Sprintf("probe-%s-%s", userID, probeID)
 	probeQueueURL, err := cr.service.GetQueueUrl(&sqs.GetQueueUrlInput{
 		QueueName: aws.String(probeQueueName),
@@ -204,7 +199,7 @@ func (cr *sqsControlRouter) Handle(ctx context.Context, probeID string, req xfer
 		return xfer.Response{}, err
 	}
 
-	// Wait for a response befor we send the request, to prevent races
+	// Add a response channel before we send the request, to prevent races
 	id := fmt.Sprintf("request-%s-%d", userID, rand.Int63())
 	waiter := make(chan xfer.Response, 1)
 	cr.mtx.Lock()
@@ -220,7 +215,7 @@ func (cr *sqsControlRouter) Handle(ctx context.Context, probeID string, req xfer
 	if err := cr.sendMessage(probeQueueURL.QueueUrl, sqsRequestMessage{
 		ID:               id,
 		Request:          req,
-		ResponseQueueURL: *queueURL,
+		ResponseQueueURL: *responseQueueURL,
 	}); err != nil {
 		return xfer.Response{}, err
 	}
@@ -248,10 +243,10 @@ func (cr *sqsControlRouter) Register(ctx context.Context, probeID string, handle
 
 	pwID := rand.Int63()
 	pw := &probeWorker{
-		router:   cr,
-		queueURL: queueURL,
-		handler:  handler,
-		quit:     make(chan struct{}),
+		router:          cr,
+		requestQueueURL: queueURL,
+		handler:         handler,
+		quit:            make(chan struct{}),
 	}
 	pw.done.Add(1)
 	go pw.loop()
@@ -273,12 +268,13 @@ func (cr *sqsControlRouter) Deregister(_ context.Context, probeID string, id int
 	return nil
 }
 
+// a probeWorker encapsulates a goroutine serving a probe's websocket connection.
 type probeWorker struct {
-	router   *sqsControlRouter
-	queueURL *string
-	handler  xfer.ControlHandlerFunc
-	quit     chan struct{}
-	done     sync.WaitGroup
+	router          *sqsControlRouter
+	requestQueueURL *string
+	handler         xfer.ControlHandlerFunc
+	quit            chan struct{}
+	done            sync.WaitGroup
 }
 
 func (pw *probeWorker) stop() {
@@ -296,7 +292,7 @@ func (pw *probeWorker) loop() {
 		}
 
 		res, err := pw.router.service.ReceiveMessage(&sqs.ReceiveMessageInput{
-			QueueUrl:        pw.queueURL,
+			QueueUrl:        pw.requestQueueURL,
 			WaitTimeSeconds: longPollTime,
 		})
 		if err != nil {
@@ -306,8 +302,10 @@ func (pw *probeWorker) loop() {
 		if len(res.Messages) == 0 {
 			continue
 		}
+		if err := pw.router.deleteMessages(pw.requestQueueURL, res.Messages); err != nil {
+			log.Errorf("Error deleting message from %s: %v", *pw.requestQueueURL, err)
+		}
 
-		// TODO do we need to parallelise the handling of requests?
 		for _, message := range res.Messages {
 			var sqsRequest sqsRequestMessage
 			if err := json.NewDecoder(bytes.NewBufferString(*message.Body)).Decode(&sqsRequest); err != nil {
@@ -315,16 +313,14 @@ func (pw *probeWorker) loop() {
 				continue
 			}
 
+			response := pw.handler(sqsRequest.Request)
+
 			if err := pw.router.sendMessage(&sqsRequest.ResponseQueueURL, sqsResponseMessage{
 				ID:       sqsRequest.ID,
-				Response: pw.handler(sqsRequest.Request),
+				Response: response,
 			}); err != nil {
 				log.Errorf("Error sending response: %v", err)
 			}
 		}
-
-		if err := pw.router.deleteMessages(pw.queueURL, res.Messages); err != nil {
-			log.Errorf("Error deleting message from %s: %v", *pw.queueURL, err)
-		}
 	}
 }

app/multitenant/user.go

@@ -9,9 +9,9 @@ import (
 	"github.com/weaveworks/scope/app"
 )
 
-// ErrNotFound should be returned by a UserIDer when it fails to ID the
+// ErrUserIDNotFound should be returned by a UserIDer when it fails to ID the
 // user for a request.
-var ErrNotFound = fmt.Errorf("User ID not found")
+var ErrUserIDNotFound = fmt.Errorf("User ID not found")
 
 // UserIDer identifies users given a request context.
 type UserIDer func(context.Context) (string, error)
@@ -21,11 +21,11 @@ func UserIDHeader(headerName string) UserIDer {
 	return func(ctx context.Context) (string, error) {
 		request, ok := ctx.Value(app.RequestCtxKey).(*http.Request)
 		if !ok || request == nil {
-			return "", ErrNotFound
+			return "", ErrUserIDNotFound
 		}
 		userID := request.Header.Get(headerName)
 		if userID == "" {
-			return "", ErrNotFound
+			return "", ErrUserIDNotFound
 		}
 		return userID, nil
 	}

common/network/interface.go

@@ -5,7 +5,7 @@ import (
 	"net"
 )
 
-// GetFirstAddressOf returns the first address of the supplied interface name.
+// GetFirstAddressOf returns the first IPv4 address of the supplied interface name.
 func GetFirstAddressOf(name string) (string, error) {
 	inf, err := net.InterfaceByName(name)
 	if err != nil {
@@ -20,10 +20,14 @@ func GetFirstAddressOf(name string) (string, error) {
 		return "", fmt.Errorf("No address found for %s", name)
 	}
 
-	switch v := addrs[0].(type) {
-	case *net.IPNet:
-		return v.IP.String(), nil
-	default:
-		return "", fmt.Errorf("No address found for %s", name)
+	for _, addr := range addrs {
+		switch v := addr.(type) {
+		case *net.IPNet:
+			if ip := v.IP.To4(); ip != nil {
+				return v.IP.String(), nil
+			}
+		}
 	}
+
+	return "", fmt.Errorf("No address found for %s", name)
 }

common/xfer/pipes.go

@@ -93,6 +93,9 @@ func (p *pipe) CopyToWebsocket(end io.ReadWriter, conn Websocket) error {
 	p.mtx.Unlock()
 	defer p.wg.Done()
 
+	// The goroutines below both post their errors to the channel, but if you close()
+	// the pipe before any errors then the pipe may not get read from. Therefore it
+	// needs up to 2 slots free.
 	errors := make(chan error, 2)
 
 	// Read-from-UI loop

prog/app.go

@@ -106,8 +106,11 @@ func appMain() {
 	if *pipeRouterType == "local" {
 		pipeRouter = app.NewLocalPipeRouter()
 	} else if *pipeRouterType == "consul" {
-		var err error
-		pipeRouter, err = multitenant.NewConsulPipeRouter(*consulAddr, *consulPrefix, *consulInf, userIDer)
+		consulClient, err := multitenant.NewConsulClient(*consulAddr)
+		if err != nil {
+			log.Fatalf("Error createing consul client: %v", err)
+		}
+		pipeRouter, err = multitenant.NewConsulPipeRouter(consulClient, *consulPrefix, *consulInf, userIDer)
 		if err != nil {
 			log.Fatalf("Error createing consul pipe router: %v", err)
 		}