feat: add GetNodeMetricSpec in statesInformer interface

add func GetNodeMetricSpec, so we can get nodeMetricSpec conveniently by statesInformer interface Signed-off-by: j4ckstraw <[email protected]>
koordinator-sh · Nov 27, 2024 · 53221fd · 53221fd
1 parent ceff9cd
commit 53221fd
Show file tree

Hide file tree

Showing 5 changed files with 108 additions and 0 deletions.
diff --git a/pkg/koordlet/prediction/prediction_test.go b/pkg/koordlet/prediction/prediction_test.go
@@ -43,10 +43,15 @@ func (m *mockStatesInformer) HasSynced() bool {
 func (m *mockStatesInformer) GetNode() *v1.Node {
 	return m.node
 }
+
 func (m *mockStatesInformer) GetNodeSLO() *slov1alpha1.NodeSLO {
 	return nil
 }
 
+func (m *mockStatesInformer) GetNodeMetricSpec() *slov1alpha1.NodeMetricSpec {
+	return nil
+}
+
 func (m *mockStatesInformer) GetAllPods() []*statesinformer.PodMeta {
 	return m.pods
 }

diff --git a/pkg/koordlet/statesinformer/api.go b/pkg/koordlet/statesinformer/api.go
@@ -120,6 +120,7 @@ type StatesInformer interface {
 
 	GetNode() *corev1.Node
 	GetNodeSLO() *slov1alpha1.NodeSLO
+	GetNodeMetricSpec() *slov1alpha1.NodeMetricSpec
 
 	GetAllPods() []*PodMeta
 

diff --git a/pkg/koordlet/statesinformer/impl/states_informer.go b/pkg/koordlet/statesinformer/impl/states_informer.go
@@ -51,6 +51,7 @@ type StatesInformer interface {
 
 	GetNode() *corev1.Node
 	GetNodeSLO() *slov1alpha1.NodeSLO
+	GetNodeMetricSpec() *slov1alpha1.NodeMetricSpec
 
 	GetAllPods() []*statesinformer.PodMeta
 
@@ -101,6 +102,8 @@ type informerPlugin interface {
 	HasSynced() bool
 }
 
+var _ StatesInformer = &statesInformer{}
+
 // TODO merge all clients into one struct
 func NewStatesInformer(config *Config, kubeClient clientset.Interface, crdClient koordclientset.Interface, topologyClient topologyclientset.Interface,
 	metricsCache metriccache.MetricCache, nodeName string, schedulingClient schedv1alpha1.SchedulingV1alpha1Interface, predictorFactory prediction.PredictorFactory) StatesInformer {
@@ -225,6 +228,16 @@ func (s *statesInformer) GetNodeSLO() *slov1alpha1.NodeSLO {
 	return nodeSLOInformer.GetNodeSLO()
 }
 
+func (s *statesInformer) GetNodeMetricSpec() *slov1alpha1.NodeMetricSpec {
+	nodeMetricInformerIf := s.states.informerPlugins[nodeMetricInformerName]
+	nodeMetricInformer, ok := nodeMetricInformerIf.(*nodeMetricInformer)
+	if !ok {
+		klog.Errorf("node metric informer format error")
+		return nil
+	}
+	return nodeMetricInformer.getNodeMetricSpec()
+}
+
 func (s *statesInformer) GetNodeTopo() *topov1alpha1.NodeResourceTopology {
 	nodeTopoInformerIf := s.states.informerPlugins[nodeTopoInformerName]
 	nodeTopoInformer, ok := nodeTopoInformerIf.(*nodeTopoInformer)

diff --git a/pkg/koordlet/statesinformer/impl/states_informer_test.go b/pkg/koordlet/statesinformer/impl/states_informer_test.go
@@ -130,6 +130,81 @@ func Test_statesInformer_GetNodeSLO(t *testing.T) {
 	}
 }
 
+func Test_statesInformer_GetNodeMetricSpec(t *testing.T) {
+	type fields struct {
+		nodeMetric *slov1alpha1.NodeMetric
+	}
+	collectPolicy := slov1alpha1.UsageWithoutPageCache
+	tests := []struct {
+		name   string
+		fields fields
+		want   *slov1alpha1.NodeMetricSpec
+	}{
+		{
+			name: "get node metric spec with default",
+			fields: fields{
+				nodeMetric: nil,
+			},
+			want: &defaultNodeMetricSpec,
+		},
+		{
+			name: "get node metric spec",
+			fields: fields{
+				nodeMetric: &slov1alpha1.NodeMetric{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test-node-slo-name",
+						UID:  "test-node-slo-uid",
+					},
+					Spec: slov1alpha1.NodeMetricSpec{
+						CollectPolicy: &slov1alpha1.NodeMetricCollectPolicy{},
+					},
+				},
+			},
+			want: &slov1alpha1.NodeMetricSpec{
+				CollectPolicy: &slov1alpha1.NodeMetricCollectPolicy{},
+			},
+		},
+		{
+			name: "get node metric spec",
+			fields: fields{
+				nodeMetric: &slov1alpha1.NodeMetric{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test-node-slo-name",
+						UID:  "test-node-slo-uid",
+					},
+					Spec: slov1alpha1.NodeMetricSpec{
+						CollectPolicy: &slov1alpha1.NodeMetricCollectPolicy{
+							NodeMemoryCollectPolicy: &collectPolicy,
+						},
+					},
+				},
+			},
+			want: &slov1alpha1.NodeMetricSpec{
+				CollectPolicy: &slov1alpha1.NodeMetricCollectPolicy{
+					NodeMemoryCollectPolicy: &collectPolicy,
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			nodeMetricInformer := &nodeMetricInformer{
+				nodeMetric: tt.fields.nodeMetric,
+			}
+			s := &statesInformer{
+				states: &PluginState{
+					informerPlugins: map[PluginName]informerPlugin{
+						nodeMetricInformerName: nodeMetricInformer,
+					},
+				},
+			}
+			if got := s.GetNodeMetricSpec(); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("GetNodeMetricSpec() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func Test_statesInformer_GetNodeTopo(t *testing.T) {
 	type fields struct {
 		nodeTopo *topov1alpha1.NodeResourceTopology

diff --git a/pkg/koordlet/statesinformer/mockstatesinformer/mock.go b/pkg/koordlet/statesinformer/mockstatesinformer/mock.go