Improvements

- `kubectl-p` Replace go routines with the use of the `errgroup` module. This allows me to run multiple routines with a wait group still and handle the errors. I no longer need to muck around with channels myself, nor having to add the count of go routines to a wait group beforehand. It's cleaner and leads to simpler code. Removed the error handling in the defer statements since trying to return values from deferred functions can lead to strange results. If they error, just display it at the time. - `ssm` Use `strings.Builder` for a more efficient way of building up strings. Move the line wrapping logic out of the `WrapTextToWidth` function and into its own `WrapLine` function. Minor fix around a width calculation when tabs are involved. Created a `lineVisualWidth` function to make things a little easier.
jim-barber-he · Jul 21, 2024 · 3076707 · 3076707
1 parent 433dc4a
commit 3076707
Show file tree

Hide file tree

Showing 4 changed files with 107 additions and 99 deletions.
diff --git a/go.mod b/go.mod
@@ -13,6 +13,7 @@ require (
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/pflag v1.0.5
+	golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9
 	golang.org/x/term v0.18.0
 	k8s.io/api v0.30.2
 	k8s.io/apimachinery v0.30.2

diff --git a/go.sum b/go.sum
@@ -134,6 +134,7 @@ golang.org/x/oauth2 v0.10.0 h1:zHCpF2Khkwy4mMB4bv0U37YtJdTGW8jI0glAApi0Kh8=
 golang.org/x/oauth2 v0.10.0/go.mod h1:kTpgurOux7LqtuxjuyZa4Gj2gdezIt/jQtGnNFfypQI=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9 h1:SQFwaSi55rU7vdNs9Yr0Z324VNlrF+0wMqRXT4St8ck=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

diff --git a/kubectl-plugins/kubectl-p/main.go b/kubectl-plugins/kubectl-p/main.go
@@ -14,12 +14,12 @@ import (
 	"runtime"
 	"runtime/pprof"
 	"slices"
-	"sync"
 
 	"github.com/jim-barber-he/go/k8s"
 	"github.com/jim-barber-he/go/texttable"
 	"github.com/jim-barber-he/go/util"
 	flag "github.com/spf13/pflag"
+	"golang.org/x/sync/errgroup"
 	v1 "k8s.io/api/core/v1"
 )
 
@@ -88,9 +88,6 @@ func main() {
 // Error handling isn't perfect here, and not sure how to do it better.
 // If an error is returned early, then I guess any errors from the defer functions will be lost.
 func run(opts options) error {
-	// If the defer anonymous functions encounter an error, they can set this var to be returned to the calling function.
-	var deferErrors error
-
 	// CPU profiling.
 	if opts.profileCPU != "" {
 		fp, err := os.Create(opts.profileCPU)
@@ -99,7 +96,7 @@ func run(opts options) error {
 		}
 		defer func(fp *os.File) {
 			if err := fp.Close(); err != nil {
-				deferErrors = errors.Join(err, deferErrors)
+				log.Println(err)
 			}
 		}(fp)
 		if err := pprof.StartCPUProfile(fp); err != nil {
@@ -122,52 +119,35 @@ func run(opts options) error {
 	}
 
 	// Fetch the log of nodes and pods in parallel.
-	var wg sync.WaitGroup
-	wg.Add(2)
+	g := new(errgroup.Group)
 
-	errGetNodes := make(chan error)
-	defer close(errGetNodes)
 	nodes := make(map[string]*v1.Node)
-	go func(errChan chan<- error) {
-		defer wg.Done()
-
+	g.Go(func() error {
 		nodeList, err := k8s.ListNodes(clientset)
 		if err != nil {
-			errChan <- err
-			return
+			return err
 		}
 		for i, node := range nodeList.Items {
 			nodes[node.Name] = &nodeList.Items[i]
 		}
-	}(errGetNodes)
+		return nil
+	})
 
-	errGetPods := make(chan error)
-	defer close(errGetPods)
 	var pods *v1.PodList = &v1.PodList{}
-	go func(pods *v1.PodList, errChan chan<- error) {
-		defer wg.Done()
-
+	g.Go(func() error {
 		listPods, err := k8s.ListPods(clientset, namespace, opts.labelSelector)
 		if err != nil {
-			errChan <- err
-			return
+			return err
 		}
 		if len(listPods.Items) == 0 {
-			errChan <- errNoPodsFound
-			return
+			return errNoPodsFound
 		}
 		*pods = *listPods
-	}(pods, errGetPods)
-
-	wg.Wait()
+		return nil
+	})
 
-	select {
-	case err := <-errGetNodes:
+	if err := g.Wait(); err != nil {
 		return err
-	case err := <-errGetPods:
-		return err
-	default:
-		// No errors.
 	}
 
 	var tbl texttable.Table[*tableRow]
@@ -227,7 +207,7 @@ func run(opts options) error {
 		}
 		defer func(fp *os.File) {
 			if err := fp.Close(); err != nil {
-				deferErrors = errors.Join(err, deferErrors)
+				log.Println(err)
 			}
 		}(fp)
 		// Get up-to-date statistics.
@@ -237,5 +217,5 @@ func run(opts options) error {
 		}
 	}
 
-	return deferErrors
+	return nil
 }
diff --git a/util/util.go b/util/util.go
@@ -18,7 +18,7 @@ const (
 	numSecondsPerDay    = 24 * numSecondsPerHour
 	numSecondsPerWeek   = 7 * numSecondsPerDay
 
-	tabWidth = 8
+	tabStopWidth = 8
 )
 
 // FormatAge returns the age in a human readable format of the first 2 non-zero time units from weeks to seconds,
@@ -104,95 +104,121 @@ func TimeTaken(start time.Time, name string) {
 	fmt.Fprintf(os.Stderr, "%s took %s\n", name, time.Since(start))
 }
 
+// lineVisualWidth returns the visual width of a line with a string added taking tab stops into account.
+// The line position of where the string is to be written is passed in since it affects the tabstop width at that point.
+func lineVisualWidth(linePos int, str string) int {
+	width := linePos
+	for _, r := range str {
+		if r == '\t' {
+			width += tabStopWidth - (width % tabStopWidth)
+		} else {
+			width++
+		}
+	}
+	return width
+}
+
+// WrapLine takes a string representing a single line and wraps it to a specified width.
+// Any tab characters are handled based on the tabstop they'd pad out to.
+func WrapLine(str string, width int) string {
+	if len(str) == 0 {
+		return ""
+	}
+
+	var currentLine, wrappedLine strings.Builder
+
+	pos := 0
+	prevWord := ""
+	for _, word := range strings.Split(str, " ") {
+		wordLength := len(word)
+
+		// Preserve leading spaces which end up as empty words when split on a space.
+		if wordLength == 0 {
+			word = " "
+		}
+
+		// If the length of the current line + a space + the length of the word doesn't fit in the width, then
+		// write out the current line, so that the word just read will start on a new line.
+		// lineVisualWidth() is used since if a word contains tabs it will likely be visually wider.
+		if pos > 0 && lineVisualWidth(pos+1, word) > width {
+			wrappedLine.WriteString(currentLine.String() + "\n")
+			currentLine.Reset()
+			pos = 0
+		}
+
+		// Pad the new word with a leading space unless it is a space itself, or the previous word was a space.
+		if pos > 0 && prevWord != " " && word != " " {
+			currentLine.WriteString(" ")
+			pos++
+		}
+
+		// If the word contains tabs, write it character by character incrementing pos as we go, except when we
+		// get to a tab character where we increment pos by the number of characters to reach the next tabstop.
+		if strings.Contains(word, `\t`) {
+			for _, r := range word {
+				currentLine.WriteRune(r)
+				if r == '\t' {
+					pos += tabStopWidth - (pos % tabStopWidth)
+				} else {
+					pos++
+				}
+			}
+		} else {
+			currentLine.WriteString(word)
+			pos += wordLength
+		}
+		prevWord = word
+	}
+
+	// Add the last line we were working on.
+	if pos > 0 {
+		wrappedLine.WriteString(currentLine.String())
+	}
+
+	return wrappedLine.String()
+}
+
 // WrapTextToWidth takes paragraphs in a multi-line string and writes them into a string, wrapping before the specified
 // width when possible, without breaking up words.
 // It's not possible to wrap in time if a single word is longer than the width.
 // Multiple lines of the paragraph are joined into the longest line that fits within the width before wrapping.
 // Each paragraph in the multi-line string is represented by a blank line between them.
 // The constructed string is then returned.
 func WrapTextToWidth(width int, str string) string {
-	var finalStr string
-	var paragraphs []string
-
 	// Get rid of an initial and final newlines on the string if it has them.
 	str = strings.TrimPrefix(str, string('\n'))
 	str = strings.TrimRight(str, string('\n'))
 
 	// Produce an array of paragraphs joining multi-line paragraphs into a single line.
-	var paragraph string
+	var paragraph strings.Builder
+	var paragraphs []string
 	for _, line := range strings.Split(str, "\n") {
 		if line != "" {
 			// Add the line just read to the existing line since it is part of the same paragraph.
-			if paragraph != "" {
-				paragraph += " "
+			if paragraph.Len() != 0 {
+				paragraph.WriteString(" ")
 			}
-			paragraph += line
+			paragraph.WriteString(line)
 		} else {
 			// When we encounter a blank line, we have completed a paragraph.
-			paragraphs = append(paragraphs, paragraph)
-			paragraph = ""
+			paragraphs = append(paragraphs, paragraph.String())
+			paragraph.Reset()
 		}
 	}
 	// Add the last paragraph.
-	if paragraph != "" {
-		paragraphs = append(paragraphs, paragraph)
+	if paragraph.Len() != 0 {
+		paragraphs = append(paragraphs, paragraph.String())
 	}
 
-	// Process each paragraph, and break it up into words by splitting on the space character.
-	// Write out words on a line until they are going to exceed the width, at which point start a new line.
-	// At the end of each paragraph append a blank line (except for the final paragraph).
+	// Format the paragraphs to the width we want.
+	var finalStr strings.Builder
 	lastParagraph := len(paragraphs) - 1
 	for i, line := range paragraphs {
-		pos := 0
-		prevWord := ""
-		for _, word := range strings.Split(line, " ") {
-			wordLength := len(word)
-			// Preserve leading spaces which end up as empty words when split on a space.
-			if wordLength == 0 {
-				word = " "
-				wordLength++
-			}
-			if pos+1+wordLength > width {
-				// If characters written on the line + an added space + the word > width,
-				// then will need to move to a new line before writing the word.
-				if pos != 0 {
-					finalStr += fmt.Sprintln()
-					pos = 0
-				}
-			} else {
-				// The word will fit with what has already been written plus a space to separate them.
-				// Add a space, but only if we're not at the start of the line; that last thing written
-				// to the line wasn't a space; and that the current word isn't a space.
-				if pos != 0 && prevWord != " " && word != " " {
-					finalStr += " "
-					pos++
-				}
-			}
-
-			// If the word contains tabs, write it character by character incrementing pos as we go, except
-			// when we get to a tab character where we increment pos by the number of characters to reach
-			// the next tab-stop.
-			if strings.Contains(word, `\t`) {
-				for _, char := range word {
-					finalStr += string(char)
-					if char == '\t' {
-						// Add how many characters to the next tab-stop.
-						pos += tabWidth - (pos % tabWidth)
-					} else {
-						pos++
-					}
-				}
-			} else {
-				finalStr += word
-				pos += wordLength
-			}
-			prevWord = word
-		}
-		finalStr += fmt.Sprintln()
+		finalStr.WriteString(WrapLine(line, width) + "\n")
+		// Insert a newline between paragraphs to separate them.
 		if i < lastParagraph {
-			finalStr += fmt.Sprintln()
+			finalStr.WriteString("\n")
 		}
 	}
-
-	return finalStr
+	return finalStr.String()
 }