Fix(Alpha): MASA: Make Alpha Shutdown Again (#6313)

This PR removes the usage of context.Background() in groups.go and ensures that all the various goroutines exit as intended.

Changes:
* Shutdown SubscribeForUpdates correctly.
* Fix up all the closing conditions.
* Consolidate updaters into one closer
* Update Badger to master
* fix(build): Update ResetAcl args for OSS build.
* chore: Remove TODO comment.

Co-authored-by: Daniel Mai <[email protected]>

Author: manishrjain

dgraph/cmd/alpha/run.go

@@ -753,43 +753,52 @@ func run() {
 		}
 	}()
 
-	// Setup external communication.
-	aclCloser := y.NewCloser(1)
+	updaters := y.NewCloser(4)
 	go func() {
 		worker.StartRaftNodes(worker.State.WALstore, bindall)
 		// initialization of the admin account can only be done after raft nodes are running
 		// and health check passes
-		edgraph.ResetAcl()
-		edgraph.RefreshAcls(aclCloser)
-		edgraph.ResetCors()
+		edgraph.ResetAcl(updaters)
+		edgraph.RefreshAcls(updaters)
+		edgraph.ResetCors(updaters)
 		// Update the accepted cors origins.
-		for {
-			origins, err := edgraph.GetCorsOrigins(context.TODO())
+		for updaters.Ctx().Err() == nil {
+			origins, err := edgraph.GetCorsOrigins(updaters.Ctx())
 			if err != nil {
 				glog.Errorf("Error while retriving cors origins: %s", err.Error())
 				continue
 			}
 			x.UpdateCorsOrigins(origins)
-			break
+			return
 		}
 	}()
 	// Listen for any new cors origin update.
-	corsCloser := y.NewCloser(1)
-	go listenForCorsUpdate(corsCloser)
+	go listenForCorsUpdate(updaters)
 
 	// Graphql subscribes to alpha to get schema updates. We need to close that before we
 	// close alpha. This closer is for closing and waiting that subscription.
 	adminCloser := y.NewCloser(1)
 
 	setupServer(adminCloser)
 	glog.Infoln("GRPC and HTTP stopped.")
-	aclCloser.SignalAndWait()
-	corsCloser.SignalAndWait()
+
+	// This might not close until group is given the signal to close. So, only signal here,
+	// wait for it after group is closed.
+	updaters.Signal()
+
 	worker.BlockingStop()
+	glog.Infoln("worker stopped.")
+
 	adminCloser.SignalAndWait()
-	glog.Info("Disposing server state.")
+	glog.Infoln("adminCloser closed.")
+
 	worker.State.Dispose()
 	x.RemoveCidFile()
+	glog.Info("worker.State disposed.")
+
+	updaters.Wait()
+	glog.Infoln("updaters closed.")
+
 	glog.Infoln("Server shutdown. Bye!")
 }
 

dgraph/cmd/zero/run.go

@@ -336,5 +336,10 @@ func run() {
 
 	glog.Infoln("Running Dgraph Zero...")
 	st.zero.closer.Wait()
-	glog.Infoln("All done.")
+	glog.Infoln("Closer closed.")
+
+	err = kv.Close()
+	glog.Infof("Badger closed with err: %v\n", err)
+
+	glog.Infoln("All done. Goodbye!")
 }

edgraph/access.go

@@ -42,7 +42,7 @@ func (s *Server) Login(ctx context.Context,
 }
 
 // ResetAcl is an empty method since ACL is only supported in the enterprise version.
-func ResetAcl() {
+func ResetAcl(closer *y.Closer) {
 	// do nothing
 }
 

edgraph/access_ee.go

@@ -306,15 +306,15 @@ func authorizeUser(ctx context.Context, userid string, password string) (
 
 // RefreshAcls queries for the ACL triples and refreshes the ACLs accordingly.
 func RefreshAcls(closer *y.Closer) {
-	defer closer.Done()
+	defer func() {
+		glog.Infoln("RefreshAcls closed")
+		closer.Done()
+	}()
 	if len(worker.Config.HmacSecret) == 0 {
 		// the acl feature is not turned on
 		return
 	}
 
-	ticker := time.NewTicker(worker.Config.AclRefreshInterval)
-	defer ticker.Stop()
-
 	// retrieve the full data set of ACLs from the corresponding alpha server, and update the
 	// aclCachePtr
 	retrieveAcls := func() error {
@@ -324,9 +324,7 @@ func RefreshAcls(closer *y.Closer) {
 			ReadOnly: true,
 		}
 
-		ctx := context.Background()
-		var err error
-		queryResp, err := (&Server{}).doQuery(ctx, &queryRequest, NoAuthorize)
+		queryResp, err := (&Server{}).doQuery(closer.Ctx(), &queryRequest, NoAuthorize)
 		if err != nil {
 			return errors.Errorf("unable to retrieve acls: %v", err)
 		}
@@ -340,14 +338,16 @@ func RefreshAcls(closer *y.Closer) {
 		return nil
 	}
 
+	ticker := time.NewTicker(worker.Config.AclRefreshInterval)
+	defer ticker.Stop()
 	for {
 		select {
-		case <-closer.HasBeenClosed():
-			return
 		case <-ticker.C:
 			if err := retrieveAcls(); err != nil {
-				glog.Errorf("Error while retrieving acls:%v", err)
+				glog.Errorf("Error while retrieving acls: %v", err)
 			}
+		case <-closer.HasBeenClosed():
+			return
 		}
 	}
 }
@@ -368,7 +368,12 @@ const queryAcls = `
 `
 
 // ResetAcl clears the aclCachePtr and upserts the Groot account.
-func ResetAcl() {
+func ResetAcl(closer *y.Closer) {
+	defer func() {
+		glog.Infof("ResetAcl closed")
+		closer.Done()
+	}()
+
 	if len(worker.Config.HmacSecret) == 0 {
 		// The acl feature is not turned on.
 		return
@@ -435,8 +440,8 @@ func ResetAcl() {
 		return nil
 	}
 
-	for {
-		ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	for closer.Ctx().Err() == nil {
+		ctx, cancel := context.WithTimeout(closer.Ctx(), time.Minute)
 		defer cancel()
 		if err := upsertGuardians(ctx); err != nil {
 			glog.Infof("Unable to upsert the guardian group. Error: %v", err)
@@ -446,8 +451,8 @@ func ResetAcl() {
 		break
 	}
 
-	for {
-		ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	for closer.Ctx().Err() == nil {
+		ctx, cancel := context.WithTimeout(closer.Ctx(), time.Minute)
 		defer cancel()
 		if err := upsertGroot(ctx); err != nil {
 			glog.Infof("Unable to upsert the groot account. Error: %v", err)

edgraph/server.go

@@ -42,6 +42,7 @@ import (
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 
+	"github.com/dgraph-io/badger/v2/y"
 	"github.com/dgraph-io/dgo/v200"
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	"github.com/dgraph-io/dgraph/chunker"
@@ -328,8 +329,8 @@ func (s *Server) Alter(ctx context.Context, op *api.Operation) (*api.Payload, er
 		// reset their in-memory GraphQL schema
 		_, err = UpdateGQLSchema(ctx, "", "")
 		// recreate the admin account after a drop all operation
-		ResetAcl()
-		ResetCors()
+		ResetAcl(nil)
+		ResetCors(nil)
 		return empty, err
 	}
 
@@ -353,8 +354,8 @@ func (s *Server) Alter(ctx context.Context, op *api.Operation) (*api.Payload, er
 		// just reinsert the GraphQL schema, no need to alter dgraph schema as this was drop_data
 		_, err = UpdateGQLSchema(ctx, graphQLSchema, "")
 		// recreate the admin account after a drop data operation
-		ResetAcl()
-		ResetCors()
+		ResetAcl(nil)
+		ResetCors(nil)
 		return empty, err
 	}
 
@@ -1541,7 +1542,12 @@ func isDropAll(op *api.Operation) bool {
 
 // ResetCors make the dgraph to accept all the origins if no origins were given
 // by the users.
-func ResetCors() {
+func ResetCors(closer *y.Closer) {
+	defer func() {
+		glog.Infof("ResetCors closed")
+		closer.Done()
+	}()
+
 	req := &api.Request{
 		Query: `query{
 			cors as var(func: has(dgraph.cors))
@@ -1561,8 +1567,8 @@ func ResetCors() {
 		CommitNow: true,
 	}
 
-	for {
-		ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	for closer.Ctx().Err() == nil {
+		ctx, cancel := context.WithTimeout(closer.Ctx(), time.Minute)
 		defer cancel()
 		if _, err := (&Server{}).doQuery(ctx, req, CorsMutationAllowed); err != nil {
 			glog.Infof("Unable to upsert cors. Error: %v", err)
@@ -1623,9 +1629,8 @@ func GetCorsOrigins(ctx context.Context) ([]string, error) {
 	if err = json.Unmarshal(res.Json, corsRes); err != nil {
 		return nil, err
 	}
-	if len(corsRes.Me) > 1 {
-		glog.Errorf("Something went wrong in cors predicate, expected 1 predicate but got %d",
-			len(corsRes.Me))
+	if len(corsRes.Me) != 1 {
+		return []string{}, fmt.Errorf("GetCorsOrigins returned %d results", len(corsRes.Me))
 	}
 	return corsRes.Me[0].DgraphCors, nil
 }

go.mod

@@ -17,10 +17,10 @@ require (
 	github.com/blevesearch/segment v0.0.0-20160915185041-762005e7a34f // indirect
 	github.com/blevesearch/snowballstem v0.0.0-20180110192139-26b06a2c243d // indirect
 	github.com/codahale/hdrhistogram v0.0.0-20161010025455-3a0bb77429bd
-	github.com/dgraph-io/badger/v2 v2.2007.2-0.20200826122734-bc243f38bfe1
+	github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200828220306-806a325a0462
 	github.com/dgraph-io/dgo/v200 v200.0.0-20200805103119-a3544c464dd6
 	github.com/dgraph-io/graphql-transport-ws v0.0.0-20200715131837-c0460019ead2
-	github.com/dgraph-io/ristretto v0.0.3-0.20200630154024-f66de99634de
+	github.com/dgraph-io/ristretto v0.0.4-0.20200820164438-623d8ef1614b
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/dgrijalva/jwt-go/v4 v4.0.0-preview1
 	github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13

go.sum

@@ -84,14 +84,14 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgraph-io/badger v1.6.0 h1:DshxFxZWXUcO0xX476VJC07Xsr6ZCBVRHKZ93Oh7Evo=
 github.com/dgraph-io/badger v1.6.0/go.mod h1:zwt7syl517jmP8s94KqSxTlM6IMsdhYy6psNgSztDR4=
-github.com/dgraph-io/badger/v2 v2.2007.2-0.20200826122734-bc243f38bfe1 h1:vPPlQYByX3+Z3NOaB06i7VCb0bNOznVQ9eEnqLxbrH0=
-github.com/dgraph-io/badger/v2 v2.2007.2-0.20200826122734-bc243f38bfe1/go.mod h1:26P/7fbL4kUZVEVKLAKXkBXKOydDmM2p1e+NhhnBCAE=
+github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200828220306-806a325a0462 h1:5A3xbCP8emF/lJ1btFdvMZ/oBaVpF9Ncldbop5MFcwM=
+github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200828220306-806a325a0462/go.mod h1:EdVb0qoPSOGlrBrHSG4ArQTao3rVHe14qaVx1qqO7Vk=
 github.com/dgraph-io/dgo/v200 v200.0.0-20200805103119-a3544c464dd6 h1:toHzMCdCUgYsjM0cW9+wafnKFXfp1HizIJUyzihN+vk=
 github.com/dgraph-io/dgo/v200 v200.0.0-20200805103119-a3544c464dd6/go.mod h1:rHa+h3kI4M8ASOirxyIyNeXBfHFgeskVUum2OrDMN3U=
 github.com/dgraph-io/graphql-transport-ws v0.0.0-20200715131837-c0460019ead2 h1:NSl3XXyON9bgmBJSAvr5FPrgILAovtoTs7FwdtaZZq0=
 github.com/dgraph-io/graphql-transport-ws v0.0.0-20200715131837-c0460019ead2/go.mod h1:7z3c/5w0sMYYZF5bHsrh8IH4fKwG5O5Y70cPH1ZLLRQ=
-github.com/dgraph-io/ristretto v0.0.3-0.20200630154024-f66de99634de h1:t0UHb5vdojIDUqktM6+xJAfScFBsVpXZmqC9dsgJmeA=
-github.com/dgraph-io/ristretto v0.0.3-0.20200630154024-f66de99634de/go.mod h1:KPxhHT9ZxKefz+PCeOGsrHpl1qZ7i70dGTu2u+Ahh6E=
+github.com/dgraph-io/ristretto v0.0.4-0.20200820164438-623d8ef1614b h1:/g8jOqvD1UzHTOwENtkqcLmMLzTcN18P3ut8aSUZ45g=
+github.com/dgraph-io/ristretto v0.0.4-0.20200820164438-623d8ef1614b/go.mod h1:KPxhHT9ZxKefz+PCeOGsrHpl1qZ7i70dGTu2u+Ahh6E=
 github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumCAMpl/TFQ4/5kLM=
 github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
 github.com/dgrijalva/jwt-go/v4 v4.0.0-preview1 h1:CaO/zOnF8VvUfEbhRatPcwKVWamvbYd8tQGRWacE9kU=