gistfile1.txt

package main

import (
	"flag"
	"fmt"
	v1 "k8s.io/api/core/v1"
	"os"
	"path/filepath"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
	"golang.org/x/sync/errgroup"
	"github.com/urfave/cli"
)


const loadTestLabel = "load-test"

func main() {
	app := cli.NewApp()
	app.Name = "flood"
	app.Commands = []cli.Command{
		UpCommand(),
		DownCommand(),
	}
	app.Run(os.Args)
}


func UpCommand() cli.Command {
	upFlags := []cli.Flag{
		cli.IntFlag{
			Name:  "worker-threads",
			Usage: "Number of threads to use for resources creation",
			Value: 10,
		},
		cli.IntFlag{
			Name:  "namespaces",
			Usage: "Number of namespaces to create",
			Value: 10,
		},
		cli.IntFlag{
			Name:  "pods",
			Usage: "Number of pods per namespace",
			Value: 1,
		},
	}

	upFlags = append(upFlags)
	return cli.Command{
		Name:   "up",
		Usage:  "Bring resources up",
		Action: createResources,
		Flags:  upFlags,
	}
}


func DownCommand() cli.Command {
	downFlags := []cli.Flag{}
	downFlags = append(downFlags)

	return cli.Command{
		Name:   "down",
		Usage:  "Tear down resources",
		Action: deleteResources,
		Flags:  downFlags,
	}
}

func deleteResources(ctx *cli.Context) error {
	clientset, err := getClient()
	if err != nil {
		return err
	}

	opts := metav1.ListOptions{
		LabelSelector: fmt.Sprintf("%s=true", loadTestLabel),
	}
	namespaces, err := clientset.CoreV1().Namespaces().List(opts)
	if err != nil {
		return err
	}
	fmt.Printf("Found %v namespaces to delete\n", len(namespaces.Items))
	for _, ns := range namespaces.Items {
		if err := clientset.CoreV1().Namespaces().Delete(ns.Name, &metav1.DeleteOptions{}); err != nil {
			fmt.Printf("Failed to cleanup namespace %s: [%v]", ns.Name, err)
		}
		fmt.Printf("Removed namespace %s\n", ns.Name)
	}
	return nil
}

func createResources(ctx *cli.Context) error {
	clientset, err := getClient()
	if err != nil {
		return err
	}

	fmt.Printf("Creating namespaces\n")
	labels := map[string]string{loadTestLabel: "true"}
	var errgrp errgroup.Group
	workerThreads :=  ctx.Int("worker-threads")
	for w := 0; w < workerThreads; w++ {
		errgrp.Go(func() error {
			return createResource(labels, clientset)
		})
	}

	if err := errgrp.Wait(); err != nil {
		return err
	}

	return nil
}

func createResource(labels map[string]string, clientset *kubernetes.Clientset) error {
	var errList []error
	namespace := &v1.Namespace{
		TypeMeta: metav1.TypeMeta{},
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "load-test-",
			Labels:       labels,
		},
	}
	_, err := clientset.CoreV1().Namespaces().Create(namespace)
	if err != nil {
		errList = append(errList, err)
	}
	return listErr(errList)
}

func getClient() (*kubernetes.Clientset, error) {
	var kubeconfig *string
	if home := os.Getenv("HOME"); home != "" {
		kubeconfig = flag.String("kubeconfig", filepath.Join(home, ".kube", "config"), "(optional) absolute path to the kubeconfig file")
	} else {
		kubeconfig = flag.String("kubeconfig", "", "absolute path to the kubeconfig file")
	}
	flag.Parse()

	config, err := clientcmd.BuildConfigFromFlags("", *kubeconfig)
	if err != nil {
		return nil, err
	}

	return kubernetes.NewForConfig(config)
}


func listErr(e []error) error {
	if len(e) > 0 {
		return fmt.Errorf("%v", e)
	}
	return nil
}