From 63df3adab121b7e48fd86b7e889a8874b0e0bfd2 Mon Sep 17 00:00:00 2001
From: Alex Yatskov <alex@foosoft.net>
Date: Fri, 31 Jul 2015 15:33:09 +0900
Subject: [PATCH] Adding some profiling code

---
 server.go | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/server.go b/server.go
index dc0c204..5b9ea60 100644
--- a/server.go
+++ b/server.go
@@ -30,9 +30,13 @@ import (
 	"log"
 	"math"
 	"net/http"
+	"os"
+	"os/signal"
 	"runtime"
+	"runtime/pprof"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/GaryBoone/GoStats/stats"
@@ -309,6 +313,7 @@ func main() {
 	staticDir := flag.String("static", "static", "path to static files")
 	portNum := flag.Int("port", 8080, "port to serve content on")
 	dataSrc := flag.String("data", "hscd@/hscd", "data source for database")
+	profile := flag.String("profile", "", "write cpu profile to file")
 	flag.Parse()
 
 	var err error
@@ -317,6 +322,25 @@ func main() {
 	}
 	defer db.Close()
 
+	if *profile != "" {
+		f, err := os.Create(*profile)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		pprof.StartCPUProfile(f)
+
+		c := make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+
+		go func() {
+			<-c
+			log.Print("interrupted")
+			pprof.StopCPUProfile()
+			os.Exit(1)
+		}()
+	}
+
 	http.HandleFunc("/query", executeQuery)
 	http.HandleFunc("/categories", getCategories)
 	http.HandleFunc("/learn", addCategory)