Avoid deadlock in etcd.Close when stopping during bootstrapping

Signed-off-by: Joshua Zhang <[email protected]>

Author: joshuazh-x

server/embed/etcd.go

@@ -86,6 +86,7 @@ type Etcd struct {
 	errc  chan error
 
 	closeOnce sync.Once
+	wg        sync.WaitGroup
 }
 
 type peerListener struct {
@@ -111,7 +112,7 @@ func StartEtcd(inCfg *Config) (e *Etcd, err error) {
 		if !serving {
 			// errored before starting gRPC server for serveCtx.serversC
 			for _, sctx := range e.sctxs {
-				close(sctx.serversC)
+				sctx.close()
 			}
 		}
 		e.Close()
@@ -436,6 +437,7 @@ func (e *Etcd) Close() {
 		}
 	}
 	if e.errc != nil {
+		e.wg.Wait()
 		close(e.errc)
 	}
 }
@@ -880,6 +882,9 @@ func (e *Etcd) serveMetrics() (err error) {
 }
 
 func (e *Etcd) errHandler(err error) {
+	e.wg.Add(1)
+	defer e.wg.Done()
+
 	select {
 	case <-e.stopc:
 		return

server/embed/serve.go

@@ -23,6 +23,7 @@ import (
 	"net"
 	"net/http"
 	"strings"
+	"sync"
 
 	etcdservergw "go.etcd.io/etcd/api/v3/etcdserverpb/gw"
 	"go.etcd.io/etcd/client/pkg/v3/transport"
@@ -65,6 +66,7 @@ type serveCtx struct {
 	userHandlers    map[string]http.Handler
 	serviceRegister func(*grpc.Server)
 	serversC        chan *servers
+	closeOnce       sync.Once
 }
 
 type servers struct {
@@ -100,6 +102,9 @@ func (sctx *serveCtx) serve(
 	gopts ...grpc.ServerOption) (err error) {
 	logger := defaultLog.New(ioutil.Discard, "etcdhttp", 0)
 
+	// Make sure serversC is closed even if we prematurely exit the function.
+	defer sctx.close()
+
 	select {
 	case <-s.StoppingNotify():
 		return errors.New("server is stopping")
@@ -119,8 +124,6 @@ func (sctx *serveCtx) serve(
 	servElection := v3election.NewElectionServer(v3c)
 	servLock := v3lock.NewLockServer(v3c)
 
-	// Make sure serversC is closed even if we prematurely exit the function.
-	defer close(sctx.serversC)
 	var gwmux *gw.ServeMux
 	if s.Cfg.EnableGRPCGateway {
 		// GRPC gateway connects to grpc server via connection provided by grpc dial.
@@ -503,3 +506,9 @@ func (sctx *serveCtx) registerTrace() {
 	evf := func(w http.ResponseWriter, r *http.Request) { trace.RenderEvents(w, r, true) }
 	sctx.registerUserHandler("/debug/events", http.HandlerFunc(evf))
 }
+
+func (sctx *serveCtx) close() {
+	sctx.closeOnce.Do(func() {
+		close(sctx.serversC)
+	})
+}

server/etcdserver/server.go

@@ -2130,6 +2130,7 @@ func (s *EtcdServer) publish(timeout time.Duration) {
 		Val:    string(b),
 	}
 
+	// gofail: var beforePublishing struct{}
 	for {
 		ctx, cancel := context.WithTimeout(s.ctx, timeout)
 		_, err := s.Do(ctx, req)

tests/integration/embed/embed_test.go

@@ -30,11 +30,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
 	"go.etcd.io/etcd/client/pkg/v3/testutil"
 	"go.etcd.io/etcd/client/pkg/v3/transport"
-	"go.etcd.io/etcd/client/v3"
+	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.etcd.io/etcd/server/v3/embed"
 	"go.etcd.io/etcd/tests/v3/integration"
+	gofail "go.etcd.io/gofail/runtime"
 )
 
 var (
@@ -210,3 +212,56 @@ func setupEmbedCfg(cfg *embed.Config, curls []url.URL, purls []url.URL) {
 	}
 	cfg.InitialCluster = cfg.InitialCluster[1:]
 }
+
+func TestEmbedEtcdStopDuringBootstrapping(t *testing.T) {
+	if len(gofail.List()) == 0 {
+		t.Skip("please run 'make gofail-enable' before running the test")
+	}
+
+	fpName := "beforePublishing"
+	require.NoError(t, gofail.Enable(fpName, `sleep("2s")`))
+	t.Cleanup(func() {
+		terr := gofail.Disable(fpName)
+		if terr != nil && terr != gofail.ErrDisabled {
+			t.Fatalf("failed to disable %s: %v", fpName, terr)
+		}
+	})
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+
+		cfg := embed.NewConfig()
+		urls := newEmbedURLs(false, 2)
+		setupEmbedCfg(cfg, []url.URL{urls[0]}, []url.URL{urls[1]})
+		cfg.Dir = filepath.Join(t.TempDir(), "embed-etcd")
+
+		e, err := embed.StartEtcd(cfg)
+		if err != nil {
+			t.Errorf("Failed to start etcd, got error %v", err)
+		}
+		defer e.Close()
+
+		go func() {
+			time.Sleep(time.Second)
+			e.Server.Stop()
+			t.Log("Stopped server during bootstrapping")
+		}()
+
+		select {
+		case <-e.Server.ReadyNotify():
+			t.Log("Server is ready!")
+		case <-e.Server.StopNotify():
+			t.Log("Server is stopped")
+		case <-time.After(20 * time.Second):
+			e.Server.Stop() // trigger a shutdown
+			t.Error("Server took too long to start!")
+		}
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(10 * time.Second):
+		t.Error("timeout in bootstrapping etcd")
+	}
+}