refactor: a new graph client to ease the DAG manipulation (#3999)

1. a new graph client to ease DAG manipulation
2. package model coverage from 9.7% to 86.4%

internal/controller/model/graph_client.go

@@ -0,0 +1,138 @@
+/*
+Copyright (C) 2022-2023 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package model
+
+import (
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"github.com/apecloud/kubeblocks/internal/controller/graph"
+)
+
+type GraphWriter interface {
+	// Root setups the given obj as root vertex of the underlying DAG.
+	// this func should be called once before any others.
+	Root(dag *graph.DAG, objOld, objNew client.Object)
+
+	// Create saves the object obj in the underlying DAG.
+	Create(dag *graph.DAG, obj client.Object)
+
+	// Delete deletes the given obj from the underlying DAG.
+	Delete(dag *graph.DAG, obj client.Object)
+
+	// Update updates the given obj in the underlying DAG.
+	Update(dag *graph.DAG, objOld, objNew client.Object)
+
+	// Status updates the given obj's status in the underlying DAG.
+	Status(dag *graph.DAG, objOld, objNew client.Object)
+
+	// DependOn setups dependencies between 'object' and 'dependency',
+	// which will guarantee the Write Order to the K8s cluster of these objects.
+	DependOn(dag *graph.DAG, object client.Object, dependency ...client.Object)
+}
+
+type GraphClient interface {
+	client.Reader
+	GraphWriter
+}
+
+// TODO(free6om): make DAG a member of realGraphClient
+type realGraphClient struct {
+	client.Client
+}
+
+func (r *realGraphClient) Root(dag *graph.DAG, objOld, objNew client.Object) {
+	vertex := &ObjectVertex{
+		Obj:    objNew,
+		OriObj: objOld,
+		Action: ActionPtr(STATUS),
+	}
+	dag.AddVertex(vertex)
+}
+
+func (r *realGraphClient) Create(dag *graph.DAG, obj client.Object) {
+	r.doWrite(dag, nil, obj, ActionPtr(CREATE))
+}
+
+func (r *realGraphClient) Update(dag *graph.DAG, objOld, objNew client.Object) {
+	r.doWrite(dag, objOld, objNew, ActionPtr(UPDATE))
+}
+
+func (r *realGraphClient) Delete(dag *graph.DAG, obj client.Object) {
+	r.doWrite(dag, nil, obj, ActionPtr(DELETE))
+}
+
+func (r *realGraphClient) Status(dag *graph.DAG, objOld, objNew client.Object) {
+	r.doWrite(dag, objOld, objNew, ActionPtr(STATUS))
+}
+
+func (r *realGraphClient) DependOn(dag *graph.DAG, object client.Object, dependency ...client.Object) {
+	objectVertex := r.findMatchedVertex(dag, object)
+	if objectVertex == nil {
+		return
+	}
+	for _, d := range dependency {
+		v := r.findMatchedVertex(dag, d)
+		if v != nil {
+			dag.Connect(objectVertex, v)
+		}
+	}
+}
+
+func (r *realGraphClient) doWrite(dag *graph.DAG, objOld, objNew client.Object, action *Action) {
+	vertex := r.findMatchedVertex(dag, objNew)
+	switch {
+	case vertex != nil:
+		objVertex, _ := vertex.(*ObjectVertex)
+		objVertex.Action = action
+	default:
+		vertex = &ObjectVertex{
+			Obj:    objNew,
+			OriObj: objOld,
+			Action: action,
+		}
+		dag.AddConnectRoot(vertex)
+	}
+}
+
+func (r *realGraphClient) findMatchedVertex(dag *graph.DAG, object client.Object) graph.Vertex {
+	keyLookfor, err := GetGVKName(object)
+	if err != nil {
+		return nil
+	}
+	for _, vertex := range dag.Vertices() {
+		v, _ := vertex.(*ObjectVertex)
+		key, err := GetGVKName(v.Obj)
+		if err != nil {
+			return nil
+		}
+		if *keyLookfor == *key {
+			return vertex
+		}
+	}
+	return nil
+}
+
+var _ GraphClient = &realGraphClient{}
+
+func NewGraphClient(cli client.Client) GraphClient {
+	return &realGraphClient{
+		Client: cli,
+	}
+}

internal/controller/model/graph_client_test.go

@@ -0,0 +1,76 @@
+/*
+Copyright (C) 2022-2023 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package model
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/apecloud/kubeblocks/internal/controller/builder"
+	"github.com/apecloud/kubeblocks/internal/controller/graph"
+)
+
+var _ = Describe("graph client test.", func() {
+	Context("GraphWriter", func() {
+		It("should work well", func() {
+			graphCli := NewGraphClient(nil)
+			dag := graph.NewDAG()
+
+			By("create without root vertex")
+			namespace := "foo"
+			name := "bar"
+			root := builder.NewStatefulSetBuilder(namespace, name).GetObject()
+			graphCli.Create(dag, root)
+			dagExpected := graph.NewDAG()
+			Expect(dag.Equals(dagExpected, DefaultLess)).Should(BeTrue())
+
+			By("init root vertex")
+			graphCli.Root(dag, root, root.DeepCopy())
+			dagExpected.AddVertex(&ObjectVertex{Obj: root, OriObj: root, Action: ActionPtr(STATUS)})
+			Expect(dag.Equals(dagExpected, DefaultLess)).Should(BeTrue())
+
+			By("create object")
+			obj0 := builder.NewPodBuilder(namespace, name+"0").GetObject()
+			obj1 := builder.NewPodBuilder(namespace, name+"1").GetObject()
+			obj2 := builder.NewPodBuilder(namespace, name+"2").GetObject()
+			graphCli.Create(dag, obj0)
+			graphCli.Create(dag, obj1)
+			graphCli.Create(dag, obj2)
+			graphCli.DependOn(dag, obj1, obj2)
+			v0 := &ObjectVertex{Obj: obj0, Action: ActionPtr(CREATE)}
+			v1 := &ObjectVertex{Obj: obj1, Action: ActionPtr(CREATE)}
+			v2 := &ObjectVertex{Obj: obj2, Action: ActionPtr(CREATE)}
+			dagExpected.AddConnectRoot(v0)
+			dagExpected.AddConnectRoot(v1)
+			dagExpected.AddConnectRoot(v2)
+			dagExpected.Connect(v1, v2)
+			Expect(dag.Equals(dagExpected, DefaultLess)).Should(BeTrue())
+
+			By("update&delete&status object")
+			graphCli.Status(dag, obj0, obj0.DeepCopy())
+			graphCli.Update(dag, obj1, obj1.DeepCopy())
+			graphCli.Delete(dag, obj2)
+			v0.Action = ActionPtr(STATUS)
+			v1.Action = ActionPtr(UPDATE)
+			v2.Action = ActionPtr(DELETE)
+			Expect(dag.Equals(dagExpected, DefaultLess)).Should(BeTrue())
+		})
+	})
+})

internal/controller/model/parallel_transformer_test.go

@@ -0,0 +1,60 @@
+/*
+Copyright (C) 2022-2023 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package model
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/apecloud/kubeblocks/internal/controller/graph"
+)
+
+var _ = Describe("parallel transformer test", func() {
+	Context("Transform function", func() {
+		It("should work well", func() {
+			id1, id2 := 1, 2
+			transformer := &ParallelTransformer{
+				Transformers: []graph.Transformer{
+					&testTransformer{id: id1},
+					&testTransformer{id: id2},
+				},
+			}
+			dag := graph.NewDAG()
+			Expect(transformer.Transform(nil, dag)).Should(Succeed())
+			dag.Connect(id1, id2)
+			dagExpected := graph.NewDAG()
+			dagExpected.AddVertex(id1)
+			dagExpected.AddVertex(id2)
+			dagExpected.Connect(id1, id2)
+			Expect(dag.Equals(dagExpected, DefaultLess)).Should(BeTrue())
+		})
+	})
+})
+
+type testTransformer struct {
+	id int
+}
+
+var _ graph.Transformer = &testTransformer{}
+
+func (t *testTransformer) Transform(ctx graph.TransformContext, dag *graph.DAG) error {
+	dag.AddVertex(t.id)
+	return nil
+}

internal/controller/model/transform_types.go

@@ -49,9 +49,9 @@ const (
 	STATUS = Action("STATUS")
 )
 
-type GVKName struct {
-	gvk      schema.GroupVersionKind
-	ns, name string
+type GVKNObjKey struct {
+	schema.GroupVersionKind
+	client.ObjectKey
 }
 
 // ObjectVertex describes expected object spec and how to reach it
@@ -70,7 +70,7 @@ type ObjectVertex struct {
 	Action    *Action
 }
 
-func (v ObjectVertex) String() string {
+func (v *ObjectVertex) String() string {
 	if v.Action == nil {
 		return fmt.Sprintf("{obj:%T, name: %s, immutable: %v, orphan: %v, action: nil}",
 			v.Obj, v.Obj.GetName(), v.Immutable, v.IsOrphan)
@@ -79,7 +79,7 @@ func (v ObjectVertex) String() string {
 		v.Obj, v.Obj.GetName(), v.Immutable, v.IsOrphan, *v.Action)
 }
 
-type ObjectSnapshot map[GVKName]client.Object
+type ObjectSnapshot map[GVKNObjKey]client.Object
 
 type RequeueError interface {
 	RequeueAfter() time.Duration

internal/controller/model/transform_types_test.go

@@ -0,0 +1,44 @@
+/*
+Copyright (C) 2022-2023 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package model
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/apecloud/kubeblocks/internal/controller/builder"
+)
+
+var _ = Describe("transform types test", func() {
+	const (
+		namespace = "foo"
+		name      = "bar"
+	)
+
+	Context("FindX function", func() {
+		It("should work well", func() {
+			root := builder.NewStatefulSetBuilder(namespace, name).GetObject()
+			vertex := &ObjectVertex{Obj: root}
+			Expect(vertex.String()).Should(Equal("{obj:*v1.StatefulSet, name: bar, immutable: false, orphan: false, action: nil}"))
+			vertex.Action = ActionPtr(CREATE)
+			Expect(vertex.String()).Should(Equal("{obj:*v1.StatefulSet, name: bar, immutable: false, orphan: false, action: CREATE}"))
+		})
+	})
+})