From 722489db3670bc6b4f4d9d5dfe6b76bd87920c0c Mon Sep 17 00:00:00 2001
From: kerthcet <kerthcet@gmail.com>
Date: Fri, 8 Oct 2021 18:54:31 +0800
Subject: [PATCH] fix kubemark flushlogs not working together with os.Exit()

Signed-off-by: kerthcet <kerthcet@gmail.com>
---
 cmd/kubemark/hollow-node.go  | 30 ++++++++++++++++--------------
 pkg/kubemark/hollow_proxy.go |  5 +++--
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/cmd/kubemark/hollow-node.go b/cmd/kubemark/hollow-node.go
index de254acb964..ee6690ff82e 100644
--- a/cmd/kubemark/hollow-node.go
+++ b/cmd/kubemark/hollow-node.go
@@ -145,9 +145,9 @@ func newHollowNodeCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:  "kubemark",
 		Long: "kubemark",
-		Run: func(cmd *cobra.Command, args []string) {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			verflag.PrintAndExitIfRequested()
-			run(cmd, s)
+			return run(cmd, s)
 		},
 		Args: func(cmd *cobra.Command, args []string) error {
 			for _, arg := range args {
@@ -166,24 +166,24 @@ func newHollowNodeCommand() *cobra.Command {
 	return cmd
 }
 
-func run(cmd *cobra.Command, config *hollowNodeConfig) {
+func run(cmd *cobra.Command, config *hollowNodeConfig) error {
 	// To help debugging, immediately log version and print flags.
 	klog.Infof("Version: %+v", version.Get())
 	cliflag.PrintFlags(cmd.Flags())
 
 	if !knownMorphs.Has(config.Morph) {
-		klog.Fatalf("Unknown morph: %v. Allowed values: %v", config.Morph, knownMorphs.List())
+		return fmt.Errorf("Unknown morph: %v. allowed values: %v", config.Morph, knownMorphs.List())
 	}
 
 	// create a client to communicate with API server.
 	clientConfig, err := config.createClientConfigFromFile()
 	if err != nil {
-		klog.Fatalf("Failed to create a ClientConfig: %v. Exiting.", err)
+		return fmt.Errorf("Failed to create a ClientConfig, error: %w. Exiting", err)
 	}
 
 	client, err := clientset.NewForConfig(clientConfig)
 	if err != nil {
-		klog.Fatalf("Failed to create a ClientSet: %v. Exiting.", err)
+		return fmt.Errorf("Failed to create a ClientSet, error: %w. Exiting", err)
 	}
 
 	if config.Morph == "kubelet" {
@@ -200,7 +200,7 @@ func run(cmd *cobra.Command, config *hollowNodeConfig) {
 		heartbeatClientConfig.QPS = float32(-1)
 		heartbeatClient, err := clientset.NewForConfig(&heartbeatClientConfig)
 		if err != nil {
-			klog.Fatalf("Failed to create a ClientSet: %v. Exiting.", err)
+			return fmt.Errorf("Failed to create a ClientSet, error: %w. Exiting", err)
 		}
 
 		cadvisorInterface := &cadvisortest.Fake{
@@ -221,23 +221,23 @@ func run(cmd *cobra.Command, config *hollowNodeConfig) {
 
 		endpoint, err := fakeremote.GenerateEndpoint()
 		if err != nil {
-			klog.Fatalf("Failed to generate fake endpoint %v.", err)
+			return fmt.Errorf("Failed to generate fake endpoint, error: %w", err)
 		}
 		fakeRemoteRuntime := fakeremote.NewFakeRemoteRuntime()
 		if err = fakeRemoteRuntime.Start(endpoint); err != nil {
-			klog.Fatalf("Failed to start fake runtime %v.", err)
+			return fmt.Errorf("Failed to start fake runtime, error: %w", err)
 		}
 		defer fakeRemoteRuntime.Stop()
 		runtimeService, err := remote.NewRemoteRuntimeService(endpoint, 15*time.Second)
 		if err != nil {
-			klog.Fatalf("Failed to init runtime service %v.", err)
+			return fmt.Errorf("Failed to init runtime service, error: %w", err)
 		}
 
 		var imageService internalapi.ImageManagerService = fakeRemoteRuntime.ImageService
 		if config.UseHostImageService {
 			imageService, err = remote.NewRemoteImageService(f.RemoteImageEndpoint, 15*time.Second)
 			if err != nil {
-				klog.Fatalf("Failed to init image service %v.", err)
+				return fmt.Errorf("Failed to init image service, error: %w", err)
 			}
 		}
 
@@ -256,7 +256,7 @@ func run(cmd *cobra.Command, config *hollowNodeConfig) {
 	if config.Morph == "proxy" {
 		client, err := clientset.NewForConfig(clientConfig)
 		if err != nil {
-			klog.Fatalf("Failed to create API Server client: %v", err)
+			return fmt.Errorf("Failed to create API Server client, error: %w", err)
 		}
 		iptInterface := fakeiptables.NewFake()
 		sysctl := fakesysctl.NewFake()
@@ -280,8 +280,10 @@ func run(cmd *cobra.Command, config *hollowNodeConfig) {
 			config.ProxierMinSyncPeriod,
 		)
 		if err != nil {
-			klog.Fatalf("Failed to create hollowProxy instance: %v", err)
+			return fmt.Errorf("Failed to create hollowProxy instance, error: %w", err)
 		}
-		hollowProxy.Run()
+		return hollowProxy.Run()
 	}
+
+	return nil
 }
diff --git a/pkg/kubemark/hollow_proxy.go b/pkg/kubemark/hollow_proxy.go
index 40e36035b28..390204cc478 100644
--- a/pkg/kubemark/hollow_proxy.go
+++ b/pkg/kubemark/hollow_proxy.go
@@ -132,8 +132,9 @@ func NewHollowProxyOrDie(
 	}, nil
 }
 
-func (hp *HollowProxy) Run() {
+func (hp *HollowProxy) Run() error {
 	if err := hp.ProxyServer.Run(); err != nil {
-		klog.Fatalf("Error while running proxy: %v\n", err)
+		return fmt.Errorf("Error while running proxy: %w", err)
 	}
+	return nil
 }