Save runtime profile data periodically.

Arvados-DCO-1.1-Signed-off-by: Tom Clegg <tom@curii.com>

diff --git a/export.go b/export.go
index 1668cae007842d2fca2dc3b86bba1b3a16904a3a..7ad24cb915091272c45cd2256c3c216a9e7578f3 100644 (file)
--- a/export.go
+++ b/export.go
@@ -52,6 +52,7 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
        flags := flag.NewFlagSet("", flag.ContinueOnError)
        flags.SetOutput(stderr)
        pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
+       pprofdir := flags.String("pprof-dir", "", "write Go profile data to `directory` periodically")
        runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
        projectUUID := flags.String("project", "", "project `UUID` for output data")
        priority := flags.Int("priority", 500, "container request priority")
@@ -81,6 +82,9 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
                        log.Println(http.ListenAndServe(*pprof, nil))
                }()
        }
+       if *pprofdir != "" {
+               go writeProfilesPeriodically(*pprofdir)
+       }
 
        if !*runlocal {
                if *outputFilename != "-" {
@@ -107,6 +111,8 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
                        *outputBed = "/mnt/output/" + *outputBed
                }
                runner.Args = []string{"export", "-local=true",
+                       "-pprof", ":6000",
+                       "-pprof-dir", "/mnt/output",
                        "-ref", *refname,
                        "-output-format", *outputFormatStr,
                        "-output-bed", *outputBed,

diff --git a/profile.go b/profile.go
new file mode 100644 (file)
index 0000000..159c271
--- /dev/null
+++ b/profile.go
@@ -0,0 +1,67 @@
+package lightning
+
+import (
+       "os"
+       "runtime"
+       "runtime/pprof"
+       "time"
+
+       log "github.com/sirupsen/logrus"
+)
+
+func writeProfilesPeriodically(outdir string) {
+       for range time.NewTicker(time.Minute).C {
+               writeMemProfile(outdir)
+               writeCPUProfile(outdir)
+       }
+}
+
+func writeCPUProfile(outdir string) {
+       f, err := os.OpenFile(outdir+"/cpu.prof~", os.O_CREATE|os.O_WRONLY, 0666)
+       if err != nil {
+               log.Print(err)
+               return
+       }
+       defer f.Close()
+       runtime.GC()
+       if err := pprof.StartCPUProfile(f); err != nil {
+               log.Print(err)
+               return
+       }
+       time.Sleep(time.Second)
+       pprof.StopCPUProfile()
+       err = f.Close()
+       if err != nil {
+               log.Print(err)
+               return
+       }
+       err = os.Rename(outdir+"/cpu.prof~", outdir+"/cpu.prof")
+       if err != nil {
+               log.Print(err)
+       }
+       return
+}
+
+func writeMemProfile(outdir string) {
+       f, err := os.OpenFile(outdir+"/mem.prof~", os.O_CREATE|os.O_WRONLY, 0666)
+       if err != nil {
+               log.Print(err)
+               return
+       }
+       defer f.Close()
+       runtime.GC()
+       if err := pprof.WriteHeapProfile(f); err != nil {
+               log.Print(err)
+               return
+       }
+       err = f.Close()
+       if err != nil {
+               log.Print(err)
+               return
+       }
+       err = os.Rename(outdir+"/mem.prof~", outdir+"/mem.prof")
+       if err != nil {
+               log.Print(err)
+       }
+       return
+}