From f428feab09f20e309a6513cb30f6f6a2aa3beb34 Mon Sep 17 00:00:00 2001
From: Jingyi Hu <jingyih@google.com>
Date: Tue, 29 Oct 2019 20:07:15 -0700
Subject: [PATCH] etcdserver: wait purge file loop during shutdown

---
 etcdserver/server.go  | 16 +++++++++++++---
 pkg/fileutil/purge.go |  8 +++++---
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/etcdserver/server.go b/etcdserver/server.go
index 976acd684c17..c6beeb7ff00d 100644
--- a/etcdserver/server.go
+++ b/etcdserver/server.go
@@ -807,12 +807,13 @@ func (s *EtcdServer) start() {
 
 func (s *EtcdServer) purgeFile() {
 	var dberrc, serrc, werrc <-chan error
+	var dbdonec, sdonec, wdonec <-chan struct{}
 	if s.Cfg.MaxSnapFiles > 0 {
-		dberrc = fileutil.PurgeFile(s.getLogger(), s.Cfg.SnapDir(), "snap.db", s.Cfg.MaxSnapFiles, purgeFileInterval, s.done)
-		serrc = fileutil.PurgeFile(s.getLogger(), s.Cfg.SnapDir(), "snap", s.Cfg.MaxSnapFiles, purgeFileInterval, s.done)
+		dbdonec, dberrc = fileutil.PurgeFile(s.getLogger(), s.Cfg.SnapDir(), "snap.db", s.Cfg.MaxSnapFiles, purgeFileInterval, s.stopping)
+		sdonec, serrc = fileutil.PurgeFile(s.getLogger(), s.Cfg.SnapDir(), "snap", s.Cfg.MaxSnapFiles, purgeFileInterval, s.stopping)
 	}
 	if s.Cfg.MaxWALFiles > 0 {
-		werrc = fileutil.PurgeFile(s.getLogger(), s.Cfg.WALDir(), "wal", s.Cfg.MaxWALFiles, purgeFileInterval, s.done)
+		wdonec, werrc = fileutil.PurgeFile(s.getLogger(), s.Cfg.WALDir(), "wal", s.Cfg.MaxWALFiles, purgeFileInterval, s.stopping)
 	}
 
 	lg := s.getLogger()
@@ -836,6 +837,15 @@ func (s *EtcdServer) purgeFile() {
 			plog.Fatalf("failed to purge wal file %v", e)
 		}
 	case <-s.stopping:
+		if dbdonec != nil {
+			<-dbdonec
+		}
+		if sdonec != nil {
+			<-sdonec
+		}
+		if wdonec != nil {
+			<-wdonec
+		}
 		return
 	}
 }
diff --git a/pkg/fileutil/purge.go b/pkg/fileutil/purge.go
index fda96c371143..fb214338fa31 100644
--- a/pkg/fileutil/purge.go
+++ b/pkg/fileutil/purge.go
@@ -24,14 +24,16 @@ import (
 	"go.uber.org/zap"
 )
 
-func PurgeFile(lg *zap.Logger, dirname string, suffix string, max uint, interval time.Duration, stop <-chan struct{}) <-chan error {
+func PurgeFile(lg *zap.Logger, dirname string, suffix string, max uint, interval time.Duration, stop <-chan struct{}) (<-chan struct{}, <-chan error) {
 	return purgeFile(lg, dirname, suffix, max, interval, stop, nil)
 }
 
 // purgeFile is the internal implementation for PurgeFile which can post purged files to purgec if non-nil.
-func purgeFile(lg *zap.Logger, dirname string, suffix string, max uint, interval time.Duration, stop <-chan struct{}, purgec chan<- string) <-chan error {
+func purgeFile(lg *zap.Logger, dirname string, suffix string, max uint, interval time.Duration, stop <-chan struct{}, purgec chan<- string) (<-chan struct{}, <-chan error) {
 	errC := make(chan error, 1)
+	doneC := make(chan struct{})
 	go func() {
+		defer close(doneC)
 		for {
 			fnames, err := ReadDir(dirname)
 			if err != nil {
@@ -84,5 +86,5 @@ func purgeFile(lg *zap.Logger, dirname string, suffix string, max uint, interval
 			}
 		}
 	}()
-	return errC
+	return doneC, errC
 }