Make Profile copying and encoding thread-safe. (#232)

Added some multi-threaded tests that allow the race detector
to detect race conditions in pprof internals.

internal/driver/webui_test.go

@@ -23,6 +23,7 @@ import (
 	"net/url"
 	"os/exec"
 	"regexp"
+	"sync"
 	"testing"
 
 	"github.com/google/pprof/internal/plugin"
@@ -101,6 +102,23 @@ func TestWebInterface(t *testing.T) {
 		}
 	}
 
+	// Also fetch all the test case URLs in parallel to test thread
+	// safety when run under the race detector.
+	var wg sync.WaitGroup
+	for _, c := range testcases {
+		if c.needDot && !haveDot {
+			continue
+		}
+		path := server.URL + c.path
+		for count := 0; count < 2; count++ {
+			wg.Add(1)
+			go func() {
+				http.Get(path)
+				wg.Done()
+			}()
+		}
+	}
+	wg.Wait()
 }
 
 // Implement fake object file support.

profile/profile.go

@@ -26,6 +26,7 @@ import (
 	"regexp"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -47,6 +48,10 @@ type Profile struct {
 	PeriodType    *ValueType
 	Period        int64
 
+	// The following fields are modified during encoding and copying,
+	// so are protected by a Mutex.
+	encodeMu sync.Mutex
+
 	commentX           []int64
 	dropFramesX        int64
 	keepFramesX        int64
@@ -296,21 +301,25 @@ func (p *Profile) updateLocationMapping(from, to *Mapping) {
 	}
 }
 
-// Write writes the profile as a gzip-compressed marshaled protobuf.
-func (p *Profile) Write(w io.Writer) error {
+func serialize(p *Profile) []byte {
+	p.encodeMu.Lock()
 	p.preEncode()
 	b := marshal(p)
+	p.encodeMu.Unlock()
+	return b
+}
+
+// Write writes the profile as a gzip-compressed marshaled protobuf.
+func (p *Profile) Write(w io.Writer) error {
 	zw := gzip.NewWriter(w)
 	defer zw.Close()
-	_, err := zw.Write(b)
+	_, err := zw.Write(serialize(p))
 	return err
 }
 
 // WriteUncompressed writes the profile as a marshaled protobuf.
 func (p *Profile) WriteUncompressed(w io.Writer) error {
-	p.preEncode()
-	b := marshal(p)
-	_, err := w.Write(b)
+	_, err := w.Write(serialize(p))
 	return err
 }
 
@@ -605,11 +614,8 @@ func (m *Mapping) Unsymbolizable() bool {
 
 // Copy makes a fully independent copy of a profile.
 func (p *Profile) Copy() *Profile {
-	p.preEncode()
-	b := marshal(p)
-
 	pp := &Profile{}
-	if err := unmarshal(b, pp); err != nil {
+	if err := unmarshal(serialize(p), pp); err != nil {
 		panic(err)
 	}
 	if err := pp.postDecode(); err != nil {

profile/profile_test.go

@@ -21,6 +21,7 @@ import (
 	"path/filepath"
 	"regexp"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/google/pprof/internal/proftest"
@@ -694,3 +695,29 @@ func TestSetMain(t *testing.T) {
 		t.Errorf("got %s for main", testProfile1.Mapping[0].File)
 	}
 }
+
+// parallel runs n copies of fn in parallel.
+func parallel(n int, fn func()) {
+	var wg sync.WaitGroup
+	wg.Add(n)
+	for i := 0; i < n; i++ {
+		go func() {
+			fn()
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+}
+
+func TestThreadSafety(t *testing.T) {
+	src := testProfile1.Copy()
+	parallel(4, func() { src.Copy() })
+	parallel(4, func() {
+		var b bytes.Buffer
+		src.WriteUncompressed(&b)
+	})
+	parallel(4, func() {
+		var b bytes.Buffer
+		src.Write(&b)
+	})
+}