Zero server shutdown endpoint

dgraph/cmd/zero/assign.go

@@ -26,10 +26,7 @@ import (
 	"github.com/golang/glog"
 )
 
-var (
-	emptyNum         pb.Num
-	emptyAssignedIds pb.AssignedIds
-)
+var emptyAssignedIds pb.AssignedIds
 
 const (
 	leaseBandwidth = uint64(10000)

dgraph/cmd/zero/http.go

@@ -22,7 +22,6 @@ import (
 	"net"
 	"net/http"
 	"strconv"
-	"sync"
 	"time"
 
 	"github.com/dgraph-io/dgraph/protos/pb"
@@ -224,15 +223,30 @@ func (st *state) getState(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (st *state) serveHTTP(l net.Listener, wg *sync.WaitGroup) {
+func (st *state) shutdown(w http.ResponseWriter, r *http.Request) {
+	x.AddCorsHeaders(w)
+	if r.Method == "OPTIONS" {
+		return
+	}
+	if r.Method != http.MethodGet {
+		w.WriteHeader(http.StatusBadRequest)
+		x.SetStatus(w, x.ErrorInvalidMethod, "Invalid method")
+		return
+	}
+
+	st.zero.closer.Signal()
+	w.Write([]byte("Server is shutting down...\n"))
+}
+
+func (st *state) serveHTTP(l net.Listener) {
 	srv := &http.Server{
 		ReadTimeout:  10 * time.Second,
 		WriteTimeout: 600 * time.Second,
 		IdleTimeout:  2 * time.Minute,
 	}
 
 	go func() {
-		defer wg.Done()
+		defer st.zero.closer.Done()
 		err := srv.Serve(l)
 		glog.Errorf("Stopped taking more http(s) requests. Err: %v", err)
 		ctx, cancel := context.WithTimeout(context.Background(), 630*time.Second)

dgraph/cmd/zero/oracle.go

@@ -448,7 +448,7 @@ func (s *Server) Oracle(unused *api.Payload, server pb.Zero_OracleServer) error
 			}
 		case <-ctx.Done():
 			return ctx.Err()
-		case <-s.shutDownCh:
+		case <-s.closer.HasBeenClosed():
 			return errServerShutDown
 		}
 	}

dgraph/cmd/zero/run.go

@@ -23,7 +23,6 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
-	"sync"
 	"syscall"
 	"time"
 
@@ -105,7 +104,7 @@ type state struct {
 	zero *Server
 }
 
-func (st *state) serveGRPC(l net.Listener, wg *sync.WaitGroup, store *raftwal.DiskStorage) {
+func (st *state) serveGRPC(l net.Listener, store *raftwal.DiskStorage) {
 	if collector := Zero.Conf.GetString("jaeger.collector"); len(collector) > 0 {
 		// Port details: https://www.jaegertracing.io/docs/getting-started/
 		// Default collectorEndpointURI := "http://localhost:14268"
@@ -148,7 +147,7 @@ func (st *state) serveGRPC(l net.Listener, wg *sync.WaitGroup, store *raftwal.Di
 	pb.RegisterRaftServer(s, st.rs)
 
 	go func() {
-		defer wg.Done()
+		defer st.zero.closer.Done()
 		err := s.Serve(l)
 		glog.Infof("gRpc server stopped : %v", err)
 		st.node.stop <- struct{}{}
@@ -228,17 +227,16 @@ func run() {
 	defer kv.Close()
 	store := raftwal.Init(kv, opts.nodeId, 0)
 
-	var wg sync.WaitGroup
-	wg.Add(3)
 	// Initialize the servers.
 	var st state
-	st.serveGRPC(grpcListener, &wg, store)
-	st.serveHTTP(httpListener, &wg)
+	st.serveGRPC(grpcListener, store)
+	st.serveHTTP(httpListener)
 
 	http.HandleFunc("/state", st.getState)
 	http.HandleFunc("/removeNode", st.removeNode)
 	http.HandleFunc("/moveTablet", st.moveTablet)
 	http.HandleFunc("/assign", st.assign)
+	http.HandleFunc("/shutdown", st.shutdown)
 	zpages.Handle(http.DefaultServeMux, "/z")
 
 	// This must be here. It does not work if placed before Grpc init.
@@ -251,18 +249,35 @@ func run() {
 	sdCh := make(chan os.Signal, 1)
 	signal.Notify(sdCh, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 
+	// handle signals
 	go func() {
-		defer wg.Done()
-		<-sdCh
-		glog.Infof("Shutting down...")
+		for {
+			select {
+			case sig, ok := <-sdCh:
+				if !ok {
+					return
+				}
+				glog.Infof("--- Received %s signal", sig)
+				signal.Stop(sdCh)
+				st.zero.closer.Signal()
+			}
+		}
+	}()
+
+	st.zero.closer.AddRunning(1)
+
+	go func() {
+		defer st.zero.closer.Done()
+		<-st.zero.closer.HasBeenClosed()
+		glog.Infoln("Shutting down...")
+		close(sdCh)
 		// Close doesn't close already opened connections.
 		httpListener.Close()
 		grpcListener.Close()
-		close(st.zero.shutDownCh)
 		st.node.trySnapshot(0)
 	}()
 
-	glog.Infof("Running Dgraph Zero...")
-	wg.Wait()
-	glog.Infof("All done.")
+	glog.Infoln("Running Dgraph Zero...")
+	st.zero.closer.Wait()
+	glog.Infoln("All done.")
 }

dgraph/cmd/zero/zero.go

@@ -25,6 +25,7 @@ import (
 	otrace "go.opencensus.io/trace"
 	"golang.org/x/net/context"
 
+	"github.com/dgraph-io/badger/y"
 	"github.com/dgraph-io/dgo/protos/api"
 	"github.com/dgraph-io/dgraph/conn"
 	"github.com/dgraph-io/dgraph/protos/pb"
@@ -56,8 +57,8 @@ type Server struct {
 	// groupMap    map[uint32]*Group
 	nextGroup      uint32
 	leaderChangeCh chan struct{}
-	shutDownCh     chan struct{} // Used to tell stream to close.
-	connectLock    sync.Mutex    // Used to serialize connect requests from servers.
+	closer         *y.Closer  // Used to tell stream to close.
+	connectLock    sync.Mutex // Used to serialize connect requests from servers.
 
 	blockCommitsOn map[string]struct{}
 }
@@ -76,7 +77,7 @@ func (s *Server) Init() {
 	s.nextTxnTs = 1
 	s.nextGroup = 1
 	s.leaderChangeCh = make(chan struct{}, 1)
-	s.shutDownCh = make(chan struct{}, 1)
+	s.closer = y.NewCloser(2) // grpc and http
 	s.blockCommitsOn = make(map[string]struct{})
 	go s.rebalanceTablets()
 }
@@ -614,7 +615,7 @@ func (s *Server) StreamMembership(_ *api.Payload, stream pb.Zero_StreamMembershi
 			}
 		case <-ctx.Done():
 			return ctx.Err()
-		case <-s.shutDownCh:
+		case <-s.closer.HasBeenClosed():
 			return errServerShutDown
 		}
 	}