github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/adhoc/writer/writer.go (about)

     1  package writer
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/sirupsen/logrus"
     9  
    10  	"github.com/pyroscope-io/pyroscope/pkg/config"
    11  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    12  	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
    13  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    14  	"github.com/pyroscope-io/pyroscope/pkg/structs/flamebearer"
    15  )
    16  
    17  type Writer struct {
    18  	maxNodesRender int
    19  	outputFormat   string
    20  	outputJSON     bool
    21  	logger         *logrus.Logger
    22  	storage        *storage.Storage
    23  
    24  	adhocDataDirWriter *AdhocDataDirWriter
    25  	stripTimestamp     bool
    26  }
    27  
    28  func NewWriter(cfg *config.Adhoc, st *storage.Storage, logger *logrus.Logger) Writer {
    29  	return Writer{
    30  		maxNodesRender:     cfg.MaxNodesRender,
    31  		outputFormat:       cfg.OutputFormat,
    32  		outputJSON:         !cfg.NoJSONOutput,
    33  		logger:             logger,
    34  		storage:            st,
    35  		adhocDataDirWriter: NewAdhocDataDirWriter(cfg.DataPath),
    36  		stripTimestamp:     cfg.StripTimestamp,
    37  	}
    38  }
    39  
    40  func (w Writer) Write(t0, t1 time.Time) error {
    41  	err := w.adhocDataDirWriter.EnsureExists()
    42  	if err != nil {
    43  		return err
    44  	}
    45  	ew, err := newExternalWriter(w.outputFormat, w.maxNodesRender, t0)
    46  	if err != nil {
    47  		return fmt.Errorf("could not create the external writer: %w", err)
    48  	}
    49  	defer ew.close() // It's fine to call this multiple times
    50  
    51  	profiles := 0
    52  
    53  	// The assumption is that these were the only ingested apps
    54  	for _, name := range w.storage.GetAppNames(context.TODO()) {
    55  		skey, err := segment.ParseKey(name)
    56  		if err != nil {
    57  			w.logger.WithError(err).Error("parsing storage key")
    58  			continue
    59  		}
    60  		gi := &storage.GetInput{
    61  			StartTime: t0,
    62  			EndTime:   t1,
    63  			Key:       skey,
    64  		}
    65  		out, err := w.storage.Get(context.TODO(), gi)
    66  
    67  		if err != nil {
    68  			w.logger.WithError(err).Error("retrieving storage key")
    69  			continue
    70  		}
    71  		if out == nil {
    72  			w.logger.Warn("no data retrieved")
    73  			continue
    74  		}
    75  
    76  		// TODO: Remove stripTimestamp flag and instead receive a formatter
    77  		if err := ew.write(name, out, w.stripTimestamp); err != nil {
    78  			w.logger.WithError(err).Error("saving output file")
    79  			continue
    80  		}
    81  
    82  		if w.outputJSON {
    83  			flame := flamebearer.NewProfile(flamebearer.ProfileConfig{
    84  				Name:      name,
    85  				MaxNodes:  w.maxNodesRender,
    86  				Tree:      out.Tree,
    87  				Timeline:  out.Timeline,
    88  				Groups:    out.Groups,
    89  				Telemetry: out.Telemetry,
    90  				Metadata: metadata.Metadata{
    91  					SpyName:         out.SpyName,
    92  					SampleRate:      out.SampleRate,
    93  					Units:           out.Units,
    94  					AggregationType: out.AggregationType,
    95  				},
    96  			})
    97  			filename := fmt.Sprintf("%s-%s.json", name, t0.Format("2006-01-02-15-04-05"))
    98  			path, err := w.adhocDataDirWriter.Write(filename, flame)
    99  			if err != nil {
   100  				w.logger.WithError(err).Error("saving to AdhocDir")
   101  				continue
   102  			}
   103  
   104  			w.logger.Infof("profiling data has been saved to %s", path)
   105  		}
   106  
   107  		profiles++
   108  	}
   109  
   110  	path, err := ew.close()
   111  	if err != nil {
   112  		w.logger.WithError(err).Error("closing external writer")
   113  	}
   114  	if path == "" {
   115  		if profiles == 0 {
   116  			w.logger.Warning("no profiling data was saved, maybe the profiled process didn't run long enough?")
   117  		} else {
   118  			w.logger.Info("you can now run `pyroscope server` and see the profiling data at http://localhost:4040/adhoc-single")
   119  		}
   120  	} else {
   121  		if profiles == 0 {
   122  			w.logger.Infof("profiling data was saved in '%s'", path)
   123  		} else {
   124  			w.logger.Infof(
   125  				"profiling data was saved in '%s' and you can also run `pyroscope server` to see the profiling data at http://localhost:4040/adhoc-single",
   126  				path,
   127  			)
   128  		}
   129  	}
   130  	return nil
   131  }