migrate from github
This commit is contained in:
@@ -0,0 +1,46 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/XingfenD/yoresee_doc/internal/bootstrap"
|
||||
"github.com/XingfenD/yoresee_doc/internal/service/mq_service"
|
||||
"github.com/XingfenD/yoresee_doc/internal/utils"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/constant"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/key"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/mq"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func main() {
|
||||
initializer := bootstrap.NewInitializer().
|
||||
InitConfig().
|
||||
InitPostgres().
|
||||
InitRedis().
|
||||
InitElasticsearchAllowFail().
|
||||
InitMQ().
|
||||
InitRepository().
|
||||
InitService()
|
||||
if err := initializer.Err(); err != nil {
|
||||
logrus.Fatalf("Init snapshot-worker failed: %v", err)
|
||||
}
|
||||
|
||||
w := &Worker{
|
||||
collabCoreHTTP: utils.GetEnvVar("COLLAB_CORE_HTTP", "http://collab-core:1234"),
|
||||
mqGroup: utils.GetEnvVar("SNAPSHOT_MQ_GROUP", "snapshot-worker"),
|
||||
topic: constant.DirtyDocTopicDefault,
|
||||
mqBackend: mq.BackendRabbitMQ,
|
||||
client: &http.Client{Timeout: 10 * time.Second},
|
||||
dirtySetKey: key.KeyCollabDirtyDocSet(),
|
||||
inFlight: &sync.Map{},
|
||||
}
|
||||
|
||||
logrus.Infof("Snapshot worker started: topic=%s group=%s collabCore=%s", w.topic, w.mqGroup, w.collabCoreHTTP)
|
||||
|
||||
go w.runMQConsumer(mq_service.MQSvc)
|
||||
go w.runScanLoop()
|
||||
|
||||
initializer.ShutdownOnSignal(0)
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/bytedance/sonic"
|
||||
)
|
||||
|
||||
type dirtyDocMessage struct {
|
||||
DocID string `json:"doc_id"`
|
||||
DocId string `json:"docId"`
|
||||
}
|
||||
|
||||
func parseDocID(data []byte) string {
|
||||
payload := strings.TrimSpace(string(data))
|
||||
if payload == "" {
|
||||
return ""
|
||||
}
|
||||
var msg dirtyDocMessage
|
||||
if err := sonic.Unmarshal(data, &msg); err == nil {
|
||||
if msg.DocID != "" {
|
||||
return msg.DocID
|
||||
}
|
||||
if msg.DocId != "" {
|
||||
return msg.DocId
|
||||
}
|
||||
}
|
||||
return payload
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/XingfenD/yoresee_doc/internal/domain_event"
|
||||
"github.com/XingfenD/yoresee_doc/internal/service/document_service"
|
||||
"github.com/XingfenD/yoresee_doc/internal/status"
|
||||
"github.com/XingfenD/yoresee_doc/internal/utils"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/key"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/storage"
|
||||
"github.com/bytedance/sonic"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func (w *Worker) snapshotDoc(ctx context.Context, docID string, force bool) error {
|
||||
if _, loaded := w.inFlight.LoadOrStore(docID, struct{}{}); loaded {
|
||||
return nil
|
||||
}
|
||||
defer w.inFlight.Delete(docID)
|
||||
|
||||
logrus.Infof("Snapshot start docId=%s force=%v", docID, force)
|
||||
|
||||
docType, err := document_service.DocumentSvc.GetDocumentTypeByExternalID(ctx, docID)
|
||||
if err != nil {
|
||||
if errors.Is(err, status.StatusDocumentNotFound) {
|
||||
logrus.Infof("Snapshot doc not found docId=%s, cleaning up Redis", docID)
|
||||
cleanupCollabKeys(ctx, docID)
|
||||
return nil
|
||||
}
|
||||
logrus.Errorf("Snapshot get doc type failed docId=%s err=%v", docID, err)
|
||||
return err
|
||||
}
|
||||
|
||||
state, content, err := w.fetchSnapshot(ctx, docID, string(docType))
|
||||
if err != nil {
|
||||
logrus.Errorf("Snapshot fetch failed docId=%s err=%v", docID, err)
|
||||
return err
|
||||
}
|
||||
if len(state) == 0 {
|
||||
logrus.Infof("Snapshot empty docId=%s", docID)
|
||||
return nil
|
||||
}
|
||||
logrus.Infof("Snapshot fetched docId=%s bytes=%d", docID, len(state))
|
||||
|
||||
contentChanged, err := document_service.DocumentSvc.SaveDocumentSnapshotAndContent(ctx, docID, state, content)
|
||||
if err != nil {
|
||||
logrus.Errorf("Snapshot save failed docId=%s err=%v", docID, err)
|
||||
return err
|
||||
}
|
||||
if contentChanged {
|
||||
if pubErr := domain_event.PublishDocumentUpsertEvent(ctx, docID); pubErr != nil {
|
||||
logrus.Warnf("Snapshot publish search sync event failed docId=%s err=%v", docID, pubErr)
|
||||
}
|
||||
}
|
||||
|
||||
if err := flushCollabKeys(ctx, docID, state); err != nil {
|
||||
logrus.Errorf("Snapshot redis flush failed docId=%s err=%v", docID, err)
|
||||
return err
|
||||
}
|
||||
|
||||
if force {
|
||||
logrus.Infof("Snapshot saved (mq) for %s", docID)
|
||||
} else {
|
||||
logrus.Infof("Snapshot saved (scan) for %s", docID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Worker) fetchSnapshot(ctx context.Context, docID, docType string) (state []byte, content string, err error) {
|
||||
url := fmt.Sprintf("%s/internal/yjs/doc-snapshot/%s?type=%s", strings.TrimRight(w.collabCoreHTTP, "/"), docID, docType)
|
||||
resp, err := w.client.Get(url)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode == http.StatusNotFound {
|
||||
return nil, "", nil
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, "", fmt.Errorf("unexpected status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
var payload struct {
|
||||
State string `json:"state"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
if err := sonic.Unmarshal(body, &payload); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if payload.State == "" {
|
||||
return nil, payload.Content, nil
|
||||
}
|
||||
|
||||
state, err = utils.DecodeBase64(payload.State)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return state, payload.Content, nil
|
||||
}
|
||||
|
||||
func flushCollabKeys(ctx context.Context, docID string, state []byte) error {
|
||||
rdb := storage.GetRedis()
|
||||
if rdb == nil {
|
||||
return nil
|
||||
}
|
||||
updatesKey := key.KeyCollabDocUpdates(docID)
|
||||
pipe := rdb.TxPipeline()
|
||||
pipe.Del(ctx, updatesKey)
|
||||
pipe.RPush(ctx, updatesKey, state)
|
||||
pipe.SRem(ctx, key.KeyCollabDirtyDocSet(), docID)
|
||||
_, err := pipe.Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func cleanupCollabKeys(ctx context.Context, docID string) {
|
||||
rdb := storage.GetRedis()
|
||||
if rdb == nil {
|
||||
return
|
||||
}
|
||||
pipe := rdb.TxPipeline()
|
||||
pipe.SRem(ctx, key.KeyCollabDirtyDocSet(), docID)
|
||||
pipe.Del(ctx, key.KeyCollabDocUpdates(docID))
|
||||
pipe.Del(ctx, key.KeyCollabRoom(docID))
|
||||
if _, err := pipe.Exec(ctx); err != nil {
|
||||
logrus.Warnf("Snapshot cleanup Redis failed docId=%s err=%v", docID, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/XingfenD/yoresee_doc/internal/service/mq_service"
|
||||
"github.com/XingfenD/yoresee_doc/internal/utils"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/key"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/mq"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/storage"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type Worker struct {
|
||||
collabCoreHTTP string
|
||||
mqGroup string
|
||||
topic string
|
||||
mqBackend mq.Backend
|
||||
client *http.Client
|
||||
dirtySetKey string
|
||||
inFlight *sync.Map
|
||||
}
|
||||
|
||||
func (w *Worker) runMQConsumer(mqSvc *mq_service.MQService) {
|
||||
err := mqSvc.Consume(
|
||||
context.Background(),
|
||||
w.mqBackend,
|
||||
mq.ConsumeOptions{
|
||||
Topic: w.topic,
|
||||
Mode: mq.ConsumeModeGroup,
|
||||
Group: w.mqGroup,
|
||||
AutoAck: false,
|
||||
OnError: mq.ErrorActionRequeue,
|
||||
},
|
||||
func(ctx context.Context, message mq.Message) error {
|
||||
docID := parseDocID(message.Body)
|
||||
if docID == "" {
|
||||
return nil
|
||||
}
|
||||
return w.snapshotDoc(ctx, docID, true)
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
logrus.Fatalf("MQ consumer failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Worker) runScanLoop() {
|
||||
ticker := time.NewTicker(2 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
w.scanAndSnapshot()
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Worker) scanAndSnapshot() {
|
||||
ctx := context.Background()
|
||||
rdb := storage.GetRedis()
|
||||
if rdb == nil {
|
||||
return
|
||||
}
|
||||
|
||||
docIDs, err := rdb.SMembers(ctx, w.dirtySetKey).Result()
|
||||
if err != nil {
|
||||
logrus.Errorf("Scan dirty docs failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
candidates := filterCandidates(ctx, docIDs)
|
||||
logrus.Infof("Scan dirty docs: candidates=%d", len(candidates))
|
||||
if len(candidates) > 0 {
|
||||
logrus.Infof("Dirty doc candidates: %v", candidates)
|
||||
}
|
||||
for _, docID := range candidates {
|
||||
_ = w.snapshotDoc(ctx, docID, false)
|
||||
}
|
||||
}
|
||||
|
||||
func filterCandidates(ctx context.Context, docIDs []string) []string {
|
||||
now := time.Now().UnixMilli()
|
||||
candidates := make([]string, 0, len(docIDs))
|
||||
for _, docID := range docIDs {
|
||||
if docID == "" {
|
||||
continue
|
||||
}
|
||||
lastStr, err := storage.GetRedis().Get(ctx, key.KeyCollabRoom(docID)).Result()
|
||||
if err != nil {
|
||||
logrus.Infof("Dirty doc %s skipped: room key missing", docID)
|
||||
continue
|
||||
}
|
||||
last, err := utils.ParseInt64(lastStr)
|
||||
if err != nil {
|
||||
logrus.Infof("Dirty doc %s skipped: invalid room timestamp %s", docID, lastStr)
|
||||
continue
|
||||
}
|
||||
if now-last < 10_000 {
|
||||
logrus.Infof("Dirty doc %s skipped: lastEditAgo=%dms", docID, now-last)
|
||||
continue
|
||||
}
|
||||
candidates = append(candidates, docID)
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
Reference in New Issue
Block a user