Merge pull request #55 from WithLin/cleanup-test-database

Fix  capture a Ctrl+C ,kill, Interrupt signal and run a cleanup test environment

cmd/soar/soar.go

@@ -80,17 +80,24 @@ func main() {
 	// 环境初始化，连接检查线上环境+构建测试环境
 	vEnv, rEnv := env.BuildEnv()
 
-	// 如果使用到测试环境，在这里环境清理
-	if common.Config.DropTestTemporary {
-		defer vEnv.CleanUp()
-	}
-
 	// 使用CleanupTestDatabase命令手动清理残余的`optimizer_xxx`数据库
 	// 为了保护当前正在评审的SQL，只清除前1小时前的残余
 	if common.Config.CleanupTestDatabase {
 		vEnv.CleanupTestDatabase()
 	}
 
+	// 如果使用到测试环境，在这里环境清理
+	if common.Config.DropTestTemporary {
+		defer vEnv.CleanUp()
+	}
+
+	//当程序卡死的时候，或者由于某些原因程序没有退出，可以通过捕获信号量的形式让程序优雅退出并且清理测试环境
+	common.HandleSignal(func() {
+		if common.Config.DropTestTemporary {
+			vEnv.CleanUp()
+		}
+	})
+
 	// 对指定的库表进行索引重复检查
 	if common.Config.ReportType == "duplicate-key-checker" {
 		dupKeySuggest := advisor.DuplicateKeyChecker(rEnv)

common/signal.go

+/*
+ * Copyright 2018 Xiaomi, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package common
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+)
+
+//当程序卡死的时候，或者由于某些原因程序没有退出，可以通过捕获信号量的形式让程序优雅退出并且清理测试环境
+func HandleSignal(f func()) {
+	sc := make(chan os.Signal, 1)
+	signal.Notify(sc,
+		syscall.SIGHUP,
+		syscall.SIGINT,
+		syscall.SIGTERM,
+		syscall.SIGQUIT)
+
+	go func() {
+
+		select {
+		case n := <-sc:
+			Log.Info("receive signal %v, closing", n)
+			f()
+			time.Sleep(250 * time.Millisecond)
+			os.Exit(0)
+		}
+	}()
+}