rpc: add the ability to write out a memory profile during testing.

R=rsc, dsymonds
CC=golang-dev
https://golang.org/cl/4290047

diff --git a/src/pkg/rpc/server_test.go b/src/pkg/rpc/server_test.go
index 9e32b740f1e8b37a4cd54dabbcc55c2f4e7d3161..71c283ed844c8d56438b02b48b2e28df804a46ce 100644 (file)
--- a/src/pkg/rpc/server_test.go
+++ b/src/pkg/rpc/server_test.go
@@ -5,12 +5,14 @@
 package rpc
 
 import (
+       "flag"
        "fmt"
        "http/httptest"
        "log"
        "net"
        "os"
        "runtime"
+       "runtime/pprof"
        "strings"
        "sync"
        "testing"
@@ -23,6 +25,8 @@ var (
        once, newOnce, httpOnce   sync.Once
 )
 
+var memprofile = flag.String("memprofile", "", "write the memory profile in TestCountMallocs to the named file")
+
 const (
        second      = 1e9
        newHttpPath = "/foo"
@@ -352,6 +356,7 @@ func testSendDeadlock(client *Client) {
 }
 
 func TestCountMallocs(t *testing.T) {
+       runtime.MemProfileRate = 1
        once.Do(startServer)
        client, err := Dial("tcp", serverAddr)
        if err != nil {
@@ -360,7 +365,7 @@ func TestCountMallocs(t *testing.T) {
        args := &Args{7, 8}
        reply := new(Reply)
        mallocs := 0 - runtime.MemStats.Mallocs
-       const count = 100
+       const count = 10000
        for i := 0; i < count; i++ {
                err = client.Call("Arith.Add", args, reply)
                if err != nil {
@@ -371,6 +376,16 @@ func TestCountMallocs(t *testing.T) {
                }
        }
        mallocs += runtime.MemStats.Mallocs
+       if *memprofile != "" {
+               if fd, err := os.Open(*memprofile, os.O_WRONLY|os.O_CREAT|os.O_TRUNC, 0666); err != nil {
+                       t.Errorf("can't open %s: %s", *memprofile, err)
+               } else {
+                       if err = pprof.WriteHeapProfile(fd); err != nil {
+                               t.Errorf("can't write %s: %s", *memprofile, err)
+                       }
+                       fd.Close()
+               }
+       }
        fmt.Printf("mallocs per rpc round trip: %d\n", mallocs/count)
 }