Simple StoreFinder and new capacity gossiping

gossip/infostore.go

@@ -210,7 +210,7 @@ func (is *infoStore) addInfo(i *info) error {
 	if existingInfo, ok := is.Infos[i.Key]; ok {
 		if i.Timestamp < existingInfo.Timestamp ||
 			(i.Timestamp == existingInfo.Timestamp && i.Hops >= existingInfo.Hops) {
-			return util.Errorf("info %+v older than current group info %+v", i, existingInfo)
+			return util.Errorf("info %+v older than current info %+v", i, existingInfo)
 		}
 		contentsChanged = !reflect.DeepEqual(existingInfo.Val, i.Val)
 	} else {

gossip/keys.go

@@ -35,9 +35,8 @@ const (
 	KeyConfigZone = "zones"
 
 	// KeyMaxAvailCapacityPrefix is the key prefix for gossiping available
-	// store capacity. The suffix is composed of:
-	// <datacenter>-<hex node ID>-<hex store ID>. The value is a
-	// storage.StoreDescriptor struct.
+	// store capacity. The suffix is composed of: <hex node ID>-<hex store ID>.
+	// The value is a storage.StoreDescriptor struct.
 	KeyMaxAvailCapacityPrefix = "max-avail-capacity-"
 
 	// KeyNodeCount is the count of gossip nodes in the network. The

proto/config.go

@@ -25,8 +25,8 @@ import (
 	"strings"
 )
 
-// IsSubset returns whether attributes list b is a subset of
-// attributes list a.
+// IsSubset returns whether attributes list a is a subset of
+// attributes list b.
 func (a Attributes) IsSubset(b Attributes) bool {
 	m := map[string]struct{}{}
 	for _, s := range b.Attrs {

server/node.go

@@ -368,11 +368,10 @@ func (n *Node) gossipCapacities() {
 			log.Warningf("problem getting store descriptor for store %+v: %v", s.Ident, err)
 			return nil
 		}
-		gossipPrefix := gossip.KeyMaxAvailCapacityPrefix + storeDesc.CombinedAttrs().SortedString()
-		keyMaxCapacity := gossipPrefix + strconv.FormatInt(int64(storeDesc.Node.NodeID), 10) + "-" +
-			strconv.FormatInt(int64(storeDesc.StoreID), 10)
-		// Register gossip group.
-		n.gossip.RegisterGroup(gossipPrefix, gossipGroupLimit, gossip.MaxGroup)
+		// Unique gossip key per store.
+		keyMaxCapacity := gossip.KeyMaxAvailCapacityPrefix +
+			strconv.FormatInt(int64(storeDesc.Node.NodeID), 16) + "-" +
+			strconv.FormatInt(int64(storeDesc.StoreID), 16)
 		// Gossip store descriptor.
 		n.gossip.AddInfo(keyMaxCapacity, *storeDesc, ttlCapacityGossip)
 		return nil

storage/allocator.go

@@ -25,9 +25,6 @@ import (
 	"github.com/cockroachdb/cockroach/util"
 )
 
-// StoreFinder finds the disks in a datacenter with the most available capacity.
-type StoreFinder func(proto.Attributes) ([]*StoreDescriptor, error)
-
 // allocator makes allocation decisions based on a zone configuration,
 // existing range metadata and available stores. Configuration
 // settings and range metadata information is stored directly in the

storage/allocator_test.go

@@ -51,9 +51,8 @@ var multiDCConfig = proto.ZoneConfig{
 func filterStores(a proto.Attributes, stores []*StoreDescriptor) ([]*StoreDescriptor, error) {
 	var filtered []*StoreDescriptor
 	for _, s := range stores {
-		b := s.Attrs.Attrs
-		b = append(b, s.Node.Attrs.Attrs...)
-		if a.IsSubset(proto.Attributes{Attrs: b}) {
+		sAttrs := s.CombinedAttrs()
+		if a.IsSubset(*sAttrs) {
 			filtered = append(filtered, s)
 		}
 	}

storage/store.go

@@ -164,17 +164,18 @@ func (s StoreDescriptor) Less(b util.Ordered) bool {
 // A Store maintains a map of ranges by start key. A Store corresponds
 // to one physical device.
 type Store struct {
-	Ident        proto.StoreIdent
-	clock        *hlc.Clock
-	engine       engine.Engine  // The underlying key-value store
-	db           *client.KV     // Cockroach KV DB
-	allocator    *allocator     // Makes allocation decisions
-	gossip       *gossip.Gossip // Configs and store capacities
-	raftIDAlloc  *IDAllocator   // Raft ID allocator
-	rangeIDAlloc *IDAllocator   // Range ID allocator
-	configMu     sync.Mutex     // Limit config update processing
-	raft         raft
-	closer       chan struct{}
+	Ident         proto.StoreIdent
+	clock         *hlc.Clock
+	engine        engine.Engine      // The underlying key-value store
+	db            *client.KV         // Cockroach KV DB
+	allocator     *allocator         // Makes allocation decisions
+	gossip        *gossip.Gossip     // Configs and store capacities
+	raftIDAlloc   *IDAllocator       // Raft ID allocator
+	rangeIDAlloc  *IDAllocator       // Range ID allocator
+	configMu      sync.Mutex         // Limit config update processing
+	finderContext storeFinderContext // Context to find stores for allocation
+	raft          raft
+	closer        chan struct{}
 
 	mu             sync.RWMutex     // Protects variables below...
 	ranges         map[int64]*Range // Map of ranges by range ID
@@ -184,7 +185,7 @@ type Store struct {
 
 // NewStore returns a new instance of a store.
 func NewStore(clock *hlc.Clock, eng engine.Engine, db *client.KV, gossip *gossip.Gossip) *Store {
-	return &Store{
+	s := &Store{
 		clock:          clock,
 		engine:         eng,
 		db:             db,
@@ -195,6 +196,8 @@ func NewStore(clock *hlc.Clock, eng engine.Engine, db *client.KV, gossip *gossip
 		ranges:         map[int64]*Range{},
 		rangesByRaftID: map[int64]*Range{},
 	}
+	s.allocator.storeFinder = s.findStores
+	return s
 }
 
 // Stop calls Range.Stop() on all active ranges.
@@ -286,6 +289,9 @@ func (s *Store) Start() error {
 	if s.gossip != nil {
 		s.gossip.RegisterCallback(gossip.KeyConfigAccounting, s.configGossipUpdate)
 		s.gossip.RegisterCallback(gossip.KeyConfigZone, s.configGossipUpdate)
+		// Callback triggers on capacity gossip from all stores.
+		capacityRegex := fmt.Sprintf("%s.*", gossip.KeyMaxAvailCapacityPrefix)
+		s.gossip.RegisterCallback(capacityRegex, s.capacityGossipUpdate)
 	}
 
 	return nil

storage/store_finder.go

@@ -0,0 +1,88 @@
+// Copyright 2014 The Cockroach Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+// implied. See the License for the specific language governing
+// permissions and limitations under the License. See the AUTHORS file
+// for names of contributors.
+//
+// Author: Kathy Spradlin ([email protected])
+
+package storage
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/cockroachdb/cockroach/gossip"
+	"github.com/cockroachdb/cockroach/proto"
+)
+
+// StoreFinder finds the disks in a datacenter that have the requested
+// attributes.
+type StoreFinder func(proto.Attributes) ([]*StoreDescriptor, error)
+
+type stringSet map[string]struct{}
+
+type storeFinderContext struct {
+	sync.Mutex
+	capacityKeys stringSet
+}
+
+// capacityGossipUpdate is a gossip callback triggered whenever capacity
+// information is gossiped. It just tracks keys used for capacity gossip.
+func (s *Store) capacityGossipUpdate(key string, contentsChanged bool) {
+	s.finderContext.Lock()
+	defer s.finderContext.Unlock()
+
+	if s.finderContext.capacityKeys == nil {
+		s.finderContext.capacityKeys = stringSet{}
+	}
+	s.finderContext.capacityKeys[key] = struct{}{}
+}
+
+// findStores is the Store's implementation of a StoreFinder. It returns a list
+// of stores with attributes that are a superset of the required attributes. It
+// never returns an error.
+//
+// If it cannot retrieve a StoreDescriptor from the Store's gossip, it garbage
+// collects the failed key.
+func (s *Store) findStores(required proto.Attributes) ([]*StoreDescriptor, error) {
+	s.finderContext.Lock()
+	defer s.finderContext.Unlock()
+	var stores []*StoreDescriptor
+	for key := range s.finderContext.capacityKeys {
+		storeDesc, err := storeDescFromGossip(key, s.gossip)
+		if err != nil {
+			// We can no longer retrieve this key from the gossip store,
+			// perhaps it expired.
+			delete(s.finderContext.capacityKeys, key)
+		} else if required.IsSubset(storeDesc.Attrs) {
+			stores = append(stores, storeDesc)
+		}
+	}
+	return stores, nil
+}
+
+// storeDescFromGossip retrieves a StoreDescriptor from the specified capacity
+// gossip key. Returns an error if the gossip doesn't exist or is not
+// a StoreDescriptor.
+func storeDescFromGossip(key string, g *gossip.Gossip) (*StoreDescriptor, error) {
+	info, err := g.GetInfo(key)
+
+	if err != nil {
+		return nil, err
+	}
+	storeDesc, ok := info.(StoreDescriptor)
+	if !ok {
+		return nil, fmt.Errorf("gossiped info is not a StoreDescriptor: %+v", info)
+	}
+	return &storeDesc, nil
+}

storage/store_finder_test.go

@@ -0,0 +1,121 @@
+// Copyright 2014 The Cockroach Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+// implied. See the License for the specific language governing
+// permissions and limitations under the License. See the AUTHORS file
+// for names of contributors.
+//
+// Author: Kathy Spradlin ([email protected])
+
+package storage
+
+import (
+	"reflect"
+	"sort"
+	"testing"
+	"time"
+
+	"github.com/cockroachdb/cockroach/proto"
+)
+
+func TestCapacityGossipUpdate(t *testing.T) {
+	s, _ := createTestStore(t)
+	defer s.Stop()
+	key := "testkey"
+
+	// Order and value of contentsChanged shouldn't matter.
+	s.capacityGossipUpdate(key, true)
+	s.capacityGossipUpdate(key, false)
+
+	expectedKeys := stringSet{key: struct{}{}}
+	s.finderContext.Lock()
+	actualKeys := s.finderContext.capacityKeys
+	s.finderContext.Unlock()
+
+	if !reflect.DeepEqual(expectedKeys, actualKeys) {
+		t.Errorf("expected to fetch %+v, instead %+v", expectedKeys, actualKeys)
+	}
+}
+
+func TestStoreFinder(t *testing.T) {
+	s, _ := createTestStore(t)
+	defer s.Stop()
+	required := []string{"ssd", "dc"}
+	// Nothing yet.
+	if stores, _ := s.findStores(proto.Attributes{Attrs: required}); stores != nil {
+		t.Errorf("expected no stores, instead %+v", stores)
+	}
+
+	matchingStore := StoreDescriptor{
+		Attrs: proto.Attributes{Attrs: required},
+	}
+	supersetStore := StoreDescriptor{
+		Attrs: proto.Attributes{Attrs: append(required, "db")},
+	}
+	unmatchingStore := StoreDescriptor{
+		Attrs: proto.Attributes{Attrs: []string{"ssd", "otherdc"}},
+	}
+	emptyStore := StoreDescriptor{Attrs: proto.Attributes{}}
+
+	// Explicitly add keys rather than registering a gossip callback to avoid
+	// waiting for the goroutine callback to finish.
+	s.finderContext.capacityKeys = stringSet{
+		"k1": struct{}{},
+		"k2": struct{}{},
+		"k3": struct{}{},
+		"k4": struct{}{},
+	}
+	s.gossip.AddInfo("k1", matchingStore, time.Hour)
+	s.gossip.AddInfo("k2", supersetStore, time.Hour)
+	s.gossip.AddInfo("k3", unmatchingStore, time.Hour)
+	s.gossip.AddInfo("k4", emptyStore, time.Hour)
+
+	expected := []string{matchingStore.Attrs.SortedString(), supersetStore.Attrs.SortedString()}
+	stores, err := s.findStores(proto.Attributes{Attrs: required})
+	if err != nil {
+		t.Errorf("expected no err, got %s", err)
+	}
+	var actual []string
+	for _, store := range stores {
+		actual = append(actual, store.Attrs.SortedString())
+	}
+	sort.Strings(expected)
+	sort.Strings(actual)
+	if !reflect.DeepEqual(expected, actual) {
+		t.Errorf("expected %+v Attrs, instead %+v", expected, actual)
+	}
+}
+
+// TestStoreFinderGarbageCollection ensures removal of capacity gossip keys in
+// the map, if their gossip does not exist when we try to retrieve them.
+func TestStoreFinderGarbageCollection(t *testing.T) {
+	s, _ := createTestStore(t)
+	defer s.Stop()
+
+	s.finderContext.capacityKeys = stringSet{
+		"key0": struct{}{},
+		"key1": struct{}{},
+	}
+	required := []string{}
+
+	// No gossip added for either key, so they should be removed.
+	stores, err := s.findStores(proto.Attributes{Attrs: required})
+	if err != nil {
+		t.Errorf("unexpected error retrieving stores %s", err)
+	} else if len(stores) != 0 {
+		t.Errorf("expected no stores found, instead %+v", stores)
+	}
+
+	if len(s.finderContext.capacityKeys) != 0 {
+		t.Errorf("expected keys to be cleared, instead are %+v",
+			s.finderContext.capacityKeys)
+	}
+}