github.com/alibaba/ilogtail/pkg@v0.0.0-20250526110833-c53b480d046c/protocol/decoder/pyroscope/decoder.go (about)

     1  // Copyright 2023 iLogtail Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package pyroscope
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"net/http"
    21  	"strconv"
    22  	"strings"
    23  	"time"
    24  
    25  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    26  	"github.com/pyroscope-io/pyroscope/pkg/util/attime"
    27  
    28  	"github.com/alibaba/ilogtail/pkg/helper/profile"
    29  	"github.com/alibaba/ilogtail/pkg/helper/profile/pyroscope/jfr"
    30  	"github.com/alibaba/ilogtail/pkg/helper/profile/pyroscope/pprof"
    31  	"github.com/alibaba/ilogtail/pkg/helper/profile/pyroscope/raw"
    32  	"github.com/alibaba/ilogtail/pkg/logger"
    33  	"github.com/alibaba/ilogtail/pkg/models"
    34  	"github.com/alibaba/ilogtail/pkg/protocol"
    35  	"github.com/alibaba/ilogtail/pkg/protocol/decoder/common"
    36  )
    37  
    38  const AlarmType = "PYROSCOPE_ALARM"
    39  
    40  type Decoder struct {
    41  }
    42  
    43  func (d *Decoder) DecodeV2(data []byte, req *http.Request) (groups []*models.PipelineGroupEvents, err error) {
    44  	// do nothing
    45  	return nil, nil
    46  }
    47  
    48  func (d *Decoder) Decode(data []byte, req *http.Request, tags map[string]string) (logs []*protocol.Log, err error) {
    49  	in, err := d.extractRawInput(data, req)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	return in.Profile.Parse(context.Background(), &in.Metadata, tags)
    54  }
    55  
    56  func (d *Decoder) extractRawInput(data []byte, req *http.Request) (*profile.Input, error) {
    57  	in, ft, err := d.parseInputMeta(req)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  	ct := req.Header.Get("Content-Type")
    62  	var category string
    63  	switch {
    64  	case ft == profile.FormatPprof:
    65  		in.Profile = pprof.NewRawProfile(data, "")
    66  		category = "pprof"
    67  	case ft == profile.FormatJFR:
    68  		in.Profile = jfr.NewRawProfile(data, ct)
    69  		category = "JFR"
    70  	case strings.Contains(ct, "multipart/form-data"):
    71  		in.Profile = pprof.NewRawProfile(data, ct)
    72  		category = "pprof"
    73  	case ft == profile.FormatTrie, ct == "binary/octet-stream+trie":
    74  		in.Profile = raw.NewRawProfile(data, profile.FormatTrie)
    75  		category = "tire"
    76  	default:
    77  		in.Profile = raw.NewRawProfile(data, profile.FormatGroups)
    78  		category = "groups"
    79  	}
    80  	if logger.DebugFlag() {
    81  		var h string
    82  		for k, v := range req.Header {
    83  			h += "key: " + k + " val: " + strings.Join(v, ",")
    84  		}
    85  		logger.Debug(context.Background(), "CATEGORY", category, "URL", req.URL.Query().Encode(), "Header", h)
    86  	}
    87  	return in, nil
    88  }
    89  
    90  func (d *Decoder) ParseRequest(res http.ResponseWriter, req *http.Request, maxBodySize int64) (data []byte, statusCode int, err error) {
    91  	return common.CollectBody(res, req, maxBodySize)
    92  }
    93  
    94  func (d *Decoder) parseInputMeta(req *http.Request) (*profile.Input, profile.Format, error) {
    95  	var input profile.Input
    96  	q := req.URL.Query()
    97  
    98  	n := q.Get("name")
    99  	key, err := segment.ParseKey(n)
   100  	if err != nil {
   101  		logger.Error(context.Background(), AlarmType, "invalid name", n)
   102  		return nil, "", fmt.Errorf("pyroscope protocol get name err: %w", err)
   103  	}
   104  	name := key.AppName()
   105  	if strings.HasSuffix(name, ".cpu") {
   106  		key.Add("__name__", name[:len(name)-4])
   107  	}
   108  	input.Metadata.Tags = key.Labels()
   109  
   110  	if f := q.Get("from"); f != "" {
   111  		input.Metadata.StartTime = attime.Parse(f)
   112  	}
   113  	if input.Metadata.StartTime.IsZero() {
   114  		input.Metadata.StartTime = time.Now()
   115  	}
   116  
   117  	if f := q.Get("until"); f != "" {
   118  		input.Metadata.EndTime = attime.Parse(f)
   119  	}
   120  	if input.Metadata.StartTime.IsZero() {
   121  		input.Metadata.EndTime = time.Now()
   122  	}
   123  
   124  	input.Metadata.SampleRate = 100
   125  	if sr := q.Get("sampleRate"); sr != "" {
   126  		sampleRate, err := strconv.Atoi(sr)
   127  		if err != nil {
   128  			logger.Error(context.Background(), AlarmType, "invalid sampleRate", sr)
   129  		} else {
   130  			input.Metadata.SampleRate = uint32(sampleRate)
   131  		}
   132  	}
   133  
   134  	if sn := q.Get("spyName"); sn != "" {
   135  		sn = strings.TrimPrefix(sn, "pyroscope-")
   136  		sn = strings.TrimSuffix(sn, "spy")
   137  		input.Metadata.SpyName = sn
   138  	} else {
   139  		input.Metadata.SpyName = "unknown"
   140  	}
   141  
   142  	if u := q.Get("units"); u != "" {
   143  		input.Metadata.Units = profile.Units(u)
   144  	} else {
   145  		input.Metadata.Units = profile.SamplesUnits
   146  	}
   147  
   148  	if at := q.Get("aggregationType"); at != "" {
   149  		input.Metadata.AggregationType = profile.AggType(at)
   150  	} else {
   151  		input.Metadata.AggregationType = profile.SumAggType
   152  	}
   153  
   154  	format := q.Get("format")
   155  	return &input, profile.Format(format), nil
   156  }