volcano.sh/volcano@v1.9.0/pkg/controllers/apis/job_info.go (about)

     1  /*
     2  Copyright 2019 The Volcano Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package apis
    18  
    19  import (
    20  	"fmt"
    21  
    22  	v1 "k8s.io/api/core/v1"
    23  
    24  	batch "volcano.sh/apis/pkg/apis/batch/v1alpha1"
    25  )
    26  
    27  // JobInfo struct.
    28  type JobInfo struct {
    29  	Namespace string
    30  	Name      string
    31  
    32  	Job  *batch.Job
    33  	Pods map[string]map[string]*v1.Pod
    34  }
    35  
    36  // Clone function clones the k8s pod values to the JobInfo struct.
    37  func (ji *JobInfo) Clone() *JobInfo {
    38  	job := &JobInfo{
    39  		Namespace: ji.Namespace,
    40  		Name:      ji.Name,
    41  		Job:       ji.Job,
    42  
    43  		Pods: make(map[string]map[string]*v1.Pod, len(ji.Pods)),
    44  	}
    45  
    46  	for key, pods := range ji.Pods {
    47  		job.Pods[key] = make(map[string]*v1.Pod, len(pods))
    48  		for pn, pod := range pods {
    49  			job.Pods[key][pn] = pod
    50  		}
    51  	}
    52  
    53  	return job
    54  }
    55  
    56  // SetJob sets the volcano jobs values to the JobInfo struct.
    57  func (ji *JobInfo) SetJob(job *batch.Job) {
    58  	ji.Name = job.Name
    59  	ji.Namespace = job.Namespace
    60  	ji.Job = job
    61  }
    62  
    63  // AddPod adds the k8s pod object values to the Pods field
    64  // of JobStruct if it doesn't exist. Otherwise it throws error.
    65  func (ji *JobInfo) AddPod(pod *v1.Pod) error {
    66  	taskName, found := pod.Annotations[batch.TaskSpecKey]
    67  	if !found {
    68  		return fmt.Errorf("failed to find taskName of Pod <%s/%s>",
    69  			pod.Namespace, pod.Name)
    70  	}
    71  
    72  	_, found = pod.Annotations[batch.JobVersion]
    73  	if !found {
    74  		return fmt.Errorf("failed to find jobVersion of Pod <%s/%s>",
    75  			pod.Namespace, pod.Name)
    76  	}
    77  
    78  	if _, found := ji.Pods[taskName]; !found {
    79  		ji.Pods[taskName] = make(map[string]*v1.Pod)
    80  	}
    81  	if _, found := ji.Pods[taskName][pod.Name]; found {
    82  		return fmt.Errorf("duplicated pod")
    83  	}
    84  	ji.Pods[taskName][pod.Name] = pod
    85  
    86  	return nil
    87  }
    88  
    89  // UpdatePod updates the k8s pod object values to the existing pod.
    90  func (ji *JobInfo) UpdatePod(pod *v1.Pod) error {
    91  	taskName, found := pod.Annotations[batch.TaskSpecKey]
    92  	if !found {
    93  		return fmt.Errorf("failed to find taskName of Pod <%s/%s>",
    94  			pod.Namespace, pod.Name)
    95  	}
    96  	_, found = pod.Annotations[batch.JobVersion]
    97  	if !found {
    98  		return fmt.Errorf("failed to find jobVersion of Pod <%s/%s>",
    99  			pod.Namespace, pod.Name)
   100  	}
   101  
   102  	if _, found := ji.Pods[taskName]; !found {
   103  		return fmt.Errorf("can not find task %s in cache", taskName)
   104  	}
   105  	if _, found := ji.Pods[taskName][pod.Name]; !found {
   106  		return fmt.Errorf("can not find pod <%s/%s> in cache",
   107  			pod.Namespace, pod.Name)
   108  	}
   109  	ji.Pods[taskName][pod.Name] = pod
   110  
   111  	return nil
   112  }
   113  
   114  // DeletePod deletes the given k8s pod from the JobInfo struct.
   115  func (ji *JobInfo) DeletePod(pod *v1.Pod) error {
   116  	taskName, found := pod.Annotations[batch.TaskSpecKey]
   117  	if !found {
   118  		return fmt.Errorf("failed to find taskName of Pod <%s/%s>",
   119  			pod.Namespace, pod.Name)
   120  	}
   121  	_, found = pod.Annotations[batch.JobVersion]
   122  	if !found {
   123  		return fmt.Errorf("failed to find jobVersion of Pod <%s/%s>",
   124  			pod.Namespace, pod.Name)
   125  	}
   126  
   127  	if pods, found := ji.Pods[taskName]; found {
   128  		delete(pods, pod.Name)
   129  		if len(pods) == 0 {
   130  			delete(ji.Pods, taskName)
   131  		}
   132  	}
   133  
   134  	return nil
   135  }