138 lines
3.4 KiB
Go
138 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
|
|
"tyapi-server/internal/config"
|
|
"tyapi-server/internal/domains/article/entities"
|
|
|
|
"github.com/hibiken/asynq"
|
|
"go.uber.org/zap"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
const (
|
|
TaskTypeArticlePublish = "article:publish"
|
|
)
|
|
|
|
func main() {
|
|
// 加载配置
|
|
cfg, err := config.LoadConfig()
|
|
if err != nil {
|
|
log.Fatal("加载配置失败:", err)
|
|
}
|
|
|
|
// 创建日志器
|
|
logger, err := zap.NewProduction()
|
|
if err != nil {
|
|
log.Fatal("创建日志器失败:", err)
|
|
}
|
|
defer logger.Sync()
|
|
|
|
// 连接数据库
|
|
// 在 Docker 环境中使用容器名
|
|
dbHost := os.Getenv("DB_HOST")
|
|
if dbHost == "" {
|
|
dbHost = cfg.Database.Host
|
|
}
|
|
|
|
// 使用默认端口 5432
|
|
dbPort := 5432
|
|
|
|
dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%d sslmode=disable TimeZone=Asia/Shanghai",
|
|
dbHost, cfg.Database.User, cfg.Database.Password, cfg.Database.Name, dbPort)
|
|
fmt.Printf("dsn: %s\n", dsn)
|
|
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
|
if err != nil {
|
|
logger.Fatal("连接数据库失败", zap.Error(err))
|
|
}
|
|
|
|
// 从环境变量获取 Redis 地址
|
|
redisAddr := os.Getenv("REDIS_ADDR")
|
|
if redisAddr == "" {
|
|
redisAddr = fmt.Sprintf("%s:%d", cfg.Redis.Host, cfg.Redis.Port)
|
|
}
|
|
|
|
// 创建 Asynq Server
|
|
server := asynq.NewServer(
|
|
asynq.RedisClientOpt{Addr: redisAddr},
|
|
asynq.Config{
|
|
Concurrency: 10,
|
|
Queues: map[string]int{
|
|
"critical": 6,
|
|
"default": 3,
|
|
"low": 1,
|
|
},
|
|
},
|
|
)
|
|
|
|
// 创建任务处理器
|
|
mux := asynq.NewServeMux()
|
|
mux.HandleFunc(TaskTypeArticlePublish, func(ctx context.Context, t *asynq.Task) error {
|
|
return handleArticlePublish(ctx, t, db, logger)
|
|
})
|
|
|
|
// 启动 Worker
|
|
go func() {
|
|
logger.Info("启动 Asynq Worker", zap.String("redis_addr", redisAddr))
|
|
if err := server.Run(mux); err != nil {
|
|
logger.Fatal("启动 Worker 失败", zap.Error(err))
|
|
}
|
|
}()
|
|
|
|
// 等待信号
|
|
quit := make(chan os.Signal, 1)
|
|
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
|
<-quit
|
|
|
|
// 优雅关闭
|
|
logger.Info("正在关闭 Worker...")
|
|
server.Stop()
|
|
server.Shutdown()
|
|
logger.Info("Worker 已关闭")
|
|
}
|
|
|
|
// handleArticlePublish 处理文章定时发布任务
|
|
func handleArticlePublish(ctx context.Context, t *asynq.Task, db *gorm.DB, logger *zap.Logger) error {
|
|
var payload map[string]interface{}
|
|
if err := json.Unmarshal(t.Payload(), &payload); err != nil {
|
|
logger.Error("解析任务载荷失败", zap.Error(err))
|
|
return fmt.Errorf("解析任务载荷失败: %w", err)
|
|
}
|
|
|
|
articleID, ok := payload["article_id"].(string)
|
|
if !ok {
|
|
logger.Error("任务载荷中缺少文章ID")
|
|
return fmt.Errorf("任务载荷中缺少文章ID")
|
|
}
|
|
|
|
// 获取文章
|
|
var article entities.Article
|
|
if err := db.WithContext(ctx).First(&article, "id = ?", articleID).Error; err != nil {
|
|
logger.Error("获取文章失败", zap.String("article_id", articleID), zap.Error(err))
|
|
return fmt.Errorf("获取文章失败: %w", err)
|
|
}
|
|
|
|
// 发布文章
|
|
if err := article.Publish(); err != nil {
|
|
logger.Error("发布文章失败", zap.String("article_id", articleID), zap.Error(err))
|
|
return fmt.Errorf("发布文章失败: %w", err)
|
|
}
|
|
|
|
// 保存更新
|
|
if err := db.WithContext(ctx).Save(&article).Error; err != nil {
|
|
logger.Error("保存文章失败", zap.String("article_id", articleID), zap.Error(err))
|
|
return fmt.Errorf("保存文章失败: %w", err)
|
|
}
|
|
|
|
logger.Info("定时发布文章成功", zap.String("article_id", articleID))
|
|
return nil
|
|
}
|