Add interface for tokens with a SimpleListTokens type

Signed-off-by: Ganesh Vernekar <[email protected]>

Author: codesome

pkg/ring/lifecycler.go

@@ -2,8 +2,6 @@ package ring
 
 import (
 	"context"
-	"encoding/binary"
-	"errors"
 	"flag"
 	"fmt"
 	"io/ioutil"
@@ -133,7 +131,7 @@ type Lifecycler struct {
 	// back empty.  And it changes during lifecycle of ingester.
 	stateMtx sync.Mutex
 	state    IngesterState
-	tokens   []uint32
+	tokens   Tokens
 
 	// Controls the ready-reporting
 	readyLock sync.Mutex
@@ -207,7 +205,7 @@ func (i *Lifecycler) CheckReady(ctx context.Context) error {
 		return fmt.Errorf("error talking to consul: %s", err)
 	}
 
-	if len(i.getTokens()) == 0 {
+	if i.getTokens().Len() == 0 {
 		return fmt.Errorf("this ingester owns no tokens")
 	}
 
@@ -246,9 +244,14 @@ func (i *Lifecycler) ChangeState(ctx context.Context, state IngesterState) error
 	return <-err
 }
 
-func (i *Lifecycler) getTokens() []uint32 {
+var emptyTokens Tokens = &SimpleListTokens{}
+
+func (i *Lifecycler) getTokens() Tokens {
 	i.stateMtx.Lock()
 	defer i.stateMtx.Unlock()
+	if i.tokens == nil {
+		return emptyTokens
+	}
 	return i.tokens
 }
 
@@ -263,18 +266,14 @@ func (i *Lifecycler) flushTokensToFile() {
 		return
 	}
 
-	b := make([]byte, 4*len(i.tokens))
-	for idx, token := range i.tokens {
-		binary.BigEndian.PutUint32(b[idx*4:], token)
-	}
-
+	b := i.tokens.Marshal(nil)
 	if _, err = f.Write(b); err != nil {
 		level.Error(util.Logger).Log("msg", "error in writing token file", "err", err)
 		return
 	}
 }
 
-func (i *Lifecycler) getTokensFromFile() ([]uint32, error) {
+func (i *Lifecycler) getTokensFromFile() (Tokens, error) {
 	tokenFilePath := path.Join(i.cfg.TokensFileDir, tokensFileName)
 	b, err := ioutil.ReadFile(tokenFilePath)
 	if err != nil {
@@ -283,23 +282,12 @@ func (i *Lifecycler) getTokensFromFile() ([]uint32, error) {
 		}
 		return nil, err
 	}
-	if len(b) == 0 {
-		return nil, nil
-	} else if len(b)%4 != 0 {
-		return nil, errors.New("token data is not 4 byte aligned")
-	}
-
-	numTokens := len(b) >> 2
-	tokens := make([]uint32, 0, numTokens)
-	for i := 0; i < numTokens; i++ {
-		tokens = append(tokens, binary.BigEndian.Uint32(b[i<<2:]))
-	}
 
-	return tokens, nil
+	return UnmarshalTokens(b)
 }
 
-func (i *Lifecycler) setTokens(tokens []uint32) {
-	tokensOwned.WithLabelValues(i.RingName).Set(float64(len(tokens)))
+func (i *Lifecycler) setTokens(tokens Tokens) {
+	tokensOwned.WithLabelValues(i.RingName).Set(float64(tokens.Len()))
 
 	i.stateMtx.Lock()
 	i.tokens = tokens
@@ -313,7 +301,7 @@ func (i *Lifecycler) ClaimTokensFor(ctx context.Context, ingesterID string) erro
 	err := make(chan error)
 
 	i.actorChan <- func() {
-		var tokens []uint32
+		var tokens Tokens
 
 		claimTokens := func(in interface{}) (out interface{}, retry bool, err error) {
 			ringDesc, ok := in.(*Desc)
@@ -451,12 +439,12 @@ func (i *Lifecycler) initRing(ctx context.Context) error {
 			tokens, err := i.getTokensFromFile()
 			if err != nil {
 				level.Error(util.Logger).Log("msg", "error in getting tokens from file", "err", err)
-			} else if len(tokens) > 0 {
-				level.Info(util.Logger).Log("msg", "adding tokens from file", "num_tokens", len(tokens))
-				if len(tokens) == i.cfg.NumTokens {
+			} else if tokens != nil && tokens.Len() > 0 {
+				level.Info(util.Logger).Log("msg", "adding tokens from file", "num_tokens", tokens.Len())
+				if tokens.Len() == i.cfg.NumTokens {
 					i.setState(ACTIVE)
 				}
-				ringDesc.AddIngester(i.ID, i.Addr, tokens, i.GetState(), i.cfg.NormaliseTokens)
+				ringDesc.AddIngester(i.ID, i.Addr, tokens.Tokens(), i.GetState(), i.cfg.NormaliseTokens)
 				i.setTokens(tokens)
 				return ringDesc, true, nil
 			}
@@ -472,7 +460,7 @@ func (i *Lifecycler) initRing(ctx context.Context) error {
 		tokens, _ := ringDesc.TokensFor(i.ID)
 		i.setTokens(tokens)
 
-		level.Info(util.Logger).Log("msg", "existing entry found in ring", "state", i.GetState(), "tokens", len(tokens))
+		level.Info(util.Logger).Log("msg", "existing entry found in ring", "state", i.GetState(), "tokens", tokens.Len())
 		return ringDesc, true, nil
 	})
 }
@@ -489,17 +477,17 @@ func (i *Lifecycler) autoJoin(ctx context.Context) error {
 
 		// At this point, we should not have any tokens, and we should be in PENDING state.
 		myTokens, takenTokens := ringDesc.TokensFor(i.ID)
-		if len(myTokens) > 0 {
-			level.Error(util.Logger).Log("msg", "tokens already exist for this ingester - wasn't expecting any!", "num_tokens", len(myTokens))
+		if myTokens.Len() > 0 {
+			level.Error(util.Logger).Log("msg", "tokens already exist for this ingester - wasn't expecting any!", "num_tokens", myTokens.Len())
 		}
 
-		newTokens := GenerateTokens(i.cfg.NumTokens-len(myTokens), takenTokens)
+		newTokens := GenerateTokens(i.cfg.NumTokens-myTokens.Len(), takenTokens.Tokens())
 		i.setState(ACTIVE)
 		ringDesc.AddIngester(i.ID, i.Addr, newTokens, i.GetState(), i.cfg.NormaliseTokens)
 
-		tokens := append(myTokens, newTokens...)
-		sort.Sort(sortableUint32(tokens))
-		i.setTokens(tokens)
+		myTokens.Add(newTokens...)
+		sort.Sort(sortableUint32(myTokens.Tokens()))
+		i.setTokens(myTokens)
 
 		return ringDesc, true, nil
 	})
@@ -520,7 +508,7 @@ func (i *Lifecycler) updateConsul(ctx context.Context) error {
 		if !ok {
 			// consul must have restarted
 			level.Info(util.Logger).Log("msg", "found empty ring, inserting tokens")
-			ringDesc.AddIngester(i.ID, i.Addr, i.getTokens(), i.GetState(), i.cfg.NormaliseTokens)
+			ringDesc.AddIngester(i.ID, i.Addr, i.getTokens().Tokens(), i.GetState(), i.cfg.NormaliseTokens)
 		} else {
 			ingesterDesc.Timestamp = time.Now().Unix()
 			ingesterDesc.State = i.GetState()

pkg/ring/lifecycler_test.go

@@ -85,7 +85,7 @@ func TestRingNormaliseMigration(t *testing.T) {
 		return checkDenormalised(d, "ing1")
 	})
 
-	token := l1.tokens[0]
+	token := l1.tokens.Tokens()[0]
 
 	// Add a second ingester with normalised tokens.
 	var lifecyclerConfig2 = testLifecyclerConfig(t, ringConfig, "ing2")
@@ -139,7 +139,7 @@ func TestRingRestart(t *testing.T) {
 		return checkNormalised(d, "ing1")
 	})
 
-	token := l1.tokens[0]
+	token := l1.tokens.Tokens()[0]
 
 	// Add a second ingester with the same settings, so it will think it has restarted
 	l2, err := NewLifecycler(lifecyclerConfig1, &nopFlushTransferer{}, "ingester")
@@ -151,8 +151,8 @@ func TestRingRestart(t *testing.T) {
 		require.NoError(t, err)
 		l2Tokens := l2.getTokens()
 		return checkNormalised(d, "ing1") &&
-			len(l2Tokens) == 1 &&
-			l2Tokens[0] == token
+			l2Tokens.Len() == 1 &&
+			l2Tokens.Tokens()[0] == token
 	})
 }
 
@@ -197,13 +197,14 @@ func TestCheckReady(t *testing.T) {
 	flagext.DefaultValues(&ringConfig)
 	ringConfig.KVStore.Mock = &MockClient{}
 
+	tokens := SimpleListTokens([]uint32{1})
 	r, err := New(ringConfig, "ingester")
 	require.NoError(t, err)
 	defer r.Stop()
 	cfg := testLifecyclerConfig(t, ringConfig, "ring1")
 	cfg.MinReadyDuration = 1 * time.Nanosecond
 	l1, err := NewLifecycler(cfg, &nopFlushTransferer{}, "ingester")
-	l1.setTokens([]uint32{1})
+	l1.setTokens(&tokens)
 	require.NoError(t, err)
 
 	// Delete the ring key before checking ready

pkg/ring/model.go

@@ -75,14 +75,14 @@ func (d *Desc) RemoveIngester(id string) {
 
 // ClaimTokens transfers all the tokens from one ingester to another,
 // returning the claimed token.
-func (d *Desc) ClaimTokens(from, to string, normaliseTokens bool) []uint32 {
-	var result []uint32
+func (d *Desc) ClaimTokens(from, to string, normaliseTokens bool) Tokens {
+	result := &SimpleListTokens{}
 
 	if normaliseTokens {
 
 		// If the ingester we are claiming from is normalising, get its tokens then erase them from the ring.
 		if fromDesc, found := d.Ingesters[from]; found {
-			result = fromDesc.Tokens
+			result.Add(fromDesc.Tokens...)
 			fromDesc.Tokens = nil
 			d.Ingesters[from] = fromDesc
 		}
@@ -93,24 +93,24 @@ func (d *Desc) ClaimTokens(from, to string, normaliseTokens bool) []uint32 {
 		// When all ingesters are in normalised mode, d.Tokens is empty here
 		for i := 0; i < len(d.Tokens); {
 			if d.Tokens[i].Ingester == from {
-				result = append(result, d.Tokens[i].Token)
+				result.Add(d.Tokens[i].Token)
 				d.Tokens = d.Tokens[:i+copy(d.Tokens[i:], d.Tokens[i+1:])]
 				continue
 			}
 			i++
 		}
 
-		sort.Sort(uint32s(result))
+		sort.Sort(uint32s(result.Tokens()))
 		ing := d.Ingesters[to]
-		ing.Tokens = result
+		ing.Tokens = result.Tokens()
 		d.Ingesters[to] = ing
 
 	} else {
 
 		for i := 0; i < len(d.Tokens); i++ {
 			if d.Tokens[i].Ingester == from {
 				d.Tokens[i].Ingester = to
-				result = append(result, d.Tokens[i].Token)
+				result.Add(d.Tokens[i].Token)
 			}
 		}
 	}
@@ -148,12 +148,12 @@ func (d *Desc) Ready(heartbeatTimeout time.Duration) error {
 }
 
 // TokensFor partitions the tokens into those for the given ID, and those for others.
-func (d *Desc) TokensFor(id string) (tokens, other []uint32) {
-	var takenTokens, myTokens []uint32
+func (d *Desc) TokensFor(id string) (tokens, other Tokens) {
+	takenTokens, myTokens := &SimpleListTokens{}, &SimpleListTokens{}
 	for _, token := range migrateRing(d) {
-		takenTokens = append(takenTokens, token.Token)
+		takenTokens.Add(token.Token)
 		if token.Ingester == id {
-			myTokens = append(myTokens, token.Token)
+			myTokens.Add(token.Token)
 		}
 	}
 	return myTokens, takenTokens

pkg/ring/tokens.go

@@ -0,0 +1,100 @@
+package ring
+
+import (
+	"encoding/binary"
+	"errors"
+)
+
+// Tokens is the interface to store Token.
+type Tokens interface {
+	// Encoding return the type of the token implementation.
+	Encoding() TokenType
+	// Add adds a new token.
+	Add(...uint32)
+	// Tokens returns a slice of uint32 tokens.
+	Tokens() []uint32
+	// Len returns the number of tokens.
+	Len() int
+	// Marshal marshalls all the tokens into byte stream.
+	Marshal([]byte) []byte
+	// Unmarshal reads the tokens from the byte stream.
+	Unmarshal([]byte) error
+}
+
+// TokenType is the type identifier for the token.
+type TokenType byte
+
+const (
+	// SimpleList is the type for Type1Tokens.
+	SimpleList TokenType = 1
+)
+
+// SimpleListTokens is a simple list of tokens.
+type SimpleListTokens []uint32
+
+// Encoding implements the Tokens interface.
+func (slt SimpleListTokens) Encoding() TokenType {
+	return SimpleList
+}
+
+// Tokens implements the Tokens interface.
+func (slt SimpleListTokens) Tokens() []uint32 {
+	return slt
+}
+
+// Len implements the Tokens interface.
+func (slt SimpleListTokens) Len() int {
+	return len(slt)
+}
+
+// Add implements the Tokens interface.
+func (slt *SimpleListTokens) Add(tokens ...uint32) {
+	*slt = append(*slt, tokens...)
+}
+
+// Marshal implements the Tokens interface.
+func (slt SimpleListTokens) Marshal(b []byte) []byte {
+	if cap(b) < 1+(4*len(slt)) {
+		b = make([]byte, 1+(4*len(slt)))
+	}
+	b = b[:1+(4*len(slt))]
+	b[0] = byte(SimpleList)
+	for idx, token := range slt {
+		binary.BigEndian.PutUint32(b[1+(idx*4):], uint32(token))
+	}
+	return b
+}
+
+// Unmarshal implements the Tokens interface.
+func (slt *SimpleListTokens) Unmarshal(b []byte) error {
+	*slt = (*slt)[:0]
+	if len(b) == 0 {
+		return nil
+	} else if len(b)%4 != 0 {
+		return errors.New("token data is not 4 byte aligned")
+	}
+
+	numTokens := len(b) >> 2
+	for i := 0; i < numTokens; i++ {
+		*slt = append(*slt, binary.BigEndian.Uint32(b[i<<2:]))
+	}
+
+	return nil
+}
+
+// UnmarshalTokens converts the byte stream into Tokens.
+// The first byte should be the token type follwed by the
+// actual token data.
+func UnmarshalTokens(b []byte) (Tokens, error) {
+	if len(b) == 0 {
+		return nil, errors.New("empty bytes")
+	}
+	switch TokenType(b[0]) {
+	case SimpleList:
+		tokens := &SimpleListTokens{}
+		err := tokens.Unmarshal(b[1:])
+		return tokens, err
+	default:
+		return nil, errors.New("invalid token type")
+	}
+}