Allow mem profiling without shutting down the server

gadget-inc · Aug 1, 2024 · f4b85a6 · f4b85a6
1 parent aa72f12
commit f4b85a6
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 16 deletions.
diff --git a/Makefile b/Makefile
@@ -125,7 +125,7 @@ server: internal/pb/fs.pb.go internal/pb/fs_grpc.pb.go
 
 server-profile: export DL_ENV=dev
 server-profile: internal/pb/fs.pb.go internal/pb/fs_grpc.pb.go
-	go run cmd/server/main.go --dburi $(DB_URI) --port $(GRPC_PORT) --profile cpu.prof --memprofile mem.pb.gz --log-level info
+	go run cmd/server/main.go --dburi $(DB_URI) --port $(GRPC_PORT) --profile cpu.prof --log-level info
 
 cached: export DL_ENV=dev
 cached: export DL_TOKEN=$(DEV_SHARED_READER_TOKEN)

diff --git a/pkg/cli/server.go b/pkg/cli/server.go
@@ -130,23 +130,35 @@ func NewServerCommand() *cobra.Command {
 			signal.Notify(osSignals, os.Interrupt, syscall.SIGTERM)
 			go func() {
 				<-osSignals
+				s.Grpc.GracefulStop()
+			}()
 
-				if memProfilePath != "" {
-					memProfile, err := os.Create(memProfilePath)
-					if err != nil {
-						logger.Error(ctx, "cannot create heap profile file", zap.Error(err), zap.String("path", memProfilePath))
-					}
-					defer memProfile.Close()
-					runtime.GC()
+			if memProfilePath != "" {
+				memSnapshotSignals := make(chan os.Signal, 1)
+				signal.Notify(memSnapshotSignals, syscall.SIGUSR2)
 
-					err = pprof.WriteHeapProfile(memProfile)
-					if err != nil {
-						logger.Error(ctx, "cannot write heap profile", zap.Error(err), zap.String("path", memProfilePath))
-					}
-				}
+				go func() {
+					for {
+						<-memSnapshotSignals
 
-				s.Grpc.GracefulStop()
-			}()
+						logger.Info(ctx, "SIGUSR2 received, building heap profile", zap.String("path", memProfilePath))
+
+						memProfile, err := os.Create(memProfilePath)
+						if err != nil {
+							logger.Error(ctx, "cannot create heap profile file", zap.Error(err), zap.String("path", memProfilePath))
+						}
+
+						runtime.GC()
+
+						err = pprof.WriteHeapProfile(memProfile)
+						if err != nil {
+							logger.Error(ctx, "cannot write heap profile", zap.Error(err), zap.String("path", memProfilePath))
+						}
+
+						memProfile.Close()
+					}
+				}()
+			}
 
 			logger.Info(ctx, "start fs server", key.Port.Field(port), key.Environment.Field(env.String()))
 			return s.Serve(listen)
@@ -171,7 +183,7 @@ func NewServerCommand() *cobra.Command {
 	flags.StringVar(&encoding, "log-encoding", "console", "Log encoding (console | json)")
 	flags.BoolVar(&tracing, "tracing", false, "Whether tracing is enabled")
 	flags.StringVar(&profilePath, "profile", "", "CPU profile output path (CPU profiling enabled if set)")
-	flags.StringVar(&memProfilePath, "memprofile", "", "Memory profile output path")
+	flags.StringVar(&memProfilePath, "memprofile", "mem.pb.gz", "Memory profile output path")
 
 	flags.IntVar(&port, "port", 5051, "GRPC server port")
 	flags.StringVar(&dbUri, "dburi", "postgres://postgres@127.0.0.1:5432/dl", "Postgres URI")