Merge pull request #330 from shangjin92/feature/support-imagepullpoli…

…cy-setting
alibaba · Dec 13, 2024 · 05c2d69 · 05c2d69
2 parents 0caffc7 + 616d5d0
commit 05c2d69
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 3 deletions.
diff --git a/pkg/skoop/collector/manager/config.go b/pkg/skoop/collector/manager/config.go
@@ -20,6 +20,7 @@ func init() {
 
 type SimplePodCollectorConfig struct {
 	Image                string
+	ImagePullPolicy      string
 	CollectorNamespace   string
 	RuntimeAPIAddress    string
 	WaitInterval         time.Duration
@@ -29,6 +30,7 @@ type SimplePodCollectorConfig struct {
 
 func (cc *SimplePodCollectorConfig) BindFlags(fs *pflag.FlagSet) {
 	fs.StringVarP(&cc.Image, "collector-image", "", "kubeskoop/agent:v1.0.0", "Image used for collector.")
+	fs.StringVarP(&cc.ImagePullPolicy, "collector-image-pull-policy", "", "Always", "The image pull policy for collector.")
 	fs.StringVarP(&cc.CollectorNamespace, "collector-namespace", "", "skoop", "Namespace where collector pods in.")
 	fs.StringVarP(&cc.RuntimeAPIAddress, "collector-cri-address", "", "", "Runtime CRI API endpoint address.")
 	fs.DurationVarP(&cc.WaitInterval, "collector-pod-wait-interval", "", 2*time.Second, "Collector pod running check interval.")

diff --git a/pkg/skoop/collector/manager/manager.go b/pkg/skoop/collector/manager/manager.go
@@ -45,6 +45,7 @@ type SimplePodCollectorManagerOptions struct {
 
 type simplePodCollectorManager struct {
 	image                string
+	imagePullPolicy      v1.PullPolicy
 	namespace            string
 	runtimeAPIAddress    string
 	client               *kubernetes.Clientset
@@ -75,8 +76,14 @@ func NewSimplePodCollectorManager(ctx *ctx.Context) (collector.Manager, error) {
 		Config.SimplePodCollectorConfig.WaitTimeout = defaultWaitTimeout * time.Second
 	}
 
+	pullPolicy, err := utils.ConvertToImagePullPolicy(Config.SimplePodCollectorConfig.ImagePullPolicy)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create pod collector manager: %w", err)
+	}
+
 	return &simplePodCollectorManager{
 		image:                Config.SimplePodCollectorConfig.Image,
+		imagePullPolicy:      pullPolicy,
 		namespace:            Config.SimplePodCollectorConfig.CollectorNamespace,
 		client:               ctx.KubernetesClient(),
 		restConfig:           ctx.KubernetesRestClient(),
@@ -332,7 +339,7 @@ func (m *simplePodCollectorManager) createCollectorPod(nodeName string) (*v1.Pod
 				{
 					Name:            "collector",
 					Image:           m.image,
-					ImagePullPolicy: "Always",
+					ImagePullPolicy: m.imagePullPolicy,
 					SecurityContext: &v1.SecurityContext{
 						Privileged: pointer.Bool(true),
 					},
@@ -359,8 +366,9 @@ func (m *simplePodCollectorManager) createCollectorPod(nodeName string) (*v1.Pod
 			},
 			Containers: []v1.Container{
 				{
-					Name:  "alive",
-					Image: m.image,
+					Name:            "alive",
+					Image:           m.image,
+					ImagePullPolicy: m.imagePullPolicy,
 					Command: []string{
 						"/bin/sh",
 						"-c",

diff --git a/pkg/skoop/utils/k8s.go b/pkg/skoop/utils/k8s.go
@@ -154,3 +154,16 @@ func ContainsLoadBalancerIP(svc *v1.Service, ip string) bool {
 	}
 	return false
 }
+
+func ConvertToImagePullPolicy(policy string) (v1.PullPolicy, error) {
+	policyMap := map[string]v1.PullPolicy{
+		"Always":       v1.PullAlways,
+		"IfNotPresent": v1.PullIfNotPresent,
+		"Never":        v1.PullNever,
+	}
+
+	if pullPolicy, exists := policyMap[policy]; exists {
+		return pullPolicy, nil
+	}
+	return "", fmt.Errorf("invalid image pull policy: %s, valid options are: Always, IfNotPresent, Never", policy)
+}