add cpu profiling support

1 files changed, 21 insertions, 0 deletions

diff --git a/main.go b/main.go
index 6f97efa..fa99806 100644
--- a/main.go
+++ b/main.go
@@ -5,6 +5,7 @@ import (
         "fmt"
         "net/http"
         "os"
+        "runtime/pprof"
 )
 
 var base_link string
@@ -22,6 +23,7 @@ func usage() {
         fmt.Println("  -F, --listfav\tlists favourited canvas courses with name and ID")
         fmt.Println("  -v, --verbose\tprints with verbose output")
         fmt.Println("  -h, --help\tdisplays this help message")
+        fmt.Println("  -p, --profile\twrites cpu profile to a file")
 }
 
 func main() {
@@ -32,6 +34,7 @@ func main() {
         var listall bool
         var listfav bool
         var help bool
+        var cpuprofile string
 
         flag.StringVar(&code, "code", "", "5 character attendance code")
         flag.StringVar(&code, "c", "", "5 character attendance code")
@@ -45,9 +48,27 @@ func main() {
         flag.BoolVar(&verbose, "v", false, "verbose output")
         flag.BoolVar(&help, "help", false, "view usage message")
         flag.BoolVar(&help, "h", false, "view usage message")
+        flag.StringVar(&cpuprofile, "p", "", "CPU profiling report location")
+        flag.StringVar(&cpuprofile, "profile", "", "CPU profiling report location")
 
         flag.Parse()
 
+        if cpuprofile != "" {
+                cpuf, err := os.Create(cpuprofile)
+                if err != nil {
+                        fmt.Println(err)
+                        os.Exit(1)
+                }
+                defer cpuf.Close()
+
+                err = pprof.StartCPUProfile(cpuf)
+                if err != nil {
+                        fmt.Println(err)
+                        os.Exit(1)
+                }
+                defer pprof.StopCPUProfile()
+        }
+
         if code == "" && course == -1 && !listall && !listfav && !help {
                 usage()
                 os.Exit(1)