customize bfs

Signed-off-by: Mikhail Scherba <[email protected]>

pkg/module_manager/scheduler/node/node.go

@@ -33,8 +33,9 @@ type Node struct {
 }
 
 const (
-	ModuleType NodeType = "module"
-	WeightType NodeType = "weight"
+	ModuleType    NodeType = "module"
+	WeightType    NodeType = "weight"
+	TypeAttribute string   = "type"
 )
 
 func NewNode() *Node {

pkg/module_manager/scheduler/scheduler.go

@@ -143,7 +143,7 @@ func (s *Scheduler) AddModuleVertex(module node.ModuleInterface) error {
 	defer s.l.Unlock()
 	vertex := node.NewNode().WithName(module.GetName()).WithWeight(module.GetOrder()).WithType(node.ModuleType).WithModule(module)
 	// add module vertex
-	if err := s.dag.AddVertex(vertex, graph.VertexAttribute("colorscheme", "greens3"), graph.VertexAttribute("style", "filled"), graph.VertexAttribute("color", "2"), graph.VertexAttribute("fillcolor", "1"), graph.VertexAttribute("type", string(node.ModuleType))); err != nil {
+	if err := s.dag.AddVertex(vertex, graph.VertexAttribute("colorscheme", "greens3"), graph.VertexAttribute("style", "filled"), graph.VertexAttribute("color", "2"), graph.VertexAttribute("fillcolor", "1"), graph.VertexAttribute(node.TypeAttribute, string(node.ModuleType))); err != nil {
 		return err
 	}
 
@@ -218,7 +218,7 @@ func (s *Scheduler) AddModuleVertex(module node.ModuleInterface) error {
 
 // addWeightVertex adds a new vertex of type Weight to the graph
 func (s *Scheduler) addWeightVertex(vertex *node.Node) error {
-	if err := s.dag.AddVertex(vertex, graph.VertexWeight(int(vertex.GetWeight())), graph.VertexAttribute("colorscheme", "blues3"), graph.VertexAttribute("style", "filled"), graph.VertexAttribute("color", "2"), graph.VertexAttribute("fillcolor", "1"), graph.VertexAttribute("type", string(node.WeightType))); err != nil {
+	if err := s.dag.AddVertex(vertex, graph.VertexWeight(int(vertex.GetWeight())), graph.VertexAttribute("colorscheme", "blues3"), graph.VertexAttribute("style", "filled"), graph.VertexAttribute("color", "2"), graph.VertexAttribute("fillcolor", "1"), graph.VertexAttribute(node.TypeAttribute, string(node.WeightType))); err != nil {
 		return err
 	}
 
@@ -240,7 +240,7 @@ func (s *Scheduler) addWeightVertex(vertex *node.Node) error {
 		}
 
 		// filter out module vertices
-		if props.Attributes["type"] == string(node.ModuleType) {
+		if props.Attributes[node.TypeAttribute] == string(node.ModuleType) {
 			return false
 		}
 
@@ -403,18 +403,19 @@ func (s *Scheduler) getModuleVertices() ([]*node.Node, error) {
 		bfsErr error
 	)
 
-	err := s.customBFS(s.root.GetName(), func(name string, _ int) bool {
+	err := s.customBFS(s.root.GetName(), func(name, _ string, _ int) (bool, int) {
+		var dummyDepth int
 		vertex, props, err := s.dag.VertexWithProperties(name)
 		if err != nil {
 			bfsErr = fmt.Errorf("couldn't get %s vertex from the graph: %w", name, err)
-			return true
+			return true, dummyDepth
 		}
 
-		if props.Attributes["type"] == string(node.ModuleType) {
+		if props.Attributes[node.TypeAttribute] == string(node.ModuleType) {
 			nodes = append(nodes, vertex)
 		}
 
-		return false
+		return false, dummyDepth
 	})
 
 	if bfsErr != nil {
@@ -425,7 +426,7 @@ func (s *Scheduler) getModuleVertices() ([]*node.Node, error) {
 }
 
 // customBFS reimplements BFS but allows visiting vertices multiple times and vertices at the same level are traversed in a sorted way
-func (s *Scheduler) customBFS(root string, visit func(node string, depth int) bool) error {
+func (s *Scheduler) customBFS(root string, visit func(node, prevNode string, prevDepth int) (bool, int)) error {
 	adjacencyMap, err := s.dag.AdjacencyMap()
 	if err != nil {
 		return fmt.Errorf("could not get adjacency map: %w", err)
@@ -435,37 +436,42 @@ func (s *Scheduler) customBFS(root string, visit func(node string, depth int) bo
 		return fmt.Errorf("could not find root vertex with hash %v", root)
 	}
 
-	queue := make([]string, 0)
-	queue = append(queue, root)
-	depth := 1
-	currLvl := 1
-	nextLvl := 0
+	// ad-hoc type to save additional meta about vertices
+	type vertex struct {
+		name       string
+		prevVertex string
+		prevDepth  int
+	}
+
+	queue := make([]vertex, 0)
+	queue = append(queue, vertex{name: root})
 
 	for len(queue) > 0 {
-		currLvl -= 1
-		currentHash := queue[0]
+		currentHash := queue[0].name
+		prevHash := queue[0].prevVertex
+		prevDepth := queue[0].prevDepth
+		depth := prevDepth + 1
 
 		queue = queue[1:]
 
 		// Stop traversing the graph if the visit function returns true.
-		if stop := visit(currentHash, depth); stop {
+		stop, updatedDepth := visit(currentHash, prevHash, prevDepth)
+		if stop {
 			break
 		}
 
+		if updatedDepth != 0 {
+			depth = updatedDepth
+		}
+
 		sliceOfAdjacencies := make([]string, 0)
 		for adjacency := range adjacencyMap[currentHash] {
-			nextLvl += 1
 			sliceOfAdjacencies = append(sliceOfAdjacencies, adjacency)
 		}
 		slices.Sort(sliceOfAdjacencies)
-		queue = append(queue, sliceOfAdjacencies...)
-
-		if currLvl == 0 {
-			currLvl = nextLvl
-			nextLvl = 0
-			depth += 1
+		for _, adj := range sliceOfAdjacencies {
+			queue = append(queue, vertex{name: adj, prevVertex: currentHash, prevDepth: depth})
 		}
-
 	}
 
 	return nil
@@ -561,21 +567,38 @@ func (s *Scheduler) getEnabledModuleNamesByOrder() ([][]string, error) {
 		bfsErr   error
 	)
 
-	if err := s.customBFS(s.root.GetName(), func(name string, depth int) bool {
+	if err := s.customBFS(s.root.GetName(), func(name, prevNodeName string, prevDepth int) (bool, int) {
 		vertex, props, err := s.dag.VertexWithProperties(name)
-
 		if err != nil {
 			bfsErr = fmt.Errorf("couldn't get %s vertex from the graph: %w", name, err)
-			return true
+			return true, 0
 		}
 
-		if props.Attributes["type"] == string(node.ModuleType) && vertex.GetState() {
+		var depth int
+		// the root node
+		if len(prevNodeName) == 0 {
+			depth = 1
+		} else {
+			_, prevProps, err := s.dag.VertexWithProperties(prevNodeName)
+			if err != nil {
+				bfsErr = fmt.Errorf("couldn't get %s vertex from the graph: %w", prevNodeName, err)
+				return true, 0
+			}
+
+			if prevProps.Attributes[node.TypeAttribute] == string(node.WeightType) {
+				depth = prevDepth
+			} else {
+				depth = prevDepth + 1
+			}
+		}
+
+		if props.Attributes[node.TypeAttribute] == string(node.ModuleType) && vertex.GetState() {
 			if depth > vertices[name] {
 				vertices[name] = depth
 			}
 		}
 
-		return false
+		return false, depth
 	}); err != nil {
 		return nil, err
 	}

pkg/module_manager/scheduler/scheduler_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/deckhouse/deckhouse/pkg/log"
 	"github.com/dominikbraun/graph"
+	"github.com/dominikbraun/graph/draw"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -234,6 +235,9 @@ operatorTrivyEnabled: true
 
 	_, _ = s.RecalculateGraph(logLabels)
 
+	file, _ := os.Create("./mygraph.gv")
+	_ = draw.DOT(s.dag, file)
+
 	// get all enabled modules by order case
 	enabledModules, err := s.getEnabledModuleNamesByOrder()
 	for _, v := range enabledModules {