Skip to content

Commit f998d37

Browse files
kraukrau
committed
feat: graceful shutdown
1 parent ee2e6bb commit f998d37

File tree

12 files changed

+94
-72
lines changed

12 files changed

+94
-72
lines changed

‎api/restful/api.go‎

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,24 @@
11
package restful
22

33
import (
4+
"context"
45
"fmt"
6+
"net/http"
57
"os"
8+
"time"
69

710
"github.com/krau/ManyACG/api/restful/middleware"
811
"github.com/krau/ManyACG/api/restful/routers"
12+
"github.com/krau/ManyACG/common"
913
"github.com/krau/ManyACG/config"
1014

1115
"github.com/penglongli/gin-metrics/ginmetrics"
1216

1317
"github.com/gin-contrib/cors"
1418
"github.com/gin-gonic/gin"
15-
"github.com/krau/ManyACG/common"
1619
)
1720

18-
func Run() {
21+
func Run(ctx context.Context) {
1922
if config.Cfg.Debug {
2023
gin.SetMode(gin.DebugMode)
2124
} else {
@@ -46,8 +49,26 @@ func Run() {
4649
v1 := r.Group("/api/v1")
4750
routers.RegisterAllRouters(v1, middleware.JWTAuthMiddleware)
4851

49-
if err := r.Run(config.Cfg.API.Address); err != nil {
50-
common.Logger.Fatalf("Failed to start server: %v", err)
51-
os.Exit(1)
52+
server := &http.Server{
53+
Addr: config.Cfg.API.Address,
54+
Handler: r,
5255
}
56+
57+
go func() {
58+
<-ctx.Done()
59+
common.Logger.Info("Shutting down api server...")
60+
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
61+
defer cancel()
62+
if err := server.Shutdown(shutdownCtx); err != nil {
63+
common.Logger.Fatalf("Failed to shutdown server: %v", err)
64+
}
65+
common.Logger.Info("API server stopped")
66+
}()
67+
68+
go func() {
69+
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
70+
common.Logger.Fatalf("Failed to start server: %v", err)
71+
os.Exit(1)
72+
}
73+
}()
5374
}

‎cmd/run.go‎

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -51,34 +51,40 @@ func Run() {
5151
}()
5252
}
5353

54-
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
55-
defer cancel()
56-
common.Logger.Info("Start running")
54+
ctx, stop := signal.NotifyContext(context.TODO(), syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
55+
defer stop()
56+
common.Logger.Info("Starting...")
5757
dao.InitDB(ctx)
5858
defer func() {
5959
if err := dao.Client.Disconnect(ctx); err != nil {
6060
common.Logger.Fatal(err)
61-
os.Exit(1)
6261
}
6362
}()
64-
service.InitService()
63+
service.InitService(ctx)
64+
sources.InitSources(service.NewService())
65+
storage.InitStorage(ctx)
6566
if config.Cfg.Telegram.Token != "" {
66-
go telegram.RunPolling()
67+
telegram.RunPolling(ctx)
6768
}
68-
storage.InitStorage()
69-
sources.InitSources(service.NewService())
70-
go fetcher.StartScheduler(context.TODO())
69+
70+
go fetcher.StartScheduler(ctx)
7171
if config.Cfg.API.Enable {
72-
go restful.Run()
72+
restful.Run(ctx)
7373
}
74-
quit := make(chan os.Signal, 1)
75-
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
76-
sig := <-quit
77-
common.Logger.Info(sig, " Exiting...")
74+
75+
common.Logger.Info("ManyACG is running !")
76+
7877
defer common.Logger.Info("Exited.")
79-
if err := service.Cleanup(context.TODO()); err != nil {
78+
<-ctx.Done()
79+
cleanCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
80+
defer cancel()
81+
if err := service.Cleanup(cleanCtx); err != nil {
8082
common.Logger.Error(err)
8183
}
84+
cleanCacheDir()
85+
}
86+
87+
func cleanCacheDir() {
8288
if config.Cfg.Storage.CacheDir != "" && !config.Cfg.Debug {
8389
for _, path := range []string{"/", ".", "\\", ".."} {
8490
if filepath.Clean(config.Cfg.Storage.CacheDir) == path {

‎service/change_stream.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,7 @@ func GetArtworkChangeStream(ctx context.Context) (*mongo.ChangeStream, error) {
2626
return changeStream, nil
2727
}
2828

29-
func syncArtworkToSearchEngine() {
30-
ctx := context.Background()
29+
func syncArtworkToSearchEngine(ctx context.Context) {
3130
changeStream, err := GetArtworkChangeStream(ctx)
3231
if err != nil {
3332
common.Logger.Fatalf("get artwork change stream error: %s", err)

‎service/service.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@ import (
77
"github.com/krau/ManyACG/types"
88
)
99

10-
func InitService() {
10+
func InitService(ctx context.Context) {
1111
go listenProcessPictureTask()
1212
if config.Cfg.Search.Enable {
13-
go syncArtworkToSearchEngine()
13+
go syncArtworkToSearchEngine(ctx)
1414
}
1515
if config.Cfg.Tagger.Enable {
1616
go listenPredictArtworkTagsTask()

‎storage/alist/alist.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ var (
3333
loginReq *loginRequset
3434
)
3535

36-
func (a *Alist) Init() {
36+
func (a *Alist) Init(ctx context.Context) {
3737
alistConfig := config.Cfg.Storage.Alist
3838
basePath = strings.TrimSuffix(alistConfig.Path, "/")
3939
baseUrl = strings.TrimSuffix(alistConfig.URL, "/")
@@ -45,7 +45,7 @@ func (a *Alist) Init() {
4545
Username: alistConfig.Username,
4646
Password: alistConfig.Password,
4747
}
48-
token, err := getJwtToken()
48+
token, err := getJwtToken(ctx)
4949
if err != nil {
5050
common.Logger.Errorf("Failed to login to Alist: %v", err)
5151
os.Exit(1)

‎storage/alist/auth.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package alist
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"net/http"
@@ -12,8 +13,8 @@ import (
1213
"github.com/imroc/req/v3"
1314
)
1415

15-
func getJwtToken() (string, error) {
16-
resp, err := reqClient.R().SetBodyJsonMarshal(loginReq).Post("/api/auth/login")
16+
func getJwtToken(ctx context.Context) (string, error) {
17+
resp, err := reqClient.R().SetContext(ctx).SetBodyJsonMarshal(loginReq).Post("/api/auth/login")
1718
if err != nil {
1819
return "", err
1920
}
@@ -30,7 +31,7 @@ func getJwtToken() (string, error) {
3031
func refreshJwtToken(client *req.Client) {
3132
for {
3233
time.Sleep(time.Duration(config.Cfg.Storage.Alist.TokenExpire) * time.Second)
33-
token, err := getJwtToken()
34+
token, err := getJwtToken(context.Background())
3435
if err != nil {
3536
common.Logger.Errorf("Failed to refresh jwt token: %v", err)
3637
continue

‎storage/interface.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import (
77
)
88

99
type Storage interface {
10-
Init()
10+
Init(ctx context.Context)
1111
// filePath 本地文件路径, storagePath 存储路径.
1212
//
1313
// 存储实现可能会对传入的存储路径进行其他处理 (如添加前缀), 因此返回的 StorageDetail 中的 Path 可能与传入的 storagePath 不同.

‎storage/local/local.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ var (
1919
basePath string
2020
)
2121

22-
func (l *Local) Init() {
22+
func (l *Local) Init(ctx context.Context) {
2323
basePath = strings.TrimSuffix(config.Cfg.Storage.Local.Path, "/")
2424
if basePath == "" {
2525
common.Logger.Fatalf("Local storage path not set,for example: manyacg/storage")

‎storage/storage.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,23 +24,23 @@ import (
2424

2525
var Storages = make(map[types.StorageType]Storage)
2626

27-
func InitStorage() {
27+
func InitStorage(ctx context.Context) {
2828
common.Logger.Info("Initializing storage")
2929
if config.Cfg.Storage.Local.Enable {
3030
Storages[types.StorageTypeLocal] = new(local.Local)
31-
Storages[types.StorageTypeLocal].Init()
31+
Storages[types.StorageTypeLocal].Init(ctx)
3232
}
3333
if config.Cfg.Storage.Webdav.Enable {
3434
Storages[types.StorageTypeWebdav] = new(webdav.Webdav)
35-
Storages[types.StorageTypeWebdav].Init()
35+
Storages[types.StorageTypeWebdav].Init(ctx)
3636
}
3737
if config.Cfg.Storage.Alist.Enable {
3838
Storages[types.StorageTypeAlist] = new(alist.Alist)
39-
Storages[types.StorageTypeAlist].Init()
39+
Storages[types.StorageTypeAlist].Init(ctx)
4040
}
4141
if config.Cfg.Storage.Telegram.Enable {
4242
Storages[types.StorageTypeTelegram] = new(telegram.TelegramStorage)
43-
Storages[types.StorageTypeTelegram].Init()
43+
Storages[types.StorageTypeTelegram].Init(ctx)
4444
}
4545
}
4646

‎storage/telegram/telegram.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ var (
2222
ChatID telego.ChatID
2323
)
2424

25-
func (t *TelegramStorage) Init() {
25+
func (t *TelegramStorage) Init(ctx context.Context) {
2626
common.Logger.Infof("Initializing telegram storage")
2727
ChatID = telegoutil.ID(config.Cfg.Storage.Telegram.ChatID)
2828
var err error
@@ -31,8 +31,6 @@ func (t *TelegramStorage) Init() {
3131
common.Logger.Fatalf("failed to create telegram bot: %s", err)
3232
os.Exit(1)
3333
}
34-
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
35-
defer cancel()
3634
botInfo, err := Bot.GetMe(ctx)
3735
if err != nil {
3836
common.Logger.Fatalf("failed to get bot info: %s", err)

0 commit comments

Comments
 (0)