[v16] Add ability to capture pprof profiles to the event-handler

integrations/event-handler/app.go

@@ -18,6 +18,7 @@ package main
 
 import (
 	"context"
+	"path/filepath"
 	"time"
 
 	"github.com/gravitational/trace"
@@ -27,6 +28,7 @@ import (
 	"github.com/gravitational/teleport/integrations/lib"
 	"github.com/gravitational/teleport/integrations/lib/backoff"
 	"github.com/gravitational/teleport/integrations/lib/logger"
+	"github.com/gravitational/teleport/lib/integrations/diagnostics"
 )
 
 // App is the app structure
@@ -247,3 +249,9 @@ func (a *App) RegisterSession(ctx context.Context, e *TeleportEvent) {
 		log.Error("Registering session: ", err)
 	}
 }
+
+func (a *App) Profile() {
+	if err := diagnostics.Profile(filepath.Join(a.Config.StorageDir, "profiles")); err != nil {
+		logrus.WithError(err).Warn("failed to capture profiles")
+	}
+}

integrations/lib/signals.go

@@ -41,6 +41,7 @@ func ServeSignals(app Terminable, shutdownTimeout time.Duration) {
 	signal.Notify(sigC,
 		syscall.SIGTERM, // graceful shutdown
 		syscall.SIGINT,  // graceful-then-fast shutdown
+		syscall.SIGUSR1, // capture pprof profiles
 	)
 	defer signal.Stop(sigC)
 
@@ -67,6 +68,10 @@ func ServeSignals(app Terminable, shutdownTimeout time.Duration) {
 			}
 			go gracefulShutdown()
 			alreadyInterrupted = true
+		case syscall.SIGUSR1:
+			if p, ok := app.(interface{ Profile() }); ok {
+				go p.Profile()
+			}
 		}
 	}
 }

lib/integrations/diagnostics/profile.go

@@ -0,0 +1,111 @@
+// Teleport
+// Copyright (C) 2024 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package diagnostics
+
+import (
+	"os"
+	"path/filepath"
+	"runtime"
+	"runtime/pprof"
+	runtimetrace "runtime/trace"
+	"strconv"
+	"time"
+
+	"github.com/gravitational/trace"
+	"github.com/sirupsen/logrus"
+)
+
+// Profile captures various Go pprof profiles and writes
+// them to the profivided directory. All profiles are prefixed
+// with the same epoch time so that profiles can easily be associated
+// as being captured from the same call.
+func Profile(dir string) error {
+	if err := os.MkdirAll(dir, 0o755); err != nil {
+		return trace.Wrap(err, "creating profile directory %v", dir)
+	}
+
+	timestamp := strconv.FormatInt(time.Now().UnixMilli(), 10)
+	traceFile, err := os.Create(filepath.Join(dir, timestamp+"-trace.profile"))
+	if err != nil {
+		return trace.Wrap(err, "creating trace proile file")
+	}
+	defer traceFile.Close()
+
+	cpuFile, err := os.Create(filepath.Join(dir, timestamp+"-cpu.profile"))
+	if err != nil {
+		return trace.Wrap(err, "creating cpu proile file")
+	}
+	defer cpuFile.Close()
+
+	heapFile, err := os.Create(filepath.Join(dir, timestamp+"-heap.profile"))
+	if err != nil {
+		return trace.Wrap(err, "creating heap proile file")
+	}
+	defer heapFile.Close()
+
+	goroutineFile, err := os.Create(filepath.Join(dir, timestamp+"-goroutine.profile"))
+	if err != nil {
+		return trace.Wrap(err, "creating goroutine proile file")
+	}
+	defer goroutineFile.Close()
+
+	blockFile, err := os.Create(filepath.Join(dir, timestamp+"-block.profile"))
+	if err != nil {
+		return trace.Wrap(err, "creating block proile file")
+	}
+	defer blockFile.Close()
+
+	logrus.Debugf("capturing trace profile to %s", traceFile.Name())
+
+	if err := runtimetrace.Start(traceFile); err != nil {
+		return trace.Wrap(err, "capturing trace profile")
+	}
+
+	logrus.Debugf("capturing cpu profile to %s", cpuFile.Name())
+
+	if err := pprof.StartCPUProfile(cpuFile); err != nil {
+		return trace.Wrap(err, "capturing cpu profile")
+	}
+
+	defer func() {
+		logrus.Debugf("capturing goroutine profile to %s", cpuFile.Name())
+
+		if err := pprof.Lookup("goroutine").WriteTo(goroutineFile, 0); err != nil {
+			logrus.WithError(err).Warn("failed to capture goroutine profile")
+		}
+
+		logrus.Debugf("capturing block profile to %s", cpuFile.Name())
+
+		if err := pprof.Lookup("block").WriteTo(blockFile, 0); err != nil {
+			logrus.WithError(err).Warn("failed to capture block profile")
+		}
+
+		runtime.GC()
+
+		logrus.Debugf("capturing heap profile to %s", cpuFile.Name())
+
+		if err := pprof.WriteHeapProfile(heapFile); err != nil {
+			logrus.WithError(err).Warn("failed to capture heap profile")
+		}
+
+		pprof.StopCPUProfile()
+		runtimetrace.Stop()
+	}()
+
+	<-time.After(30 * time.Second)
+	return nil
+}