github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/sigx/signal.go (about)

     1  package sigx
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io/ioutil"
     7  	"log"
     8  	"os"
     9  	"os/signal"
    10  	"runtime/pprof"
    11  	"syscall"
    12  	"time"
    13  
    14  	"github.com/bingoohuang/gg/pkg/iox"
    15  	"github.com/bingoohuang/gg/pkg/osx"
    16  	"github.com/bingoohuang/gg/pkg/profile"
    17  )
    18  
    19  // RegisterSignals registers signal handlers.
    20  func RegisterSignals(c context.Context, signals ...os.Signal) (context.Context, context.CancelFunc) {
    21  	if c == nil {
    22  		c = context.Background()
    23  	}
    24  	cc, cancel := context.WithCancel(c)
    25  	sig := make(chan os.Signal, 1)
    26  	if len(signals) == 0 {
    27  		// syscall.SIGINT: ctl + c, syscall.SIGTERM: kill pid
    28  		signals = []os.Signal{syscall.SIGINT, syscall.SIGTERM}
    29  	}
    30  	signal.Notify(sig, signals...)
    31  	go func() {
    32  		<-sig
    33  		cancel()
    34  	}()
    35  
    36  	return cc, cancel
    37  }
    38  
    39  func RegisterSignalCallback(f func(), signals ...os.Signal) {
    40  	sig := make(chan os.Signal, 1)
    41  	signal.Notify(sig, signals...)
    42  	go func() {
    43  		for range sig {
    44  			f()
    45  		}
    46  	}()
    47  }
    48  
    49  func RegisterSignalProfile(signals ...os.Signal) {
    50  	if len(signals) == 0 {
    51  		signals = defaultSignals
    52  	}
    53  
    54  	RegisterSignalCallback(func() {
    55  		val, ok := ReadCmdOK("jj.cpu", false)
    56  		if ok {
    57  			collectCpuProfile()
    58  			if val = bytes.TrimSpace(val); len(val) > 0 {
    59  				if duration, err := time.ParseDuration(string(val)); err != nil {
    60  					log.Printf("ignore duration %s in jj.cpu, parse failed: %v", val, err)
    61  				} else if duration > 0 {
    62  					log.Printf("after %s, cpu.profile will be generated ", val)
    63  					go func() {
    64  						time.Sleep(duration)
    65  						collectCpuProfile()
    66  					}()
    67  				}
    68  			}
    69  		}
    70  
    71  		if HasCmd("jj.mem", true) {
    72  			if err := CollectMemProfile("mem.profile"); err != nil {
    73  				log.Printf("failed to collect profile: %v", err)
    74  			}
    75  		}
    76  		if v := ReadCmd("jj.profile"); len(v) > 0 {
    77  			go profile.Start(profile.Specs(string(v)))
    78  		}
    79  	}, signals...)
    80  }
    81  
    82  func collectCpuProfile() {
    83  	if completed, err := CollectCpuProfile("cpu.profile"); err != nil {
    84  		log.Printf("failed to collect profile: %v", err)
    85  	} else if completed {
    86  		osx.Remove("jj.cpu")
    87  	}
    88  }
    89  
    90  var cpuProfileFile *os.File
    91  
    92  func HasCmd(f string, remove bool) bool {
    93  	s, err := os.Stat(f)
    94  	if err == nil && !s.IsDir() {
    95  		if remove {
    96  			osx.Remove(f)
    97  		}
    98  		return true
    99  	}
   100  
   101  	return false
   102  }
   103  
   104  func ReadCmdOK(f string, remove bool) ([]byte, bool) {
   105  	s, err := os.Stat(f)
   106  	if err == nil && !s.IsDir() {
   107  		data, _ := ioutil.ReadFile(f)
   108  		if remove {
   109  			osx.Remove(f)
   110  		}
   111  		return data, true
   112  	}
   113  
   114  	return nil, false
   115  }
   116  
   117  func ReadCmd(f string) []byte {
   118  	s, err := os.Stat(f)
   119  	if err == nil && !s.IsDir() {
   120  		data, _ := ioutil.ReadFile(f)
   121  		osx.Remove(f)
   122  		return data
   123  	}
   124  
   125  	return nil
   126  }
   127  
   128  func CollectCpuProfile(cpuProfile string) (bool, error) {
   129  	if cpuProfile == "" {
   130  		return false, nil
   131  	}
   132  
   133  	if cpuProfileFile != nil {
   134  		pprof.StopCPUProfile()
   135  		iox.Close(cpuProfileFile)
   136  
   137  		log.Printf("%s collected", cpuProfileFile.Name())
   138  		cpuProfileFile = nil
   139  		return true, nil
   140  	}
   141  
   142  	f, err := os.Create(cpuProfile)
   143  	if err != nil {
   144  		return false, err
   145  	}
   146  	cpuProfileFile = f
   147  
   148  	if err := pprof.StartCPUProfile(f); err != nil {
   149  		return false, err
   150  	}
   151  
   152  	log.Printf("%s started", cpuProfile)
   153  	return false, nil
   154  }
   155  
   156  func CollectMemProfile(memProfile string) error {
   157  	if memProfile == "" {
   158  		return nil
   159  	}
   160  
   161  	f, err := os.Create(memProfile)
   162  	if err != nil {
   163  		return err
   164  	}
   165  	defer f.Close()
   166  
   167  	return pprof.WriteHeapProfile(f)
   168  }