From 986d90ae046f28e2f7377715a4feb721b57b52eb Mon Sep 17 00:00:00 2001
From: feitianbubu <feitianbubu@qq.com>
Date: Thu, 2 Jul 2026 21:53:09 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E6=9C=8D=E5=8A=A1=E4=BC=98?=
 =?UTF-8?q?=E9=9B=85=E5=85=B3=E9=97=AD,=E9=81=BF=E5=85=8D=E9=87=8D?=
 =?UTF-8?q?=E5=90=AF=E5=9B=9E=E5=A4=8D=E4=B8=AD=E6=96=AD=E5=92=8C=E9=9D=A2?=
 =?UTF-8?q?=E6=9D=BF=E7=BC=93=E5=AD=98=E6=95=B0=E6=8D=AE=E4=B8=A2=E5=A4=B1?=
 =?UTF-8?q?=20(#4258)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: add graceful shutdown with configurable timeout

* fix: flush quota dashboard cache on graceful shutdown

Persist the in-memory CacheQuotaData aggregation to the quota_data table
before process exit, so a restart no longer drops up to one DataExportInterval
window of dashboard data (issue #5679).
---
 main.go | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/main.go b/main.go
index c157bc0a..adc9ddec 100644
--- a/main.go
+++ b/main.go
@@ -2,13 +2,17 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"embed"
+	"errors"
 	"fmt"
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
 	"strconv"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/QuantumNous/new-api/common"
@@ -209,10 +213,34 @@ func main() {
 	// Log startup success message
 	common.LogStartupSuccess(startTime, port)
 
-	err = server.Run(":" + port)
-	if err != nil {
-		common.FatalLog("failed to start HTTP server: " + err.Error())
+	srv := &http.Server{
+		Addr:    ":" + port,
+		Handler: server,
 	}
+
+	go func() {
+		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			common.FatalLog("failed to start HTTP server: " + err.Error())
+		}
+	}()
+
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
+	sig := <-quit
+	common.SysLog(fmt.Sprintf("received signal: %v, shutting down...", sig))
+
+	// SSE streams may run for minutes; give them time to finish before forced exit
+	shutdownTimeout := time.Duration(common.GetEnvOrDefault("SHUTDOWN_TIMEOUT_SECONDS", 120)) * time.Second
+	ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
+	defer cancel()
+	if err := srv.Shutdown(ctx); err != nil {
+		common.SysError(fmt.Sprintf("server forced to shutdown: %v", err))
+	}
+	// 内存中的看板数据保存入库，避免重启丢失未落库数据 (issue #5679)
+	if common.DataExportEnabled {
+		model.SaveQuotaDataCache()
+	}
+	common.SysLog("server exited")
 }
 
 func InjectUmamiAnalytics() {