profiler

cmd/wal-g/main.go

@@ -9,22 +9,27 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"runtime/pprof"
 	"time"
 )
 
+var profile bool
+var mem bool
 var help bool
 var l *log.Logger
-var helpMsg = "\tbackup-fetch\tfetch a backup from S3\n" +
-	"\tbackup-push\tstarts and uploads a finished backup to S3\n" +
-	"\twal-fetch\tfetch a WAL file from S3\n" +
-	"\twal-push\tupload a WAL file to S3\n"
+var helpMsg = "  backup-fetch\tfetch a backup from S3\n" +
+	"  backup-push\tstarts and uploads a finished backup to S3\n" +
+	"  wal-fetch\tfetch a WAL file from S3\n" +
+	"  wal-push\tupload a WAL file to S3\n"
 
 func init() {
 	flag.Usage = func() {
 		fmt.Fprintf(os.Stderr, "Usage of WAL-G:\n")
 		fmt.Fprintf(os.Stderr, "%s", helpMsg)
 		flag.PrintDefaults()
 	}
+	flag.BoolVar(&profile, "p", false, "\tProfiler (false on default)")
+	flag.BoolVar(&mem, "m", false, "\tMemory profiler (false on default)")
 	l = log.New(os.Stderr, "", 0)
 }
 
@@ -44,6 +49,16 @@ func main() {
 		backupName = all[2]
 	}
 
+	// Various profiling options
+	if profile {
+		f, err := os.Create("cpu.prof")
+		if err != nil {
+			log.Fatal(err)
+		}
+		pprof.StartCPUProfile(f)
+		defer pprof.StopCPUProfile()
+	}
+
 	tu, pre, err := walg.Configure()
 	if err != nil {
 		log.Fatalf("FATAL: %+v\n", err)
@@ -81,7 +96,7 @@ func main() {
 				log.Fatalf("Backup '%s' does not exist.\n", *bk.Name)
 			}
 
-		// Find the LATEST valid backup (checks against JSON file and grabs backup name) and extract to DIRARC.
+			// Find the LATEST valid backup (checks against JSON file and grabs backup name) and extract to DIRARC.
 		} else if backupName == "LATEST" {
 			bk = &walg.Backup{
 				Prefix: pre,
@@ -158,6 +173,16 @@ func main() {
 			}
 		}
 
+		if mem {
+			f, err := os.Create("mem.prof")
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			pprof.WriteHeapProfile(f)
+			defer f.Close()
+		}
+
 	} else if command == "wal-fetch" {
 		// Fetch and decompress a WAL file from S3.
 		a := &walg.Archive{

extract.go

@@ -133,7 +133,6 @@ func ExtractAll(ti TarInterpreter, files []ReaderMaker) error {
 			pr, tempW := io.Pipe()
 			pw := &EmptyWriteIgnorer{tempW}
 
-
 			// Collect errors returned by tarHandler or parsing.
 			collectLow := make(chan error)