From 2000908baca361b41bad7632dcf124b261dde793 Mon Sep 17 00:00:00 2001
From: ryan <kuyoeku@gmail.com>
Date: Sat, 15 Aug 2026 01:28:51 +0800
Subject: [PATCH] feat: enhance Safew watcher initialization and token
 management

- Added functionality to list and start Safew tokens in the main server function, ensuring proper initialization of the Safew watcher.
- Introduced `ensureSafewWatcher` method in the ChannelHandler to manage Safew tokens during channel creation and updates.
- Implemented `StartTokens` method in the Watcher to handle multiple tokens efficiently.
- Enhanced error handling and logging for Safew watcher operations to improve observability.
---
 cmd/server/main.go          |  6 ++++++
 internal/handler/channel.go | 15 +++++++++++++++
 internal/safew/watcher.go   | 11 ++++++++++-
 internal/store/channel.go   | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/cmd/server/main.go b/cmd/server/main.go
index ec72faa..503e0f8 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -139,6 +139,12 @@ func main() {
 	}
 	safewWatcher := safew.NewWatcher(watcherStore, poller)
 	defer safewWatcher.Stop()
+	if tokens, err := st.ListSafewTokens(context.Background()); err != nil {
+		slog.Warn("list safew tokens", "error", err)
+	} else {
+		safewWatcher.StartTokens(tokens)
+		slog.Info("safew watchers ensured", "count", len(tokens))
+	}
 
 	channelH := handler.NewChannelHandler(st, redisCache, safewWatcher)
 	ruleH := handler.NewRuleHandler(st, redisCache)
diff --git a/internal/handler/channel.go b/internal/handler/channel.go
index 230f488..f77d1c8 100644
--- a/internal/handler/channel.go
+++ b/internal/handler/channel.go
@@ -49,6 +49,7 @@ func (h *ChannelHandler) Create(c *gin.Context) {
 		c.JSON(http.StatusConflict, gin.H{"error": err.Error()})
 		return
 	}
+	h.ensureSafewWatcher(req.Type, req.Config)
 	c.JSON(http.StatusCreated, ch)
 }
 
@@ -91,6 +92,7 @@ func (h *ChannelHandler) Update(c *gin.Context) {
 		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
 		return
 	}
+	h.ensureSafewWatcher(req.Type, req.Config)
 	c.JSON(http.StatusOK, gin.H{"ok": true})
 }
 
@@ -160,6 +162,19 @@ func (h *ChannelHandler) respondSafewChats(c *gin.Context, token, q string) {
 	c.JSON(http.StatusOK, gin.H{"data": list, "total": len(list)})
 }
 
+func (h *ChannelHandler) ensureSafewWatcher(typ string, config json.RawMessage) {
+	if h.chats == nil || typ != "safew" {
+		return
+	}
+	raw := config
+	ch := &model.Channel{Type: typ, Config: &raw}
+	tok, err := safewTokenFromChannel(ch)
+	if err != nil {
+		return
+	}
+	h.chats.Ensure(tok)
+}
+
 func safewTokenFromChannel(ch *model.Channel) (string, error) {
 	if ch.Type != "safew" {
 		return "", fmt.Errorf("channel is not safew")
diff --git a/internal/safew/watcher.go b/internal/safew/watcher.go
index 60087cc..6672c07 100644
--- a/internal/safew/watcher.go
+++ b/internal/safew/watcher.go
@@ -67,6 +67,12 @@ func (w *Watcher) Ensure(token string) {
 	w.ensure(token)
 }
 
+func (w *Watcher) StartTokens(tokens []string) {
+	for _, token := range tokens {
+		w.Ensure(strings.TrimSpace(token))
+	}
+}
+
 func (w *Watcher) ensure(token string) *runner {
 	if token == "" || w == nil {
 		return nil
@@ -82,6 +88,7 @@ func (w *Watcher) ensure(token string) *runner {
 	ctx, cancel := context.WithCancel(context.Background())
 	r := &runner{cancel: cancel, firstDone: make(chan struct{})}
 	w.running[token] = r
+	slog.Info("safew watcher started")
 	go w.loop(ctx, token, r)
 	return r
 }
@@ -144,7 +151,7 @@ func (w *Watcher) loop(ctx context.Context, token string, r *runner) {
 				return
 			}
 			if isGetUpdatesConflict(err) {
-				slog.Debug("safew watcher poll conflict", "error", err)
+				slog.Info("safew watcher poll conflict", "error", err)
 				if !w.sleep(ctx, conflictBackoff(w.bgIdle)) {
 					return
 				}
@@ -176,6 +183,7 @@ func (w *Watcher) pollOnce(ctx context.Context, token string, timeout int) error
 		return err
 	}
 	if !ok {
+		slog.Info("safew poll skipped", "reason", "lock held", "timeout", timeout)
 		return nil
 	}
 	defer func() { _ = w.store.UnlockSafewPoll(ctx, token) }()
@@ -188,6 +196,7 @@ func (w *Watcher) pollOnce(ctx context.Context, token string, timeout int) error
 	if err != nil {
 		return err
 	}
+	slog.Info("safew poll", "timeout", timeout, "groups", len(chats), "offset", next)
 	if err := w.mergeAndLog(ctx, token, chats); err != nil {
 		return err
 	}
diff --git a/internal/store/channel.go b/internal/store/channel.go
index fe9eaab..59ded69 100644
--- a/internal/store/channel.go
+++ b/internal/store/channel.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	"aiaa-notification-service/internal/model"
 )
@@ -94,3 +95,38 @@ func (s *Store) DeleteChannel(ctx context.Context, id int) error {
 	}
 	return nil
 }
+
+func (s *Store) ListSafewTokens(ctx context.Context) ([]string, error) {
+	rows, err := s.DB.QueryContext(ctx, `SELECT config FROM notification_channel WHERE type = 'safew'`)
+	if err != nil {
+		return nil, fmt.Errorf("list safew tokens: %w", err)
+	}
+	defer rows.Close()
+	seen := map[string]struct{}{}
+	var tokens []string
+	for rows.Next() {
+		var configBytes []byte
+		if err := rows.Scan(&configBytes); err != nil {
+			return nil, fmt.Errorf("scan safew config: %w", err)
+		}
+		var cfg struct {
+			Token string `json:"token"`
+		}
+		if err := json.Unmarshal(configBytes, &cfg); err != nil {
+			continue
+		}
+		tok := strings.TrimSpace(cfg.Token)
+		if tok == "" {
+			continue
+		}
+		if _, ok := seen[tok]; ok {
+			continue
+		}
+		seen[tok] = struct{}{}
+		tokens = append(tokens, tok)
+	}
+	if tokens == nil {
+		tokens = []string{}
+	}
+	return tokens, rows.Err()
+}