2 files changed, 184 insertions, 104 deletions
diff --git a/pkg/sentry/control/pprof.go b/pkg/sentry/control/pprof.go
index 2bf3c45e1..2f3664c57 100644
--- a/pkg/sentry/control/pprof.go
+++ b/pkg/sentry/control/pprof.go
@@ -15,10 +15,10 @@
 package control
 
 import (
-	"errors"
 	"runtime"
 	"runtime/pprof"
 	"runtime/trace"
+	"time"
 
 	"gvisor.dev/gvisor/pkg/fd"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
@@ -26,184 +26,263 @@ import (
 	"gvisor.dev/gvisor/pkg/urpc"
 )
 
-var errNoOutput = errors.New("no output writer provided")
+// Profile includes profile-related RPC stubs. It provides a way to
+// control the built-in runtime profiling facilities.
+//
+// The profile object must be instantied via NewProfile.
+type Profile struct {
+	// kernel is the kernel under profile. It's immutable.
+	kernel *kernel.Kernel
 
-// ProfileOpts contains options for the StartCPUProfile/Goroutine RPC call.
-type ProfileOpts struct {
-	// File is the filesystem path for the profile.
-	File string `json:"path"`
+	// cpuMu protects CPU profiling.
+	cpuMu sync.Mutex
 
-	// FilePayload is the destination for the profiling output.
-	urpc.FilePayload
+	// blockMu protects block profiling.
+	blockMu sync.Mutex
+
+	// mutexMu protects mutex profiling.
+	mutexMu sync.Mutex
+
+	// traceMu protects trace profiling.
+	traceMu sync.Mutex
+
+	// done is closed when profiling is done.
+	done chan struct{}
 }
 
-// Profile includes profile-related RPC stubs. It provides a way to
-// control the built-in pprof facility in sentry via sentryctl.
-//
-// The following options to sentryctl are added:
-//
-// - collect CPU profile on-demand.
-//   sentryctl -pid <pid> pprof-cpu-start
-//   sentryctl -pid <pid> pprof-cpu-stop
-//
-// - dump out the stack trace of current go routines.
-//   sentryctl -pid <pid> pprof-goroutine
-type Profile struct {
-	// Kernel is the kernel under profile. It's immutable.
-	Kernel *kernel.Kernel
+// NewProfile returns a new Profile object.
+func NewProfile(k *kernel.Kernel) *Profile {
+	return &Profile{
+		kernel: k,
+		done:   make(chan struct{}),
+	}
+}
 
-	// mu protects the fields below.
-	mu sync.Mutex
+// Stop implements urpc.Stopper.Stop.
+func (p *Profile) Stop() {
+	close(p.done)
+}
 
-	// cpuFile is the current CPU profile output file.
-	cpuFile *fd.FD
+// CPUProfileOpts contains options specifically for CPU profiles.
+type CPUProfileOpts struct {
+	// FilePayload is the destination for the profiling output.
+	urpc.FilePayload
 
-	// traceFile is the current execution trace output file.
-	traceFile *fd.FD
+	// Duration is the duration of the profile.
+	Duration time.Duration `json:"duration"`
 }
 
-// StartCPUProfile is an RPC stub which starts recording the CPU profile in a
-// file.
-func (p *Profile) StartCPUProfile(o *ProfileOpts, _ *struct{}) error {
+// CPU is an RPC stub which collects a CPU profile.
+func (p *Profile) CPU(o *CPUProfileOpts, _ *struct{}) error {
 	if len(o.FilePayload.Files) < 1 {
-		return errNoOutput
+		return nil // Allowed.
 	}
 
-	output, err := fd.NewFromFile(o.FilePayload.Files[0])
-	if err != nil {
-		return err
-	}
+	output := o.FilePayload.Files[0]
+	defer output.Close()
 
-	p.mu.Lock()
-	defer p.mu.Unlock()
+	p.cpuMu.Lock()
+	defer p.cpuMu.Unlock()
 
 	// Returns an error if profiling is already started.
 	if err := pprof.StartCPUProfile(output); err != nil {
-		output.Close()
 		return err
 	}
+	defer pprof.StopCPUProfile()
+
+	// Collect the profile.
+	select {
+	case <-time.After(o.Duration):
+	case <-p.done:
+	}
 
-	p.cpuFile = output
 	return nil
 }
 
-// StopCPUProfile is an RPC stub which stops the CPU profiling and flush out the
-// profile data. It takes no argument.
-func (p *Profile) StopCPUProfile(_, _ *struct{}) error {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-
-	if p.cpuFile == nil {
-		return errors.New("CPU profiling not started")
-	}
+// HeapProfileOpts contains options specifically for heap profiles.
+type HeapProfileOpts struct {
+	// FilePayload is the destination for the profiling output.
+	urpc.FilePayload
 
-	pprof.StopCPUProfile()
-	p.cpuFile.Close()
-	p.cpuFile = nil
-	return nil
+	// Delay is the sleep time, similar to Duration. This may
+	// not affect the data collected however, as the heap will
+	// continue only the memory associated with the last alloc.
+	Delay time.Duration `json:"delay"`
 }
 
-// HeapProfile generates a heap profile for the sentry.
-func (p *Profile) HeapProfile(o *ProfileOpts, _ *struct{}) error {
+// Heap generates a heap profile.
+func (p *Profile) Heap(o *HeapProfileOpts, _ *struct{}) error {
 	if len(o.FilePayload.Files) < 1 {
-		return errNoOutput
+		return nil // Allowed.
 	}
+
 	output := o.FilePayload.Files[0]
 	defer output.Close()
-	runtime.GC() // Get up-to-date statistics.
-	if err := pprof.WriteHeapProfile(output); err != nil {
-		return err
+
+	// Wait for the given delay.
+	select {
+	case <-time.After(o.Delay):
+	case <-p.done:
 	}
-	return nil
+
+	// Get up-to-date statistics.
+	runtime.GC()
+
+	// Write the given profile.
+	return pprof.WriteHeapProfile(output)
+}
+
+// GoroutineProfileOpts contains options specifically for goroutine profiles.
+type GoroutineProfileOpts struct {
+	// FilePayload is the destination for the profiling output.
+	urpc.FilePayload
 }
 
-// GoroutineProfile is an RPC stub which dumps out the stack trace for all
-// running goroutines.
-func (p *Profile) GoroutineProfile(o *ProfileOpts, _ *struct{}) error {
+// Goroutine dumps out the stack trace for all running goroutines.
+func (p *Profile) Goroutine(o *GoroutineProfileOpts, _ *struct{}) error {
 	if len(o.FilePayload.Files) < 1 {
-		return errNoOutput
+		return nil // Allowed.
 	}
+
 	output := o.FilePayload.Files[0]
 	defer output.Close()
-	if err := pprof.Lookup("goroutine").WriteTo(output, 2); err != nil {
-		return err
-	}
-	return nil
+
+	return pprof.Lookup("goroutine").WriteTo(output, 2)
+}
+
+// BlockProfileOpts contains options specifically for block profiles.
+type BlockProfileOpts struct {
+	// FilePayload is the destination for the profiling output.
+	urpc.FilePayload
+
+	// Duration is the duration of the profile.
+	Duration time.Duration `json:"duration"`
+
+	// Rate is the block profile rate.
+	Rate int `json:"rate"`
 }
 
-// BlockProfile is an RPC stub which dumps out the stack trace that led to
-// blocking on synchronization primitives.
-func (p *Profile) BlockProfile(o *ProfileOpts, _ *struct{}) error {
+// Block dumps a blocking profile.
+func (p *Profile) Block(o *BlockProfileOpts, _ *struct{}) error {
 	if len(o.FilePayload.Files) < 1 {
-		return errNoOutput
+		return nil // Allowed.
 	}
+
 	output := o.FilePayload.Files[0]
 	defer output.Close()
-	if err := pprof.Lookup("block").WriteTo(output, 0); err != nil {
-		return err
+
+	p.blockMu.Lock()
+	defer p.blockMu.Unlock()
+
+	// Always set the rate. We then wait to collect a profile at this rate,
+	// and disable when we're done. Note that the default here is 10%, which
+	// will record a stacktrace 10% of the time when blocking occurs. Since
+	// these events should not be super frequent, we expect this to achieve
+	// a reasonable balance between collecting the data we need and imposing
+	// a high performance cost (e.g. skewing even the CPU profile).
+	rate := 10
+	if o.Rate != 0 {
+		rate = o.Rate
 	}
-	return nil
+	runtime.SetBlockProfileRate(rate)
+	defer runtime.SetBlockProfileRate(0)
+
+	// Collect the profile.
+	select {
+	case <-time.After(o.Duration):
+	case <-p.done:
+	}
+
+	return pprof.Lookup("block").WriteTo(output, 0)
+}
+
+// MutexProfileOpts contains options specifically for mutex profiles.
+type MutexProfileOpts struct {
+	// FilePayload is the destination for the profiling output.
+	urpc.FilePayload
+
+	// Duration is the duration of the profile.
+	Duration time.Duration `json:"duration"`
+
+	// Fraction is the mutex profile fraction.
+	Fraction int `json:"fraction"`
 }
 
-// MutexProfile is an RPC stub which dumps out the stack trace of holders of
-// contended mutexes.
-func (p *Profile) MutexProfile(o *ProfileOpts, _ *struct{}) error {
+// Mutex dumps a mutex profile.
+func (p *Profile) Mutex(o *MutexProfileOpts, _ *struct{}) error {
 	if len(o.FilePayload.Files) < 1 {
-		return errNoOutput
+		return nil // Allowed.
 	}
+
 	output := o.FilePayload.Files[0]
 	defer output.Close()
-	if err := pprof.Lookup("mutex").WriteTo(output, 0); err != nil {
-		return err
+
+	p.mutexMu.Lock()
+	defer p.mutexMu.Unlock()
+
+	// Always set the fraction. Like the block rate above, we use
+	// a default rate of 10% for the same reasons.
+	fraction := 10
+	if o.Fraction != 0 {
+		fraction = o.Fraction
 	}
-	return nil
+	runtime.SetMutexProfileFraction(fraction)
+	defer runtime.SetMutexProfileFraction(0)
+
+	// Collect the profile.
+	select {
+	case <-time.After(o.Duration):
+	case <-p.done:
+	}
+
+	return pprof.Lookup("mutex").WriteTo(output, 0)
 }
 
-// StartTrace is an RPC stub which starts collection of an execution trace.
-func (p *Profile) StartTrace(o *ProfileOpts, _ *struct{}) error {
+// TraceProfileOpts contains options specifically for traces.
+type TraceProfileOpts struct {
+	// FilePayload is the destination for the profiling output.
+	urpc.FilePayload
+
+	// Duration is the duration of the profile.
+	Duration time.Duration `json:"duration"`
+}
+
+// Trace is an RPC stub which starts collection of an execution trace.
+func (p *Profile) Trace(o *TraceProfileOpts, _ *struct{}) error {
 	if len(o.FilePayload.Files) < 1 {
-		return errNoOutput
+		return nil // Allowed.
 	}
 
 	output, err := fd.NewFromFile(o.FilePayload.Files[0])
 	if err != nil {
 		return err
 	}
+	defer output.Close()
 
-	p.mu.Lock()
-	defer p.mu.Unlock()
+	p.traceMu.Lock()
+	defer p.traceMu.Unlock()
 
 	// Returns an error if profiling is already started.
 	if err := trace.Start(output); err != nil {
 		output.Close()
 		return err
 	}
+	defer trace.Stop()
 
 	// Ensure all trace contexts are registered.
-	p.Kernel.RebuildTraceContexts()
-
-	p.traceFile = output
-	return nil
-}
-
-// StopTrace is an RPC stub which stops collection of an ongoing execution
-// trace and flushes the trace data. It takes no argument.
-func (p *Profile) StopTrace(_, _ *struct{}) error {
-	p.mu.Lock()
-	defer p.mu.Unlock()
+	p.kernel.RebuildTraceContexts()
 
-	if p.traceFile == nil {
-		return errors.New("Execution tracing not started")
+	// Wait for the trace.
+	select {
+	case <-time.After(o.Duration):
+	case <-p.done:
 	}
 
 	// Similarly to the case above, if tasks have not ended traces, we will
 	// lose information. Thus we need to rebuild the tasks in order to have
 	// complete information. This will not lose information if multiple
 	// traces are overlapping.
-	p.Kernel.RebuildTraceContexts()
+	p.kernel.RebuildTraceContexts()
 
-	trace.Stop()
-	p.traceFile.Close()
-	p.traceFile = nil
 	return nil
 }
diff --git a/pkg/sentry/control/state.go b/pkg/sentry/control/state.go
index d800f2c85..62eaca965 100644
--- a/pkg/sentry/control/state.go
+++ b/pkg/sentry/control/state.go
@@ -62,6 +62,7 @@ func (s *State) Save(o *SaveOpts, _ *struct{}) error {
 		Callback: func(err error) {
 			if err == nil {
 				log.Infof("Save succeeded: exiting...")
+				s.Kernel.SetSaveSuccess(false /* autosave */)
 			} else {
 				log.Warningf("Save failed: exiting...")
 				s.Kernel.SetSaveError(err)