From 80391ed27965c739b65ca60b5589536cba0516e2 Mon Sep 17 00:00:00 2001
From: Michael McLoughlin <mmcloughlin@gmail.com>
Date: Sun, 25 Apr 2021 22:20:20 -0700
Subject: [PATCH] cmd/addchain: cpu profile option (#88)

Update #25
Update #60
---
 cmd/addchain/search.go | 35 +++++++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/cmd/addchain/search.go b/cmd/addchain/search.go
index 92b6534..8c8f55a 100644
--- a/cmd/addchain/search.go
+++ b/cmd/addchain/search.go
@@ -3,7 +3,10 @@ package main
 import (
 	"context"
 	"flag"
+	"os"
+	"os/signal"
 	"runtime"
+	"runtime/pprof"
 
 	"github.com/google/subcommands"
 
@@ -21,12 +24,13 @@ type search struct {
 
 	concurrency int
 	verbose     bool
+	cpuprofile  string
 }
 
 func (*search) Name() string     { return "search" }
 func (*search) Synopsis() string { return "search for an addition chain." }
 func (*search) Usage() string {
-	return `Usage: search [-v] [-p <N>] <expr>
+	return `Usage: search [-v] [-p <N>] [-cpuprofile <file>] <expr>
 
 Search for an addition chain for <expr>.
 
@@ -36,9 +40,10 @@ Search for an addition chain for <expr>.
 func (cmd *search) SetFlags(f *flag.FlagSet) {
 	f.IntVar(&cmd.concurrency, "p", runtime.NumCPU(), "number of algorithms to run in parallel")
 	f.BoolVar(&cmd.verbose, "v", false, "verbose output")
+	f.StringVar(&cmd.cpuprofile, "cpuprofile", "", "write cpu profile to `file`")
 }
 
-func (cmd *search) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
+func (cmd *search) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{}) (status subcommands.ExitStatus) {
 	if f.NArg() < 1 {
 		return cmd.UsageError("missing expression")
 	}
@@ -55,6 +60,32 @@ func (cmd *search) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{})
 	cmd.Log.Printf("hex: %x", n)
 	cmd.Log.Printf("dec: %s", n)
 
+	// Start profiling.
+	if cmd.cpuprofile != "" {
+		f, err := os.Create(cmd.cpuprofile)
+		if err != nil {
+			return cmd.Fail("could not create cpu profile: %v", err)
+		}
+		defer cmd.CheckClose(&status, f)
+
+		if err := pprof.StartCPUProfile(f); err != nil {
+			return cmd.Fail("could not start cpu profile: %v", err)
+		}
+
+		go func() {
+			c := make(chan os.Signal, 1)
+			signal.Notify(c, os.Interrupt)
+			s := <-c
+
+			cmd.Log.Printf("caught %s: stopping cpu profile", s)
+			pprof.StopCPUProfile()
+
+			os.Exit(0)
+		}()
+
+		defer pprof.StopCPUProfile()
+	}
+
 	// Execute an ensemble of algorithms.
 	ex := exec.NewParallel()
 	if cmd.verbose {