Merge release-20201208.0-109-gb06e5bc5b (automated)

3 files changed, 55 insertions, 25 deletions

diff --git a/pkg/control/server/server.go b/pkg/control/server/server.go
index 41abe1f2d..629dae8f4 100644
--- a/pkg/control/server/server.go
+++ b/pkg/control/server/server.go
@@ -67,9 +67,10 @@ func (s *Server) Wait() {
 // and the server should not be used afterwards.
 func (s *Server) Stop() {
 	s.socket.Close()
-	s.wg.Wait()
+	s.Wait()
 
-	// This will cause existing clients to be terminated safely.
+	// This will cause existing clients to be terminated safely. If the
+	// registered handlers have a Stop callback, it will be called.
 	s.server.Stop()
 }
 
diff --git a/pkg/sentry/control/pprof.go b/pkg/sentry/control/pprof.go
index b78e29416..2f3664c57 100644
--- a/pkg/sentry/control/pprof.go
+++ b/pkg/sentry/control/pprof.go
@@ -50,17 +50,17 @@ type Profile struct {
 	done chan struct{}
 }
 
-// NewProfile returns a new Profile object, and a stop callback.
-//
-// The stop callback should be used at most once.
-func NewProfile(k *kernel.Kernel) (*Profile, func()) {
-	p := &Profile{
+// NewProfile returns a new Profile object.
+func NewProfile(k *kernel.Kernel) *Profile {
+	return &Profile{
 		kernel: k,
 		done:   make(chan struct{}),
 	}
-	return p, func() {
-		close(p.done)
-	}
+}
+
+// Stop implements urpc.Stopper.Stop.
+func (p *Profile) Stop() {
+	close(p.done)
 }
 
 // CPUProfileOpts contains options specifically for CPU profiles.
@@ -70,9 +70,6 @@ type CPUProfileOpts struct {
 
 	// Duration is the duration of the profile.
 	Duration time.Duration `json:"duration"`
-
-	// Hz is the rate, which may be zero.
-	Hz int `json:"hz"`
 }
 
 // CPU is an RPC stub which collects a CPU profile.
@@ -81,19 +78,13 @@ func (p *Profile) CPU(o *CPUProfileOpts, _ *struct{}) error {
 		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.cpuMu.Lock()
 	defer p.cpuMu.Unlock()
 
 	// Returns an error if profiling is already started.
-	if o.Hz != 0 {
-		runtime.SetCPUProfileRate(o.Hz)
-	}
 	if err := pprof.StartCPUProfile(output); err != nil {
 		return err
 	}
@@ -112,6 +103,11 @@ func (p *Profile) CPU(o *CPUProfileOpts, _ *struct{}) error {
 type HeapProfileOpts struct {
 	// FilePayload is the destination for the profiling output.
 	urpc.FilePayload
+
+	// 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"`
 }
 
 // Heap generates a heap profile.
@@ -123,7 +119,16 @@ func (p *Profile) Heap(o *HeapProfileOpts, _ *struct{}) error {
 	output := o.FilePayload.Files[0]
 	defer output.Close()
 
-	runtime.GC() // Get up-to-date statistics.
+	// Wait for the given delay.
+	select {
+	case <-time.After(o.Delay):
+	case <-p.done:
+	}
+
+	// Get up-to-date statistics.
+	runtime.GC()
+
+	// Write the given profile.
 	return pprof.WriteHeapProfile(output)
 }
 
@@ -170,8 +175,12 @@ func (p *Profile) Block(o *BlockProfileOpts, _ *struct{}) error {
 	defer p.blockMu.Unlock()
 
 	// Always set the rate. We then wait to collect a profile at this rate,
-	// and disable when we're done.
-	rate := 1
+	// 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
 	}
@@ -211,8 +220,9 @@ func (p *Profile) Mutex(o *MutexProfileOpts, _ *struct{}) error {
 	p.mutexMu.Lock()
 	defer p.mutexMu.Unlock()
 
-	// Always set the fraction.
-	fraction := 1
+	// 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
 	}
diff --git a/pkg/urpc/urpc.go b/pkg/urpc/urpc.go
index dfd23032c..0e9a829f6 100644
--- a/pkg/urpc/urpc.go
+++ b/pkg/urpc/urpc.go
@@ -170,6 +170,9 @@ type Server struct {
 	// methods is the set of server methods.
 	methods map[string]registeredMethod
 
+	// stoppers are all registered stoppers.
+	stoppers []Stopper
+
 	// clients is a map of clients.
 	clients map[*unet.Socket]clientState
 
@@ -195,6 +198,12 @@ func NewServerWithCallback(afterRPCCallback func()) *Server {
 	}
 }
 
+// Stopper is an optional interface, that when implemented, allows an object
+// to have a callback executed when the server is shutting down.
+type Stopper interface {
+	Stop()
+}
+
 // Register registers the given object as an RPC receiver.
 //
 // This functions is the same way as the built-in RPC package, but it does not
@@ -206,6 +215,7 @@ func (s *Server) Register(obj interface{}) {
 	defer s.mu.Unlock()
 
 	typ := reflect.TypeOf(obj)
+	stopper, hasStop := obj.(Stopper)
 
 	// If we got a pointer, deref it to the underlying object. We need this to
 	// obtain the name of the underlying type.
@@ -221,6 +231,10 @@ func (s *Server) Register(obj interface{}) {
 			// Can't be anonymous.
 			panic("type not named.")
 		}
+		if hasStop && method.Name == "Stop" {
+			s.stoppers = append(s.stoppers, stopper)
+			continue // Legal stop method.
+		}
 
 		prettyName := typDeref.Name() + "." + method.Name
 		if _, ok := s.methods[prettyName]; ok {
@@ -448,6 +462,11 @@ func (s *Server) Stop() {
 	// Wait for all outstanding requests.
 	defer s.wg.Wait()
 
+	// Call any Stop callbacks.
+	for _, stopper := range s.stoppers {
+		stopper.Stop()
+	}
+
 	// Close all known clients.
 	s.mu.Lock()
 	defer s.mu.Unlock()