diff --git a/pkg/appmgmt/general/general.go b/pkg/appmgmt/general/general.go index 8a04d1310..93255806d 100644 --- a/pkg/appmgmt/general/general.go +++ b/pkg/appmgmt/general/general.go @@ -25,7 +25,6 @@ import ( v1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/labels" - "k8s.io/apimachinery/pkg/selection" k8sCache "k8s.io/client-go/tools/cache" "github.com/apache/incubator-yunikorn-k8shim/pkg/appmgmt/interfaces" @@ -293,15 +292,8 @@ func (os *Manager) deletePod(obj interface{}) { } func (os *Manager) ListApplications() (map[string]interfaces.ApplicationMetadata, error) { - // selector: applicationID exist - slt := labels.NewSelector() - req, err := labels.NewRequirement(constants.LabelApplicationID, selection.Exists, nil) - if err != nil { - return nil, err - } - slt = slt.Add(*req) - // list all pods on this cluster + slt := labels.NewSelector() appPods, err := os.apiProvider.GetAPIs().PodInformer.Lister().List(slt) if err != nil { return nil, err diff --git a/pkg/appmgmt/general/general_test.go b/pkg/appmgmt/general/general_test.go index bcd59deeb..3e0f57333 100644 --- a/pkg/appmgmt/general/general_test.go +++ b/pkg/appmgmt/general/general_test.go @@ -30,6 +30,7 @@ import ( "github.com/apache/incubator-yunikorn-k8shim/pkg/client" "github.com/apache/incubator-yunikorn-k8shim/pkg/common/constants" "github.com/apache/incubator-yunikorn-k8shim/pkg/common/events" + "github.com/apache/incubator-yunikorn-k8shim/pkg/common/test" ) const taskGroupInfo = ` @@ -526,3 +527,269 @@ func TestGetOwnerReferences(t *testing.T) { assert.Equal(t, returnedOwnerRefs[0].Kind, "Pod", "Unexpected owner reference Kind") assert.Equal(t, returnedOwnerRefs[0].APIVersion, v1.SchemeGroupVersion.String(), "Unexpected owner reference Kind") } + +// nolint: funlen +func TestListApplication(t *testing.T) { + var app01, app02, app03, app04, app05, app06 = "app00001", + "app00002", "app00003", "app00004", "app00005", "app00006" + var queue01, queue02 = "root.queue01", "root.queue02" + var ns01, ns02 = "namespace01", "namespace02" + + // mock the pod lister for this test + mockedAPIProvider := client.NewMockedAPIProvider() + mockedPodLister := test.NewPodListerMock() + mockedAPIProvider.SetPodLister(mockedPodLister) + + // app01 pods, running in namespace01 and queue01 + // all pods are having applicationID and queue name specified + // 2 pods have assigned nodes, 1 pod is pending for scheduling + // allocated pod + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app01pod00001", + Namespace: ns01, + Labels: map[string]string{ + constants.LabelApplicationID: app01, + constants.LabelQueueName: queue01, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + NodeName: "allocated-node", + }, + }) + + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app01pod00002", + Namespace: ns01, + Labels: map[string]string{ + constants.LabelApplicationID: app01, + constants.LabelQueueName: queue01, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + NodeName: "allocated-node", + }, + }) + + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app01pod00003", + Namespace: ns01, + Labels: map[string]string{ + constants.LabelApplicationID: app01, + constants.LabelQueueName: queue01, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + }, + }) + + // app02 pods, running in queue02 and namespace02 + // 2 pods are having applicationID and queue name specified + // both 2 pods are pending + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app02pod0001", + Namespace: ns02, + Labels: map[string]string{ + constants.LabelApplicationID: app02, + constants.LabelQueueName: queue02, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + }, + }) + + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app02pod0002", + Namespace: ns02, + Labels: map[string]string{ + constants.LabelApplicationID: app02, + constants.LabelQueueName: queue02, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + }, + }) + + // app03 pods, running in queue02 and namespace02 + // 2 pods do not have label applicationID specified, but have spark-app-selector + // both 2 pods are allocated + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app03pod0001", + Namespace: ns01, + Labels: map[string]string{ + constants.SparkLabelAppID: app03, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + NodeName: "some-node", + }, + }) + + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app03pod0002", + Namespace: ns02, + Labels: map[string]string{ + constants.SparkLabelAppID: app03, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + NodeName: "some-node", + }, + }) + + // app04 pods, running in queue01 and namespace01 + // app04 has 2 pods which only has annotation yunikorn.apache.org/app-id specified + // both 2 pods are allocated + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app04pod0001", + Namespace: ns01, + Annotations: map[string]string{ + constants.AnnotationApplicationID: app04, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + NodeName: "some-node", + }, + }) + + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app04pod0002", + Namespace: ns01, + Annotations: map[string]string{ + constants.AnnotationApplicationID: app04, + }, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + NodeName: "some-node", + }, + }) + + // app05 pods, running in queue01 and namespace01 + // app05 pod has no label or annotation specified + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app05pod0001", + Namespace: ns01, + }, + Spec: v1.PodSpec{ + SchedulerName: constants.SchedulerName, + NodeName: "some-node", + }, + }) + + // app06 pods, running in queue01 and namespace01 + // pod has spark-app-selector set and it is allocated but not scheduled by yunikorn + mockedPodLister.AddPod(&v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "app06pod0001", + Namespace: ns01, + Labels: map[string]string{ + constants.SparkLabelAppID: app06, + }, + }, + Spec: v1.PodSpec{ + NodeName: "some-node", + }, + }) + + // init the app manager and run listApp + am := NewManager(cache.NewMockedAMProtocol(), mockedAPIProvider) + apps, err := am.ListApplications() + assert.NilError(t, err) + assert.Equal(t, len(apps), 3) + _, exist := apps[app01] + assert.Equal(t, exist, true, + "app01 should be included in the list because "+ + "it has applicationID and queue namespace specified in the"+ + "queue and it has 2 pods allocated.") + + _, exist = apps[app02] + assert.Equal(t, exist, false, + "app02 should be excluded from the list because"+ + " it has no allocated pods found.") + + _, exist = apps[app03] + assert.Equal(t, exist, true, + "app03 should be included in the list because"+ + " it has 2 pods allocated and both pods have "+ + "spark-app-selector set.") + + _, exist = apps[app04] + assert.Equal(t, exist, true, + "app04 should be included in the list because"+ + " it has 2 pods allocated and both pods have "+ + "annotation yunikorn.apache.org/app-id specified set.") + + _, exist = apps[app05] + assert.Equal(t, exist, false, + "app05 should be excluded in the list because"+ + " pods have no appID set in annotation/label and "+ + "spark-app-selector doesn't exist either") + + _, exist = apps[app06] + assert.Equal(t, exist, false, + "app06 should be excluded in the list because"+ + " pods have spark-app-selector but the schedulerName "+ + "is not yunikorn, which is not scheduled by yunikorn.") +} diff --git a/pkg/client/apifactory_mock.go b/pkg/client/apifactory_mock.go index 5bb6f122c..1b1ca2848 100644 --- a/pkg/client/apifactory_mock.go +++ b/pkg/client/apifactory_mock.go @@ -58,7 +58,7 @@ func NewMockedAPIProvider() *MockedAPIProvider { KubeClient: NewKubeClientMock(), SchedulerAPI: test.NewSchedulerAPIMock(), AppClient: fake.NewSimpleClientset(), - PodInformer: nil, + PodInformer: test.NewMockedPodInformer(), NodeInformer: test.NewMockedNodeInformer(), ConfigMapInformer: test.NewMockedConfigMapInformer(), PVInformer: &MockedPersistentVolumeInformer{}, @@ -116,6 +116,13 @@ func (m *MockedAPIProvider) SetNodeLister(lister corev1.NodeLister) { } } +func (m *MockedAPIProvider) SetPodLister(lister corev1.PodLister) { + informer := m.clients.PodInformer + if i, ok := informer.(*test.MockedPodInformer); ok { + i.SetLister(lister) + } +} + func (m *MockedAPIProvider) GetAPIs() *Clients { return m.clients } diff --git a/pkg/common/test/podinformer_mock.go b/pkg/common/test/podinformer_mock.go new file mode 100644 index 000000000..0862697f5 --- /dev/null +++ b/pkg/common/test/podinformer_mock.go @@ -0,0 +1,46 @@ +/* + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +*/ + +package test + +import ( + v1 "k8s.io/client-go/listers/core/v1" + "k8s.io/client-go/tools/cache" +) + +type MockedPodInformer struct { + podLister v1.PodLister +} + +func NewMockedPodInformer() *MockedPodInformer { + return &MockedPodInformer{ + podLister: NewPodListerMock(), + } +} + +func (m *MockedPodInformer) Informer() cache.SharedIndexInformer { + return nil +} + +func (m *MockedPodInformer) Lister() v1.PodLister { + return m.podLister +} + +func (m *MockedPodInformer) SetLister(lister v1.PodLister) { + m.podLister = lister +} diff --git a/pkg/common/test/podlister_mock.go b/pkg/common/test/podlister_mock.go index cde4529fc..cdd50266b 100644 --- a/pkg/common/test/podlister_mock.go +++ b/pkg/common/test/podlister_mock.go @@ -41,7 +41,13 @@ func (n *PodListerMock) AddPod(pod *v1.Pod) { } func (n *PodListerMock) List(selector labels.Selector) (ret []*v1.Pod, err error) { - return n.allPods, nil + result := make([]*v1.Pod, 0) + for _, pod := range n.allPods { + if selector.Matches(labels.Set(pod.Labels)) { + result = append(result, pod) + } + } + return result, nil } func (n *PodListerMock) Get(name string) (*v1.Pod, error) {