change prune strategy to separate resource removal from the strategy implementations (#76)

Co-authored-by: churromechanic <info@churrodata.com>

Author: Jeff McCormick

prune/maxage.go

@@ -20,24 +20,26 @@ import (
 )
 
 // maxAge looks for and prunes resources, currently jobs and pods,
-// that exceed a user specified age (e.g. 3d)
-func pruneByMaxAge(ctx context.Context, config Config, resources []ResourceInfo) (err error) {
+// that exceed a user specified age (e.g. 3d), resources to be removed
+// are returned
+func pruneByMaxAge(_ context.Context, config Config, resources []ResourceInfo) (resourcesToRemove []ResourceInfo, err error) {
 	config.log.V(1).Info("maxAge running", "setting", config.Strategy.MaxAgeSetting)
 
-	maxAgeDuration, _ := time.ParseDuration(config.Strategy.MaxAgeSetting)
+	maxAgeDuration, e := time.ParseDuration(config.Strategy.MaxAgeSetting)
+	if e != nil {
+		return resourcesToRemove, e
+	}
+
 	maxAgeTime := time.Now().Add(-maxAgeDuration)
 
 	for i := 0; i < len(resources); i++ {
 		config.log.V(1).Info("age of pod ", "age", time.Since(resources[i].StartTime), "maxage", maxAgeTime)
 		if resources[i].StartTime.Before(maxAgeTime) {
 			config.log.V(1).Info("pruning ", "kind", resources[i].GVK, "name", resources[i].Name)
 
-			err := config.removeResource(ctx, resources[i])
-			if err != nil {
-				return err
-			}
+			resourcesToRemove = append(resourcesToRemove, resources[i])
 		}
 	}
 
-	return nil
+	return resourcesToRemove, nil
 }

prune/maxcount.go

@@ -16,24 +16,25 @@ package prune
 
 import (
 	"context"
+	"fmt"
 	"time"
 )
 
 // pruneByMaxCount looks for and prunes resources, currently jobs and pods,
 // that exceed a user specified count (e.g. 3), the oldest resources
-// are pruned
-func pruneByMaxCount(ctx context.Context, config Config, resources []ResourceInfo) (err error) {
+// are pruned, resources to remove are returned
+func pruneByMaxCount(_ context.Context, config Config, resources []ResourceInfo) (resourcesToRemove []ResourceInfo, err error) {
 	config.log.V(1).Info("pruneByMaxCount running ", "max count", config.Strategy.MaxCountSetting, "resource count", len(resources))
+	if config.Strategy.MaxCountSetting < 0 {
+		return resourcesToRemove, fmt.Errorf("max count setting less than zero")
+	}
 
 	if len(resources) > config.Strategy.MaxCountSetting {
 		removeCount := len(resources) - config.Strategy.MaxCountSetting
 		for i := len(resources) - 1; i >= 0; i-- {
 			config.log.V(1).Info("pruning pod ", "pod name", resources[i].Name, "age", time.Since(resources[i].StartTime))
 
-			err := config.removeResource(ctx, resources[i])
-			if err != nil {
-				return err
-			}
+			resourcesToRemove = append(resourcesToRemove, resources[i])
 
 			removeCount--
 			if removeCount == 0 {
@@ -42,5 +43,5 @@ func pruneByMaxCount(ctx context.Context, config Config, resources []ResourceInf
 		}
 	}
 
-	return nil
+	return resourcesToRemove, nil
 }

prune/prune.go

@@ -54,7 +54,7 @@ type StrategyConfig struct {
 
 // StrategyFunc function allows a means to specify
 // custom prune strategies
-type StrategyFunc func(cfg Config, resources []ResourceInfo) error
+type StrategyFunc func(cfg Config, resources []ResourceInfo) ([]ResourceInfo, error)
 
 // PreDelete function is called before a resource is pruned
 type PreDelete func(cfg Config, something ResourceInfo) error
@@ -101,19 +101,26 @@ func (config Config) Execute(ctx context.Context) error {
 			config.log.V(1).Info("jobs ", "count", len(resourceList))
 		}
 
+		var resourcesToRemove []ResourceInfo
+
 		switch config.Strategy.Mode {
 		case MaxAgeStrategy:
-			err = pruneByMaxAge(ctx, config, resourceList)
+			resourcesToRemove, err = pruneByMaxAge(ctx, config, resourceList)
 		case MaxCountStrategy:
-			err = pruneByMaxCount(ctx, config, resourceList)
+			resourcesToRemove, err = pruneByMaxCount(ctx, config, resourceList)
 		case CustomStrategy:
-			err = config.CustomStrategy(config, resourceList)
+			resourcesToRemove, err = config.CustomStrategy(config, resourceList)
 		default:
 			return fmt.Errorf("unknown strategy")
 		}
 		if err != nil {
 			return err
 		}
+
+		err = config.removeResources(ctx, resourcesToRemove)
+		if err != nil {
+			return err
+		}
 	}
 
 	config.log.V(1).Info("Prune completed")

prune/remove.go

@@ -21,32 +21,36 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-func (config Config) removeResource(ctx context.Context, resource ResourceInfo) (err error) {
+func (config Config) removeResources(ctx context.Context, resources []ResourceInfo) (err error) {
 
 	if config.DryRun {
 		return nil
 	}
 
-	if config.PreDeleteHook != nil {
-		err = config.PreDeleteHook(config, resource)
-		if err != nil {
-			return err
-		}
-	}
+	for i := 0; i < len(resources); i++ {
+		r := resources[i]
 
-	switch resource.GVK.Kind {
-	case PodKind:
-		err := config.Clientset.CoreV1().Pods(resource.Namespace).Delete(ctx, resource.Name, metav1.DeleteOptions{})
-		if err != nil {
-			return err
+		if config.PreDeleteHook != nil {
+			err = config.PreDeleteHook(config, r)
+			if err != nil {
+				return err
+			}
 		}
-	case JobKind:
-		err := config.Clientset.BatchV1().Jobs(resource.Namespace).Delete(ctx, resource.Name, metav1.DeleteOptions{})
-		if err != nil {
-			return err
+
+		switch resources[i].GVK.Kind {
+		case PodKind:
+			err := config.Clientset.CoreV1().Pods(r.Namespace).Delete(ctx, r.Name, metav1.DeleteOptions{})
+			if err != nil {
+				return err
+			}
+		case JobKind:
+			err := config.Clientset.BatchV1().Jobs(r.Namespace).Delete(ctx, r.Name, metav1.DeleteOptions{})
+			if err != nil {
+				return err
+			}
+		default:
+			return fmt.Errorf("unsupported resource kind")
 		}
-	default:
-		return fmt.Errorf("unsupported resource kind")
 	}
 
 	return nil

prune/resource_test.go

@@ -325,11 +325,12 @@ func myhook(cfg Config, x ResourceInfo) error {
 
 // myStrategy shows how you can write your own strategy, in this
 // example, the strategy doesn't really do another other than count
-// the number of resources
-func myStrategy(cfg Config, resources []ResourceInfo) error {
+// the number of resources, returning a list of resources to delete in
+// this case zero.
+func myStrategy(cfg Config, resources []ResourceInfo) (resourcesToRemove []ResourceInfo, err error) {
 	fmt.Printf("myStrategy is called with resources %v config %v\n", resources, cfg)
 	if len(resources) != 3 {
-		return fmt.Errorf("count of resources did not equal our expectation")
+		return resourcesToRemove, fmt.Errorf("count of resources did not equal our expectation")
 	}
-	return nil
+	return resourcesToRemove, nil
 }