Skip to content

Commit 62cb6dd

Browse files
kraukrau
committed
更好的缓存与退出
1 parent 0c28b81 commit 62cb6dd

11 files changed

+123
-44
lines changed

‎bot/handlers_admin.go

+35-14
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ func deletePicture(ctx context.Context, bot *telego.Bot, message telego.Message)
142142
channelMessageID = originChannel.MessageID
143143
}
144144
if cmd == "del" {
145-
if !service.CheckAdminPermission(ctx, message.From.ID, types.DeletePicture) {
145+
if !service.CheckAdminPermission(ctx, message.From.ID, types.PermissionDeleteArtwork) {
146146
telegram.ReplyMessage(bot, message, "你没有删除图片的权限")
147147
return
148148
}
@@ -164,7 +164,7 @@ func deletePicture(ctx context.Context, bot *telego.Bot, message telego.Message)
164164
}
165165
return
166166
}
167-
if !service.CheckAdminPermission(ctx, message.From.ID, types.DeleteArtwork) {
167+
if !service.CheckAdminPermission(ctx, message.From.ID, types.PermissionDeleteArtwork) {
168168
telegram.ReplyMessage(bot, message, "你没有删除作品的权限")
169169
return
170170
}
@@ -196,7 +196,7 @@ func deletePicture(ctx context.Context, bot *telego.Bot, message telego.Message)
196196
}
197197

198198
func fetchArtwork(ctx context.Context, bot *telego.Bot, message telego.Message) {
199-
if !service.CheckAdminPermission(ctx, message.From.ID, types.FetchArtwork) {
199+
if !service.CheckAdminPermission(ctx, message.From.ID, types.PermissionFetchArtwork) {
200200
telegram.ReplyMessage(bot, message, "你没有拉取作品的权限")
201201
return
202202
}
@@ -206,7 +206,7 @@ func fetchArtwork(ctx context.Context, bot *telego.Bot, message telego.Message)
206206
}
207207

208208
func postArtwork(ctx context.Context, bot *telego.Bot, query telego.CallbackQuery) {
209-
if !service.CheckAdminPermission(ctx, query.From.ID, types.PostArtwork) {
209+
if !service.CheckAdminPermission(ctx, query.From.ID, types.PermissionPostArtwork) {
210210
bot.AnswerCallbackQuery(&telego.AnswerCallbackQueryParams{
211211
CallbackQueryID: query.ID,
212212
Text: "你没有发布作品的权限",
@@ -227,18 +227,39 @@ func postArtwork(ctx context.Context, bot *telego.Bot, query telego.CallbackQuer
227227
return
228228
}
229229
Logger.Infof("posting artwork: %s", sourceURL)
230-
artwork, err := service.GetCachedArtworkByURL(ctx, sourceURL)
230+
231+
var artwork *types.Artwork
232+
cachedArtwork, err := service.GetCachedArtworkByURL(ctx, sourceURL)
231233
if err != nil {
232234
artwork, err = sources.GetArtworkInfo(sourceURL)
233-
}
234-
if err != nil {
235-
bot.AnswerCallbackQuery(&telego.AnswerCallbackQueryParams{
236-
CallbackQueryID: query.ID,
237-
Text: "获取作品信息失败" + err.Error(),
238-
ShowAlert: true,
239-
CacheTime: 60,
240-
})
241-
return
235+
if err != nil {
236+
bot.AnswerCallbackQuery(&telego.AnswerCallbackQueryParams{
237+
CallbackQueryID: query.ID,
238+
Text: "获取作品信息失败" + err.Error(),
239+
ShowAlert: true,
240+
CacheTime: 60,
241+
})
242+
return
243+
}
244+
} else {
245+
if cachedArtwork.Status == types.ArtworkStatusPosting {
246+
bot.AnswerCallbackQuery(&telego.AnswerCallbackQueryParams{
247+
CallbackQueryID: query.ID,
248+
Text: "该作品正在发布中",
249+
ShowAlert: true,
250+
CacheTime: 60,
251+
})
252+
return
253+
}
254+
if err := service.UpdateCachedArtworkByURL(ctx, sourceURL, types.ArtworkStatusPosting); err != nil {
255+
Logger.Errorf("更新缓存作品状态失败: %s", err)
256+
}
257+
artwork = cachedArtwork.Artwork
258+
defer func() {
259+
if err := service.UpdateCachedArtworkByURL(ctx, sourceURL, types.ArtworkStatusCached); err != nil {
260+
Logger.Errorf("更新缓存作品状态失败: %s", err)
261+
}
262+
}()
242263
}
243264
go bot.EditMessageCaption(&telego.EditMessageCaptionParams{
244265
ChatID: telegoutil.ID(query.Message.GetChat().ID),

‎bot/handlers_common.go

+9-5
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"math/rand"
1313
"strconv"
1414
"strings"
15+
"time"
1516

1617
. "ManyACG/logger"
1718

@@ -151,7 +152,7 @@ func randomPicture(ctx context.Context, bot *telego.Bot, message telego.Message)
151152
}
152153

153154
func getArtworkInfo(ctx context.Context, bot *telego.Bot, message telego.Message) {
154-
hasPermission := CheckPermissionInGroup(ctx, message, types.GetArtworkInfo)
155+
hasPermission := CheckPermissionInGroup(ctx, message, types.PermissionGetArtworkInfo)
155156
sourceURL := FindSourceURLForMessage(&message)
156157
var waitMessageID int
157158
if hasPermission {
@@ -165,6 +166,7 @@ func getArtworkInfo(ctx context.Context, bot *telego.Bot, message telego.Message
165166
}()
166167
}
167168
defer func() {
169+
time.Sleep(1 * time.Second)
168170
if waitMessageID != 0 {
169171
bot.DeleteMessage(telegoutil.Delete(message.Chat.ChatID(), waitMessageID))
170172
}
@@ -173,20 +175,22 @@ func getArtworkInfo(ctx context.Context, bot *telego.Bot, message telego.Message
173175
isAlreadyPosted := true
174176
artwork, err := service.GetArtworkByURL(ctx, sourceURL)
175177
if err != nil {
176-
artwork, err = service.GetCachedArtworkByURL(ctx, sourceURL)
178+
cachedArtwork, err := service.GetCachedArtworkByURL(ctx, sourceURL)
177179
if err != nil && !hasPermission {
178180
return
179181
}
180182
isAlreadyPosted = false
181-
if artwork == nil {
183+
if cachedArtwork == nil {
182184
artwork, err = sources.GetArtworkInfo(sourceURL)
183185
if err != nil {
184186
telegram.ReplyMessage(bot, message, "获取作品信息失败: "+err.Error())
185187
return
186188
}
187-
if err := service.CreateCachedArtwork(ctx, artwork); err != nil {
189+
if err := service.CreateCachedArtwork(ctx, artwork, types.ArtworkStatusCached); err != nil {
188190
Logger.Warnf("缓存作品失败: %s", err)
189191
}
192+
} else {
193+
artwork = cachedArtwork.Artwork
190194
}
191195
}
192196

@@ -237,7 +241,7 @@ func getArtworkInfo(ctx context.Context, bot *telego.Bot, message telego.Message
237241
}
238242

239243
func searchPicture(ctx context.Context, bot *telego.Bot, message telego.Message) {
240-
hasPermission := CheckPermissionInGroup(ctx, message, types.SearchPicture)
244+
hasPermission := CheckPermissionInGroup(ctx, message, types.PermissionSearchPicture)
241245
if message.ReplyToMessage == nil {
242246
telegram.ReplyMessage(bot, message, "请使用该命令回复一条图片消息")
243247
return

‎cmd/run.go

+11-1
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,12 @@ import (
77
"ManyACG/dao"
88
"ManyACG/fetcher"
99
"ManyACG/logger"
10+
"ManyACG/service"
1011
"ManyACG/sources"
1112
"context"
1213
"os"
14+
"os/signal"
15+
"syscall"
1316
"time"
1417
)
1518

@@ -32,5 +35,12 @@ func Run() {
3235
if config.Cfg.API.Enable {
3336
go restful.Run()
3437
}
35-
select {}
38+
quit := make(chan os.Signal, 1)
39+
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
40+
sig := <-quit
41+
logger.Logger.Info(sig, " Exiting...")
42+
if err := service.Cleanup(context.TODO()); err != nil {
43+
logger.Logger.Error(err)
44+
}
45+
logger.Logger.Info("See you next time.")
3646
}

‎cmd/version.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import (
77
)
88

99
const (
10-
Version string = "0.4.8"
10+
Version string = "0.4.9"
1111
)
1212

1313
var VersionCmd = &cobra.Command{

‎dao/cached_artwork.go

+13-1
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,12 @@ import (
1313

1414
var cachedArtworkCollection *mongo.Collection
1515

16-
func CreateCachedArtwork(ctx context.Context, artwork *types.Artwork) (*mongo.InsertOneResult, error) {
16+
func CreateCachedArtwork(ctx context.Context, artwork *types.Artwork, status types.ArtworkStatus) (*mongo.InsertOneResult, error) {
1717
cachedArtwork := &model.CachedArtworksModel{
1818
SourceURL: artwork.SourceURL,
1919
Artwork: artwork,
2020
CreatedAt: primitive.NewDateTimeFromTime(time.Now()),
21+
Status: status,
2122
}
2223
return cachedArtworkCollection.InsertOne(ctx, cachedArtwork)
2324
}
@@ -30,3 +31,14 @@ func GetCachedArtworkByURL(ctx context.Context, url string) (*model.CachedArtwor
3031
}
3132
return &cachedArtwork, err
3233
}
34+
35+
func UpdateCachedArtworkByURL(ctx context.Context, url string, status types.ArtworkStatus) (*mongo.UpdateResult, error) {
36+
filter := bson.M{"source_url": url}
37+
update := bson.M{"$set": bson.M{"status": status}}
38+
return cachedArtworkCollection.UpdateOne(ctx, filter, update)
39+
}
40+
41+
func CleanPostingCachedArtwork(ctx context.Context) (*mongo.DeleteResult, error) {
42+
filter := bson.M{"status": types.ArtworkStatusPosting}
43+
return cachedArtworkCollection.DeleteMany(ctx, filter)
44+
}

‎dao/model/model.go

+5-4
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,9 @@ type CallbackDataModel struct {
7474
}
7575

7676
type CachedArtworksModel struct {
77-
ID primitive.ObjectID `bson:"_id,omitempty"`
78-
SourceURL string `bson:"source_url"`
79-
CreatedAt primitive.DateTime `bson:"created_at"`
80-
Artwork *types.Artwork `bson:"artwork"`
77+
ID primitive.ObjectID `bson:"_id,omitempty"`
78+
SourceURL string `bson:"source_url"`
79+
CreatedAt primitive.DateTime `bson:"created_at"`
80+
Artwork *types.Artwork `bson:"artwork"`
81+
Status types.ArtworkStatus `bson:"status"`
8182
}

‎fetcher/service.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,8 @@ func PostAndCreateArtwork(ctx context.Context, artwork *types.Artwork, bot *tele
3434
for i, picture := range artwork.Pictures {
3535
info, err := storage.SavePicture(artwork, picture)
3636
if err != nil {
37-
Logger.Errorf("saving picture %s: %s", picture.Original, err)
38-
return fmt.Errorf("saving picture %s: %w", picture.Original, err)
37+
Logger.Errorf("saving picture %d of artwork %s: %s", i, artwork.Title, err)
38+
return fmt.Errorf("saving picture %d of artwork %s: %w", i, artwork.Title, err)
3939
}
4040
artwork.Pictures[i].StorageInfo = info
4141
}

‎service/artwork.go

+9-4
Original file line numberDiff line numberDiff line change
@@ -386,15 +386,20 @@ func GetArtworkByID(ctx context.Context, id primitive.ObjectID) (*types.Artwork,
386386
}, nil
387387
}
388388

389-
func CreateCachedArtwork(ctx context.Context, artwork *types.Artwork) error {
390-
_, err := dao.CreateCachedArtwork(ctx, artwork)
389+
func CreateCachedArtwork(ctx context.Context, artwork *types.Artwork, status types.ArtworkStatus) error {
390+
_, err := dao.CreateCachedArtwork(ctx, artwork, status)
391391
return err
392392
}
393393

394-
func GetCachedArtworkByURL(ctx context.Context, sourceURL string) (*types.Artwork, error) {
394+
func GetCachedArtworkByURL(ctx context.Context, sourceURL string) (*model.CachedArtworksModel, error) {
395395
cachedArtwork, err := dao.GetCachedArtworkByURL(ctx, sourceURL)
396396
if err != nil {
397397
return nil, err
398398
}
399-
return cachedArtwork.Artwork, nil
399+
return cachedArtwork, nil
400+
}
401+
402+
func UpdateCachedArtworkByURL(ctx context.Context, sourceURL string, status types.ArtworkStatus) error {
403+
_, err := dao.UpdateCachedArtworkByURL(ctx, sourceURL, status)
404+
return err
400405
}

‎service/cleanup.go

+18
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package service
2+
3+
import (
4+
"ManyACG/dao"
5+
"context"
6+
"fmt"
7+
)
8+
9+
func Cleanup(ctx context.Context) error {
10+
var errs []error
11+
if _, err := dao.CleanPostingCachedArtwork(ctx); err != nil {
12+
errs = append(errs, err)
13+
}
14+
if len(errs) > 0 {
15+
return fmt.Errorf("failed to cleanup: %v", errs)
16+
}
17+
return nil
18+
}

‎types/permissions.go

+12-12
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,19 @@ package types
33
type Permission string
44

55
const (
6-
PostArtwork Permission = "post_artwork"
7-
DeleteArtwork Permission = "delete_artwork"
8-
DeletePicture Permission = "delete_picture"
9-
FetchArtwork Permission = "fetch_artwork"
10-
GetArtworkInfo Permission = "get_artwork_info"
11-
SearchPicture Permission = "search_picture"
6+
PermissionPostArtwork Permission = "post_artwork"
7+
PermissionDeleteArtwork Permission = "delete_artwork"
8+
PermissionDeletePicture Permission = "delete_picture"
9+
PermissionFetchArtwork Permission = "fetch_artwork"
10+
PermissionGetArtworkInfo Permission = "get_artwork_info"
11+
PermissionSearchPicture Permission = "search_picture"
1212
)
1313

1414
var AllPermissions = []Permission{
15-
PostArtwork,
16-
DeleteArtwork,
17-
DeletePicture,
18-
FetchArtwork,
19-
GetArtworkInfo,
20-
SearchPicture,
15+
PermissionPostArtwork,
16+
PermissionDeleteArtwork,
17+
PermissionDeletePicture,
18+
PermissionFetchArtwork,
19+
PermissionGetArtworkInfo,
20+
PermissionSearchPicture,
2121
}

‎types/status.go

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
package types
2+
3+
type ArtworkStatus string
4+
5+
const (
6+
ArtworkStatusCached ArtworkStatus = "cached"
7+
ArtworkStatusPosting ArtworkStatus = "posting"
8+
)

0 commit comments

Comments
 (0)