Harden shutdown logic

cmd/ipfs/daemon.go

@@ -80,6 +80,15 @@ func daemonFunc(req cmds.Request, res cmds.Response) {
  // let the user know we're going.
  fmt.Printf("Initializing daemon...\n")
 
+ ctx := req.Context()
+
+ go func() {
+ select {
+ case <-ctx.Context.Done():
+ fmt.Println("Received interrupt signal, shutting down...")
+ }
+ }()
+
  // first, whether user has provided the initialization flag. we may be
  // running in an uninitialized state.
  initialize, _, err := req.Option(initOptionKwd).Bool()
@@ -111,7 +120,6 @@ func daemonFunc(req cmds.Request, res cmds.Response) {
  return
  }
 
- ctx := req.Context()
  cfg, err := ctx.GetConfig()
  if err != nil {
  res.SetError(err, cmds.ErrNormal)
@@ -149,7 +157,19 @@ func daemonFunc(req cmds.Request, res cmds.Response) {
  res.SetError(err, cmds.ErrNormal)
  return
  }
- defer node.Close()
+
+ defer func() {
+ // We wait for the node to close first, as the node has children
+ // that it will wait for before closing, such as the API server.
+ node.Close()
+
+ select {
+ case <-ctx.Context.Done():
+ log.Info("Gracefully shut down daemon")
+ default:
+ }
+ }()
+
  req.Context().ConstructNode = func() (*core.IpfsNode, error) {
  return node, nil
  }
@@ -262,9 +282,6 @@ func daemonFunc(req cmds.Request, res cmds.Response) {
  corehttp.VersionOption(),
  }
 
- // our global interrupt handler can now try to stop the daemon
- close(req.Context().InitDone)
-
  if rootRedirect != nil {
  opts = append(opts, rootRedirect)
  }

cmd/ipfs/main.go

@@ -11,6 +11,7 @@ import (
  "runtime"
  "runtime/pprof"
  "strings"
+ "sync"
  "syscall"
  "time"
 
@@ -39,7 +40,6 @@ const (
  cpuProfile = "ipfs.cpuprof"
  heapProfile = "ipfs.memprof"
  errorFormat = "ERROR: %v\n\n"
- shutdownMessage = "Received interrupt signal, shutting down..."
 )
 
 type cmdInvocation struct {
@@ -132,15 +132,10 @@ func main() {
  os.Exit(1)
  }
 
- // our global interrupt handler may try to stop the daemon
- // before the daemon is ready to be stopped; this dirty
- // workaround is for the daemon only; other commands are always
- // ready to be stopped
- if invoc.cmd != daemonCmd {
- close(invoc.req.Context().InitDone)
- }
-
  // ok, finally, run the command invocation.
+ intrh, ctx := invoc.SetupInterruptHandler(ctx)
+ defer intrh.Close()
+
  output, err := invoc.Run(ctx)
  if err != nil {
  printErr(err)
@@ -157,8 +152,6 @@ func main() {
 }
 
 func (i *cmdInvocation) Run(ctx context.Context) (output io.Reader, err error) {
- // setup our global interrupt handler.
- i.setupInterruptHandler()
 
  // check if user wants to debug. option OR env var.
  debug, _, err := i.req.Option("debug").Bool()
@@ -226,7 +219,6 @@ func (i *cmdInvocation) Parse(ctx context.Context, args []string) error {
  if err != nil {
  return err
  }
- i.req.Context().Context = ctx
 
  repoPath, err := getRepoPath(i.req)
  if err != nil {
@@ -279,6 +271,8 @@ func callCommand(ctx context.Context, req cmds.Request, root *cmds.Command, cmd
  log.Info(config.EnvDir, " ", req.Context().ConfigRoot)
  var res cmds.Response
 
+ req.Context().Context = ctx
+
  details, err := commandDetails(req.Path(), root)
  if err != nil {
  return nil, err
@@ -474,59 +468,70 @@ func writeHeapProfileToFile() error {
  return pprof.WriteHeapProfile(mprof)
 }
 
-// listen for and handle SIGTERM
-func (i *cmdInvocation) setupInterruptHandler() {
+// IntrHandler helps set up an interrupt handler that can
+// be cleanly shut down through the io.Closer interface.
+type IntrHandler struct {
+ sig chan os.Signal
+ wg sync.WaitGroup
+}
+
+func NewIntrHandler() *IntrHandler {
+ ih := &IntrHandler{}
+ ih.sig = make(chan os.Signal, 1)
+ return ih
+}
+
+func (ih *IntrHandler) Close() error {
+ close(ih.sig)
+ ih.wg.Wait()
+ return nil
+}
 
- ctx := i.req.Context()
- sig := allInterruptSignals()
 
+// Handle starts handling the given signals, and will call the handler
+// callback function each time a signal is catched. The function is passed
+// the number of times the handler has been triggered in total, as
+// well as the handler itself, so that the handling logic can use the
+// handler's wait group to ensure clean shutdown when Close() is called.
+func (ih *IntrHandler) Handle(handler func(count int, ih *IntrHandler), sigs ...os.Signal) {
+ signal.Notify(ih.sig, sigs...)
+ ih.wg.Add(1)
  go func() {
- // first time, try to shut down.
-
- // loop because we may be
- for count := 0; ; count++ {
- <-sig
-
- // if we're still initializing, cannot use `ctx.GetNode()`
- select {
- default: // initialization not done
- fmt.Println(shutdownMessage)
- os.Exit(-1)
- case <-ctx.InitDone:
- }
-
- switch count {
- case 0:
- fmt.Println(shutdownMessage)
- if ctx.Online {
- go func() {
- // TODO cancel the command context instead
- n, err := ctx.GetNode()
- if err != nil {
- log.Error(err)
- fmt.Println(shutdownMessage)
- os.Exit(-1)
- }
- n.Close()
- log.Info("Gracefully shut down.")
- }()
- } else {
- os.Exit(0)
- }
-
- default:
- fmt.Println("Received another interrupt before graceful shutdown, terminating...")
- os.Exit(-1)
- }
+ defer ih.wg.Done()
+ count := 0
+ for _ = range ih.sig {
+ count++
+ handler(count, ih)
  }
+ signal.Stop(ih.sig)
  }()
 }
 
-func allInterruptSignals() chan os.Signal {
- sigc := make(chan os.Signal, 1)
- signal.Notify(sigc, syscall.SIGHUP, syscall.SIGINT,
- syscall.SIGTERM)
- return sigc
+func (i *cmdInvocation) SetupInterruptHandler(ctx context.Context) (io.Closer, context.Context) {
+
+ intrh := NewIntrHandler()
+ ctx, cancelFunc := context.WithCancel(ctx)
+
+ handlerFunc := func(count int, ih *IntrHandler) {
+ switch count {
+ case 1:
+ fmt.Println() // Prevent un-terminated ^C character in terminal
+
+ ih.wg.Add(1)
+ go func() {
+ defer ih.wg.Done()
+ cancelFunc()
+ }()
+
+ default:
+ fmt.Println("Received another interrupt before graceful shutdown, terminating...")
+ os.Exit(-1)
+ }
+ }
+
+ intrh.Handle(handlerFunc, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM)
+
+ return intrh, ctx
 }
 
 func profileIfEnabled() (func(), error) {

commands/http/client.go

@@ -82,25 +82,44 @@ func (c *client) Send(req cmds.Request) (cmds.Response, error) {
  version := config.CurrentVersionNumber
  httpReq.Header.Set("User-Agent", fmt.Sprintf("/go-ipfs/%s/", version))
 
- httpRes, err := http.DefaultClient.Do(httpReq)
- if err != nil {
- return nil, err
- }
+ ec := make(chan error, 1)
+ rc := make(chan cmds.Response, 1)
+ dc := req.Context().Context.Done()
 
- // using the overridden JSON encoding in request
- res, err := getResponse(httpRes, req)
- if err != nil {
- return nil, err
- }
-
- if found && len(previousUserProvidedEncoding) > 0 {
- // reset to user provided encoding after sending request
- // NB: if user has provided an encoding but it is the empty string,
- // still leave it as JSON.
- req.SetOption(cmds.EncShort, previousUserProvidedEncoding)
+ go func() {
+ httpRes, err := http.DefaultClient.Do(httpReq)
+ if err != nil {
+ ec <- err
+ return
+ }
+ // using the overridden JSON encoding in request
+ res, err := getResponse(httpRes, req)
+ if err != nil {
+ ec <- err
+ return
+ }
+ rc <- res
+ }()
+
+ for {
+ select {
+ case <-dc:
+ log.Debug("Context cancelled, cancelling HTTP request...")
+ tr := http.DefaultTransport.(*http.Transport)
+ tr.CancelRequest(httpReq)
+ dc = nil // Wait for ec or rc
+ case err := <-ec:
+ return nil, err
+ case res := <-rc:
+ if found && len(previousUserProvidedEncoding) > 0 {
+ // reset to user provided encoding after sending request
+ // NB: if user has provided an encoding but it is the empty string,
+ // still leave it as JSON.
+ req.SetOption(cmds.EncShort, previousUserProvidedEncoding)
+ }
+ return res, nil
+ }
  }
-
- return res, nil
 }
 
 func getQuery(req cmds.Request) (string, error) {
@@ -162,6 +181,8 @@ func getResponse(httpRes *http.Response, req cmds.Request) (cmds.Response, error
  dec := json.NewDecoder(httpRes.Body)
  outputType := reflect.TypeOf(req.Command().Type)
 
+ ctx := req.Context().Context
+
  for {
  var v interface{}
  var err error
@@ -175,6 +196,14 @@ func getResponse(httpRes *http.Response, req cmds.Request) (cmds.Response, error
  fmt.Println(err.Error())
  return
  }
+
+ select {
+ case <-ctx.Done():
+ close(outChan)
+ return
+ default:
+ }
+
  if err == io.EOF {
  close(outChan)
  return

commands/request.go

@@ -30,7 +30,6 @@ type Context struct {
 
  node *core.IpfsNode
  ConstructNode func() (*core.IpfsNode, error)
- InitDone chan bool
 }
 
 // GetConfig returns the config of the current Command exection
@@ -288,7 +287,7 @@ func NewRequest(path []string, opts OptMap, args []string, file files.File, cmd
  optDefs = make(map[string]Option)
  }
 
- ctx := Context{Context: context.TODO(), InitDone: make(chan bool)}
+ ctx := Context{Context: context.TODO()}
  values := make(map[string]interface{})
  req := &request{path, opts, args, file, cmd, ctx, optDefs, values, os.Stdin}
  err := req.ConvertOptions()

core/corehttp/corehttp.go

@@ -2,6 +2,7 @@ package corehttp
 
 import (
  "net/http"
+ "time"
 
  manners "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/braintree/manners"
  ma "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
@@ -63,6 +64,9 @@ func listenAndServe(node *core.IpfsNode, addr ma.Multiaddr, handler http.Handler
  var serverError error
  serverExited := make(chan struct{})
 
+ node.Children().Add(1)
+ defer node.Children().Done()
+
  go func() {
  serverError = server.ListenAndServe(host, handler)
  close(serverExited)
@@ -75,8 +79,22 @@ func listenAndServe(node *core.IpfsNode, addr ma.Multiaddr, handler http.Handler
  // if node being closed before server exits, close server
  case <-node.Closing():
  log.Infof("server at %s terminating...", addr)
+
+ // make sure keep-alive connections do not keep the server running
+ server.InnerServer.SetKeepAlivesEnabled(false)
+
  server.Shutdown <- true
- <-serverExited // now, DO wait until server exit
+
+ outer:
+ for {
+ // wait until server exits
+ select {
+ case <-serverExited:
+ break outer
+ case <-time.After(5 * time.Second):
+ log.Infof("waiting for server at %s to terminate...", addr)
+ }
+ }
  }
 
  log.Infof("server at %s terminated", addr)