feat(cli): enhance signal handling for graceful shutdown

macrat · Dec 10, 2022 · 817ace5 · 817ace5
1 parent 739e935
commit 817ace5
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 11 deletions.
diff --git a/cmd/ayd/main.go b/cmd/ayd/main.go
@@ -6,8 +6,8 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"os/signal"
 	"text/template"
+	"time"
 
 	"github.com/macrat/ayd/internal/scheme"
 	"github.com/macrat/ayd/internal/store"
@@ -38,7 +38,8 @@ type AydCommand struct {
 	ShowVersion bool
 	ShowHelp    bool
 
-	Tasks []Task
+	Tasks     []Task
+	StartedAt time.Time
 }
 
 var defaultAydCommand = &AydCommand{
@@ -142,8 +143,8 @@ func (cmd *AydCommand) Run(args []string) (exitCode int) {
 		return 1
 	}
 
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
-	defer stop()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	if len(cmd.AlertURLs) > 0 {
 		alert, err := scheme.NewAlerterSet(cmd.AlertURLs)

diff --git a/cmd/ayd/oneshot.go b/cmd/ayd/oneshot.go
@@ -2,8 +2,10 @@ package main
 
 import (
 	"context"
+	"os/signal"
 	"sync"
 	"sync/atomic"
+	"syscall"
 
 	"github.com/macrat/ayd/internal/store"
 	api "github.com/macrat/ayd/lib-ayd"
@@ -12,6 +14,9 @@ import (
 func (cmd *AydCommand) RunOneshot(ctx context.Context, s *store.Store) (exitCode int) {
 	var unhealthy atomic.Value
 
+	ctx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGHUP)
+	defer stop()
+
 	s.OnStatusChanged = append(s.OnStatusChanged, func(r api.Record) {
 		if r.Status != api.StatusHealthy {
 			unhealthy.Store(true)

diff --git a/cmd/ayd/server.go b/cmd/ayd/server.go
@@ -5,7 +5,9 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"os/signal"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/macrat/ayd/internal/endpoint"
@@ -14,7 +16,7 @@ import (
 	"github.com/robfig/cron/v3"
 )
 
-func (cmd *AydCommand) reportStartLog(s *store.Store, protocol, listen string) {
+func (cmd *AydCommand) reportStartServer(s *store.Store, protocol, listen string) {
 	var tasks [][]string
 
 	for _, t := range cmd.Tasks {
@@ -24,12 +26,14 @@ func (cmd *AydCommand) reportStartLog(s *store.Store, protocol, listen string) {
 		})
 	}
 
+	cmd.StartedAt = time.Now()
+
 	u := &api.URL{Scheme: "ayd", Opaque: "server"}
 	s.Report(u, api.Record{
-		Time:    time.Now(),
+		Time:    cmd.StartedAt,
 		Status:  api.StatusHealthy,
 		Target:  u,
-		Message: fmt.Sprintf("start Ayd server"),
+		Message: "start Ayd server",
 		Extra: map[string]interface{}{
 			"url":     fmt.Sprintf("%s://%s", protocol, listen),
 			"targets": tasks,
@@ -38,6 +42,21 @@ func (cmd *AydCommand) reportStartLog(s *store.Store, protocol, listen string) {
 	})
 }
 
+func (cmd *AydCommand) reportStopServer(s *store.Store, protocol, listen string) {
+	u := &api.URL{Scheme: "ayd", Opaque: "server"}
+	s.Report(u, api.Record{
+		Time:    time.Now(),
+		Status:  api.StatusHealthy,
+		Target:  u,
+		Message: "stop Ayd server",
+		Extra: map[string]interface{}{
+			"url":     fmt.Sprintf("%s://%s", protocol, listen),
+			"version": fmt.Sprintf("%s (%s)", version, commit),
+			"since":   cmd.StartedAt.Format(time.RFC3339),
+		},
+	})
+}
+
 func (cmd *AydCommand) RunServer(ctx context.Context, s *store.Store) (exitCode int) {
 	startDebugLogger(s)
 
@@ -54,9 +73,6 @@ func (cmd *AydCommand) RunServer(ctx context.Context, s *store.Store) (exitCode
 		}
 	}
 
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
 	scheduler := cron.New()
 
 	if err := s.Restore(); err != nil {
@@ -65,7 +81,20 @@ func (cmd *AydCommand) RunServer(ctx context.Context, s *store.Store) (exitCode
 	}
 
 	listen := fmt.Sprintf("0.0.0.0:%d", cmd.ListenPort)
-	cmd.reportStartLog(s, protocol, listen)
+	cmd.reportStartServer(s, protocol, listen)
+
+	ctx, cancel := context.WithCancel(ctx)
+	go func() {
+		ch := make(chan os.Signal)
+		signal.Notify(ch, syscall.SIGINT, syscall.SIGHUP)
+		select {
+		case <-ctx.Done():
+		case <-ch:
+			cmd.reportStopServer(s, protocol, listen)
+			cancel()
+		}
+	}()
+	defer cancel()
 
 	wg := &sync.WaitGroup{}
 	for _, t := range cmd.Tasks {