feat: add auto shutdown (#2)

* feat: add auto shutdown

Signed-off-by: cbluebird <crk2778355976@163.com>

* feat: add gracefully shutting down

Signed-off-by: cbluebird <crk2778355976@163.com>

* bugfix: fix the graceful shutdown and http close

Signed-off-by: cbluebird <crk2778355976@163.com>

* bugfix: move http close after err check

Signed-off-by: cbluebird <crk2778355976@163.com>

---------

Signed-off-by: cbluebird <crk2778355976@163.com>

Author: cbluebird

.gitignore

@@ -1,2 +1,3 @@
 wst
 **/*.local
+.idea

go.work.sum

@@ -1,4 +1,8 @@
+golang.org/x/crypto v0.30.0 h1:RwoQn3GkWiMkzlX562cLB7OxWvjH1L8xutO2WoJcRoY=
 golang.org/x/crypto v0.30.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
+golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
 golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/term v0.27.0 h1:WP60Sv1nlK1T6SupCHbXzSaN0b9wUmsPoRS9b61A23Q=
 golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
+golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
 golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=

server/healthy.go

@@ -0,0 +1,57 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"log"
+	"net/http"
+	"os"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	interval      = os.Getenv("AUTO_SHUTDOWN_INTERVAL")
+	serviceTarget = os.Getenv("AUTO_SHUTDOWN_SERVICE_URL")
+	jwtToken      = os.Getenv("JWT_TOKEN")
+)
+
+var ActiveNum int32 = 0
+
+func HealthyCheck(ctx context.Context) {
+	shutdownDuration, _ := time.ParseDuration(interval)
+	ticker := time.NewTicker(1 * time.Minute)
+	defer ticker.Stop()
+
+	zeroDuration := 0 * time.Minute
+	for {
+		select {
+		case <-ticker.C:
+			if atomic.LoadInt32(&ActiveNum) == 0 {
+				zeroDuration += 1 * time.Minute
+				if zeroDuration >= shutdownDuration {
+					sendShutdownRequest()
+				}
+			} else {
+				zeroDuration = 0
+			}
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func sendShutdownRequest() {
+	url := serviceTarget + "/opsrequest"
+	data := map[string]string{
+		"operation": "shutdown",
+		"jwt_token": jwtToken,
+	}
+	jsonData, _ := json.Marshal(data)
+	resp, err := http.Post(url, "application/json", bytes.NewBuffer(jsonData))
+	if err != nil {
+		log.Println("error:fail to send shutdown request", err)
+	}
+	defer resp.Body.Close()
+}

server/main.go

@@ -1,19 +1,50 @@
 package main
 
-import "os"
+import (
+	"context"
+	"log"
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+)
 
 var (
 	listen = os.Getenv("LISTEN")
 	target = os.Getenv("TARGET")
+	flag   = os.Getenv("ENABLE_AUTO_SHUTDOWN")
 )
 
 func main() {
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
+
+	var wg sync.WaitGroup
+
+	if flag == "true" {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			HealthyCheck(ctx)
+		}()
+	}
+
 	if listen == "" || target == "" {
 		panic("LISTEN or TARGET is not set")
 	}
-	_ = NewServer(
-		listen,
-		"/",
-		NewHandler(target),
-	).Serve()
+
+	go func() {
+		err := NewServer(
+			listen,
+			"/",
+			NewHandler(target),
+		).Serve()
+		if err != nil {
+			log.Fatalf("Failed to ws serve: %v", err)
+		}
+	}()
+
+	<-ctx.Done()
+	log.Println("Shutting down gracefully")
+	wg.Wait()
 }

server/wsh.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/http"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"golang.org/x/net/websocket"
@@ -95,7 +96,9 @@ func (h *Handler) putBuffer(buffer *[]byte) {
 }
 
 func (h *Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	atomic.AddInt32(&ActiveNum, 1)
 	h.wsServer.ServeHTTP(w, req)
+	atomic.AddInt32(&ActiveNum, -1)
 }
 
 var pingCodec = websocket.Codec{