Add multi-layered caching with LRU for symbols, Ristretto for debuginfo, and optimized DWARF parsing

Author: marcsanmi

.mockery.yaml

@@ -57,3 +57,6 @@ packages:
     interfaces:
       Exporter:
       Ruler:
+  github.com/grafana/pyroscope/pkg/experiment/symbolizer:
+    interfaces:
+      DebuginfodClient:

cmd/symbolization/main.go

@@ -21,7 +21,7 @@ func main() {
 	// Alternatively, use a local debug info file:
 	//client := &localDebuginfodClient{debugFilePath: "/path/to/your/debug/file"}
 
-	s := symbolizer.NewProfileSymbolizer(client, nil, symbolizer.NewMetrics(nil))
+	s := symbolizer.NewProfileSymbolizer(client, nil, symbolizer.NewMetrics(nil), 0)
 	ctx := context.Background()
 
 	_, err := client.FetchDebuginfo(ctx, buildID)

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/colega/zeropool v0.0.0-20230505084239-6fb4a4f75381
 	github.com/dennwc/varint v1.0.0
+	github.com/dgraph-io/ristretto v0.2.0
 	github.com/dgryski/go-groupvarint v0.0.0-20230630160417-2bfb7969fb3c
 	github.com/dolthub/swiss v0.2.1
 	github.com/drone/envsubst v1.0.3

go.sum

@@ -197,6 +197,10 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dennwc/varint v1.0.0 h1:kGNFFSSw8ToIy3obO/kKr8U9GZYUAxQEVuix4zfDWzE=
 github.com/dennwc/varint v1.0.0/go.mod h1:hnItb35rvZvJrbTALZtY/iQfDs48JKRG1RPpgziApxA=
+github.com/dgraph-io/ristretto v0.2.0 h1:XAfl+7cmoUDWW/2Lx8TGZQjjxIQ2Ley9DSf52dru4WE=
+github.com/dgraph-io/ristretto v0.2.0/go.mod h1:8uBHCU/PBV4Ag0CJrP47b9Ofby5dqWNh4FicAdoqFNU=
+github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 h1:fAjc9m62+UWV/WAFKLNi6ZS0675eEUC9y3AlwSbQu1Y=
+github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/dgryski/go-groupvarint v0.0.0-20230630160417-2bfb7969fb3c h1:cHaw4wmusVzAZLEPWOCCGCfu6UvFXx9UboCHQCnjvxY=
 github.com/dgryski/go-groupvarint v0.0.0-20230630160417-2bfb7969fb3c/go.mod h1:MlkUQveSLEDbIgq2r1e++tSf0zfzU9mQpa9Qkczl+9Y=
 github.com/digitalocean/godo v1.109.0 h1:4W97RJLJSUQ3veRZDNbp1Ol3Rbn6Lmt9bKGvfqYI5SU=

pkg/experiment/symbolizer/debuginfod_client.go

@@ -1,21 +1,24 @@
 package symbolizer
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"net/url"
 	"os"
-	"path/filepath"
 	"regexp"
 	"strings"
+	"sync"
 	"time"
+
+	"github.com/dgraph-io/ristretto"
 )
 
 type DebuginfodClient interface {
-	FetchDebuginfo(ctx context.Context, buildID string) (string, error)
+	FetchDebuginfo(ctx context.Context, buildID string) (io.ReadCloser, error)
 }
 
 type debuginfodClientConfig struct {
@@ -25,10 +28,23 @@ type debuginfodClientConfig struct {
 	// userAgent   string
 	httpClient *http.Client
 }
-
 type debuginfodClient struct {
 	cfg     debuginfodClientConfig
 	metrics *Metrics
+
+	// In-memory cache of build IDs to file paths
+	cache *ristretto.Cache
+
+	// Track in-flight requests to prevent duplicate fetches
+	inFlightRequests      map[string]*inFlightRequest
+	inFlightRequestsMutex sync.Mutex
+}
+
+// inFlightRequest represents an ongoing fetch operation
+type inFlightRequest struct {
+	done chan struct{}
+	data []byte
+	err  error
 }
 
 func NewDebuginfodClient(baseURL string, metrics *Metrics) DebuginfodClient {
@@ -38,6 +54,16 @@ func NewDebuginfodClient(baseURL string, metrics *Metrics) DebuginfodClient {
 		TLSHandshakeTimeout: 10 * time.Second,
 	}
 
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 1e7,     // number of keys to track frequency of (10M)
+		MaxCost:     2 << 30, // maximum cost of cache (2GB)
+		BufferItems: 64,      // number of keys per Get buffer
+	})
+
+	if err != nil {
+		cache = nil
+	}
+
 	return &debuginfodClient{
 		cfg: debuginfodClientConfig{
 			baseURL:     baseURL,
@@ -55,12 +81,14 @@ func NewDebuginfodClient(baseURL string, metrics *Metrics) DebuginfodClient {
 				},
 			},
 		},
-		metrics: metrics,
+		metrics:          metrics,
+		cache:            cache,
+		inFlightRequests: make(map[string]*inFlightRequest),
 	}
 }
 
 // FetchDebuginfo fetches the debuginfo file for a specific build ID.
-func (c *debuginfodClient) FetchDebuginfo(ctx context.Context, buildID string) (string, error) {
+func (c *debuginfodClient) FetchDebuginfo(ctx context.Context, buildID string) (io.ReadCloser, error) {
 	start := time.Now()
 	var lastErr error
 	c.metrics.debuginfodRequestsTotal.Inc()
@@ -76,23 +104,67 @@ func (c *debuginfodClient) FetchDebuginfo(ctx context.Context, buildID string) (
 	sanitizedBuildID, err := sanitizeBuildID(buildID)
 	if err != nil {
 		c.metrics.debuginfodRequestErrorsTotal.WithLabelValues("invalid_id").Inc()
-		return "", err
+		return nil, err
+	}
+
+	// Check in-memory cache first
+	if c.cache != nil {
+		if data, found := c.cache.Get(sanitizedBuildID); found {
+			return io.NopCloser(bytes.NewReader(data.([]byte))), nil
+		}
 	}
 
+	// Check if there's already a request in flight for this build ID
+	c.inFlightRequestsMutex.Lock()
+	req, inFlight := c.inFlightRequests[sanitizedBuildID]
+	if inFlight {
+		// There's already a request in flight, wait for it to complete
+		done := req.done
+		c.inFlightRequestsMutex.Unlock()
+
+		select {
+		case <-done:
+			if req.err != nil {
+				lastErr = req.err
+				return nil, req.err
+			}
+			return io.NopCloser(bytes.NewReader(req.data)), nil
+		case <-ctx.Done():
+			lastErr = ctx.Err()
+			return nil, ctx.Err()
+		}
+	}
+
+	// Create a new in-flight request
+	req = &inFlightRequest{
+		done: make(chan struct{}),
+	}
+	c.inFlightRequests[sanitizedBuildID] = req
+	c.inFlightRequestsMutex.Unlock()
+
+	// Ensure we clean up the in-flight request when we're done
+	defer func() {
+		c.inFlightRequestsMutex.Lock()
+		delete(c.inFlightRequests, sanitizedBuildID)
+		close(req.done)
+		c.inFlightRequestsMutex.Unlock()
+	}()
+
 	url := fmt.Sprintf("%s/buildid/%s/debuginfo", c.cfg.baseURL, sanitizedBuildID)
+	var data []byte
 
 	// Implement retries with exponential backoff
 	for attempt := 0; attempt < c.cfg.maxRetries; attempt++ {
 		if attempt > 0 {
 			if ctx.Err() != nil {
-				return "", ctx.Err()
+				return nil, ctx.Err()
 			}
 			time.Sleep(c.cfg.backoffTime * time.Duration(attempt))
 		}
 
-		filePath, err := c.doRequest(ctx, url, sanitizedBuildID)
+		data, err = c.doRequest(ctx, url)
 		if err == nil {
-			return filePath, nil
+			break
 		}
 
 		lastErr = err
@@ -101,50 +173,52 @@ func (c *debuginfodClient) FetchDebuginfo(ctx context.Context, buildID string) (
 		}
 	}
 
-	return "", fmt.Errorf("failed to fetch debuginfo after %d attempts: %w", c.cfg.maxRetries, lastErr)
+	if lastErr != nil {
+		req.err = fmt.Errorf("failed to fetch debuginfo after %d attempts: %w", c.cfg.maxRetries, lastErr)
+		return nil, req.err
+	}
+
+	// Store in cache
+	if c.cache != nil {
+		// The cost is the size of the data in bytes
+		c.cache.Set(sanitizedBuildID, data, int64(len(data)))
+	}
+
+	req.data = data
+	return io.NopCloser(bytes.NewReader(data)), nil
 }
 
-func (c *debuginfodClient) doRequest(ctx context.Context, url, buildID string) (string, error) {
+func (c *debuginfodClient) doRequest(ctx context.Context, url string) ([]byte, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
 	if err != nil {
-		return "", fmt.Errorf("failed to create request: %w", err)
+		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
 
 	req.Header.Set("Accept-Encoding", "gzip, deflate")
 
 	resp, err := c.cfg.httpClient.Do(req)
 	if err != nil {
-		return "", fmt.Errorf("failed to execute request: %w", err)
+		return nil, fmt.Errorf("failed to execute request: %w", err)
 	}
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return "", fmt.Errorf("unexpected HTTP status: %d %s", resp.StatusCode, resp.Status)
-	}
-
-	return c.saveDebugInfo(resp.Body, buildID, resp.ContentLength)
-}
-
-func (c *debuginfodClient) saveDebugInfo(body io.Reader, buildID string, contentLength int64) (string, error) {
-	if contentLength > 0 {
-		c.metrics.debuginfodFileSize.Observe(float64(contentLength))
+		return nil, fmt.Errorf("unexpected HTTP status: %d %s", resp.StatusCode, resp.Status)
 	}
 
-	tempDir := os.TempDir()
-	filePath := filepath.Join(tempDir, fmt.Sprintf("%s.elf", buildID))
-
-	outFile, err := os.Create(filePath)
+	// Read the entire response body into memory
+	data, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return "", fmt.Errorf("failed to create temp file: %w", err)
+		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
-	defer outFile.Close()
 
-	if _, err := io.Copy(outFile, body); err != nil {
-		os.Remove(filePath) // Clean up on error
-		return "", fmt.Errorf("failed to write debug info: %w", err)
+	if resp.ContentLength > 0 {
+		c.metrics.debuginfodFileSize.Observe(float64(resp.ContentLength))
+	} else {
+		c.metrics.debuginfodFileSize.Observe(float64(len(data)))
 	}
 
-	return filePath, nil
+	return data, nil
 }
 
 func isRetryableError(err error) bool {