Merge pull request #77977 from rimadeodhar/backport21.2-76266

release-21.2: pprofui: Increase concurrency for profiles

pkg/server/debug/pprofui/BUILD.bazel

@@ -34,6 +34,8 @@ go_test(
     deps = [
         "//pkg/build/bazel",
         "//pkg/server/serverpb",
+        "//pkg/testutils",
+        "//pkg/testutils/skip",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/server/debug/pprofui/server.go

@@ -40,6 +40,22 @@ type Profiler interface {
 	) (*serverpb.JSONResponse, error)
 }
 
+const (
+	// ProfileConcurrency governs how many concurrent profiles can be collected.
+	// This impacts the maximum number of profiles stored in memory at any given point
+	// in time. Increasing this number further will increase the number of profile
+	// requests that can be served concurrently but will also increase
+	// storage requirements and should be done with caution.
+	ProfileConcurrency = 2
+
+	// ProfileExpiry governs how long a profile is retained in memory during concurrent
+	// profile requests. A profile is considered expired once its profile expiry duration
+	// is met. However, expired profiles are only cleaned up from memory when a new profile
+	// is requested. So ProfileExpiry can be considered as a soft expiry which impacts
+	// duration for which a profile is stored only when other profile requests are received.
+	ProfileExpiry = 2 * time.Second
+)
+
 // A Server serves up the pprof web ui. A request to /<profiletype>
 // generates a profile of the desired type and redirects to the UI for
 // it at /<profiletype>/<id>. Valid profile types at the time of

pkg/server/debug/pprofui/server_test.go

@@ -14,13 +14,18 @@ import (
 	"context"
 	"fmt"
 	"io/ioutil"
+	"math/rand"
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/build/bazel"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/stretchr/testify/require"
 )
 
@@ -94,3 +99,59 @@ func TestServer(t *testing.T) {
 		require.Equal(t, "/heap/4/flamegraph?node=3", loc)
 	})
 }
+
+func TestServerConcurrentAccess(t *testing.T) {
+	expectedNodeID := "local"
+	skip.UnderRace(t, "test fails under race due to known race condition with profiles")
+	const (
+		runsPerWorker = 1
+		workers       = ProfileConcurrency
+	)
+	mockProfile := func(ctx context.Context, req *serverpb.ProfileRequest) (*serverpb.JSONResponse, error) {
+		require.Equal(t, expectedNodeID, req.NodeId)
+		fileName := "heap.profile"
+		if req.Type == serverpb.ProfileRequest_CPU {
+			fileName = "cpu.profile"
+		}
+		b, err := ioutil.ReadFile(testutils.TestDataPath(t, fileName))
+		require.NoError(t, err)
+		return &serverpb.JSONResponse{Data: b}, nil
+	}
+
+	s := NewServer(NewMemStorage(ProfileConcurrency, ProfileExpiry), ProfilerFunc(mockProfile))
+	getProfile := func(profile string, t *testing.T) {
+		t.Helper()
+
+		r := httptest.NewRequest("GET", "/heap/", nil)
+		w := httptest.NewRecorder()
+		s.ServeHTTP(w, r)
+
+		require.Equal(t, http.StatusTemporaryRedirect, w.Code)
+
+		loc := w.Result().Header.Get("Location")
+
+		r = httptest.NewRequest("GET", loc, nil)
+		w = httptest.NewRecorder()
+
+		s.ServeHTTP(w, r)
+
+		require.Equal(t, http.StatusOK, w.Code)
+		require.Contains(t, w.Body.String(), "pprof</a></h1>")
+	}
+	var wg sync.WaitGroup
+	profiles := [2]string{"/heap/", "/cpu"}
+	runWorker := func() {
+		defer wg.Done()
+		for i := 0; i < runsPerWorker; i++ {
+			time.Sleep(time.Microsecond)
+			profileID := rand.Intn(len(profiles))
+			getProfile(profiles[profileID], t)
+		}
+	}
+	// Run the workers.
+	for i := 0; i < workers; i++ {
+		wg.Add(1)
+		go runWorker()
+	}
+	wg.Wait()
+}

pkg/server/debug/pprofui/storage_mem.go

@@ -102,5 +102,8 @@ func (s *MemStorage) Get(id string, read func(io.Reader) error) error {
 			return read(bytes.NewReader(v.b))
 		}
 	}
-	return errors.Errorf("profile not found; it may have expired")
+	return errors.Errorf("profile not found; it may have expired, please regenerate the profile.\n" +
+		"To generate profile for a node, use the profile generation link from the Advanced Debug page.\n" +
+		"Attempting to generate a profile by modifying the node query parameter in the URL will not work.",
+	)
 }

pkg/server/debug/server.go

@@ -119,7 +119,7 @@ func NewServer(
 	}
 	mux.HandleFunc("/debug/logspy", spy.handleDebugLogSpy)
 
-	ps := pprofui.NewServer(pprofui.NewMemStorage(1, 0), profiler)
+	ps := pprofui.NewServer(pprofui.NewMemStorage(pprofui.ProfileConcurrency, pprofui.ProfileExpiry), profiler)
 	mux.Handle("/debug/pprof/ui/", http.StripPrefix("/debug/pprof/ui", ps))
 
 	mux.HandleFunc("/debug/pprof/goroutineui/", func(w http.ResponseWriter, req *http.Request) {