Files
aiaa-notification-server/cmd/server/main.go
T
ryan 39f3774940 feat: 配置列表分页与钉钉机器人分钟级排队限流
统一 sources/templates/channels/rules 列表为分页响应,避免配置增多时全量返回;按钉钉 access_token 限制每分钟发送并在超限时等待下一分钟,降低触发官方封禁风险。

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-10 00:34:26 +08:00

259 lines
6.9 KiB
Go

package main
import (
"context"
"fmt"
"log/slog"
"net/http"
"os"
"os/signal"
"syscall"
"time"
"aiaa-notification-service/internal/adapter"
"aiaa-notification-service/internal/cache"
"aiaa-notification-service/internal/config"
"aiaa-notification-service/internal/engine"
"aiaa-notification-service/internal/handler"
"aiaa-notification-service/internal/store"
"github.com/gin-gonic/gin"
"github.com/logbull/logbull-go/logbull"
)
func main() {
// Load config first (need logbull settings for logger setup)
cfg, err := config.Load("config/config.yaml")
if err != nil {
slog.Error("failed to load config", "error", err)
os.Exit(1)
}
// Setup dual logger: stdout JSON + Logbull remote
stdoutHandler := slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo})
var logbullHandler slog.Handler
if cfg.Logbull.Host != "" {
lbCfg := logbull.Config{
Host: cfg.Logbull.Host,
ProjectID: cfg.Logbull.ProjectID,
APIKey: cfg.Logbull.APIKey,
}
switch cfg.Logbull.LogLevel {
case "DEBUG":
lbCfg.LogLevel = logbull.DEBUG
case "WARNING":
lbCfg.LogLevel = logbull.WARNING
case "ERROR":
lbCfg.LogLevel = logbull.ERROR
default:
lbCfg.LogLevel = logbull.INFO
}
h, err := logbull.NewSlogHandler(lbCfg)
if err != nil {
slog.Warn("logbull init failed, using stdout only", "error", err)
} else {
logbullHandler = h
}
}
var baseHandler slog.Handler
if logbullHandler != nil {
baseHandler = &teeHandler{handlers: []slog.Handler{stdoutHandler, logbullHandler}}
} else {
baseHandler = stdoutHandler
}
logger := slog.New(baseHandler)
slog.SetDefault(logger)
// Flush logbull on exit
defer func() {
if lbh, ok := logbullHandler.(interface{ Flush() }); ok {
lbh.Flush()
}
if lbh, ok := logbullHandler.(interface{ Shutdown() }); ok {
lbh.Shutdown()
}
time.Sleep(2 * time.Second) // allow async flush
}()
// Connect MySQL
st, err := store.NewStore(cfg.Database)
if err != nil {
slog.Error("failed to connect to MySQL", "error", err)
os.Exit(1)
}
defer st.Close()
// Connect Redis
var redisCache *cache.Cache
if cfg.Redis.Host != "" {
redisCache, err = cache.NewCache(cfg.Redis)
if err != nil {
slog.Warn("redis connection failed, running without cache", "error", err)
redisCache = nil
}
} else {
slog.Warn("redis not configured, running without cache")
}
if redisCache != nil {
defer redisCache.Close()
}
// Build engine
matcher := engine.NewMatcher(st, redisCache)
renderer := engine.NewRenderer()
// DingTalk per-robot rate limit (queue/wait when over 18/min by default)
var dingtalkLimiter *adapter.DingTalkLimiter
limitPerMin := cfg.RateLimit.DingTalkPerMin
if redisCache != nil {
dingtalkLimiter = adapter.NewDingTalkLimiter(redisCache, limitPerMin)
} else {
dingtalkLimiter = adapter.NewMemoryDingTalkLimiter(limitPerMin)
slog.Warn("dingtalk rate limit using in-memory store (single instance only)")
}
// Build sender factory
senderFactory := func(channelType string) (adapter.ChannelSender, error) {
return adapter.NewSender(channelType, &cfg.SMTP, dingtalkLimiter)
}
router := engine.NewRouter(st, redisCache, senderFactory)
// Build handlers
notifyH := handler.NewNotifyHandler(st, redisCache, matcher, renderer, router)
sourceH := handler.NewSourceHandler(st, redisCache)
templateH := handler.NewTemplateHandler(st, redisCache)
channelH := handler.NewChannelHandler(st, redisCache)
ruleH := handler.NewRuleHandler(st, redisCache)
msgLogH := handler.NewMessageLogHandler(st)
// Setup Gin
gin.SetMode(gin.ReleaseMode)
r := gin.New()
r.Use(gin.Recovery())
// Health check
r.GET("/health", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"status": "ok"})
})
api := r.Group("/api/v1")
// Notify endpoint (source auth + rate limit)
notifyGroup := api.Group("/notify")
notifyGroup.Use(handler.SourceAuth(st))
if redisCache != nil {
notifyGroup.Use(handler.RateLimit(redisCache, cfg.RateLimit.Default))
}
notifyGroup.POST("", notifyH.Handle)
// Management endpoints (admin auth)
admin := api.Group("")
admin.Use(handler.AdminAuth(cfg.Server.AdminKey))
// Sources
admin.POST("/sources", sourceH.Create)
admin.GET("/sources", sourceH.List)
admin.GET("/sources/:id", sourceH.Get)
admin.PUT("/sources/:id", sourceH.Update)
admin.DELETE("/sources/:id", sourceH.Delete)
// Templates
admin.POST("/templates", templateH.Create)
admin.GET("/templates", templateH.List)
admin.GET("/templates/:id", templateH.Get)
admin.PUT("/templates/:id", templateH.Update)
admin.DELETE("/templates/:id", templateH.Delete)
// Channels
admin.POST("/channels", channelH.Create)
admin.GET("/channels", channelH.List)
admin.GET("/channels/:id", channelH.Get)
admin.PUT("/channels/:id", channelH.Update)
admin.DELETE("/channels/:id", channelH.Delete)
// Rules
admin.POST("/rules", ruleH.Create)
admin.GET("/rules", ruleH.List)
admin.GET("/rules/:id", ruleH.Get)
admin.PUT("/rules/:id", ruleH.Update)
admin.DELETE("/rules/:id", ruleH.Delete)
admin.PATCH("/rules/:id/enable", ruleH.Enable)
admin.PATCH("/rules/:id/disable", ruleH.Disable)
admin.PATCH("/rules/:id/channels/:channel_id/enable", ruleH.EnableChannel)
admin.PATCH("/rules/:id/channels/:channel_id/disable", ruleH.DisableChannel)
// Message Logs
admin.GET("/message-logs", msgLogH.List)
// Start server
srv := &http.Server{
Addr: fmt.Sprintf(":%d", cfg.Server.Port),
Handler: r,
}
go func() {
slog.Info("server starting", "port", cfg.Server.Port)
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
slog.Error("server error", "error", err)
os.Exit(1)
}
}()
// Graceful shutdown
quit := make(chan os.Signal, 1)
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
<-quit
slog.Info("shutting down...")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if err := srv.Shutdown(ctx); err != nil {
slog.Error("forced shutdown", "error", err)
}
slog.Info("server stopped")
}
// teeHandler fans out log records to multiple handlers (stdout + remote).
type teeHandler struct {
handlers []slog.Handler
}
func (t *teeHandler) Enabled(ctx context.Context, level slog.Level) bool {
for _, h := range t.handlers {
if h.Enabled(ctx, level) {
return true
}
}
return false
}
func (t *teeHandler) Handle(ctx context.Context, r slog.Record) error {
for _, h := range t.handlers {
if h.Enabled(ctx, r.Level) {
if err := h.Handle(ctx, r.Clone()); err != nil {
return err
}
}
}
return nil
}
func (t *teeHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
handlers := make([]slog.Handler, len(t.handlers))
for i, h := range t.handlers {
handlers[i] = h.WithAttrs(attrs)
}
return &teeHandler{handlers: handlers}
}
func (t *teeHandler) WithGroup(name string) slog.Handler {
handlers := make([]slog.Handler, len(t.handlers))
for i, h := range t.handlers {
handlers[i] = h.WithGroup(name)
}
return &teeHandler{handlers: handlers}
}