From 35b61eb2e4a593521e7d111f13ab17d08538883b Mon Sep 17 00:00:00 2001
From: truthhun <1272881215@qq.com>
Date: Sun, 5 Mar 2023 14:29:37 +0800
Subject: [PATCH] =?UTF-8?q?=E4=B8=BB=E7=A8=8B=E5=BA=8F=E9=80=80=E5=87=BA?=
 =?UTF-8?q?=E6=97=B6=EF=BC=8C=E5=85=B3=E9=97=AD=E5=AD=90=E8=BF=9B=E7=A8=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/serve.go                    | 18 ++++++++++++++++++
 main.go                         |  7 +++++--
 util/command/command.go         |  8 ++++++++
 util/command/command_darwin.go  | 22 ++++++++++++++++++++++
 util/command/command_linux.go   | 22 ++++++++++++++++++++++
 util/command/command_windows.go | 25 +++++++++++++++++++++++++
 6 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 util/command/command.go

diff --git a/cmd/serve.go b/cmd/serve.go
index 691eb21..c422960 100644
--- a/cmd/serve.go
+++ b/cmd/serve.go
@@ -16,8 +16,12 @@ limitations under the License.
 package cmd
 
 import (
+	"fmt"
 	"moredoc/service"
 	"moredoc/util"
+	"moredoc/util/command"
+	"os"
+	"os/signal"
 
 	"github.com/spf13/cobra"
 )
@@ -31,6 +35,20 @@ var serveCmd = &cobra.Command{
 		util.Version = Version
 		util.Hash = GitHash
 		util.BuildAt = BuildAt
+
+		c := make(chan os.Signal, 1)
+		signal.Notify(c)
+		go func() {
+			//阻塞直至有信号传入
+			s := <-c
+			// 收到退出信号，关闭子进程
+			fmt.Println("get signal：", s)
+			fmt.Println("close child process...")
+			command.CloseChildProccess()
+			fmt.Println("close child process done.")
+			fmt.Println("exit.")
+			os.Exit(0)
+		}()
 		service.Run(cfg, logger)
 	},
 }
diff --git a/main.go b/main.go
index fdb9806..5e59683 100644
--- a/main.go
+++ b/main.go
@@ -5,7 +5,7 @@ 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
+	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,
@@ -15,8 +15,11 @@ limitations under the License.
 */
 package main
 
-import "moredoc/cmd"
+import (
+	"moredoc/cmd"
+)
 
 func main() {
 	cmd.Execute()
+
 }
diff --git a/util/command/command.go b/util/command/command.go
new file mode 100644
index 0000000..947a4ad
--- /dev/null
+++ b/util/command/command.go
@@ -0,0 +1,8 @@
+package command
+
+import (
+	"sync"
+)
+
+// pidMap 用于存储所有的子进程pid，以便在主程序退出时，kill掉所有的相关子进程
+var pidMap sync.Map
diff --git a/util/command/command_darwin.go b/util/command/command_darwin.go
index 1a663d2..004c3aa 100644
--- a/util/command/command_darwin.go
+++ b/util/command/command_darwin.go
@@ -32,6 +32,17 @@ func ExecCommand(name string, args []string, timeout ...time.Duration) (out stri
 		return
 	}
 
+	pid := 0
+	if cmd.Process != nil && cmd.Process.Pid != 0 {
+		pid = cmd.Process.Pid
+		pidMap.Store(pid, pid)
+	}
+	defer func() {
+		if pid != 0 {
+			pidMap.Delete(pid)
+		}
+	}()
+
 	time.AfterFunc(expire, func() {
 		if cmd.Process != nil && cmd.Process.Pid != 0 {
 			errs = append(errs, fmt.Sprintf("execute timeout: %d min.", int(expire.Minutes())))
@@ -51,3 +62,14 @@ func ExecCommand(name string, args []string, timeout ...time.Duration) (out stri
 	}
 	return
 }
+
+// 当主程序退出时，从pidMap中获取所有的pid，然后kill掉
+func CloseChildProccess() {
+	pidMap.Range(func(key, value interface{}) bool {
+		if pid, ok := value.(int); ok {
+			fmt.Println("kill pid:", pid)
+			syscall.Kill(-pid, syscall.SIGKILL)
+		}
+		return true
+	})
+}
diff --git a/util/command/command_linux.go b/util/command/command_linux.go
index 1a663d2..004c3aa 100644
--- a/util/command/command_linux.go
+++ b/util/command/command_linux.go
@@ -32,6 +32,17 @@ func ExecCommand(name string, args []string, timeout ...time.Duration) (out stri
 		return
 	}
 
+	pid := 0
+	if cmd.Process != nil && cmd.Process.Pid != 0 {
+		pid = cmd.Process.Pid
+		pidMap.Store(pid, pid)
+	}
+	defer func() {
+		if pid != 0 {
+			pidMap.Delete(pid)
+		}
+	}()
+
 	time.AfterFunc(expire, func() {
 		if cmd.Process != nil && cmd.Process.Pid != 0 {
 			errs = append(errs, fmt.Sprintf("execute timeout: %d min.", int(expire.Minutes())))
@@ -51,3 +62,14 @@ func ExecCommand(name string, args []string, timeout ...time.Duration) (out stri
 	}
 	return
 }
+
+// 当主程序退出时，从pidMap中获取所有的pid，然后kill掉
+func CloseChildProccess() {
+	pidMap.Range(func(key, value interface{}) bool {
+		if pid, ok := value.(int); ok {
+			fmt.Println("kill pid:", pid)
+			syscall.Kill(-pid, syscall.SIGKILL)
+		}
+		return true
+	})
+}
diff --git a/util/command/command_windows.go b/util/command/command_windows.go
index 42ac721..9e2198b 100644
--- a/util/command/command_windows.go
+++ b/util/command/command_windows.go
@@ -35,6 +35,17 @@ func ExecCommand(name string, args []string, timeout ...time.Duration) (out stri
 		return
 	}
 
+	pid := 0
+	if cmd.Process != nil && cmd.Process.Pid != 0 {
+		pid = cmd.Process.Pid
+		pidMap.Store(pid, pid)
+	}
+	defer func() {
+		if pid != 0 {
+			pidMap.Delete(pid)
+		}
+	}()
+
 	time.AfterFunc(expire, func() {
 		if cmd.Process != nil && cmd.Process.Pid != 0 {
 			errs = append(errs, fmt.Sprintf("execute timeout: %d min.", int(expire.Minutes())))
@@ -57,3 +68,17 @@ func ExecCommand(name string, args []string, timeout ...time.Duration) (out stri
 	}
 	return
 }
+
+// 当主程序退出时，从pidMap中获取所有的pid，然后kill掉
+func CloseChildProccess() {
+	pidMap.Range(func(key, value interface{}) bool {
+		if pid, ok := value.(int); ok {
+			fmt.Println("kill pid:", pid)
+			if proc, err := os.FindProcess(pid); err == nil {
+				proc.Kill()
+				proc.Release()
+			}
+		}
+		return true
+	})
+}