Final touch

Author: ajanikow

pkg/deployment/agency/generator_collection_test.go

@@ -22,6 +22,7 @@ package agency
 
 type CollectionGeneratorInterface interface {
 	WithWriteConcern(wc int) CollectionGeneratorInterface
+	WithReplicationFactor(rf int) CollectionGeneratorInterface
 	WithShard() ShardGeneratorInterface
 	Add() DatabaseGeneratorInterface
 }
@@ -31,6 +32,7 @@ type collectionGenerator struct {
 	col string
 
 	wc     *int
+	rf     *int
 	shards map[int]shardGenerator
 }
 
@@ -55,3 +57,8 @@ func (c collectionGenerator) WithWriteConcern(wc int) CollectionGeneratorInterfa
 	c.wc = &wc
 	return c
 }
+
+func (c collectionGenerator) WithReplicationFactor(rf int) CollectionGeneratorInterface {
+	c.rf = &rf
+	return c
+}

pkg/deployment/agency/generator_database_test.go

@@ -94,9 +94,10 @@ func (d databaseGenerator) Add() StateGenerator {
 			}
 
 			planCol := StatePlanCollection{
-				Name:         util.NewString(col),
-				Shards:       planShards,
-				WriteConcern: colDet.wc,
+				Name:              util.NewString(col),
+				Shards:            planShards,
+				WriteConcern:      colDet.wc,
+				ReplicationFactor: colDet.rf,
 			}
 
 			plan[col] = planCol

pkg/deployment/agency/plan_collections.go

@@ -42,13 +42,42 @@ func (a StatePlanDBCollections) IsDBServerInCollections(name string) bool {
 	return false
 }
 
+func (a StatePlanDBCollections) CountShards() int {
+	count := 0
+
+	for _, d := range a {
+		count += len(d.Shards)
+	}
+
+	return count
+}
+
 type StatePlanCollection struct {
 	Name   *string `json:"name"`
 	Shards Shards  `json:"shards"`
 	// deprecated
 	// MinReplicationFactor is deprecated, but we have to support it for backward compatibility
 	MinReplicationFactor *int `json:"minReplicationFactor,omitempty"`
 	WriteConcern         *int `json:"writeConcern,omitempty"`
+	ReplicationFactor    *int `json:"replicationFactor,omitempty"`
+}
+
+func (a *StatePlanCollection) GetReplicationFactor(shard string) int {
+	if a == nil {
+		return 0
+	}
+
+	l := len(a.Shards[shard])
+
+	if z := a.ReplicationFactor; z == nil {
+		return l
+	} else {
+		if v := *z; v > l {
+			return v
+		} else {
+			return l
+		}
+	}
 }
 
 func (a *StatePlanCollection) GetWriteConcern(def int) int {
@@ -79,7 +108,11 @@ func (a StatePlanCollection) GetName(d string) string {
 	return *a.Name
 }
 
-func (a StatePlanCollection) IsDBServerInShards(name string) bool {
+func (a *StatePlanCollection) IsDBServerInShards(name string) bool {
+	if a == nil {
+		return false
+	}
+
 	for _, planShards := range a.Shards {
 		if planShards.Contains(name) {
 			return true

pkg/deployment/agency/shards.go

@@ -33,3 +33,15 @@ func (s ShardServers) Contains(server string) bool {
 
 	return false
 }
+
+func (s ShardServers) FilterBy(b ShardServers) ShardServers {
+	q := make(ShardServers, 0, len(s))
+
+	for _, i := range s {
+		if b.Contains(i) {
+			q = append(q, i)
+		}
+	}
+
+	return q
+}

pkg/deployment/agency/state.go

@@ -120,75 +120,151 @@ func (d *StateExists) UnmarshalJSON(bytes []byte) error {
 	return nil
 }
 
-type CollectionIterator func(db, col string, planCollection *StatePlanCollection, shard string, plan ShardServers, current ShardServers) (skipRest bool)
+func (s State) CountShards() int {
+	count := 0
+
+	for _, collections := range s.Plan.Collections {
+		count += collections.CountShards()
+	}
+
+	return count
+}
+
+func (s State) PlanServers() []string {
+	q := map[string]bool{}
+
+	for _, db := range s.Plan.Collections {
+		for _, col := range db {
+			for _, shards := range col.Shards {
+				for _, shard := range shards {
+					q[shard] = true
+				}
+			}
+		}
+	}
+
+	r := make([]string, 0, len(q))
+
+	for k := range q {
+		r = append(r, k)
+	}
+
+	return r
+}
+
+type CollectionShardDetails []CollectionShardDetail
+
+type CollectionShardDetail struct {
+	Database   string
+	Collection string
+	Shard      string
+}
+
+type StateShardFilter func(s State, db, col, shard string) bool
+
+func NegateFilter(in StateShardFilter) StateShardFilter {
+	return func(s State, db, col, shard string) bool {
+		return !in(s, db, col, shard)
+	}
+}
+
+func (s State) Filter(f StateShardFilter) CollectionShardDetails {
+	shards := make(CollectionShardDetails, s.CountShards())
+	size := 0
 
-func (s State) IterateOverCollections(i CollectionIterator) {
-dbsLoop:
 	for db, collections := range s.Plan.Collections {
 		for collection, details := range collections {
-			for shard, shardDetails := range details.Shards {
-				if currShard, ok := s.Current.Collections[db][collection][shard]; ok {
-					if skipRest := i(db, collection, &details, shard, shardDetails, currShard.Servers); skipRest {
-						break dbsLoop
+			for shard := range details.Shards {
+				if f(s, db, collection, shard) {
+					shards[size] = CollectionShardDetail{
+						Database:   db,
+						Collection: collection,
+						Shard:      shard,
 					}
+					size++
 				}
 			}
 		}
 	}
-}
 
-func (s State) IsDBServerInSync(serverID string) bool {
-	isInSync := true
-	s.IterateOverCollections(func(db, col string, planCollection *StatePlanCollection, shard string, plan ShardServers, current ShardServers) bool {
-		if !plan.Contains(serverID) {
-			return false
-		}
+	if size == 0 {
+		return nil
+	}
 
-		serverIsNotInSync := !current.Contains(serverID)
-		wc := planCollection.GetWriteConcern(1)
-		if serverIsNotInSync || wc > len(current) {
-			isInSync = false
-			return true
-		}
-		return false
-	})
+	return shards[0:size]
+}
 
-	return isInSync
+func GetDBServerBlockingRestartShards(s State, serverID string) CollectionShardDetails {
+	return s.Filter(FilterDBServerShardRestart(serverID))
 }
 
-func (s State) IsDBServerReadyToRestart(serverID string) bool {
-	readyToRestart := true
+func FilterDBServerShardRestart(serverID string) StateShardFilter {
+	return NegateFilter(func(s State, db, col, shard string) bool {
+		// Filter all shards which are not blocking restart of server
+		plan := s.Plan.Collections[db][col]
+		planShard := plan.Shards[shard]
 
-	s.IterateOverCollections(func(db, col string, planCollection *StatePlanCollection, shard string, plan ShardServers, current ShardServers) bool {
-		if !plan.Contains(serverID) {
-			return false
+		if !planShard.Contains(serverID) {
+			// This DBServer is not even in plan, restart possible
+			return true
 		}
 
-		serverInSync := current.Contains(serverID)
-		if len(plan) == 1 && serverInSync {
-			// The requested server is the only one in the plan
-			return false
+		current := s.Current.Collections[db][col][shard]
+		currentShard := current.Servers.FilterBy(planShard)
+
+		serverInSync := currentShard.Contains(serverID)
+
+		if len(planShard) == 1 && serverInSync {
+			// The requested server is the only one in the plan, restart possible
+			return true
 		}
 
-		wc := planCollection.GetWriteConcern(1)
-		if wc >= len(plan) && len(plan) != 0 {
-			wc = len(plan) - 1
+		// If WriteConcern equals replicationFactor then downtime is always there
+		wc := plan.GetWriteConcern(1)
+		if rf := plan.GetReplicationFactor(shard); wc >= rf {
+			wc = rf - 1
 		}
 
-		if len(current) >= wc && !serverInSync {
+		if len(currentShard) >= wc && !serverInSync {
 			// Current shard is not in sync, but it does not matter - we have enough replicas in sync
 			// Restart of this DBServer won't affect WC
-			return false
+			return true
 		}
 
-		if len(current) <= wc {
-			// If we restart this server, write concern won't be satisfied
-			readyToRestart = false
+		if len(currentShard) > wc {
+			// We are in plan, but restart is possible
 			return true
 		}
 
+		// If we restart this server, write concern won't be satisfied
 		return false
 	})
+}
 
-	return readyToRestart
+func GetDBServerShardsNotInSync(s State, serverID string) CollectionShardDetails {
+	return s.Filter(FilterDBServerShardsNotInSync(serverID))
+}
+
+func FilterDBServerShardsNotInSync(serverID string) StateShardFilter {
+	return NegateFilter(func(s State, db, col, shard string) bool {
+		planShard := s.Plan.Collections[db][col].Shards[shard]
+
+		if serverID != "*" && !planShard.Contains(serverID) {
+			return true
+		}
+
+		currentShard := s.Current.Collections[db][col][shard]
+
+		if len(planShard) != len(currentShard.Servers) {
+			return false
+		}
+
+		for _, s := range planShard {
+			if !currentShard.Servers.Contains(s) {
+				return false
+			}
+		}
+
+		return true
+	})
 }