Stream Consumer Group - Shard retention mechanism (#125)

Author: TomerShor

go.sum

@@ -28,6 +28,7 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/rs/xid v1.1.0 h1:9Z322kTPrDR5GpxTH+1yl7As6tEHIH9aGsRccl20ELk=
 github.com/rs/xid v1.1.0/go.mod h1:+uKXf+4Djp6Md1KODXJxgGQPKngRmWyn10oCKFzNHOQ=
+github.com/stretchr/objx v0.1.0 h1:4G4v2dO3VZwixGIRoQ5Lfboy6nUhCyYzaqnIAPPhYs4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=

pkg/dataplane/streamconsumergroup/claim.go

@@ -43,7 +43,7 @@ func newClaim(member *member, shardID int) (*claim, error) {
 }
 
 func (c *claim) start() error {
-	c.logger.DebugWith("Starting claim")
+	c.logger.DebugWith("Starting claim", "shardID", c.shardID)
 
 	go func() {
 		err := c.fetchRecordBatches(c.stopRecordBatchFetchChan,

pkg/dataplane/streamconsumergroup/member.go

@@ -16,6 +16,8 @@ type member struct {
 	sequenceNumberHandler *sequenceNumberHandler
 	handler               Handler
 	session               Session
+	retainShards          bool
+	shardGroupToRetain    []int
 }
 
 func NewMember(streamConsumerGroupInterface StreamConsumerGroup, name string) (Member, error) {
@@ -41,21 +43,14 @@ func NewMember(streamConsumerGroupInterface StreamConsumerGroup, name string) (M
 		return nil, errors.Wrap(err, "Failed creating stream consumer group state handler")
 	}
 
-	err = newMember.stateHandler.start()
-	if err != nil {
-		return nil, errors.Wrap(err, "Failed starting stream consumer group state handler")
-	}
-
-	// create & start an location handler for the stream
+	// create & start a location handler for the stream
 	newMember.sequenceNumberHandler, err = newSequenceNumberHandler(&newMember)
 	if err != nil {
 		return nil, errors.Wrap(err, "Failed creating stream consumer group location handler")
 	}
 
-	// if there's no member name, just observe
-	err = newMember.sequenceNumberHandler.start()
-	if err != nil {
-		return nil, errors.Wrap(err, "Failed starting stream consumer group state handler")
+	if err := newMember.Start(); err != nil {
+		return nil, errors.Wrap(err, "Failed starting new member")
 	}
 
 	return &newMember, nil
@@ -100,3 +95,27 @@ func (m *member) Close() error {
 
 	return nil
 }
+
+func (m *member) Start() error {
+	if err := m.stateHandler.start(); err != nil {
+		return errors.Wrap(err, "Failed starting stream consumer group state handler")
+	}
+
+	if err := m.sequenceNumberHandler.start(); err != nil {
+		return errors.Wrap(err, "Failed starting stream consumer group state handler")
+	}
+
+	return nil
+}
+
+func (m *member) GetID() string {
+	return m.id
+}
+
+func (m *member) GetRetainShardFlag() bool {
+	return m.retainShards
+}
+
+func (m *member) GetShardsToRetain() []int {
+	return m.shardGroupToRetain
+}

pkg/dataplane/streamconsumergroup/session.go

@@ -25,7 +25,7 @@ func newSession(member *member,
 }
 
 func (s *session) start() error {
-	s.logger.DebugWith("Starting session")
+	s.logger.DebugWith("Starting session", "shards", s.state.Shards)
 
 	// for each shard we need handle, create a StreamConsumerGroupClaim object and start it
 	for _, shardID := range s.state.Shards {

pkg/dataplane/streamconsumergroup/statehandler.go

@@ -12,7 +12,10 @@ import (
 
 const stateContentsAttributeKey string = "state"
 
-var errNoFreeShardGroups = errors.New("No free shard groups")
+var (
+	errNoFreeShardGroups = errors.New("No free shard groups")
+	errShardRetention    = errors.New("Could not retain shard group")
+)
 
 type stateHandler struct {
 	logger       logger.Logger
@@ -33,8 +36,16 @@ func newStateHandler(member *member) (*stateHandler, error) {
 func (sh *stateHandler) start() error {
 
 	// stops on stop()
-	go sh.refreshStatePeriodically()
+	go func() {
+		if err := sh.refreshStatePeriodically(); err != nil {
+			if errors.RootCause(err) == errShardRetention {
 
+				// signal that the Handler needs to be restarted
+				sh.logger.ErrorWith("Aborting member", "memberID", sh.member.id)
+				sh.member.handler.Abort(sh.member.session) // nolint: errcheck
+			}
+		}
+	}()
 	return nil
 }
 
@@ -77,7 +88,7 @@ func (sh *stateHandler) getSessionState(state *State, memberID string) (*Session
 	return nil, errors.Errorf("Member state not found: %s", memberID)
 }
 
-func (sh *stateHandler) refreshStatePeriodically() {
+func (sh *stateHandler) refreshStatePeriodically() error {
 	var err error
 
 	// guaranteed to only be REPLACED by a new instance - not edited. as such, once this is initialized
@@ -94,6 +105,13 @@ func (sh *stateHandler) refreshStatePeriodically() {
 			} else {
 				lastState, err = sh.refreshState()
 				if err != nil {
+
+					// in case of shard retention error we want to signal the member to restart
+					if errors.RootCause(err) == errShardRetention {
+						sh.logger.WarnWith("Failed getting state on shard retention (requested by member)",
+							"err", errors.GetErrorStackString(err, 10))
+						return errors.Wrap(err, "Failed refreshing state by demand")
+					}
 					sh.logger.WarnWith("Failed getting state", "err", errors.GetErrorStackString(err, 10))
 				}
 
@@ -105,14 +123,21 @@ func (sh *stateHandler) refreshStatePeriodically() {
 		case <-time.After(sh.member.streamConsumerGroup.config.Session.HeartbeatInterval):
 			lastState, err = sh.refreshState()
 			if err != nil {
+
+				// in case of shard retention error we want to signal the member to restart
+				if errors.RootCause(err) == errShardRetention {
+					sh.logger.WarnWith("Failed getting state on shard retention (periodic refresh)",
+						"err", errors.GetErrorStackString(err, 10))
+					return errors.Wrap(err, "Failed refreshing state periodically")
+				}
 				sh.logger.WarnWith("Failed refreshing state", "err", errors.GetErrorStackString(err, 10))
 				continue
 			}
 
 		// if we're told to stop, exit the loop
 		case <-sh.stopChan:
 			sh.logger.Debug("Stopping")
-			return
+			return nil
 		}
 	}
 }
@@ -142,6 +167,14 @@ func (sh *stateHandler) refreshState() (*State, error) {
 		}
 
 		return state, nil
+
+	}, func() error {
+
+		// set retainShards flag to true only after the new state has been saved in persistency
+		// (meaning the shards have been assigned successfully)
+		sh.member.retainShards = true
+
+		return nil
 	})
 }
 
@@ -150,16 +183,42 @@ func (sh *stateHandler) createSessionState(state *State) error {
 		state.SessionStates = []*SessionState{}
 	}
 
-	// assign shards
-	shards, err := sh.assignShards(sh.member.streamConsumerGroup.maxReplicas, sh.member.streamConsumerGroup.totalNumShards, state)
-	if err != nil {
-		return errors.Wrap(err, "Failed resolving shards for session")
+	var shards []int
+	var err error
+
+	if sh.member.retainShards {
+
+		// try to retain the originally assigned shard group
+		shards, err = sh.retainShards(sh.member.shardGroupToRetain, sh.member.id, state)
+
+		// shards were "stolen" - set retainShards flag to false and commit suicide
+		if err != nil {
+			sh.logger.ErrorWith("Failed to retain shards",
+				"memberID", sh.member.id,
+				"shardsToRetain", sh.member.shardGroupToRetain,
+				"state", state,
+				"error", err.Error())
+			sh.member.retainShards = false
+			return err
+		}
+	} else {
+
+		// assign shards
+		shards, err = sh.assignShards(sh.member.streamConsumerGroup.maxReplicas,
+			sh.member.streamConsumerGroup.totalNumShards,
+			state)
+		if err != nil {
+			return errors.Wrap(err, "Failed resolving shards for session")
+		}
 	}
 
 	sh.logger.DebugWith("Assigned shards",
 		"shards", shards,
 		"state", state)
 
+	// save shards to retain on the member itself
+	sh.member.shardGroupToRetain = shards
+
 	state.SessionStates = append(state.SessionStates, &SessionState{
 		MemberID:      sh.member.id,
 		LastHeartbeat: time.Now(),
@@ -209,6 +268,23 @@ func (sh *stateHandler) assignShards(maxReplicas int, numShards int, state *Stat
 	return nil, errNoFreeShardGroups
 }
 
+func (sh *stateHandler) retainShards(memberShardGroup []int, memberID string, state *State) ([]int, error) {
+
+	for _, sessionState := range state.SessionStates {
+		if common.IntSlicesEqual(memberShardGroup, sessionState.Shards) {
+			if sessionState.MemberID == memberID {
+				return memberShardGroup, nil
+			}
+
+			// original shard group was taken
+			return nil, errShardRetention
+		}
+	}
+
+	// shard group to retain is not taken by any member - original member can retain it
+	return memberShardGroup, nil
+}
+
 func (sh *stateHandler) getReplicaShardGroups(maxReplicas int, numShards int) ([][]int, error) {
 	var replicaShardGroups [][]int
 	shards := common.MakeRange(0, numShards)

pkg/dataplane/streamconsumergroup/statehandler_test.go

@@ -1,6 +1,7 @@
 package streamconsumergroup
 
 import (
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/suite"
@@ -77,6 +78,58 @@ func (suite *stateHandlerSuite) TestAssignShards() {
 	}
 }
 
+func (suite *stateHandlerSuite) TestRetainShards() {
+	for _, testCase := range []struct {
+		name                string
+		memberID            string
+		existingShardGroups [][]int
+		expectedShardGroup  []int
+		expectedError       bool
+	}{
+		{
+			name:                "successfulRetention",
+			memberID:            "1",
+			existingShardGroups: [][]int{{0, 1}, {2, 3}},
+			expectedShardGroup:  []int{2, 3},
+			expectedError:       false,
+		},
+		{
+			name:                "failedRetention",
+			memberID:            "2",
+			existingShardGroups: [][]int{{0, 1}, {2, 3}, {4, 5}},
+			expectedShardGroup:  []int{0, 1},
+			expectedError:       true,
+		},
+		{
+			name:                "unexpectedBehaviour",
+			memberID:            "0",
+			existingShardGroups: [][]int{{0, 1}, {2, 3}},
+			expectedShardGroup:  []int{4, 5},
+			expectedError:       true,
+		},
+	} {
+		suite.Run(testCase.name, func() {
+
+			// make state from shard groups
+			state := State{}
+			for i, existingShardGroup := range testCase.existingShardGroups {
+				state.SessionStates = append(state.SessionStates, &SessionState{
+					Shards:   existingShardGroup,
+					MemberID: strconv.Itoa(i),
+				})
+			}
+
+			shards, err := suite.stateHandler.retainShards(testCase.expectedShardGroup, testCase.memberID, &state)
+			if testCase.expectedError {
+				suite.Require().Error(err)
+			} else {
+				suite.Require().NoError(err)
+				suite.Require().ElementsMatch(testCase.expectedShardGroup, shards)
+			}
+		})
+	}
+}
+
 func TestBinaryTestSuite(t *testing.T) {
 	suite.Run(t, new(stateHandlerSuite))
 }

pkg/dataplane/streamconsumergroup/streamconsumergroup.go

@@ -86,7 +86,8 @@ func (scg *streamConsumerGroup) getTotalNumberOfShards() (int, error) {
 	return response.Output.(*v3io.DescribeStreamOutput).ShardCount, nil
 }
 
-func (scg *streamConsumerGroup) setState(modifier stateModifier) (*State, error) {
+func (scg *streamConsumerGroup) setState(modifier stateModifier,
+	handlePostSetStateInPersistency postSetStateInPersistencyHandler) (*State, error) {
 	var previousState, modifiedState *State
 
 	backoff := scg.config.State.ModifyRetry.Backoff
@@ -110,6 +111,11 @@ func (scg *streamConsumerGroup) setState(modifier stateModifier) (*State, error)
 
 		modifiedState, err = modifier(state)
 		if err != nil {
+			if errors.RootCause(err) == errShardRetention {
+
+				// if shard retention failed the member needs to be aborted, so we can stop retrying
+				return false, errors.Wrap(err, "Failed modifying state")
+			}
 			return true, errors.Wrap(err, "Failed modifying state")
 		}
 
@@ -131,6 +137,10 @@ func (scg *streamConsumerGroup) setState(modifier stateModifier) (*State, error)
 			return true, errors.Wrap(err, "Failed setting state in persistency state")
 		}
 
+		if err := handlePostSetStateInPersistency(); err != nil {
+			return false, errors.Wrap(err, "Failed handling post set state in persistency")
+		}
+
 		return false, nil
 	})
 

pkg/dataplane/streamconsumergroup/types.go

@@ -8,6 +8,8 @@ import (
 
 type stateModifier func(*State) (*State, error)
 
+type postSetStateInPersistencyHandler func() error
+
 type SessionState struct {
 	MemberID      string    `json:"member_id"`
 	LastHeartbeat time.Time `json:"last_heartbeat_time"`
@@ -27,6 +29,9 @@ type Handler interface {
 	// Once the Messages() channel is closed, the Handler must finish its processing
 	// loop and exit.
 	ConsumeClaim(Session, Claim) error
+
+	// Abort signals the handler to start abort procedure
+	Abort(Session) error
 }
 
 type RecordBatch struct {
@@ -45,6 +50,10 @@ type StreamConsumerGroup interface {
 type Member interface {
 	Consume(Handler) error
 	Close() error
+	Start() error
+	GetID() string
+	GetRetainShardFlag() bool
+	GetShardsToRetain() []int
 }
 
 type Session interface {