feat: support --abort for reload and suspend (#346)

cmd/cmd.go

@@ -5,6 +5,10 @@ import (
 	"github.com/spf13/cobra"
 )
 
+const (
+	AbortFile = "/var/run/dae.abort"
+)
+
 var (
 	Version = "unknown"
 	rootCmd = &cobra.Command{

cmd/reload.go

@@ -17,6 +17,7 @@ import (
 )
 
 var (
+	abort     bool
 	reloadCmd = &cobra.Command{
 		Use:   "reload [pid]",
 		Short: "To reload config file without interrupt connections.",
@@ -35,6 +36,11 @@ var (
 				cmd.Help()
 				os.Exit(1)
 			}
+			if abort {
+				if f, err := os.Create(AbortFile); err == nil {
+					f.Close()
+				}
+			}
 			if err = syscall.Kill(pid, syscall.SIGUSR1); err != nil {
 				fmt.Println(err)
 				os.Exit(1)
@@ -46,4 +52,5 @@ var (
 
 func init() {
 	rootCmd.AddCommand(reloadCmd)
+	reloadCmd.PersistentFlags().BoolVarP(&abort, "abort", "a", false, "Abort established connections.")
 }

cmd/run.go

@@ -108,6 +108,8 @@ var (
 )
 
 func Run(log *logrus.Logger, conf *config.Config, externGeoDataDirs []string) (err error) {
+	// Remove AbortFile at beginning.
+	_ = os.Remove(AbortFile)
 
 	// New ControlPlane.
 	c, err := newControlPlane(log, nil, nil, conf, externGeoDataDirs)
@@ -135,6 +137,7 @@ func Run(log *logrus.Logger, conf *config.Config, externGeoDataDirs []string) (e
 	}()
 	reloading := false
 	isSuspend := false
+	abortConnections := false
 loop:
 	for sig := range sigs {
 		switch sig {
@@ -174,6 +177,7 @@ loop:
 			sdnotify.Reloading()
 
 			// Load new config.
+			abortConnections = os.Remove(AbortFile) == nil
 			log.Warnln("[Reload] Load new config")
 			var newConf *config.Config
 			if isSuspend {
@@ -247,6 +251,9 @@ loop:
 			reloading = true
 
 			// Ready to close.
+			if abortConnections {
+				oldC.AbortConnections()
+			}
 			oldC.Close()
 		case syscall.SIGHUP:
 			// Ignore.

cmd/suspend.go

@@ -35,6 +35,11 @@ var (
 				cmd.Help()
 				os.Exit(1)
 			}
+			if abort {
+				if f, err := os.Create(AbortFile); err == nil {
+					f.Close()
+				}
+			}
 			if err = syscall.Kill(pid, syscall.SIGUSR2); err != nil {
 				fmt.Println(err)
 				os.Exit(1)
@@ -46,4 +51,5 @@ var (
 
 func init() {
 	rootCmd.AddCommand(suspendCmd)
+	suspendCmd.PersistentFlags().BoolVarP(&abort, "abort", "a", false, "Abort established connections.")
 }