main.go raw

   1  package interrupt
   2  
   3  import (
   4  	"fmt"
   5  	"os"
   6  	"os/exec"
   7  	"os/signal"
   8  	"runtime"
   9  	"strings"
  10  	"syscall"
  11  	
  12  	uberatomic "go.uber.org/atomic"
  13  	
  14  	"github.com/p9c/p9/pkg/qu"
  15  	
  16  	"github.com/kardianos/osext"
  17  )
  18  
  19  type HandlerWithSource struct {
  20  	Source string
  21  	Fn     func()
  22  }
  23  
  24  var (
  25  	Restart bool // = true
  26  	requested uberatomic.Bool
  27  	// ch is used to receive SIGINT (Ctrl+C) signals.
  28  	ch chan os.Signal
  29  	// signals is the list of signals that cause the interrupt
  30  	signals = []os.Signal{os.Interrupt}
  31  	// ShutdownRequestChan is a channel that can receive shutdown requests
  32  	ShutdownRequestChan = qu.T()
  33  	// addHandlerChan is used to add an interrupt handler to the list of
  34  	// handlers to be invoked on SIGINT (Ctrl+C) signals.
  35  	addHandlerChan = make(chan HandlerWithSource)
  36  	// HandlersDone is closed after all interrupt handlers run the first time
  37  	// an interrupt is signaled.
  38  	HandlersDone = make(qu.C)
  39  )
  40  
  41  var interruptCallbacks []func()
  42  var interruptCallbackSources []string
  43  
  44  // Listener listens for interrupt signals, registers interrupt callbacks,
  45  // and responds to custom shutdown signals as required
  46  func Listener() {
  47  	invokeCallbacks := func() {
  48  		D.Ln(
  49  			"running interrupt callbacks",
  50  			len(interruptCallbacks),
  51  			strings.Repeat(" ", 48),
  52  			interruptCallbackSources,
  53  		)
  54  		// run handlers in LIFO order.
  55  		for i := range interruptCallbacks {
  56  			idx := len(interruptCallbacks) - 1 - i
  57  			D.Ln("running callback", idx, interruptCallbackSources[idx])
  58  			interruptCallbacks[idx]()
  59  		}
  60  		D.Ln("interrupt handlers finished")
  61  		HandlersDone.Q()
  62  		if Restart {
  63  			var file string
  64  			var e error
  65  			file, e = osext.Executable()
  66  			if e != nil {
  67  				E.Ln(e)
  68  				return
  69  			}
  70  			D.Ln("restarting")
  71  			if runtime.GOOS != "windows" {
  72  				e = syscall.Exec(file, os.Args, os.Environ())
  73  				if e != nil {
  74  					F.Ln(e)
  75  				}
  76  			} else {
  77  				D.Ln("doing windows restart")
  78  				
  79  				// procAttr := new(os.ProcAttr)
  80  				// procAttr.Files = []*os.File{os.Stdin, os.Stdout, os.Stderr}
  81  				// os.StartProcess(os.Args[0], os.Args[1:], procAttr)
  82  				
  83  				var s []string
  84  				// s = []string{"cmd.exe", "/C", "start"}
  85  				s = append(s, os.Args[0])
  86  				// s = append(s, "--delaystart")
  87  				s = append(s, os.Args[1:]...)
  88  				cmd := exec.Command(s[0], s[1:]...)
  89  				D.Ln("windows restart done")
  90  				if e = cmd.Start(); E.Chk(e) {
  91  				}
  92  				// // select{}
  93  				// os.Exit(0)
  94  			}
  95  		}
  96  		// time.Sleep(time.Second * 3)
  97  		// os.Exit(1)
  98  		// close(HandlersDone)
  99  	}
 100  out:
 101  	for {
 102  		select {
 103  		case sig := <-ch:
 104  			// if !requested {
 105  			// 	L.Printf("\r>>> received signal (%s)\n", sig)
 106  			D.Ln("received interrupt signal", sig)
 107  			requested.Store(true)
 108  			invokeCallbacks()
 109  			// pprof.Lookup("goroutine").WriteTo(os.Stderr, 2)
 110  			// }
 111  			break out
 112  		case <-ShutdownRequestChan.Wait():
 113  			// if !requested {
 114  			W.Ln("received shutdown request - shutting down...")
 115  			requested.Store(true)
 116  			invokeCallbacks()
 117  			break out
 118  			// }
 119  		case handler := <-addHandlerChan:
 120  			// if !requested {
 121  			// D.Ln("adding handler")
 122  			interruptCallbacks = append(interruptCallbacks, handler.Fn)
 123  			interruptCallbackSources = append(interruptCallbackSources, handler.Source)
 124  			// }
 125  		case <-HandlersDone.Wait():
 126  			break out
 127  		}
 128  	}
 129  }
 130  
 131  // AddHandler adds a handler to call when a SIGINT (Ctrl+C) is received.
 132  func AddHandler(handler func()) {
 133  	// Create the channel and start the main interrupt handler which invokes all other callbacks and exits if not
 134  	// already done.
 135  	_, loc, line, _ := runtime.Caller(1)
 136  	msg := fmt.Sprintf("%s:%d", loc, line)
 137  	D.Ln("handler added by:", msg)
 138  	if ch == nil {
 139  		ch = make(chan os.Signal)
 140  		signal.Notify(ch, signals...)
 141  		go Listener()
 142  	}
 143  	addHandlerChan <- HandlerWithSource{
 144  		msg, handler,
 145  	}
 146  }
 147  
 148  // Request programmatically requests a shutdown
 149  func Request() {
 150  	_, f, l, _ := runtime.Caller(1)
 151  	D.Ln("interrupt requested", f, l, requested.Load())
 152  	if requested.Load() {
 153  		D.Ln("requested again")
 154  		return
 155  	}
 156  	requested.Store(true)
 157  	ShutdownRequestChan.Q()
 158  	// qu.PrintChanState()
 159  	var ok bool
 160  	select {
 161  	case _, ok = <-ShutdownRequestChan:
 162  	default:
 163  	}
 164  	D.Ln("shutdownrequestchan", ok)
 165  	if ok {
 166  		close(ShutdownRequestChan)
 167  	}
 168  }
 169  
 170  // GoroutineDump returns a string with the current goroutine dump in order to show what's going on in case of timeout.
 171  func GoroutineDump() string {
 172  	buf := make([]byte, 1<<18)
 173  	n := runtime.Stack(buf, true)
 174  	return string(buf[:n])
 175  }
 176  
 177  // RequestRestart sets the reset flag and requests a restart
 178  func RequestRestart() {
 179  	Restart = true
 180  	D.Ln("requesting restart")
 181  	Request()
 182  }
 183  
 184  // Requested returns true if an interrupt has been requested
 185  func Requested() bool {
 186  	return requested.Load()
 187  }
 188