pprof (#135)

Co-authored-by: Luke Lombardi <[email protected]>

internal/gateway/gateway.go

@@ -8,7 +8,10 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"runtime"
+	"runtime/pprof"
 	"syscall"
+	"time"
 
 	"github.com/beam-cloud/beta9/internal/abstractions/endpoint"
 	"github.com/beam-cloud/beta9/internal/task"
@@ -301,6 +304,35 @@ func (g *Gateway) registerServices() error {
 	return nil
 }
 
+func (g *Gateway) profile() {
+	go profileCPU(5 * time.Minute)
+	go profileMemory(5 * time.Minute)
+}
+
+func profileCPU(duration time.Duration) {
+	cpuFile, err := os.Create("cpu_profile.prof")
+	if err != nil {
+		log.Fatalf("Could not create CPU profile: %v", err)
+	}
+	defer cpuFile.Close()
+
+	pprof.StartCPUProfile(cpuFile)
+	time.Sleep(duration)
+	pprof.StopCPUProfile()
+}
+
+func profileMemory(duration time.Duration) {
+	time.Sleep(duration)
+	memFile, err := os.Create("mem_profile.prof")
+	if err != nil {
+		log.Fatalf("Could not create memory profile: %v", err)
+	}
+	defer memFile.Close()
+
+	runtime.GC()
+	pprof.WriteHeapProfile(memFile)
+}
+
 // Gateway entry point
 func (g *Gateway) Start() error {
 	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", g.config.GatewayService.GRPCPort))
@@ -330,6 +362,8 @@ func (g *Gateway) Start() error {
 		}
 	}()
 
+	g.profile() // Profile cpu/mem usage
+
 	go func() {
 		if err := g.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
 			log.Fatalf("Failed to start http server: %v", err)