github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/client/entry/wait_unix.go (about)

     1  //go:build linux || darwin
     2  
     3  package entry
     4  
     5  import (
     6  	"fmt"
     7  	"os"
     8  	"os/signal"
     9  	"path/filepath"
    10  	"runtime"
    11  	"runtime/pprof"
    12  	"syscall"
    13  	"time"
    14  
    15  	log "github.com/golang/glog"
    16  
    17  	"github.com/google/fleetspeak/fleetspeak/src/client"
    18  )
    19  
    20  // Self-imposed timeout for shutting down gracefully.
    21  var shutdownTimeout = 20 * time.Second
    22  
    23  // dumpProfile writes the given pprof profile to disk with the given debug flag.
    24  func dumpProfile(profileDir string, profileName string, pprofDebugFlag int) error {
    25  	profileDumpPath := filepath.Join(profileDir, fmt.Sprintf("fleetspeakd-%s-pprof-%d-%v", profileName, os.Getpid(), time.Now().Format("2006-01-02-15-04-05.000")))
    26  	fileWriter, err := os.Create(profileDumpPath)
    27  	if err != nil {
    28  		return fmt.Errorf("unable to create profile file [%s]: %v", profileDumpPath, err)
    29  	}
    30  	defer fileWriter.Close()
    31  	if err := pprof.Lookup(profileName).WriteTo(fileWriter, pprofDebugFlag); err != nil {
    32  		return fmt.Errorf("unable to write %s profile [%s]: %v", profileName, profileDumpPath, err)
    33  	}
    34  	log.Infof("%s profile dumped to [%s].", profileName, profileDumpPath)
    35  	return nil
    36  }
    37  
    38  func Wait(cl *client.Client, profileDir string) {
    39  	s := make(chan os.Signal, 3)
    40  	signal.Notify(s, os.Interrupt, syscall.SIGTERM, syscall.SIGUSR1)
    41  	defer signal.Stop(s)
    42  
    43  	for si := range s {
    44  		switch si {
    45  		case os.Interrupt, syscall.SIGTERM:
    46  			log.Infof("Received signal %v, cleaning up...", si)
    47  			// Give ourselves a deadline within which to shut-down gracefully.
    48  			shutdownTimer := time.AfterFunc(shutdownTimeout, func() {
    49  				if err := dumpProfile(profileDir, "goroutine", 2); err != nil {
    50  					log.Errorf("Failed to dump goroutine profile: %v", err)
    51  				}
    52  				log.Exitf("Fleetspeak received SIGTERM, but failed to shut down in %v. Exiting ungracefully.", shutdownTimeout)
    53  			})
    54  			cl.Stop()
    55  			shutdownTimer.Stop()
    56  			return
    57  		case syscall.SIGUSR1:
    58  			log.Infof("Received signal %v, writing profile.", si)
    59  			runtime.GC()
    60  			if err := dumpProfile(profileDir, "heap", 0); err != nil {
    61  				log.Errorf("Failed to dump heap profile: %v", err)
    62  			}
    63  		}
    64  	}
    65  }