From ba6920f67c6c4c8daa941f582ea192e0a45481de Mon Sep 17 00:00:00 2001 From: Mikhail Scherba Date: Thu, 12 Dec 2024 11:11:32 +0300 Subject: [PATCH] customize bfs Signed-off-by: Mikhail Scherba --- pkg/module_manager/scheduler/node/node.go | 5 +- pkg/module_manager/scheduler/scheduler.go | 81 ++++++++++++------- .../scheduler/scheduler_test.go | 4 + 3 files changed, 59 insertions(+), 31 deletions(-) diff --git a/pkg/module_manager/scheduler/node/node.go b/pkg/module_manager/scheduler/node/node.go index 9aefe86e..70bffb30 100644 --- a/pkg/module_manager/scheduler/node/node.go +++ b/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 { diff --git a/pkg/module_manager/scheduler/scheduler.go b/pkg/module_manager/scheduler/scheduler.go index 55edf890..d988790c 100644 --- a/pkg/module_manager/scheduler/scheduler.go +++ b/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 } diff --git a/pkg/module_manager/scheduler/scheduler_test.go b/pkg/module_manager/scheduler/scheduler_test.go index df55a893..b9f1e295 100644 --- a/pkg/module_manager/scheduler/scheduler_test.go +++ b/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 {