enforce maximum enum id in controller (#314)

* enforce maximum enum id in controller

* test coverage

* fix healthtracking_dynamic_test

cluster/topology/healthtracking_dynamic.go

@@ -58,8 +58,8 @@ func NewHealthTrackingDynamicTopology(opts DynamicOptions) (HealthTrackingDynami
 }
 
 func (ht *healthTrackingDynamicTopoImpl) Get() Map {
-	ht.RLock()
-	defer ht.RUnlock()
+	ht.Lock()
+	defer ht.Unlock()
 
 	// filter known unhealthy hosts from dynamic topology
 	dm := ht.dynamicTopology.Get()

controller/mutators/etcd/enum_mutator.go

@@ -15,13 +15,13 @@ package etcd
 
 import (
 	"github.com/uber/aresdb/cluster/kvstore"
+	"github.com/uber/aresdb/controller/mutators/common"
 	metaCom "github.com/uber/aresdb/metastore/common"
 	"path"
 	"strconv"
 	"sync"
 
 	"github.com/uber/aresdb/controller/generated/proto"
-	"github.com/uber/aresdb/controller/mutators/common"
 
 	"github.com/m3db/m3/src/cluster/kv"
 	"github.com/uber/aresdb/utils"
@@ -60,9 +60,11 @@ func (e *enumMutator) ExtendEnumCases(namespace, tableName, columnName string, e
 		return nil, err
 	}
 	columnID := -1
+	enumIDUpperBound := 0
 	for id, column := range schema.Columns {
 		if column.Name == columnName && !column.Deleted && column.IsEnumBasedColumn() {
 			columnID = id
+			enumIDUpperBound = metaCom.EnumCardinality(column.Type)
 			break
 		}
 	}
@@ -75,7 +77,7 @@ func (e *enumMutator) ExtendEnumCases(namespace, tableName, columnName string, e
 	if !exist {
 		e.RUnlock()
 		// fetch enum case from etcd and update cache
-		return e.extendEnumCase(namespace, tableName, schema.Incarnation, columnID, 0, enumCases)
+		return e.extendEnumCase(namespace, tableName, schema.Incarnation, columnID, 0, enumCases, enumIDUpperBound)
 	}
 
 	currentNodeID := enumCache.currentNodeID
@@ -95,7 +97,7 @@ func (e *enumMutator) ExtendEnumCases(namespace, tableName, columnName string, e
 
 	if len(newEnumCases) > 0 {
 		// fetch enum cases from etcd and update cache
-		missingIDs, err := e.extendEnumCase(namespace, tableName, schema.Incarnation, columnID, currentNodeID, newEnumCases)
+		missingIDs, err := e.extendEnumCase(namespace, tableName, schema.Incarnation, columnID, currentNodeID, newEnumCases, enumIDUpperBound)
 		if err != nil {
 			return nil, err
 		}
@@ -120,7 +122,7 @@ func getEnumID(nodeID int, innerID int) int {
 	return maxEnumCasePerNode*nodeID + innerID
 }
 
-func (e *enumMutator) extendEnumCase(namespace, tableName string, incarnation, columnID int, fromEnumNodeID int, newEnumCases []string) ([]int, error) {
+func (e *enumMutator) extendEnumCase(namespace, tableName string, incarnation, columnID int, fromEnumNodeID int, newEnumCases []string, enumIDUpperBound int) ([]int, error) {
 	// track result resolvedEnumIDs
 	resolvedEnumIDs := make([]int, len(newEnumCases))
 	// newEnumCaseDict records the resolved resolvedEnumIDs for newEnumCases
@@ -194,6 +196,10 @@ func (e *enumMutator) extendEnumCase(namespace, tableName string, incarnation, c
 				enumNodeList.NumEnumNodes++
 			}
 			enumID := getEnumID(lastEnumNodeID, len(lastEnumNode.Cases))
+			// check max enum id before creating new enum ids
+			if enumID >= enumIDUpperBound {
+				return nil, metaCom.ErrMaxEnumIDReached
+			}
 			lastEnumNode.Cases = append(lastEnumNode.Cases, newCase)
 			resolvedEnumIDs[index] = enumID
 			newEnumCaseDict[newCase] = enumID

controller/mutators/etcd/enum_mutator_test.go

@@ -34,9 +34,43 @@ func TestEnumMutator(t *testing.T) {
 				Type:    metaCom.BigEnum,
 				Deleted: false,
 			},
+			{
+				Name:    "c2",
+				Type:    metaCom.SmallEnum,
+				Deleted: false,
+			},
 		},
 	}
 
+	t.Run("Extend enum case reach maximum", func(t *testing.T) {
+		// test setup
+		txnStore := mem.NewStore()
+
+		_, err := txnStore.Set(utils.EnumNodeListKey("ns1", "test", 0, 1), &pb.EnumNodeList{
+			NumEnumNodes: 1,
+		})
+		assert.NoError(t, err)
+		_, err = txnStore.Set(utils.EnumNodeKey("ns1", "test", 0, 1, 0), &pb.EnumCases{
+			Cases: []string{},
+		})
+		assert.NoError(t, err)
+
+		schemaMutator := &mocks.TableSchemaMutator{}
+		// test
+		enumMutator := NewEnumMutator(txnStore, schemaMutator)
+		schemaMutator.On("GetTable", "ns1", "test").Return(&testTable, nil)
+
+		// test more than 1 node
+		newCases := make([]string, 0)
+		for i := 0; i < 257; i++ {
+			newCases = append(newCases, strconv.Itoa(i))
+		}
+
+		enumIDs, err := enumMutator.ExtendEnumCases("ns1", "test", "c2", newCases)
+		assert.EqualError(t, err, metaCom.ErrMaxEnumIDReached.Error())
+		assert.Empty(t, enumIDs)
+	})
+
 	t.Run("Extend and get enum cases", func(t *testing.T) {
 		// test setup
 		txnStore := mem.NewStore()

metastore/common/errors.go

@@ -81,4 +81,7 @@ var (
 	ErrInvalidPrimaryKeyBucketSize       = errors.New("Table primary key bucket size should be larger than zero")
 	ErrInvalidPrimaryKeyDataType         = errors.New("Specified data type can not be used as primary key")
 	ErrInvalidSortColumnDataType         = errors.New("Specified data type can not be used as sorting column")
+	// ErrMaxEnumIDReached indicates a column has already reached its maximum enum id
+	// eg. SmallEnum: 255, BigEnum: 65535
+	ErrMaxEnumIDReached = errors.New("Maximum enum id reached")
 )

metastore/common/model.go

@@ -176,6 +176,18 @@ func (c *Column) IsOverwriteOnlyDataType() bool {
 	}
 }
 
+// EnumCardinality returns cardinality for enum type
+func EnumCardinality(columnType string) int {
+	switch columnType {
+	case SmallEnum:
+		return 1 << 8
+	case BigEnum:
+		return 1 << 16
+	default:
+		return 0
+	}
+}
+
 // ShardOwnership defines an instruction on whether the receiving instance
 // should start to own or disown the specified table shard.
 type ShardOwnership struct {

metastore/common/model_test.go

@@ -31,4 +31,10 @@ var _ = ginkgo.Describe("model tests", func() {
 		Ω(c.IsEnumBasedColumn()).Should(BeTrue())
 		Ω(c.IsEnumArrayColumn()).Should(BeTrue())
 	})
+
+	ginkgo.It("EnumCardinality should work", func() {
+		Ω(EnumCardinality(SmallEnum)).Should(Equal(256))
+		Ω(EnumCardinality(BigEnum)).Should(Equal(65536))
+		Ω(EnumCardinality(UUID)).Should(Equal(0))
+	})
 })