Mark ephemeral pods created as part of kanister functions with JobID label

pkg/function/kube_task.go

@@ -64,6 +64,8 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		Command:      command,
 		PodOverride:  podOverride,
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddLabelsToPodOptionsFromContext(ctx, options, consts.LabelPrefix, "JobID")
 
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()

pkg/kube/pod_runner_test.go

@@ -18,11 +18,15 @@ import (
 	"context"
 	"os"
 
+	"github.com/pkg/errors"
 	. "gopkg.in/check.v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/kubernetes/fake"
 	"k8s.io/client-go/testing"
+
+	"github.com/kanisterio/kanister/pkg/consts"
+	"github.com/kanisterio/kanister/pkg/field"
 )
 
 type PodRunnerTestSuite struct{}
@@ -109,9 +113,66 @@ func (s *PodRunnerTestSuite) TestPodRunnerForSuccessCase(c *C) {
 	cancel()
 }
 
+// TestPodRunnerWithJobIDDebugLabelForSuccessCase: This test adds a debug entry (kanister.io/JobID) into the context and verifies the
+// pod got created with corresponding label using the entry or not.
+func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	randomUUID := "xyz123"
+	ctx = field.Context(ctx, consts.LabelPrefix+"JobID", randomUUID)
+	ctx = field.Context(ctx, "some-test-key", "some-test-value")
+
+	cli := fake.NewSimpleClientset()
+	cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		return false, nil, nil
+	})
+	cli.PrependReactor("get", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		p := &corev1.Pod{
+			Status: corev1.PodStatus{
+				Phase: corev1.PodRunning,
+			},
+		}
+		return true, p, nil
+	})
+	deleted := make(chan struct{})
+	cli.PrependReactor("delete", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		c.Log("Pod deleted due to Context Cancelled")
+		close(deleted)
+		return true, nil, nil
+	})
+	po := &PodOptions{
+		Namespace: podRunnerNS,
+		Name:      podName,
+		Command:   []string{"sh", "-c", "tail -f /dev/null"},
+	}
+	AddLabelsToPodOptionsFromContext(ctx, po, consts.LabelPrefix, "JobID")
+	pr := NewPodRunner(cli, po)
+	errorCh := make(chan error)
+	go func() {
+		_, err := pr.Run(ctx, makePodRunnerTestForKeyPresence(consts.LabelPrefix+"JobID", randomUUID, deleted))
+		errorCh <- err
+	}()
+	deleted <- struct{}{}
+	c.Assert(<-errorCh, IsNil)
+	cancel()
+}
+
 func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 		<-ch
 		return nil, nil
 	}
 }
+
+func makePodRunnerTestForKeyPresence(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+		<-ch
+		value, ok := pc.Pod().Labels[labelKey]
+		if !ok {
+			return nil, errors.New("Key not present")
+		}
+		if value != labelValue {
+			return nil, errors.New("Value mismatch")
+		}
+		return nil, nil
+	}
+}

pkg/kube/utils.go

@@ -17,10 +17,13 @@ package kube
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	osversioned "github.com/openshift/client-go/apps/clientset/versioned"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
+
+	"github.com/kanisterio/kanister/pkg/field"
 )
 
 const (
@@ -171,3 +174,22 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 
 	return false
 }
+
+// AddLabelsToPodOptionsFromContext adds additional label selector to `PodOptions`,
+// provided the context has a key starting with `keyPrefix`.
+func AddLabelsToPodOptionsFromContext(ctx context.Context, options *PodOptions, keyPrefix, keySuffix string) {
+	fields := field.FromContext(ctx)
+	if fields == nil {
+		return
+	}
+	for _, f := range fields.Fields() {
+		if !strings.HasPrefix(f.Key(), keyPrefix) {
+			continue
+		}
+		value := f.Value().(string)
+		if options.Labels == nil {
+			options.Labels = make(map[string]string)
+		}
+		options.Labels[keyPrefix+keySuffix] = value
+	}
+}