use a context for cancelation

Author: jehiah

apps/nsqd/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"math/rand"
@@ -95,6 +96,11 @@ func (p *program) Stop() error {
 	return nil
 }
 
+// Context returns a context that will be canceled when nsqd initiates the shutdown
+func (p program) Context() context.Context {
+	return p.nsqd.Context()
+}
+
 func logFatal(f string, args ...interface{}) {
 	lg.LogFatal("[nsqd] ", f, args...)
 }

nsqd/nsqd.go

@@ -1,6 +1,7 @@
 package nsqd
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/json"
@@ -46,6 +47,8 @@ type NSQD struct {
 	clientIDSequence int64
 
 	sync.RWMutex
+	ctx       context.Context
+	ctxCancel context.CancelFunc
 
 	opts atomic.Value
 
@@ -99,6 +102,7 @@ func New(opts *Options) (*NSQD, error) {
 		optsNotificationChan: make(chan struct{}, 1),
 		dl:                   dirlock.New(dataPath),
 	}
+	n.ctx, n.ctxCancel = context.WithCancel(context.Background())
 	httpcli := http_api.NewClient(nil, opts.HTTPClientConnectTimeout, opts.HTTPClientRequestTimeout)
 	n.ci = clusterinfo.New(n.logf, httpcli)
 
@@ -786,3 +790,8 @@ func buildTLSConfig(opts *Options) (*tls.Config, error) {
 func (n *NSQD) IsAuthEnabled() bool {
 	return len(n.getOpts().AuthHTTPAddresses) != 0
 }
+
+// Context returns a context that will be canceled when nsqd initiates the shutdown
+func (n *NSQD) Context() context.Context {
+	return n.ctx
+}