From e629e4c781171094af0c9a85ee85f17c9cc389d6 Mon Sep 17 00:00:00 2001
From: Jehiah Czebotar <jehiah@gmail.com>
Date: Thu, 26 Nov 2020 00:22:25 -0500
Subject: [PATCH] use a context for cancelation

---
 apps/nsqd/main.go | 6 ++++++
 nsqd/nsqd.go      | 9 +++++++++
 2 files changed, 15 insertions(+)

diff --git a/apps/nsqd/main.go b/apps/nsqd/main.go
index 65ceaecf6..dd2feb577 100644
--- a/apps/nsqd/main.go
+++ b/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...)
 }
diff --git a/nsqd/nsqd.go b/nsqd/nsqd.go
index c745bbf49..1dd6ecdda 100644
--- a/nsqd/nsqd.go
+++ b/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
+}