github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/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/golang/protobuf/proto"
    13  
    14  	"github.com/pyroscope-io/pyroscope/pkg/ingestion"
    15  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    16  	"github.com/pyroscope-io/pyroscope/pkg/util/form"
    17  )
    18  
    19  type RawProfile struct {
    20  	FormDataContentType string
    21  	RawData             []byte
    22  }
    23  
    24  func (p *RawProfile) Bytes() ([]byte, error) { return p.RawData, nil }
    25  
    26  func (p *RawProfile) Parse(ctx context.Context, putter storage.Putter, _ storage.MetricsExporter, md ingestion.Metadata) error {
    27  	input := storage.PutInput{
    28  		StartTime:       md.StartTime,
    29  		EndTime:         md.EndTime,
    30  		Key:             md.Key,
    31  		SpyName:         md.SpyName,
    32  		SampleRate:      md.SampleRate,
    33  		Units:           md.Units,
    34  		AggregationType: md.AggregationType,
    35  	}
    36  
    37  	labels := new(LabelsSnapshot)
    38  	var r io.Reader = bytes.NewReader(p.RawData)
    39  	var err error
    40  	if strings.Contains(p.FormDataContentType, "multipart/form-data") {
    41  		if r, labels, err = loadJFRFromForm(r, p.FormDataContentType); err != nil {
    42  			return err
    43  		}
    44  	}
    45  
    46  	return ParseJFR(ctx, putter, r, &input, labels)
    47  }
    48  
    49  func (p *RawProfile) ContentType() string {
    50  	if p.FormDataContentType == "" {
    51  		return "binary/octet-stream"
    52  	}
    53  	return p.FormDataContentType
    54  }
    55  
    56  func loadJFRFromForm(r io.Reader, contentType string) (io.Reader, *LabelsSnapshot, error) {
    57  	boundary, err := form.ParseBoundary(contentType)
    58  	if err != nil {
    59  		return nil, nil, err
    60  	}
    61  
    62  	f, err := multipart.NewReader(r, boundary).ReadForm(32 << 20)
    63  	if err != nil {
    64  		return nil, nil, err
    65  	}
    66  	defer func() {
    67  		_ = f.RemoveAll()
    68  	}()
    69  
    70  	jfrField, err := form.ReadField(f, "jfr")
    71  	if err != nil {
    72  		return nil, nil, err
    73  	}
    74  	if jfrField == nil {
    75  		return nil, nil, fmt.Errorf("jfr field is required")
    76  	}
    77  	jfrField, err = decompress(jfrField)
    78  	if err != nil {
    79  		return nil, nil, fmt.Errorf("loadJFRFromForm failed to decompress jfr: %w", err)
    80  	}
    81  
    82  	labelsField, err := form.ReadField(f, "labels")
    83  	if err != nil {
    84  		return nil, nil, err
    85  	}
    86  
    87  	var labels LabelsSnapshot
    88  	if len(labelsField) > 0 {
    89  		labelsField, err = decompress(labelsField)
    90  		if err != nil {
    91  			return nil, nil, fmt.Errorf("loadJFRFromForm failed to decompress labels: %w", err)
    92  		}
    93  		if err = proto.Unmarshal(labelsField, &labels); err != nil {
    94  			return nil, nil, fmt.Errorf("failed to parse labels form field: %w", err)
    95  		}
    96  	}
    97  
    98  	return bytes.NewReader(jfrField), &labels, nil
    99  }
   100  
   101  func decompress(bs []byte) ([]byte, error) {
   102  	var err error
   103  	if len(bs) < 2 {
   104  		return nil, fmt.Errorf("failed to read magic")
   105  	} else if bs[0] == 0x1f && bs[1] == 0x8b {
   106  		var gzipr *gzip.Reader
   107  		gzipr, err = gzip.NewReader(bytes.NewReader(bs))
   108  		defer gzipr.Close()
   109  		if err != nil {
   110  			return nil, fmt.Errorf("failed to read gzip header: %w", err)
   111  		}
   112  		buf := bytes.NewBuffer(nil)
   113  		if _, err = io.Copy(buf, gzipr); err != nil {
   114  			return nil, fmt.Errorf("failed to decompress jfr: %w", err)
   115  		}
   116  		return buf.Bytes(), nil
   117  	} else {
   118  		return bs, nil
   119  	}
   120  }