Optimize with parallelism and .CSV file support (#1)

* Refactor

* Upgrade with support to .csv

* Updated Readme.md

Co-authored-by: Matteo Gazzadi Poggioli <ITGAZZADIM@tetrapak.com>

README.md

@@ -12,4 +12,15 @@ This repository provides `kuberes` tool.
 Here is a demo of `kuberes`:
 ![kuberes demo](docs/kuberes-demo.gif)
 
+## Arguments
 
+`kuberes` accept multiple arguments allowing to produce report directly in console or to a `.csv` file.
+
+Arguments can be passed usin `-` or `--`.
+
+|    Argument   |  Type  |  Default |                Description                |
+| ------------- | ------ | -------- | ----------------------------------------- |
+| `output`      | String | `table`  |  Output type. Valid values are: table,csv |
+| `group-by-ns` |  Bool  | `true`   |  Should group statistics by namespace ?   |
+| `csv-path`    | String | ``       |  Full Path to the .CSV File to produce    |
+

cmd/kuberes/main.go

@@ -2,111 +2,93 @@ package main
 
 import (
 	"context"
-	"log"
+	"flag"
+	"fmt"
 	"sort"
+	"strings"
+	"time"
 
+	calc "github.com/matteogazzadi/kuberes/pkg/calculator"
 	domain "github.com/matteogazzadi/kuberes/pkg/domain"
-	k8shelper "github.com/matteogazzadi/kuberes/pkg/k8shelper"
+	helper "github.com/matteogazzadi/kuberes/pkg/helpers"
 	"k8s.io/client-go/kubernetes"
 	ctrl "sigs.k8s.io/controller-runtime"
-
-	"github.com/rodaine/table"
 )
 
+// Arguments Parameter
+var groupByNamespace bool
+var outputAsTable bool
+var outputAsCsv bool
+var csvOutputFilePath string
+
+// Init Function - Arguments parsing
+func init() {
+	var outputFormat string
+
+	flag.BoolVar(&groupByNamespace, "group-by-ns", true, "Should group statistics by namespace ?")
+	flag.StringVar(&outputFormat, "output", "table", "Output type. Valid values are: table,csv")
+	flag.StringVar(&csvOutputFilePath, "csv-path", "", "Full Path to the .CSV File to produce")
+	flag.Parse()
+
+	// Check if output parameter is valid
+	switch strings.ToLower(outputFormat) {
+	case "table":
+		outputAsTable = true
+		outputAsCsv = false
+	case "csv":
+		outputAsTable = false
+		outputAsCsv = true
+	default:
+		panic("Unrecognized 'output' value '" + outputFormat + "'")
+	}
+
+	// Check if CSV output Path is valid (only if output is CSV)
+	if outputAsCsv && csvOutputFilePath == "" {
+		panic("CSV Output path is not set")
+	}
+}
+
 // Main function - Application Entry Point
 func main() {
 
+	startTime := time.Now().UTC()
+
 	// 1. Initialize Kubernets Clients
 	ctx := context.Background()
 	config := ctrl.GetConfigOrDie()
 	clientset := kubernetes.NewForConfigOrDie(config)
 
-	// Resources Map object
-	resources := make(map[string]*domain.PodStats)
-
-	// Retrieve the list of namespaces in cluster
-	namespaces, err := k8shelper.GetAllNamespace(clientset, ctx)
+	// Get All Pods in cluster
+	pods, err := helper.GetAllPods(clientset, ctx)
 
 	if err != nil {
 		panic(err)
 	}
 
-	// Loop all namespace and calculate POD resources
-	for _, namespace := range namespaces {
-		pods, err := k8shelper.GetPodsByNamespace(clientset, ctx, namespace.Name)
-
-		if err != nil {
-			log.Fatal(err)
-			continue
-		}
-
-		// Check if entry for the given namespace exist in the local map.
-		_, ok := resources[namespace.Name]
-
-		// If not present, add it.
-		if !ok {
-
-			var newStats domain.PodStats
-
-			newStats.Namespace = namespace.Name
-			resources[namespace.Name] = &newStats
-		}
-
-		// Loop on pods
-		for _, pod := range pods {
-
-			stats, _ := resources[namespace.Name]
-
-			// Loop On Containers
-			for _, container := range pod.Spec.Containers {
+	var resources []domain.K8sStats
 
-				// CPU
-				cpuRequest := container.Resources.Requests.Cpu().MilliValue()
-				cpuLimit := container.Resources.Limits.Cpu().MilliValue()
+	// Calculate Resources
+	calc.CalculateResources(groupByNamespace, &pods, &resources)
 
-				// Memory
-				memRequest, _ := container.Resources.Requests.Memory().AsInt64()
-				memLimit, _ := container.Resources.Limits.Memory().AsInt64()
+	// Sort Resources by Namespace
+	sort.Slice(resources, func(i, j int) bool {
+		return resources[i].Namespace < resources[j].Namespace
+	})
 
-				// Convert MB to Mib
-				memRequest = memRequest / 1048576
-				memLimit = memLimit / 1048576
+	// =============== //
+	// Generate Output //
+	// =============== //
 
-				stats.Cpu.Limit += cpuLimit
-				stats.Cpu.Request += cpuRequest
-
-				stats.Memory.Limit += memLimit
-				stats.Memory.Request += memRequest
-			}
+	if outputAsTable {
+		// Table Output
+		helper.WriteOutputAsTable(&resources, groupByNamespace)
+	} else {
+		if outputAsCsv && csvOutputFilePath != "" {
+			helper.WriteOutputAsCsv(&resources, groupByNamespace, csvOutputFilePath)
 		}
-
-	}
-
-	// Sort the resources key alphabetically
-	keys := make([]string, 0, len(resources))
-	for k := range resources {
-		keys = append(keys, k)
 	}
-	sort.Strings(keys)
-
-	var total domain.PodStats
-
-	// Generate the on screen Table
-	tbl := table.New("Namespace", "CPU-Request (mCore)", "CPU-Limit (mCore)", "Memory-Request (Mi)", "Memory-Limit (Mi)")
-
-	for _, k := range keys {
-		stats := resources[k]
-		tbl.AddRow(stats.Namespace, stats.Cpu.Request, stats.Cpu.Limit, stats.Memory.Request, stats.Memory.Limit)
-
-		total.Cpu.Request += stats.Cpu.Request
-		total.Cpu.Limit += stats.Cpu.Limit
-		total.Memory.Request += stats.Memory.Request
-		total.Memory.Limit += stats.Memory.Limit
-	}
-
-	// Add total line
-	tbl.AddRow("-----", "-----", "-----", "-----", "-----")
-	tbl.AddRow("Total", total.Cpu.Request, total.Cpu.Limit, total.Memory.Request, total.Memory.Limit)
 
-	tbl.Print()
+	// Report Elapsed Time
+	fmt.Println()
+	fmt.Println("Elapsed: ", time.Since(startTime).Milliseconds(), "ms")
 }

pkg/calculator/statscalculator.go

@@ -0,0 +1,62 @@
+package calculator
+
+import (
+	"sync"
+
+	domain "github.com/matteogazzadi/kuberes/pkg/domain"
+	helper "github.com/matteogazzadi/kuberes/pkg/helpers"
+	v1 "k8s.io/api/core/v1"
+)
+
+// Calculate Resources
+func CalculateResources(groupByNamespace bool, pods *[]v1.Pod, resources *[]domain.K8sStats) {
+
+	// Statistics Channel
+	k8sStatsChan := make(chan domain.K8sStats, len(*pods))
+
+	// Wait Group
+	var wg sync.WaitGroup
+
+	for _, pod := range *pods {
+		wg.Add(1)
+		go helper.GetPodStats(pod, k8sStatsChan, &wg)
+	}
+
+	wg.Wait()
+	close(k8sStatsChan)
+
+	if groupByNamespace {
+		resByNs := make(map[string]*domain.K8sStats)
+
+		for stats := range k8sStatsChan {
+
+			curStats, ok := resByNs[stats.Namespace]
+
+			if ok {
+				// Update current stats
+				curStats.Cpu.Limit += stats.Cpu.Limit
+				curStats.Cpu.Request += stats.Cpu.Request
+				curStats.Memory.Limit += stats.Memory.Limit
+				curStats.Memory.Request += stats.Memory.Request
+			} else {
+				// Create new stats
+				var newStats domain.K8sStats
+
+				newStats.Namespace = stats.Namespace
+				newStats.Cpu = stats.Cpu
+				newStats.Memory = stats.Memory
+
+				resByNs[stats.Namespace] = &newStats
+			}
+		}
+
+		for _, stats := range resByNs {
+			*resources = append(*resources, *stats)
+		}
+
+	} else {
+		for stats := range k8sStatsChan {
+			*resources = append(*resources, stats)
+		}
+	}
+}

pkg/domain/stats.go

@@ -1,6 +1,7 @@
 package domain
 
-type PodStats struct {
+type K8sStats struct {
+	PodName   string
 	Namespace string
 	Cpu       Resource
 	Memory    Resource

pkg/helpers/k8shelper.go

@@ -0,0 +1,73 @@
+package helpers
+
+import (
+	"context"
+	"log"
+	"sync"
+
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes"
+)
+
+// Get all Pods for a given namespace
+func GetPodsByNamespace(clientset *kubernetes.Clientset, ctx context.Context, namespace string, podChan chan<- []v1.Pod, wg *sync.WaitGroup) {
+
+	defer wg.Done()
+
+	// Get pods
+	pods, err := clientset.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{})
+
+	if err != nil {
+		log.Fatal(err)
+	} else {
+		podChan <- pods.Items
+	}
+}
+
+// Get All Namespaces in the cluster
+func GetAllNamespace(clientset *kubernetes.Clientset, ctx context.Context) ([]v1.Namespace, error) {
+	namespaces, err := clientset.CoreV1().Namespaces().List(ctx, metav1.ListOptions{})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return namespaces.Items, nil
+}
+
+// Get all Pods
+func GetAllPods(clientset *kubernetes.Clientset, ctx context.Context) ([]v1.Pod, error) {
+
+	namespaces, err := GetAllNamespace(clientset, ctx)
+
+	if err != nil {
+		log.Fatal(err)
+		return nil, err
+	}
+
+	// Pod Channel
+	podChannel := make(chan []v1.Pod, len(namespaces))
+
+	// Wait Group
+	var wg sync.WaitGroup
+
+	// Loop on Namespace
+	for _, namespace := range namespaces {
+		wg.Add(1)
+		go GetPodsByNamespace(clientset, ctx, namespace.Name, podChannel, &wg)
+	}
+
+	wg.Wait()
+	close(podChannel)
+
+	var pods []v1.Pod
+
+	for podList := range podChannel {
+		for _, pod := range podList {
+			pods = append(pods, pod)
+		}
+	}
+
+	return pods, nil
+}