Share blocking queries between RPC requests

Only run one blocking query per request and share the result with all
RPC calls. This optimizes the blocking query process when many agents
are watching the same service.

Author: Thibault Gilles

agent/consul/health_endpoint.go

@@ -2,6 +2,7 @@ package consul
 
 import (
 	"fmt"
+	"sort"
 
 	"github.com/armon/go-metrics"
 	"github.com/hashicorp/consul/agent/consul/state"
@@ -137,52 +138,72 @@ func (h *Health) ServiceNodes(args *structs.ServiceSpecificRequest, reply *struc
 		}
 	}
 
-	err := h.srv.blockingQuery(
+	err := h.srv.sharedBlockingQuery(
+		args,
+		reply,
 		&args.QueryOptions,
 		&reply.QueryMeta,
-		func(ws memdb.WatchSet, state *state.Store) error {
+		func(ws memdb.WatchSet, state *state.Store) (uint64, func(uint64, interface{}) error, error) {
 			index, nodes, err := f(ws, state, args)
 			if err != nil {
-				return err
+				return 0, nil, err
 			}
 
-			reply.Index, reply.Nodes = index, nodes
-			if len(args.NodeMetaFilters) > 0 {
-				reply.Nodes = nodeMetaFilter(args.NodeMetaFilters, reply.Nodes)
-			}
-			if err := h.srv.filterACL(args.Token, reply); err != nil {
-				return err
-			}
-			return h.srv.sortNodesByDistanceFrom(args.Source, reply.Nodes)
+			return index, func(index uint64, rawReply interface{}) error {
+				reply := rawReply.(*structs.IndexedCheckServiceNodes)
+				// copy the slice to ensure requests filtering / sorting dont affect each others
+				replyNodes := make(structs.CheckServiceNodes, len(nodes))
+				copy(replyNodes, nodes)
+				reply.Index, reply.Nodes = index, replyNodes
+				return nil
+			}, nil
 		})
+	if err != nil {
+		return err
+	}
+
+	h.srv.setQueryMeta(&reply.QueryMeta)
+
+	if len(args.NodeMetaFilters) > 0 {
+		reply.Nodes = nodeMetaFilter(args.NodeMetaFilters, reply.Nodes)
+	}
+	if err := h.srv.filterACL(args.Token, reply); err != nil {
+		return err
+	}
+	err = h.srv.sortNodesByDistanceFrom(args.Source, reply.Nodes)
+	if err != nil {
+		return err
+	}
 
 	// Provide some metrics
-	if err == nil {
-		// For metrics, we separate Connect-based lookups from non-Connect
-		key := "service"
-		if args.Connect {
-			key = "connect"
-		}
+	// For metrics, we separate Connect-based lookups from non-Connect
+	key := "service"
+	if args.Connect {
+		key = "connect"
+	}
 
-		metrics.IncrCounterWithLabels([]string{"health", key, "query"}, 1,
-			[]metrics.Label{{Name: "service", Value: args.ServiceName}})
-		if args.ServiceTag != "" {
-			metrics.IncrCounterWithLabels([]string{"health", key, "query-tag"}, 1,
-				[]metrics.Label{{Name: "service", Value: args.ServiceName}, {Name: "tag", Value: args.ServiceTag}})
-		}
-		if len(args.ServiceTags) > 0 {
-			labels := []metrics.Label{{Name: "service", Value: args.ServiceName}}
-			for _, tag := range args.ServiceTags {
-				labels = append(labels, metrics.Label{Name: "tag", Value: tag})
-			}
-			metrics.IncrCounterWithLabels([]string{"health", key, "query-tags"}, 1, labels)
-		}
-		if len(reply.Nodes) == 0 {
-			metrics.IncrCounterWithLabels([]string{"health", key, "not-found"}, 1,
-				[]metrics.Label{{Name: "service", Value: args.ServiceName}})
+	metrics.IncrCounterWithLabels([]string{"health", key, "query"}, 1,
+		[]metrics.Label{{Name: "service", Value: args.ServiceName}})
+	if args.ServiceTag != "" {
+		metrics.IncrCounterWithLabels([]string{"health", key, "query-tag"}, 1,
+			[]metrics.Label{{Name: "service", Value: args.ServiceName}, {Name: "tag", Value: args.ServiceTag}})
+	}
+	if len(args.ServiceTags) > 0 {
+		// Sort tags so that the metric is the same even if the request
+		// tags are in a different order
+		sort.Strings(args.ServiceTags)
+
+		labels := []metrics.Label{{Name: "service", Value: args.ServiceName}}
+		for _, tag := range args.ServiceTags {
+			labels = append(labels, metrics.Label{Name: "tag", Value: tag})
 		}
+		metrics.IncrCounterWithLabels([]string{"health", key, "query-tags"}, 1, labels)
+	}
+	if len(reply.Nodes) == 0 {
+		metrics.IncrCounterWithLabels([]string{"health", key, "not-found"}, 1,
+			[]metrics.Label{{Name: "service", Value: args.ServiceName}})
 	}
-	return err
+	return nil
 }
 
 // The serviceNodes* functions below are the various lookup methods that

agent/consul/rpc.go

@@ -6,9 +6,11 @@ import (
 	"io"
 	"net"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/armon/go-metrics"
+	"github.com/hashicorp/consul/agent/cache"
 	"github.com/hashicorp/consul/agent/consul/state"
 	"github.com/hashicorp/consul/agent/metadata"
 	"github.com/hashicorp/consul/agent/pool"
@@ -443,6 +445,157 @@ RUN_QUERY:
 	return err
 }
 
+type sharedQueryFn func(memdb.WatchSet, *state.Store) (uint64, func(uint64, interface{}) error, error)
+
+func (s *Server) sharedBlockingQuery(req cache.Request, res interface{}, queryOpts *structs.QueryOptions, queryMeta *structs.QueryMeta, fn sharedQueryFn) error {
+	var timeout *time.Timer
+
+	if queryOpts.RequireConsistent {
+		if err := s.consistentRead(); err != nil {
+			return err
+		}
+	}
+
+	// Restrict the max query time, and ensure there is always one.
+	if queryOpts.MaxQueryTime > maxQueryTime {
+		queryOpts.MaxQueryTime = maxQueryTime
+	} else if queryOpts.MaxQueryTime <= 0 {
+		queryOpts.MaxQueryTime = defaultQueryTime
+	}
+
+	// Apply a small amount of jitter to the request.
+	queryOpts.MaxQueryTime += lib.RandomStagger(queryOpts.MaxQueryTime / jitterFraction)
+
+	// Setup a query timeout.
+	timeout = time.NewTimer(queryOpts.MaxQueryTime)
+	defer timeout.Stop()
+
+	cacheInfo := req.CacheInfo()
+
+	// Check if a blocking query is already runnning for this request
+	s.blockingQueriesLock.RLock()
+	queryState, alreadyInserted := s.blockingQueries[cacheInfo.Key]
+	s.blockingQueriesLock.RUnlock()
+
+	// If not, run one
+	if !alreadyInserted {
+		ws := memdb.NewWatchSet()
+		// run the func a first time to get the index
+		firstRunIndex, apply, err := fn(ws, s.fsm.State())
+
+		s.blockingQueriesLock.Lock()
+		queryState, alreadyInserted = s.blockingQueries[cacheInfo.Key]
+		if alreadyInserted {
+			// Another query raced with us and already ran a blocking query
+			s.blockingQueriesLock.Unlock()
+		} else {
+			// Add query to map
+			queryState = newBlockingQueryState(firstRunIndex, apply, err)
+			s.blockingQueries[cacheInfo.Key] = queryState
+			s.blockingQueriesLock.Unlock()
+
+			// Run the shared blocking query
+			go s.runSharedBlockingQuery(firstRunIndex, cacheInfo.Key, ws, queryMeta, queryState, fn)
+		}
+	}
+
+	stateIndex := atomic.LoadUint64(&queryState.Index)
+
+	if stateIndex <= queryOpts.MinQueryIndex {
+		// Increment the shared query watcher
+		atomic.AddInt32(&queryState.Watchers, 1)
+
+		// block on either timeout or shared query
+		select {
+		case <-timeout.C:
+			if n := atomic.AddInt32(&queryState.Watchers, -1); n == 0 {
+				s.logger.Println("[TRACE] consul: cancelling shared blocking query because there is no more watchers")
+
+				// we were the last request to wait on the shared blocking wuery and we reached MaxQueryTime, cancel the blocking query
+				close(queryState.Cancel)
+			}
+		case <-queryState.Done:
+		}
+	}
+
+	if err := queryState.Err.Load(); err != nil {
+		return err.(error)
+	}
+
+	return queryState.Apply.Load().(func(uint64, interface{}) error)(atomic.LoadUint64(&queryState.Index), res)
+}
+
+func (s *Server) runSharedBlockingQuery(index uint64, cacheKey string, ws memdb.WatchSet, queryMeta *structs.QueryMeta, queryState *blockingQueryState, fn sharedQueryFn) {
+	s.logger.Println("[TRACE] consul: running shared blocking query")
+
+	// Wait initial query watchset
+	expired := ws.Watch(queryState.Cancel)
+
+RUN_QUERY:
+
+	// If the read must be consistent we verify that we are still the leader.
+
+	// Run the query.
+	metrics.IncrCounter([]string{"rpc", "query"}, 1)
+
+	// Operate on a consistent set of state. This makes sure that the
+	// abandon channel goes with the state that the caller is using to
+	// build watches.
+	state := s.fsm.State()
+
+	// We can skip all watch tracking if this isn't a blocking query.
+	if index > 0 {
+		ws = memdb.NewWatchSet()
+
+		// This channel will be closed if a snapshot is restored and the
+		// whole state store is abandoned.
+		ws.Add(state.AbandonCh())
+	}
+
+	// Block up to the timeout if we didn't see anything fresh.
+	idx, apply, err := fn(ws, state)
+	// Note we check queryOpts.MinQueryIndex is greater than zero to determine if
+	// blocking was requested by client, NOT meta.Index since the state function
+	// might return zero if something is not initialised and care wasn't taken to
+	// handle that special case (in practice this happened a lot so fixing it
+	// systematically here beats trying to remember to add zero checks in every
+	// state method). We also need to ensure that unless there is an error, we
+	// return an index > 0 otherwise the client will never block and burn CPU and
+	// requests.
+	if err == nil && idx < 1 {
+		idx = 1
+	}
+	if !expired && err == nil && index > 0 && idx <= index {
+		if expired := ws.Watch(queryState.Cancel); !expired {
+			// If a restore may have woken us up then bail out from
+			// the query immediately. This is slightly race-ey since
+			// this might have been interrupted for other reasons,
+			// but it's OK to kick it back to the caller in either
+			// case.
+			select {
+			case <-state.AbandonCh():
+			default:
+				goto RUN_QUERY
+			}
+		}
+	}
+
+	// store results
+	s.blockingQueriesLock.Lock()
+	if err != nil {
+		queryState.Err.Store(err)
+	}
+	if apply != nil {
+		queryState.Apply.Store(apply)
+	}
+	atomic.StoreUint64(&queryState.Index, idx)
+	delete(s.blockingQueries, cacheKey)
+	s.blockingQueriesLock.Unlock()
+
+	// notify changed
+	close(queryState.Done)
+}
+
 // setQueryMeta is used to populate the QueryMeta data for an RPC call
 func (s *Server) setQueryMeta(m *structs.QueryMeta) {
 	if s.IsLeader() {

agent/consul/server.go

@@ -88,6 +88,30 @@ var (
 	ErrWANFederationDisabled = fmt.Errorf("WAN Federation is disabled")
 )
 
+type blockingQueryState struct {
+	Index    uint64
+	Cancel   chan time.Time
+	Watchers int32
+	Done     chan struct{}
+	Apply    atomic.Value
+	Err      atomic.Value
+}
+
+func newBlockingQueryState(index uint64, apply func(uint64, interface{}) error, err error) *blockingQueryState {
+	queryState := &blockingQueryState{
+		Done:   make(chan struct{}),
+		Cancel: make(chan time.Time),
+		Index:  index,
+	}
+	if apply != nil {
+		queryState.Apply.Store(apply)
+	}
+	if err != nil {
+		queryState.Err.Store(err)
+	}
+	return queryState
+}
+
 // Server is Consul server which manages the service discovery,
 // health checking, DC forwarding, Raft, and multiple Serf pools.
 type Server struct {
@@ -235,6 +259,9 @@ type Server struct {
 	shutdownCh   chan struct{}
 	shutdownLock sync.Mutex
 
+	blockingQueriesLock sync.RWMutex
+	blockingQueries     map[string]*blockingQueryState
+
 	// embedded struct to hold all the enterprise specific data
 	EnterpriseServer
 }
@@ -323,6 +350,7 @@ func NewServerLogger(config *Config, logger *log.Logger, tokens *token.Store) (*
 		sessionTimers:    NewSessionTimers(),
 		tombstoneGC:      gc,
 		serverLookup:     NewServerLookup(),
+		blockingQueries:  make(map[string]*blockingQueryState),
 		shutdownCh:       shutdownCh,
 	}
 

agent/consul/state/catalog.go

@@ -1697,7 +1697,7 @@ func (s *Store) CheckServiceTagNodes(ws memdb.WatchSet, serviceName string, tags
 	for service := iter.Next(); service != nil; service = iter.Next() {
 		svc := service.(*structs.ServiceNode)
 		if !serviceTagsFilter(svc, tags) {
-			results = append(results, svc)
+			results = append(results, service.(*structs.ServiceNode))
 		}
 	}
 	return s.parseCheckServiceNodes(tx, ws, idx, serviceName, results, err)

agent/structs/structs.go

@@ -6,6 +6,7 @@ import (
 	"math/rand"
 	"reflect"
 	"regexp"
+	"sort"
 	"strconv"
 	"strings"
 	"time"
@@ -379,10 +380,12 @@ func (r *ServiceSpecificRequest) CacheInfo() cache.RequestInfo {
 	// cached results, we need to be careful we maintain the same order of fields
 	// here. We could alternatively use `hash:set` struct tag on an anonymous
 	// struct to make it more robust if it becomes significant.
+	sort.Strings(r.ServiceTags)
 	v, err := hashstructure.Hash([]interface{}{
 		r.NodeMetaFilters,
 		r.ServiceName,
 		r.ServiceTag,
+		r.ServiceTags, // make the tests pass until http://github.com/hashicorp/consul/pull/4987 is merged
 		r.ServiceAddress,
 		r.TagFilter,
 		r.Connect,