Go R1 Day 70

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"math"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	shellescape "github.com/alessio/shellescape"
	petname "github.com/dustinkirkland/golang-petname"
	"github.com/google/uuid"
	"github.com/pterm/pterm"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
)

const (
	// exitFail is the exit code if the program
	// fails.
	exitFail = 1

	// desiredPort is the port that the app forwards traffic to.
	desiredPort = 22

	// petnameLength is the length of the petname in words to generate.
	petNameLength = 2

	// startingPort is the starting port for a new connection, and will increment up from there so each connection is unique.
	startingPort = 4000

	// maxRetryCount is the number of times to retry a connection.
	maxRetryCount = 5
)

func main() {
	if err := run(os.Args, os.Stdout); err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		os.Exit(exitFail)
	}
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

// Run handles the arguments being passed in from main, and allows us to run tests against the loading of the code much more easily than embedding all the startup logic in main().
// This is based on Matt Ryers post: https://pace.dev/blog/2020/02/12/why-you-shouldnt-use-func-main-in-golang-by-mat-ryer.html
func run(args []string, stdout io.Writer) error {
	if len(args) == 0 {
		return errors.New("no arguments")
	}
	InitLogger()
	zerolog.SetGlobalLevel(zerolog.InfoLevel)

	debug := flag.Bool("debug", false, "sets log level to debug")
	Count := flag.Int("count", 0, "number of processes to open")
	delaySec := flag.Int("delay", 0, "delay between process creation. Default is 0")
	batchSize := flag.Int("batch", 0, "number of processes to create in each batch. Default is 0 to create all at once")
	Server := flag.String("server", "", "server IP address")
	ServerPort := flag.Int("port", 3000, "server port") //nolint:gomnd

	flag.Parse()
	log.Logger.Info().Int("Count", *Count).
		Int("delaySec", *delaySec).
		Int("batchSize", *batchSize).
		Str("Server", *Server).
		Msg("input parsed")

	log.Logger.Info().
		Int("desiredPort", desiredPort).
		Int("petNameLength", petNameLength).
		Int("startingPort", startingPort).
		Msg("default constants")

	if *debug {
		zerolog.SetGlobalLevel(zerolog.DebugLevel)
	}

	RunTest(*Count, *delaySec, *batchSize, *Server, *ServerPort)
	return nil
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

// InitLogger sets up the logger magic
// By default this is only configured to do pretty console output.
// JSON structured logs are also possible, but not in my default template layout at this time.
func InitLogger() {
	output := zerolog.ConsoleWriter{Out: os.Stdout, TimeFormat: time.RFC3339}
	log.Logger = log.With().Caller().Logger().Output(zerolog.ConsoleWriter{Out: os.Stderr})

	output.FormatLevel = func(i interface{}) string {
		return strings.ToUpper(fmt.Sprintf("| %-6s|", i))
	}
	output.FormatMessage = func(i interface{}) string {
		return fmt.Sprintf("%s", i)
	}
	output.FormatFieldName = func(i interface{}) string {
		return fmt.Sprintf("%s:", i)
	}
	output.FormatFieldValue = func(i interface{}) string {
		return strings.ToUpper(fmt.Sprintf("%s", i))
	}
	log.Info().Msg("logger initialized")
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

// TestBinaryExists checks to see if the binary is found in PATH and exits with failure if can't find it.
func TestBinaryExists(binary string) string {
	p, err := exec.LookPath(binary)
	if err != nil {
		log.Logger.Error().Err(err).Str("binary",binary).Msg("binary not found")
		os.Exit(exitFail)
	}

	return p
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

// buildCliArgs is an example function of building arguments via string slices
func buildCliArgs(Server string, ServerPort int, port int) (command []string) {
	command = append(command, "server")
	command = append(command, "--header")
	command = append(command, fmt.Sprintf(`user-id: %s`, petname.Generate(petNameLength, "-")))
	command = append(command, "--header")
	command = append(command, fmt.Sprintf(`session-id: %s`, uuid.Must(uuid.NewRandom()).String()))
	command = append(command, "--max-retry-count", "5")
	command = append(command, Server+":"+fmt.Sprintf("%d", ServerPort))
	return command
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

// RunTest is the main test function that calculates the batch size and then launches the  creation using a routinue.
func RunTest(Count int, delaySec int, batchSize int, Server string, ServerPort int) {
	log.Logger.Info().Msg("RunTest startings")
	totalBatches := math.Ceil(float64(Count) / float64(batchSize))
	log.Logger.Info().Float64("totalBatches", totalBatches).Msg("batches to run")
	myBinary := TestBinaryExists("binaryname")
	port := startingPort
	var wg sync.WaitGroup

	totals := 0
	p, _ := pterm.DefaultProgressbar.WithTotal(Count).WithTitle("run s").Start()

	for i := 0; i < int(totalBatches); i++ {
		log.Debug().Int("i", i).Int("port", port).Msg("batch number")

		for j := 0; j < batchSize; j++ {
			if totals == Count {
				log.Debug().Msg("totals == Count, breaking out of loop")

				break
			}

			totals++
			log.Debug().Int("i", i).Int("", totals).Msg("")
			cmdargs := buildCliArgs(Server, ServerPort, port)
			wg.Add(1)
			go func() {
				defer wg.Done()
				buf := &bytes.Buffer{}
				cmd := exec.Command(, cmdargs...)
				cmd.Stdout = buf
				cmd.Stderr = buf
				if err := cmd.Run(); err != nil {
					log.Logger.Error().Err(err).Bytes("output", buf.Bytes()).Msg(" failed")
					os.Exit(exitFail)
				}
				log.Logger.Debug().Msgf(" %v", shellescape.QuoteCommand(cmdargs))
				log.Logger.Debug().Bytes("output", buf.Bytes()).Msg("")
			}()

			p.Title = "port: " + fmt.Sprintf("%d", port)
			p.Increment()
			port++
		}
		time.Sleep(time.Second * time.Duration(delaySec))
	}
	p.Title = "s finished"
	_, _ = p.Stop()
	wg.Wait()
}