Support timeout in Context in clean.go

Extend it to support passing in a context so the caller themselves can define a timeout.

Signed-off-by: Xie Zheng <[email protected]>

cmd_clean/main.go

@@ -4,8 +4,10 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"os"
+	"time"
 
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 
@@ -62,8 +64,16 @@ func main() {
 	cf.Cluster = cluster
 
 	logf.SetLogger(logger.ZapLogger(cf.DefaultConfig.Debug, config.LogLevel))
+	// the error roughly are:
+	// 1. failed to validate config
+	// 2. failed to get nsx client
+	// 3. failed to initialize cleanup service
+	// 4. failed to clean up specific resource
+	// 5. failed because of timeout
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*5)
+	defer cancel()
 
-	status, err := clean.Clean(cf)
+	status, err := clean.Clean(ctx, cf)
 	if err != nil {
 		log.Error(err, "failed to clean nsx resources", "status", status)
 		os.Exit(1)

pkg/clean/clean.go

@@ -4,6 +4,7 @@
 package clean
 
 import (
+	"context"
 	"fmt"
 
 	"k8s.io/client-go/util/retry"
@@ -27,39 +28,44 @@ var log = logger.Log
 // including security policy, static route, subnet, subnet port, subnet set, vpc, ip pool, nsx service account
 // it is usually used when nsx-operator is uninstalled and remove all the resources created by nsx-operator
 // return error if any, return nil if no error
-func Clean(cf *config.NSXOperatorConfig) (Status, error) {
+func Clean(ctx context.Context, cf *config.NSXOperatorConfig) (Status, error) {
 	log.Info("starting NSX cleanup")
-	if err := cf.ValidateConfigFromCmd(); err != nil {
-		return ValidationFailed, err
-	}
-	nsxClient := nsx.GetClient(cf)
-	if nsxClient == nil {
-		return GetNSXClientFailed, fmt.Errorf("failed to get nsx client")
-	}
-	if cleanupService, err := InitializeCleanupService(cf); err != nil {
-		return InitCleanupServiceFailed, fmt.Errorf("failed to initialize cleanup service: %w", err)
-	} else if cleanupService.err != nil {
-		return InitCleanupServiceFailed, cleanupService.err
-	} else {
-		for _, clean := range cleanupService.cleans {
-			if err := retry.OnError(retry.DefaultRetry, func(err error) bool {
-				if err != nil {
-					log.Info("retrying to clean up NSX resources", "error", err)
-					return true
-				}
-				return false
-			}, func() error {
-				if err := clean.Cleanup(); err != nil {
-					return err
+	select {
+	case <-ctx.Done():
+		return TimeoutFailed, ctx.Err()
+	default:
+		if err := cf.ValidateConfigFromCmd(); err != nil {
+			return ValidationFailed, err
+		}
+		nsxClient := nsx.GetClient(cf)
+		if nsxClient == nil {
+			return GetNSXClientFailed, fmt.Errorf("failed to get nsx client")
+		}
+		if cleanupService, err := InitializeCleanupService(cf); err != nil {
+			return InitCleanupServiceFailed, fmt.Errorf("failed to initialize cleanup service: %w", err)
+		} else if cleanupService.err != nil {
+			return InitCleanupServiceFailed, cleanupService.err
+		} else {
+			for _, clean := range cleanupService.cleans {
+				if err := retry.OnError(retry.DefaultRetry, func(err error) bool {
+					if err != nil {
+						log.Info("retrying to clean up NSX resources", "error", err)
+						return true
+					}
+					return false
+				}, func() error {
+					if err := clean.Cleanup(); err != nil {
+						return err
+					}
+					return nil
+				}); err != nil {
+					return CleanupResourceFailed, err
 				}
-				return nil
-			}); err != nil {
-				return CleanupResourceFailed, err
 			}
 		}
+		log.Info("cleanup NSX resources successfully")
+		return OK, nil
 	}
-	log.Info("cleanup NSX resources successfully")
-	return OK, nil
 }
 
 // InitializeCleanupService initializes all the CR services

pkg/clean/types.go

@@ -41,4 +41,5 @@ var (
 	GetNSXClientFailed       = Status{Code: 2, Message: "failed to get nsx client"}
 	InitCleanupServiceFailed = Status{Code: 3, Message: "failed to initialize cleanup service"}
 	CleanupResourceFailed    = Status{Code: 4, Message: "failed to clean up"}
+	TimeoutFailed            = Status{Code: 5, Message: "failed because of timeout"}
 )