Fix race condition (also a regression of the PR 19139)

server/embed/etcd.go

@@ -607,7 +607,9 @@
 
 	// start peer servers in a goroutine
 	for _, pl := range e.Peers {
+		e.wg.Add(1)
 		go func(l *peerListener) {
+			defer e.wg.Done()
 			u := l.Addr().String()
 			e.cfg.logger.Info(
 				"serving peer traffic",
@@ -774,7 +776,9 @@
 
 	// start client servers in each goroutine
 	for _, sctx := range e.sctxs {
+		e.wg.Add(1)
 		go func(s *serveCtx) {
+			defer e.wg.Done()
 			e.errHandler(s.serve(e.Server, &e.cfg.ClientTLSInfo, mux, e.errHandler, e.grpcGatewayDial(splitHTTP), splitHTTP, gopts...))
 		}(sctx)
 	}
@@ -859,7 +863,9 @@
 				return err
 			}
 			e.metricsListeners = append(e.metricsListeners, ml)
+			e.wg.Add(1)
 			go func(u url.URL, ln net.Listener) {
+				defer e.wg.Done()
 				e.cfg.logger.Info(
 					"serving metrics",
 					zap.String("address", u.String()),
@@ -872,9 +878,6 @@
 }
 
 func (e *Etcd) errHandler(err error) {
-	e.wg.Add(1)
-	defer e.wg.Done()
-
 	if err != nil {
 		e.GetLogger().Error("setting up serving from embedded etcd failed.", zap.Error(err))
 	}

server/embed/serve.go

@@ -68,6 +68,8 @@
 	serviceRegister func(*grpc.Server)
 	serversC        chan *servers
 	closeOnce       sync.Once
+
+	wg sync.WaitGroup
 }
 
 type servers struct {
@@ -182,13 +184,17 @@
 			server = m.Serve
 
 			httpl := m.Match(cmux.HTTP1())
+			sctx.wg.Add(1)
 			go func(srvhttp *http.Server, tlsLis net.Listener) {
+				defer sctx.wg.Done()
 				errHandler(srvhttp.Serve(tlsLis))
 			}(srv, httpl)
 
 			if grpcEnabled {
 				grpcl := m.Match(cmux.HTTP2())
+				sctx.wg.Add(1)
 				go func(gs *grpc.Server, l net.Listener) {
+					defer sctx.wg.Done()
 					errHandler(gs.Serve(l))
 				}(gs, grpcl)
 			}
@@ -237,7 +243,7 @@
 				TLSConfig: tlscfg,
 				ErrorLog:  logger, // do not log user error
 			}
-			if err := configureHTTPServer(srv, s.Cfg); err != nil {
+			if err = configureHTTPServer(srv, s.Cfg); err != nil {
 				sctx.lg.Error("Configure https server failed", zap.Error(err))
 				return err
 			}
@@ -248,11 +254,13 @@
 		} else {
 			server = m.Serve
 
-			tlsl, err := transport.NewTLSListener(m.Match(cmux.Any()), tlsinfo)
-			if err != nil {
-				return err
+			tlsl, tlsErr := transport.NewTLSListener(m.Match(cmux.Any()), tlsinfo)
+			if tlsErr != nil {
+				return tlsErr
 			}
+			sctx.wg.Add(1)
 			go func(srvhttp *http.Server, tlsl net.Listener) {
+				defer sctx.wg.Done()
 				errHandler(srvhttp.Serve(tlsl))
 			}(srv, tlsl)
 		}
@@ -265,7 +273,11 @@
 		)
 	}
 
-	return server()
+	err = server()
+	sctx.close()
+	// ensure all goroutines, which are created by this method, to complete before this method returns.
+	sctx.wg.Wait()
+	return err
 }
 
 func configureHTTPServer(srv *http.Server, cfg config.ServerConfig) error {
@@ -334,7 +346,9 @@
 			return nil, err
 		}
 	}
+	sctx.wg.Add(1)
 	go func() {
+		defer sctx.wg.Done()
 		<-ctx.Done()
 		if cerr := conn.Close(); cerr != nil {
 			sctx.lg.Warn(

tests/common/auth_test.go

@@ -97,8 +97,10 @@ func TestAuthGracefulDisable(t *testing.T) {
 		donec := make(chan struct{})
 		rootAuthClient := testutils.MustClient(clus.Client(WithAuth(rootUserName, rootPassword)))
 
+		startedC := make(chan struct{}, 1)
 		go func() {
 			defer close(donec)
+			defer close(startedC)
 			// sleep a bit to let the watcher connects while auth is still enabled
 			time.Sleep(time.Second)
 			// now disable auth...
@@ -112,10 +114,12 @@ func TestAuthGracefulDisable(t *testing.T) {
 				t.Errorf("failed to restart member %v", err)
 				return
 			}
+			startedC <- struct{}{}
 			// the watcher should still work after reconnecting
 			assert.NoErrorf(t, rootAuthClient.Put(ctx, "key", "value", config.PutOptions{}), "failed to put key value")
 		}()
 
+		<-startedC
 		wCtx, wCancel := context.WithCancel(ctx)
 		defer wCancel()