github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/benchmark/internal/server/server.go (about)

     1  package server
     2  
     3  import (
     4  	"bytes"
     5  	"compress/gzip"
     6  	"io"
     7  	"net/http"
     8  	"os"
     9  	"path/filepath"
    10  	"time"
    11  
    12  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
    13  	"github.com/sirupsen/logrus"
    14  	"google.golang.org/protobuf/proto"
    15  )
    16  
    17  func getProfiles(name string) []tree.Profile {
    18  	res := []tree.Profile{}
    19  
    20  	var p tree.Profile
    21  
    22  	paths, err := filepath.Glob("./fixtures/" + name + "-*.pprof")
    23  	if err != nil {
    24  		panic(err)
    25  	}
    26  	for _, path := range paths {
    27  		logrus.Infof("parsing profile %s", path)
    28  		gb, err := os.ReadFile(path)
    29  		if err != nil {
    30  			panic(err)
    31  		}
    32  		gr, err := gzip.NewReader(bytes.NewReader(gb))
    33  		if err != nil {
    34  			panic(err)
    35  		}
    36  		b, err := io.ReadAll(gr)
    37  		if err != nil {
    38  			panic(err)
    39  		}
    40  		if err := proto.Unmarshal(b, &p); err != nil {
    41  			panic(err)
    42  		}
    43  		res = append(res, p)
    44  	}
    45  
    46  	return res
    47  }
    48  
    49  func generateHandler(profiles []tree.Profile, sleep int) func(w http.ResponseWriter, r *http.Request) {
    50  	return func(w http.ResponseWriter, r *http.Request) {
    51  		i := int(time.Now().Unix()/10) % len(profiles)
    52  		p := profiles[i]
    53  		gw := gzip.NewWriter(w)
    54  		t := time.Now()
    55  		p.TimeNanos = t.UnixNano()
    56  
    57  		marshalled, err := proto.Marshal(&p)
    58  		if err != nil {
    59  			panic(err)
    60  		}
    61  
    62  		if sleep > 0 {
    63  			time.Sleep(time.Duration(sleep) * time.Second)
    64  		}
    65  
    66  		gw.Write(marshalled)
    67  		gw.Close()
    68  	}
    69  }
    70  
    71  func StartServer() {
    72  	m := http.NewServeMux()
    73  	m.HandleFunc("/debug/pprof/profile", generateHandler(getProfiles("cpu"), 10))
    74  	m.HandleFunc("/debug/pprof/heap", generateHandler(getProfiles("heap"), 0))
    75  
    76  	s := &http.Server{
    77  		Addr:           ":4042",
    78  		Handler:        m,
    79  		ReadTimeout:    30 * time.Second,
    80  		WriteTimeout:   30 * time.Second,
    81  		IdleTimeout:    30 * time.Second,
    82  		MaxHeaderBytes: 1 << 20,
    83  	}
    84  	s.ListenAndServe()
    85  }