From 28a490b09ca508335da6f3c89379522faaf26231 Mon Sep 17 00:00:00 2001
From: pyiyun <panyiyun@huawei.com>
Date: Fri, 9 Apr 2021 09:58:56 +0800
Subject: [PATCH] etcdserver: remove temp files in snap dir when etcdServer
 starting When etcd exits abnormally, tmp files will remain in snap dir, so
 clean up tmp files in snap dir when etcdserver starting.

Fixes #12837
---
 client/pkg/fileutil/fileutil.go      | 32 ++++++++++++++++++++
 client/pkg/fileutil/fileutil_test.go | 44 ++++++++++++++++++++++++++++
 server/etcdserver/server.go          | 12 ++++++++
 3 files changed, 88 insertions(+)

diff --git a/client/pkg/fileutil/fileutil.go b/client/pkg/fileutil/fileutil.go
index f4b198642..85a9842b0 100644
--- a/client/pkg/fileutil/fileutil.go
+++ b/client/pkg/fileutil/fileutil.go
@@ -135,3 +135,35 @@ func CheckDirPermission(dir string, perm os.FileMode) error {
 	}
 	return nil
 }
+
+// RemoveMatchFile deletes file if matchFunc is true on an existing dir
+// Returns error if the dir does not exist or remove file fail
+func RemoveMatchFile(lg *zap.Logger, dir string, matchFunc func(fileName string) bool) error {
+	if lg == nil {
+		lg = zap.NewNop()
+	}
+	if !Exist(dir) {
+		return fmt.Errorf("directory %s does not exist", dir)
+	}
+	fileNames, err := ReadDir(dir)
+	if err != nil {
+		return err
+	}
+	var removeFailedFiles []string
+	for _, fileName := range fileNames {
+		if matchFunc(fileName) {
+			file := filepath.Join(dir, fileName)
+			if err = os.Remove(file); err != nil {
+				removeFailedFiles = append(removeFailedFiles, fileName)
+				lg.Error("remove file failed",
+					zap.String("file", file),
+					zap.Error(err))
+				continue
+			}
+		}
+	}
+	if len(removeFailedFiles) != 0 {
+		return fmt.Errorf("remove file(s) %v error", removeFailedFiles)
+	}
+	return nil
+}
diff --git a/client/pkg/fileutil/fileutil_test.go b/client/pkg/fileutil/fileutil_test.go
index 33b39a23a..3a761ff9a 100644
--- a/client/pkg/fileutil/fileutil_test.go
+++ b/client/pkg/fileutil/fileutil_test.go
@@ -26,6 +26,8 @@ import (
 	"strings"
 	"testing"
 	"time"
+
+	"go.uber.org/zap/zaptest"
 )
 
 func TestIsDirWriteable(t *testing.T) {
@@ -192,3 +194,45 @@ func TestDirPermission(t *testing.T) {
 		t.Errorf("expected error, got nil")
 	}
 }
+
+func TestRemoveMatchFile(t *testing.T) {
+	tmpdir := t.TempDir()
+	defer os.RemoveAll(tmpdir)
+	f, err := ioutil.TempFile(tmpdir, "tmp")
+	if err != nil {
+		t.Fatal(err)
+	}
+	f.Close()
+	f, err = ioutil.TempFile(tmpdir, "foo.tmp")
+	if err != nil {
+		t.Fatal(err)
+	}
+	f.Close()
+
+	err = RemoveMatchFile(zaptest.NewLogger(t), tmpdir, func(fileName string) bool {
+		return strings.HasPrefix(fileName, "tmp")
+	})
+	if err != nil {
+		t.Errorf("expected nil, got error")
+	}
+	fnames, err := ReadDir(tmpdir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(fnames) != 1 {
+		t.Errorf("expected exist 1 files, got %d", len(fnames))
+	}
+
+	f, err = ioutil.TempFile(tmpdir, "tmp")
+	if err != nil {
+		t.Fatal(err)
+	}
+	f.Close()
+	err = RemoveMatchFile(zaptest.NewLogger(t), tmpdir, func(fileName string) bool {
+		os.Remove(filepath.Join(tmpdir, fileName))
+		return strings.HasPrefix(fileName, "tmp")
+	})
+	if err == nil {
+		t.Errorf("expected error, got nil")
+	}
+}
diff --git a/server/etcdserver/server.go b/server/etcdserver/server.go
index 7aee4d0b9..eae94e8b7 100644
--- a/server/etcdserver/server.go
+++ b/server/etcdserver/server.go
@@ -26,6 +26,7 @@ import (
 	"path"
 	"regexp"
 	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -338,6 +339,17 @@ func NewServer(cfg config.ServerConfig) (srv *EtcdServer, err error) {
 			zap.Error(err),
 		)
 	}
+
+	if err = fileutil.RemoveMatchFile(cfg.Logger, cfg.SnapDir(), func(fileName string) bool {
+		return strings.HasPrefix(fileName, "tmp")
+	}); err != nil {
+		cfg.Logger.Error(
+			"failed to remove temp file(s) in snapshot directory",
+			zap.String("path", cfg.SnapDir()),
+			zap.Error(err),
+		)
+	}
+
 	ss := snap.New(cfg.Logger, cfg.SnapDir())
 
 	bepath := cfg.BackendPath()