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

     1  package dbmanager
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"os/signal"
     8  	"syscall"
     9  	"time"
    10  
    11  	"github.com/cheggaaa/pb/v3"
    12  	"github.com/prometheus/client_golang/prometheus"
    13  	"github.com/pyroscope-io/pyroscope/pkg/selfprofiling"
    14  	"github.com/sirupsen/logrus"
    15  
    16  	"github.com/pyroscope-io/pyroscope/pkg/config"
    17  	"github.com/pyroscope-io/pyroscope/pkg/exporter"
    18  	"github.com/pyroscope-io/pyroscope/pkg/health"
    19  	"github.com/pyroscope-io/pyroscope/pkg/parser"
    20  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    21  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    22  )
    23  
    24  func Cli(dbCfg *config.DbManager, srvCfg *config.Server, args []string) error {
    25  	switch args[0] {
    26  	case "copy":
    27  		stCfg := storage.NewConfig(srvCfg).WithPath(dbCfg.StoragePath)
    28  		err := copyData(dbCfg, stCfg)
    29  		if err != nil {
    30  			return err
    31  		}
    32  	default:
    33  		return fmt.Errorf("unknown command %q", args[0])
    34  	}
    35  
    36  	return nil
    37  }
    38  
    39  // TODO: get this from config or something like that
    40  const resolution = 10 * time.Second
    41  
    42  // src start time, src end time, dst start time
    43  func copyData(dbCfg *config.DbManager, stCfg *storage.Config) error {
    44  	appName := dbCfg.ApplicationName
    45  	srcSt := dbCfg.SrcStartTime.Truncate(resolution)
    46  	dstSt := dbCfg.DstStartTime.Truncate(resolution)
    47  	dstEt := dbCfg.DstEndTime.Truncate(resolution)
    48  	srcEt := srcSt.Add(dstEt.Sub(dstSt))
    49  
    50  	fmt.Printf("copying %s from %s-%s to %s-%s\n",
    51  		appName,
    52  		srcSt.String(),
    53  		srcEt.String(),
    54  		dstSt.String(),
    55  		dstEt.String(),
    56  	)
    57  
    58  	// TODO: add more correctness checks
    59  	if !srcSt.Before(srcEt) || !dstSt.Before(dstEt) {
    60  		return fmt.Errorf("Incorrect time parameters. Start time has to be before end time. "+
    61  			"src start: %q end: %q, dst start: %q end: %q", srcSt, srcEt, dstSt, dstEt)
    62  	}
    63  
    64  	s, err := storage.New(stCfg, logrus.StandardLogger(), prometheus.DefaultRegisterer, new(health.Controller), storage.NoopApplicationMetadataService{})
    65  	if err != nil {
    66  		return err
    67  	}
    68  
    69  	e, _ := exporter.NewExporter(nil, nil)
    70  	logger := logrus.StandardLogger()
    71  	if dbCfg.EnableProfiling {
    72  		_ = selfprofiling.NewSession(logger, parser.New(logger, s, e), "pyroscope.dbmanager", map[string]string{}).Start()
    73  	}
    74  
    75  	sk, err := segment.ParseKey(appName)
    76  	if err != nil {
    77  		return err
    78  	}
    79  
    80  	count := int(srcEt.Sub(srcSt) / (resolution))
    81  	bar := pb.StartNew(count)
    82  
    83  	durDiff := dstSt.Sub(srcSt)
    84  
    85  	sigc := make(chan os.Signal, 1)
    86  	signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM)
    87  
    88  loop:
    89  	for srct := srcSt; srct.Before(srcEt); srct = srct.Add(resolution) {
    90  		bar.Increment()
    91  		select {
    92  		case <-sigc:
    93  			break loop
    94  		default:
    95  		}
    96  
    97  		srct2 := srct.Add(resolution)
    98  		gOut, err := s.Get(context.TODO(), &storage.GetInput{
    99  			StartTime: srct,
   100  			EndTime:   srct2,
   101  			Key:       sk,
   102  		})
   103  		if err != nil {
   104  			return err
   105  		}
   106  
   107  		if gOut.Tree != nil {
   108  			dstt := srct.Add(durDiff)
   109  			dstt2 := dstt.Add(resolution)
   110  
   111  			err = s.Put(context.TODO(), &storage.PutInput{
   112  				StartTime:  dstt,
   113  				EndTime:    dstt2,
   114  				Key:        sk,
   115  				Val:        gOut.Tree,
   116  				SpyName:    gOut.SpyName,
   117  				SampleRate: gOut.SampleRate,
   118  				Units:      gOut.Units,
   119  			})
   120  			if err != nil {
   121  				return err
   122  			}
   123  		}
   124  	}
   125  
   126  	bar.Finish()
   127  	return s.Close()
   128  }