From 1f0738c1577a55a6b7229b821ddfe762b84771d0 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josharian@gmail.com>
Date: Tue, 31 Dec 2019 21:24:20 -0800
Subject: [PATCH] runtime/pprof: speed up CPU profiling shutdown

The core CPU profiling loop contains a 100ms sleep.
This is important to reduce overhead.

However, it means that it takes 200ms to shutting down a program
with CPU profiling enabled. When trying to collect many samples
by running a short-lived program many times, this adds up.

This change cuts the shutdown penalty in half by skipping
the sleep whenever possible.

Change-Id: Ic3177f8e1a2d331fe1a1ecd7c8c06f50beb42535
Reviewed-on: https://go-review.googlesource.com/c/go/+/228886
Run-TryBot: Josh Bleecher Snyder <josharian@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Hyang-Ah Hana Kim <hyangah@gmail.com>
---
 src/runtime/pprof/pprof.go | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/src/runtime/pprof/pprof.go b/src/runtime/pprof/pprof.go
index b4f9ab8f7a..b828619a86 100644
--- a/src/runtime/pprof/pprof.go
+++ b/src/runtime/pprof/pprof.go
@@ -81,6 +81,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"text/tabwriter"
 	"time"
 	"unsafe"
@@ -714,10 +715,22 @@ func (p runtimeProfile) Stack(i int) []uintptr { return p[i].Stack() }
 
 var cpu struct {
 	sync.Mutex
-	profiling bool
+	profiling uint32 // bool, accessed atomically
 	done      chan bool
 }
 
+func cpuProfiling() bool {
+	return atomic.LoadUint32(&cpu.profiling) != 0
+}
+
+func setCPUProfiling(b bool) {
+	if b {
+		atomic.StoreUint32(&cpu.profiling, 1)
+	} else {
+		atomic.StoreUint32(&cpu.profiling, 0)
+	}
+}
+
 // StartCPUProfile enables CPU profiling for the current process.
 // While profiling, the profile will be buffered and written to w.
 // StartCPUProfile returns an error if profiling is already enabled.
@@ -747,10 +760,10 @@ func StartCPUProfile(w io.Writer) error {
 		cpu.done = make(chan bool)
 	}
 	// Double-check.
-	if cpu.profiling {
+	if cpuProfiling() {
 		return fmt.Errorf("cpu profiling already in use")
 	}
-	cpu.profiling = true
+	setCPUProfiling(true)
 	runtime.SetCPUProfileRate(hz)
 	go profileWriter(w)
 	return nil
@@ -767,7 +780,9 @@ func profileWriter(w io.Writer) {
 	b := newProfileBuilder(w)
 	var err error
 	for {
-		time.Sleep(100 * time.Millisecond)
+		if cpuProfiling() {
+			time.Sleep(100 * time.Millisecond)
+		}
 		data, tags, eof := readProfile()
 		if e := b.addCPUData(data, tags); e != nil && err == nil {
 			err = e
@@ -792,10 +807,10 @@ func StopCPUProfile() {
 	cpu.Lock()
 	defer cpu.Unlock()
 
-	if !cpu.profiling {
+	if !cpuProfiling() {
 		return
 	}
-	cpu.profiling = false
+	setCPUProfiling(false)
 	runtime.SetCPUProfileRate(0)
 	<-cpu.done
 }