github.com/grafana/pyroscope@v1.18.0/pkg/og/convert/jfr/profile.go (about)

     1  package jfr
     2  
     3  import (
     4  	"bytes"
     5  	"compress/gzip"
     6  	"context"
     7  	"fmt"
     8  	"io"
     9  	"mime/multipart"
    10  	"strings"
    11  
    12  	"github.com/grafana/dskit/tenant"
    13  	jfrPprof "github.com/grafana/jfr-parser/pprof"
    14  	jfrPprofPyroscope "github.com/grafana/jfr-parser/pprof/pyroscope"
    15  
    16  	distributormodel "github.com/grafana/pyroscope/pkg/distributor/model"
    17  	"github.com/grafana/pyroscope/pkg/pprof"
    18  
    19  	"github.com/grafana/pyroscope/pkg/og/ingestion"
    20  	"github.com/grafana/pyroscope/pkg/og/storage"
    21  	"github.com/grafana/pyroscope/pkg/og/util/form"
    22  )
    23  
    24  type RawProfile struct {
    25  	FormDataContentType string
    26  	RawData             []byte
    27  }
    28  
    29  func (p *RawProfile) Bytes() ([]byte, error) { return p.RawData, nil }
    30  
    31  func (p *RawProfile) ParseToPprof(ctx context.Context, md ingestion.Metadata, limits ingestion.Limits) (*distributormodel.PushRequest, error) {
    32  	input := jfrPprof.ParseInput{
    33  		StartTime:  md.StartTime,
    34  		EndTime:    md.EndTime,
    35  		SampleRate: int64(md.SampleRate),
    36  	}
    37  
    38  	tenantID, err := tenant.TenantID(ctx)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	maxBytes := limits.MaxProfileSizeBytes(tenantID)
    43  
    44  	labels := new(jfrPprof.LabelsSnapshot)
    45  	rawSize := len(p.RawData)
    46  	var r = p.RawData
    47  	if strings.Contains(p.FormDataContentType, "multipart/form-data") {
    48  		if r, labels, err = loadJFRFromForm(r, p.FormDataContentType, maxBytes); err != nil {
    49  			return nil, err
    50  		}
    51  	}
    52  
    53  	profiles, err := jfrPprof.ParseJFR(r, &input, labels)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	res := new(distributormodel.PushRequest)
    58  	for _, req := range profiles.Profiles {
    59  		seriesLabels := jfrPprofPyroscope.Labels(
    60  			md.LabelSet.Labels(),
    61  			profiles.JFREvent,
    62  			req.Metric,
    63  			md.LabelSet.ServiceName(),
    64  			md.SpyName,
    65  		)
    66  		res.Series = append(res.Series, &distributormodel.ProfileSeries{
    67  			Labels:  seriesLabels,
    68  			Profile: pprof.RawFromProto(req.Profile),
    69  		})
    70  	}
    71  	res.ReceivedCompressedProfileSize = rawSize
    72  	res.RawProfileType = distributormodel.RawProfileTypeJFR
    73  	return res, err
    74  }
    75  
    76  func (p *RawProfile) Parse(ctx context.Context, putter storage.Putter, _ storage.MetricsExporter, md ingestion.Metadata) error {
    77  	return fmt.Errorf("parsing to Tree/storage.Putter is no longer supported")
    78  }
    79  
    80  func (p *RawProfile) ContentType() string {
    81  	if p.FormDataContentType == "" {
    82  		return "binary/octet-stream"
    83  	}
    84  	return p.FormDataContentType
    85  }
    86  
    87  func loadJFRFromForm(r []byte, contentType string, maxBytes int) ([]byte, *jfrPprof.LabelsSnapshot, error) {
    88  	boundary, err := form.ParseBoundary(contentType)
    89  	if err != nil {
    90  		return nil, nil, err
    91  	}
    92  
    93  	f, err := multipart.NewReader(bytes.NewBuffer(r), boundary).ReadForm(32 << 20)
    94  	if err != nil {
    95  		return nil, nil, err
    96  	}
    97  	defer func() {
    98  		_ = f.RemoveAll()
    99  	}()
   100  
   101  	jfrField, err := form.ReadField(f, "jfr")
   102  	if err != nil {
   103  		return nil, nil, err
   104  	}
   105  	if jfrField == nil {
   106  		return nil, nil, fmt.Errorf("jfr field is required")
   107  	}
   108  	jfrField, err = decompress(jfrField, maxBytes)
   109  	if err != nil {
   110  		return nil, nil, fmt.Errorf("loadJFRFromForm failed to decompress jfr: %w", err)
   111  	}
   112  
   113  	labelsField, err := form.ReadField(f, "labels")
   114  	if err != nil {
   115  		return nil, nil, err
   116  	}
   117  
   118  	var labels = new(jfrPprof.LabelsSnapshot)
   119  	if len(labelsField) > 0 {
   120  		labelsField, err = decompress(labelsField, maxBytes)
   121  		if err != nil {
   122  			return nil, nil, fmt.Errorf("loadJFRFromForm failed to decompress labels: %w", err)
   123  		}
   124  		if err = labels.UnmarshalVT(labelsField); err != nil {
   125  			return nil, nil, fmt.Errorf("failed to parse labels form field: %w", err)
   126  		}
   127  	}
   128  
   129  	return jfrField, labels, nil
   130  }
   131  
   132  func decompress(bs []byte, maxBytes int) ([]byte, error) {
   133  	var err error
   134  	if len(bs) < 2 {
   135  		return nil, fmt.Errorf("failed to read magic")
   136  	} else if bs[0] == 0x1f && bs[1] == 0x8b {
   137  		var gzipr *gzip.Reader
   138  		gzipr, err = gzip.NewReader(bytes.NewReader(bs))
   139  		defer gzipr.Close()
   140  		if err != nil {
   141  			return nil, fmt.Errorf("failed to read gzip header: %w", err)
   142  		}
   143  		// Use maxBytes+1 to detect if limit is exceeded
   144  		limitReader := io.LimitReader(gzipr, int64(maxBytes+1))
   145  		buf := bytes.NewBuffer(nil)
   146  		if _, err = io.Copy(buf, limitReader); err != nil {
   147  			return nil, fmt.Errorf("failed to decompress jfr: %w", err)
   148  		}
   149  		if buf.Len() > maxBytes {
   150  			return nil, fmt.Errorf("decompressed size exceeds maximum allowed size of %d bytes", maxBytes)
   151  		}
   152  		return buf.Bytes(), nil
   153  	} else {
   154  		return bs, nil
   155  	}
   156  }