Merge pull request #6579 from prattmic:runsc_do_profile

PiperOrigin-RevId: 397114051

3 files changed, 126 insertions, 6 deletions

diff --git a/runsc/boot/BUILD b/runsc/boot/BUILD
index b8585c1e9..ff7a5a44b 100644
--- a/runsc/boot/BUILD
+++ b/runsc/boot/BUILD
@@ -15,6 +15,7 @@ go_library(
         "limits.go",
         "loader.go",
         "network.go",
+        "profile.go",
         "strace.go",
         "vfs.go",
     ],
diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index 1dd0048ac..b46d84e5a 100644
--- a/runsc/boot/loader.go
+++ b/runsc/boot/loader.go
@@ -120,6 +120,10 @@ type Loader struct {
 	// container. It should be called when a sandbox is destroyed.
 	stopSignalForwarding func()
 
+	// stopProfiling stops profiling started at container creation. It
+	// should be called when a sandbox is destroyed.
+	stopProfiling func()
+
 	// restore is set to true if we are restoring a container.
 	restore bool
 
@@ -199,6 +203,21 @@ type Args struct {
 	TotalMem uint64
 	// UserLogFD is the file descriptor to write user logs to.
 	UserLogFD int
+	// ProfileBlockFD is the file descriptor to write a block profile to.
+	// Valid if >=0.
+	ProfileBlockFD int
+	// ProfileCPUFD is the file descriptor to write a CPU profile to.
+	// Valid if >=0.
+	ProfileCPUFD int
+	// ProfileHeapFD is the file descriptor to write a heap profile to.
+	// Valid if >=0.
+	ProfileHeapFD int
+	// ProfileMutexFD is the file descriptor to write a mutex profile to.
+	// Valid if >=0.
+	ProfileMutexFD int
+	// TraceFD is the file descriptor to write a Go execution trace to.
+	// Valid if >=0.
+	TraceFD int
 }
 
 // make sure stdioFDs are always the same on initial start and on restore
@@ -207,6 +226,8 @@ const startingStdioFD = 256
 // New initializes a new kernel loader configured by spec.
 // New also handles setting up a kernel for restoring a container.
 func New(args Args) (*Loader, error) {
+	stopProfiling := startProfiling(args)
+
 	// We initialize the rand package now to make sure /dev/urandom is pre-opened
 	// on kernels that do not support getrandom(2).
 	if err := rand.Init(); err != nil {
@@ -400,12 +421,13 @@ func New(args Args) (*Loader, error) {
 
 	eid := execID{cid: args.ID}
 	l := &Loader{
-		k:          k,
-		watchdog:   dog,
-		sandboxID:  args.ID,
-		processes:  map[execID]*execProcess{eid: {}},
-		mountHints: mountHints,
-		root:       info,
+		k:             k,
+		watchdog:      dog,
+		sandboxID:     args.ID,
+		processes:     map[execID]*execProcess{eid: {}},
+		mountHints:    mountHints,
+		root:          info,
+		stopProfiling: stopProfiling,
 	}
 
 	// We don't care about child signals; some platforms can generate a
@@ -498,6 +520,8 @@ func (l *Loader) Destroy() {
 	for _, f := range l.root.goferFDs {
 		_ = f.Close()
 	}
+
+	l.stopProfiling()
 }
 
 func createPlatform(conf *config.Config, deviceFile *os.File) (platform.Platform, error) {
diff --git a/runsc/boot/profile.go b/runsc/boot/profile.go
new file mode 100644
index 000000000..3ecd3e532
--- /dev/null
+++ b/runsc/boot/profile.go
@@ -0,0 +1,95 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package boot
+
+import (
+	"os"
+	"runtime"
+	"runtime/pprof"
+	"runtime/trace"
+
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sentry/control"
+)
+
+// startProfiling initiates profiling as defined by the ProfileConfig, and
+// returns a function that should be called to stop profiling.
+func startProfiling(args Args) func() {
+	var onStopProfiling []func()
+	stopProfiling := func() {
+		for _, f := range onStopProfiling {
+			f()
+		}
+	}
+
+	if args.ProfileBlockFD >= 0 {
+		file := os.NewFile(uintptr(args.ProfileBlockFD), "profile-block")
+
+		runtime.SetBlockProfileRate(control.DefaultBlockProfileRate)
+		onStopProfiling = append(onStopProfiling, func() {
+			if err := pprof.Lookup("block").WriteTo(file, 0); err != nil {
+				log.Warningf("Error writing block profile: %v", err)
+			}
+			file.Close()
+			runtime.SetBlockProfileRate(0)
+		})
+	}
+
+	if args.ProfileCPUFD >= 0 {
+		file := os.NewFile(uintptr(args.ProfileCPUFD), "profile-cpu")
+
+		pprof.StartCPUProfile(file)
+		onStopProfiling = append(onStopProfiling, func() {
+			pprof.StopCPUProfile()
+			file.Close()
+		})
+	}
+
+	if args.ProfileHeapFD >= 0 {
+		file := os.NewFile(uintptr(args.ProfileHeapFD), "profile-heap")
+
+		onStopProfiling = append(onStopProfiling, func() {
+			if err := pprof.Lookup("heap").WriteTo(file, 0); err != nil {
+				log.Warningf("Error writing heap profile: %v", err)
+			}
+			file.Close()
+		})
+	}
+
+	if args.ProfileMutexFD >= 0 {
+		file := os.NewFile(uintptr(args.ProfileMutexFD), "profile-mutex")
+
+		prev := runtime.SetMutexProfileFraction(control.DefaultMutexProfileRate)
+		onStopProfiling = append(onStopProfiling, func() {
+			if err := pprof.Lookup("mutex").WriteTo(file, 0); err != nil {
+				log.Warningf("Error writing mutex profile: %v", err)
+			}
+			file.Close()
+			runtime.SetMutexProfileFraction(prev)
+		})
+	}
+
+	if args.TraceFD >= 0 {
+		file := os.NewFile(uintptr(args.TraceFD), "trace")
+
+		trace.Start(file)
+		onStopProfiling = append(onStopProfiling, func() {
+			trace.Stop()
+			file.Close()
+		})
+	}
+
+	return stopProfiling
+}