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

     1  package writer
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"os"
     7  	"path/filepath"
     8  	"time"
     9  
    10  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    11  	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
    12  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
    13  	"github.com/pyroscope-io/pyroscope/pkg/structs/flamebearer"
    14  	"github.com/pyroscope-io/pyroscope/webapp"
    15  	"google.golang.org/protobuf/proto"
    16  )
    17  
    18  type externalWriter struct {
    19  	format         string
    20  	maxNodesRender int
    21  	now            time.Time
    22  	dataDir        string
    23  	assetsDir      http.FileSystem
    24  	filenames      []string
    25  }
    26  
    27  // newExternalWriter creates a writer of profile trees to external formats (see isSupported for supported formats).
    28  // The writer will store all the profiles in a temporary directory
    29  // and once its closed it'll move the profiles to the current directory.
    30  // If there's a single profile, the profile file is moved instead of the directory.
    31  func newExternalWriter(format string, maxNodesRender int, now time.Time) (*externalWriter, error) {
    32  	var (
    33  		dataDir   string
    34  		assetsDir http.FileSystem
    35  		err       error
    36  	)
    37  
    38  	if format == "html" {
    39  		assetsDir, err = webapp.Assets()
    40  		if err != nil {
    41  			return nil, fmt.Errorf("could not get the asset directory: %w", err)
    42  		}
    43  	}
    44  
    45  	if format != "none" {
    46  		dataDir = fmt.Sprintf("pyroscope-adhoc-%s", now.Format("2006-01-02-15-04-05"))
    47  		if err := os.MkdirAll(dataDir, os.ModeDir|os.ModePerm); err != nil {
    48  			return nil, fmt.Errorf("could not create directory for external output: %w", err)
    49  		}
    50  	}
    51  
    52  	return &externalWriter{
    53  		format:         format,
    54  		maxNodesRender: maxNodesRender,
    55  		dataDir:        dataDir,
    56  		assetsDir:      assetsDir,
    57  		now:            now,
    58  	}, nil
    59  }
    60  
    61  func (w *externalWriter) write(name string, out *storage.GetOutput, stripTimestamp bool) error {
    62  	if w.format == "none" {
    63  		return nil
    64  	}
    65  	var ext string
    66  	if w.format == "collapsed" {
    67  		ext = "collapsed.txt"
    68  	} else {
    69  		ext = w.format
    70  	}
    71  
    72  	filename := fmt.Sprintf("%s-%s.%s", name, w.now.Format("2006-01-02-15-04-05"), ext)
    73  	if stripTimestamp {
    74  		filename = fmt.Sprintf("%s.%s", name, ext)
    75  	}
    76  
    77  	path := filepath.Join(w.dataDir, filename)
    78  
    79  	f, err := os.Create(path)
    80  	if err != nil {
    81  		return fmt.Errorf("could not create temporary path %s: %w", path, err)
    82  	}
    83  	defer f.Close()
    84  
    85  	switch w.format {
    86  	case "pprof":
    87  		pprof := out.Tree.Pprof(&tree.PprofMetadata{
    88  			// TODO(petethepig): check if this conversion always makes sense
    89  			//   e.g are these units defined in pprof somewhere?
    90  			Unit:      string(out.Units),
    91  			StartTime: w.now,
    92  		})
    93  		out, err := proto.Marshal(pprof)
    94  		if err != nil {
    95  			return fmt.Errorf("could not serialize to pprof: %w", err)
    96  		}
    97  		if _, err := f.Write(out); err != nil {
    98  			return fmt.Errorf("could not write the pprof file: %w", err)
    99  		}
   100  	case "collapsed":
   101  		if _, err := f.WriteString(out.Tree.Collapsed()); err != nil {
   102  			return fmt.Errorf("could not write the collapsed file: %w", err)
   103  		}
   104  	case "html":
   105  		fb := flamebearer.NewProfile(flamebearer.ProfileConfig{
   106  			Name:      filename,
   107  			MaxNodes:  w.maxNodesRender,
   108  			Tree:      out.Tree,
   109  			Timeline:  out.Timeline,
   110  			Groups:    out.Groups,
   111  			Telemetry: out.Telemetry,
   112  			Metadata: metadata.Metadata{
   113  				SpyName:         out.SpyName,
   114  				SampleRate:      out.SampleRate,
   115  				Units:           out.Units,
   116  				AggregationType: out.AggregationType,
   117  			},
   118  		})
   119  		if err := flamebearer.FlamebearerToStandaloneHTML(&fb, w.assetsDir, f); err != nil {
   120  			return fmt.Errorf("could not write the standalone HTML file: %w", err)
   121  		}
   122  	}
   123  
   124  	w.filenames = append(w.filenames, filename)
   125  	return nil
   126  }
   127  
   128  func (w *externalWriter) close() (string, error) {
   129  	if w.format == "none" {
   130  		return "", nil
   131  	}
   132  	w.format = "none"
   133  	switch len(w.filenames) {
   134  	case 0:
   135  		if err := os.Remove(w.dataDir); err != nil {
   136  			return "", fmt.Errorf("could not remove directory %s: %w", w.dataDir, err)
   137  		}
   138  		return "", nil
   139  	case 1:
   140  		path := filepath.Join(w.dataDir, w.filenames[0])
   141  		if err := os.Rename(path, w.filenames[0]); err != nil {
   142  			return "", fmt.Errorf("could not rename %s to %s: %w", w.filenames[0], path, err)
   143  		}
   144  		if err := os.Remove(w.dataDir); err != nil {
   145  			return "", fmt.Errorf("could not remove directory %s: %w", w.dataDir, err)
   146  		}
   147  		return w.filenames[0], nil
   148  	}
   149  	return w.dataDir, nil
   150  }