profiler.go

package profiler

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/grafana/pyroscope-go"
	"time"

	log "github.com/sirupsen/logrus"
	"net"
	"net/http"
	"runtime"
	"sync"
)

var profiler *pyroscope.Profiler
var profilerMu sync.Mutex
var server http.Server

func init() {
	runtime.SetMutexProfileFraction(5)
	runtime.SetBlockProfileRate(5)

	server = http.Server{
		Addr:              ":6061",
		ReadTimeout:       60 * time.Second,
		ReadHeaderTimeout: 60 * time.Second,
		WriteTimeout:      60 * time.Second,
		IdleTimeout:       15 * time.Second,
	}

	http.HandleFunc("/profiler/toggle", restrict(profilerToggle))

	go func() {
		if err := server.ListenAndServe(); err != nil && !errors.Is(http.ErrServerClosed, err) {
			log.Error("failed on serve profiler server")
		}
	}()
}

type TogglingRequest struct {
	CollectorAddress string            `json:"collectorAddress"`
	ApplicationName  string            `json:"applicationName"`
	Tags             map[string]string `json:"tags"`
}

// func WithCustomMux()

func profilerToggle(w http.ResponseWriter, r *http.Request) {
	profilerMu.Lock()
	defer profilerMu.Unlock()

	if profiler != nil {
		err := stopProfiler()
		if err != nil {
			log.WithError(err).Error("error in toggle OFF profiler")
			http.Error(w, fmt.Sprintf("error while stoping profiler: %v", err), http.StatusInternalServerError)
			return
		}
	} else {
		var request TogglingRequest
		err := json.NewDecoder(r.Body).Decode(&request)
		if err != nil {
			http.Error(w, fmt.Sprintf("error while toggling profiler: %v", err), http.StatusBadRequest)
			return
		}

		defer r.Body.Close()
		log.Infof("toggling profiler ON")
		err = startProfiler(request.ApplicationName, request.CollectorAddress, request.Tags)
		if err != nil {
			log.WithError(err).Error("error in toggle ON profiler")
			http.Error(w, fmt.Sprintf("error while starting profiler: %v", err), http.StatusInternalServerError)
			return
		}
	}

}

// restrict middleware to prevent request from external clients other than localhost
func restrict(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ip, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			http.Error(w, "could not parse IP address", http.StatusInternalServerError)
			return
		}

		if ip != "127.0.0.1" && ip != "::1" {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		next.ServeHTTP(w, r)
	}
}

func stopProfiler() error {
	log.Infof("toggling proflier OFF")
	profiler.Flush(false)
	err := profiler.Stop()
	profiler = nil

	return err
}

func startProfiler(applicationName, collectorAddress string, tags map[string]string) error {
	var err error

	profiler, err = pyroscope.Start(pyroscope.Config{
		ApplicationName: applicationName,
		ServerAddress:   collectorAddress,
		Tags:            tags,
		ProfileTypes: []pyroscope.ProfileType{
			// these profile types are enabled by default:
			pyroscope.ProfileCPU,
			pyroscope.ProfileAllocObjects,
			pyroscope.ProfileAllocSpace,
			pyroscope.ProfileInuseObjects,
			pyroscope.ProfileInuseSpace,

			// these profile types are optional:
			pyroscope.ProfileGoroutines,
			pyroscope.ProfileMutexCount,
			pyroscope.ProfileMutexDuration,
			pyroscope.ProfileBlockCount,
			pyroscope.ProfileBlockDuration,
		},
	})
	return err
}