From 2ab192151f735f5458101e61a56aea488856c112 Mon Sep 17 00:00:00 2001
From: "xiayu.lyt" <xiayu.lyt@alibaba-inc.com>
Date: Thu, 29 Jun 2023 10:20:27 +0800
Subject: [PATCH] add option to preserve collector pod

Signed-off-by: xiayu.lyt <xiayu.lyt@alibaba-inc.com>
---
 pkg/skoop/collector/manager/config.go  | 10 +++---
 pkg/skoop/collector/manager/manager.go | 45 ++++++++++++++------------
 2 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/pkg/skoop/collector/manager/config.go b/pkg/skoop/collector/manager/config.go
index 62a02c34..bcb1d9ac 100644
--- a/pkg/skoop/collector/manager/config.go
+++ b/pkg/skoop/collector/manager/config.go
@@ -19,10 +19,11 @@ func init() {
 }
 
 type SimplePodCollectorConfig struct {
-	Image              string
-	CollectorNamespace string
-	WaitInterval       time.Duration
-	WaitTimeout        time.Duration
+	Image                string
+	CollectorNamespace   string
+	WaitInterval         time.Duration
+	WaitTimeout          time.Duration
+	PreserveCollectorPod bool
 }
 
 func (cc *SimplePodCollectorConfig) BindFlags(fs *pflag.FlagSet) {
@@ -30,6 +31,7 @@ func (cc *SimplePodCollectorConfig) BindFlags(fs *pflag.FlagSet) {
 	fs.StringVarP(&cc.CollectorNamespace, "collector-namespace", "", "skoop", "Namespace where collector pods in.")
 	fs.DurationVarP(&cc.WaitInterval, "collector-pod-wait-interval", "", 2*time.Second, "Collector pod running check interval.")
 	fs.DurationVarP(&cc.WaitTimeout, "collector-pod-wait-timeout", "", 120*time.Second, "Collector pod running check timeout.")
+	fs.BoolVarP(&cc.PreserveCollectorPod, "preserve-collector-pod", "", false, "Preserve collector pod after diagnosis complete.")
 }
 
 func (cc *SimplePodCollectorConfig) Validate() error {
diff --git a/pkg/skoop/collector/manager/manager.go b/pkg/skoop/collector/manager/manager.go
index 54747901..bf97e36e 100644
--- a/pkg/skoop/collector/manager/manager.go
+++ b/pkg/skoop/collector/manager/manager.go
@@ -44,16 +44,17 @@ type SimplePodCollectorManagerOptions struct {
 }
 
 type simplePodCollectorManager struct {
-	image        string
-	namespace    string
-	client       *kubernetes.Clientset
-	restConfig   *rest.Config
-	ipCache      *k8s.IPCache
-	cache        map[string]*k8s.NodeNetworkStackDump
-	nodeCache    map[string]*k8s.NodeInfo
-	podCache     map[string]*k8s.Pod
-	waitInterval time.Duration
-	waitTimeout  time.Duration
+	image                string
+	namespace            string
+	client               *kubernetes.Clientset
+	restConfig           *rest.Config
+	ipCache              *k8s.IPCache
+	cache                map[string]*k8s.NodeNetworkStackDump
+	nodeCache            map[string]*k8s.NodeInfo
+	podCache             map[string]*k8s.Pod
+	waitInterval         time.Duration
+	waitTimeout          time.Duration
+	preserveCollectorPod bool
 }
 
 func NewSimplePodCollectorManager(ctx *ctx.Context) (collector.Manager, error) {
@@ -74,16 +75,17 @@ func NewSimplePodCollectorManager(ctx *ctx.Context) (collector.Manager, error) {
 	}
 
 	return &simplePodCollectorManager{
-		image:        Config.SimplePodCollectorConfig.Image,
-		namespace:    Config.SimplePodCollectorConfig.CollectorNamespace,
-		client:       ctx.KubernetesClient(),
-		restConfig:   ctx.KubernetesRestClient(),
-		ipCache:      ctx.ClusterConfig().IPCache,
-		cache:        map[string]*k8s.NodeNetworkStackDump{},
-		nodeCache:    map[string]*k8s.NodeInfo{},
-		podCache:     map[string]*k8s.Pod{},
-		waitInterval: Config.SimplePodCollectorConfig.WaitInterval,
-		waitTimeout:  Config.SimplePodCollectorConfig.WaitTimeout,
+		image:                Config.SimplePodCollectorConfig.Image,
+		namespace:            Config.SimplePodCollectorConfig.CollectorNamespace,
+		client:               ctx.KubernetesClient(),
+		restConfig:           ctx.KubernetesRestClient(),
+		ipCache:              ctx.ClusterConfig().IPCache,
+		cache:                map[string]*k8s.NodeNetworkStackDump{},
+		nodeCache:            map[string]*k8s.NodeInfo{},
+		podCache:             map[string]*k8s.Pod{},
+		waitInterval:         Config.SimplePodCollectorConfig.WaitInterval,
+		waitTimeout:          Config.SimplePodCollectorConfig.WaitTimeout,
+		preserveCollectorPod: Config.SimplePodCollectorConfig.PreserveCollectorPod,
 	}, nil
 }
 
@@ -256,6 +258,9 @@ func (m *simplePodCollectorManager) collectNodeStackDump(nodeName string) (*k8s.
 	}
 
 	defer func() {
+		if m.preserveCollectorPod {
+			return
+		}
 		err := m.deleteCollectorPod(nodeName)
 		if err != nil {
 			klog.Errorf("failed delete collector pod: %s", err)