sigs.k8s.io/prow@v0.0.0-20240503223140-c5e374dc7eb1/pkg/pjutil/pprof/pprof.go (about)

     1  /*
     2  Copyright 2021 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Package pprof contains helpers for profiling binaries.
    18  package pprof
    19  
    20  import (
    21  	"net/http"
    22  	"net/http/pprof"
    23  	"os"
    24  	"runtime"
    25  	runtimepprof "runtime/pprof"
    26  	"strconv"
    27  	"time"
    28  
    29  	"github.com/felixge/fgprof"
    30  	"github.com/sirupsen/logrus"
    31  	"sigs.k8s.io/prow/pkg/flagutil"
    32  
    33  	"sigs.k8s.io/prow/pkg/interrupts"
    34  )
    35  
    36  // Instrument implements the profiling options a user has asked for on the command line.
    37  func Instrument(opts flagutil.InstrumentationOptions) {
    38  	Serve(opts.PProfPort)
    39  	if opts.ProfileMemory {
    40  		WriteMemoryProfiles(opts.MemoryProfileInterval)
    41  	}
    42  }
    43  
    44  // Serve sets up a handler for pprof debug endpoints and starts a server for them asynchronously.
    45  // The contents of this function are identical to what the `net/http/pprof` package does on import for
    46  // the simple case where the default mux is to be used, but with a custom mux to ensure we don't serve
    47  // this data from an exposed port.
    48  func Serve(port int) {
    49  	pprofMux := http.NewServeMux()
    50  	pprofMux.HandleFunc("/debug/pprof/", pprof.Index)
    51  	pprofMux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
    52  	pprofMux.HandleFunc("/debug/pprof/profile", pprof.Profile)
    53  	pprofMux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
    54  	pprofMux.HandleFunc("/debug/pprof/trace", pprof.Trace)
    55  	pprofMux.Handle("/debug/fgprof", fgprof.Handler())
    56  	server := &http.Server{Addr: ":" + strconv.Itoa(port), Handler: pprofMux}
    57  	interrupts.ListenAndServe(server, 5*time.Second)
    58  }
    59  
    60  // WriteMemoryProfiles is a non-blocking, best-effort routine to dump memory profiles at a
    61  // pre-determined interval for future parsing and analysis.
    62  func WriteMemoryProfiles(interval time.Duration) {
    63  	logrus.Info("Writing memory profiles.")
    64  	profileDir, err := os.MkdirTemp("", "heap-profiles-")
    65  	if err != nil {
    66  		logrus.WithError(err).Warn("Could not create a directory to store memory profiles.")
    67  		return
    68  	}
    69  	interrupts.TickLiteral(func() {
    70  		profile, err := os.CreateTemp(profileDir, "heap-profile-")
    71  		if err != nil {
    72  			logrus.WithError(err).Warn("Could not create a file to store a memory profile.")
    73  			return
    74  		}
    75  		logrus.Info("Writing a memory profile.")
    76  		runtime.GC() // ensure we have up-to-date data
    77  		if err := runtimepprof.WriteHeapProfile(profile); err != nil {
    78  			logrus.WithError(err).Warn("Could not write memory profile.")
    79  		}
    80  		logrus.Infof("Wrote memory profile to %s.", profile.Name())
    81  		if err := profile.Close(); err != nil {
    82  			logrus.WithError(err).Warn("Could not close file storing memory profile.")
    83  		}
    84  	}, interval)
    85  }