Merge pull request #36 from grantseltzer/debug-log-improvements

Add color to debug printing, and add log for loading all eBPF programs

Author: grantseltzer

cmd/weaver/functions_file.go

@@ -22,9 +22,8 @@ func read_functions_file(path string) ([]functionTraceContext, error) {
 			continue
 		}
 
-		if globalDebug {
-			fmt.Println("parsing: " + funcString)
-		}
+		debugLog("parsing: %s\n", funcString)
+
 		err := parseFunctionAndArgumentTypes(&contexts[i], funcString)
 		if err != nil {
 			return nil, fmt.Errorf("could not parse function string '%s': %s", funcString, err.Error())

cmd/weaver/load_uprobe.go

@@ -8,6 +8,7 @@ import (
 	"log"
 	"math"
 	"strings"
+	"sync"
 	"text/template"
 
 	bpf "github.com/iovisor/gobpf/bcc"
@@ -104,18 +105,18 @@ func bpfText(context *functionTraceContext) string {
 	buf := new(bytes.Buffer)
 	t.Execute(buf, context)
 
-	if globalDebug {
-		// Print eBPF text
-		fmt.Println(buf.String())
-	}
+	// Print eBPF text
+	debugLog("%s\n", buf.String())
 
 	return buf.String()
 }
 
 // loadUprobeAndBPFModule will, based on the passed context, install the bpf program and attach a uprobe to the specified function
 // It then prints results to the designated output stream.
 // This blocks until Ctrl-C or error occurs.
-func loadUprobeAndBPFModule(traceContext *functionTraceContext, runtimeContext context.Context) error {
+func loadUprobeAndBPFModule(traceContext *functionTraceContext, runtimeContext context.Context, wg *sync.WaitGroup) error {
+
+	defer runtimeContext.Err()
 
 	// Load eBPF filter and uprobe
 	filterText := bpfText(traceContext)
@@ -203,6 +204,7 @@ func loadUprobeAndBPFModule(traceContext *functionTraceContext, runtimeContext c
 		}
 	}()
 
+	wg.Done()
 	perfMap.Start()
 	<-runtimeContext.Done()
 	perfMap.Stop()

cmd/weaver/main.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"sync"
 
 	"github.com/urfave/cli/v2"
 )
@@ -109,15 +110,22 @@ func entry(c *cli.Context) error {
 	runtimeContext, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
+	// wg is used to communicate back with the main thread when
+	// all uprobe/eBPFs are installed
+	var wg sync.WaitGroup
+	wg.Add(len(contexts))
+
 	// Install eBPF program for each function to trace
 	for i := range contexts {
-
 		contexts[i].binaryName = binaryFullPath
-
-		// Load uprobe and BPF code. This will block until Ctrl-C or an error occurs.
-		go loadUprobeAndBPFModule(&contexts[i], runtimeContext)
+		go loadUprobeAndBPFModule(&contexts[i], runtimeContext, &wg)
 	}
 
+	go func() {
+		wg.Wait()
+		debugLog("All probes installed\n")
+	}()
+
 	<-sig
 
 	return nil

cmd/weaver/output.go

@@ -49,3 +49,9 @@ func printOutput(o output) error {
 
 	return nil
 }
+
+func debugLog(format string, a ...interface{}) {
+	if globalDebug {
+		fmt.Fprintf(os.Stderr, "\x1b[96m"+format+"\x1b[0m", a...)
+	}
+}