From 675137e2cd6f5f174b4ac068a2cf6e3f258bbdc8 Mon Sep 17 00:00:00 2001 From: hailey Date: Wed, 04 Feb 2026 23:42:19 +0000 Subject: [PATCH] fix cursors for new relay (#64) --- models/models.go | 8 ++++++++ server/handle_sync_subscribe_repos.go | 14 +++++++++++++- server/persist.go | 243 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ server/server.go | 7 ++++++- 4 file(s) changed, 270 insertion(s)(+), 2 deletion(s)(-) diff --git a/models/models.go b/models/models.go --- a/models/models.go +++ b/models/models.go @@ -136,3 +136,11 @@ Did *string `gorm:"index"` PrivateKey []byte CreatedAt time.Time `gorm:"index"` } + +type EventRecord struct { + Seq int64 `gorm:"primaryKey;autoIncrement:false"` + CreatedAt time.Time + Did string `gorm:"index"` + Type string + Data []byte +} diff --git a/server/handle_sync_subscribe_repos.go b/server/handle_sync_subscribe_repos.go --- a/server/handle_sync_subscribe_repos.go +++ b/server/handle_sync_subscribe_repos.go @@ -2,6 +2,7 @@ package server import ( "context" + "strconv" "time" "github.com/bluesky-social/indigo/events" @@ -27,6 +28,17 @@ ident := e.RealIP() + "-" + e.Request().UserAgent() logger = logger.With("ident", ident) logger.Info("new connection established") + var since *int64 + if cursorStr := e.QueryParam("cursor"); cursorStr != "" { + cursor, err := strconv.ParseInt(cursorStr, 10, 64) + if err != nil { + logger.Warn("invalid cursor parameter", "cursor", cursorStr, "err", err) + } else { + since = &cursor + logger.Info("subscribing with cursor", "cursor", cursor) + } + } + metrics.RelaysConnected.WithLabelValues(ident).Inc() defer func() { metrics.RelaysConnected.WithLabelValues(ident).Dec() @@ -34,7 +46,7 @@ }() evts, evtManCancel, err := s.evtman.Subscribe(ctx, ident, func(evt *events.XRPCStreamEvent) bool { return true - }, nil) + }, since) if err != nil { return err } diff --git a/server/persist.go b/server/persist.go new file mode 100644 --- /dev/null +++ b/server/persist.go @@ -0,0 +1,243 @@ +package server + +import ( + "bytes" + "context" + "fmt" + "sync" + "time" + + "github.com/bluesky-social/indigo/api/atproto" + "github.com/bluesky-social/indigo/events" + indigomodels "github.com/bluesky-social/indigo/models" + cbg "github.com/whyrusleeping/cbor-gen" + "gorm.io/gorm" + + "github.com/haileyok/cocoon/models" +) + +type DbPersister struct { + Db *gorm.DB + + Lk sync.Mutex + Seq int64 + + Broadcast func(*events.XRPCStreamEvent) + + // how long do we actually want to keep these things around + Retention time.Duration +} + +func NewDbPersister(db *gorm.DB, retention time.Duration) (*DbPersister, error) { + if err := db.AutoMigrate(&models.EventRecord{}); err != nil { + return nil, fmt.Errorf("failed to migrate EventRecord: %w", err) + } + + if retention == 0 { + retention = 72 * time.Hour + } + + p := &DbPersister{ + Db: db, + Retention: retention, + } + + // kind of hacky. we will try and get the latest one from the db, but if it doesn't exist...well we have a problem + // because the relay will already have _some_ value > 0 set as a cursor, we'll want to just set this to some high value + // we'll just grab a current unix timestamp and set that as the cursor + var lastEvent models.EventRecord + if err := db.Order("seq desc").Limit(1).First(&lastEvent).Error; err != nil { + if err != gorm.ErrRecordNotFound { + return nil, fmt.Errorf("failed to get last event seq: %w", err) + } + p.Seq = time.Now().Unix() + } else { + p.Seq = lastEvent.Seq + } + + go p.cleanupRoutine() + + return p, nil +} + +func (p *DbPersister) SetEventBroadcaster(brc func(*events.XRPCStreamEvent)) { + p.Broadcast = brc +} + +func (p *DbPersister) Persist(ctx context.Context, e *events.XRPCStreamEvent) error { + p.Lk.Lock() + defer p.Lk.Unlock() + + p.Seq++ + seq := p.Seq + + var did string + var evtType string + + switch { + case e.RepoCommit != nil: + e.RepoCommit.Seq = seq + did = e.RepoCommit.Repo + evtType = "commit" + case e.RepoSync != nil: + e.RepoSync.Seq = seq + did = e.RepoSync.Did + evtType = "sync" + case e.RepoIdentity != nil: + e.RepoIdentity.Seq = seq + did = e.RepoIdentity.Did + evtType = "identity" + case e.RepoAccount != nil: + e.RepoAccount.Seq = seq + did = e.RepoAccount.Did + evtType = "account" + default: + return fmt.Errorf("unknown event type") + } + + data, err := serializeEvent(e) + if err != nil { + return fmt.Errorf("failed to serialize event: %w", err) + } + + rec := &models.EventRecord{ + Seq: seq, + CreatedAt: time.Now(), + Did: did, + Type: evtType, + Data: data, + } + + if err := p.Db.Create(rec).Error; err != nil { + return fmt.Errorf("failed to persist event: %w", err) + } + + if p.Broadcast != nil { + p.Broadcast(e) + } + + return nil +} + +func (p *DbPersister) Playback(ctx context.Context, since int64, cb func(*events.XRPCStreamEvent) error) error { + const pageSize = 500 + + cursor := since + for { + var records []models.EventRecord + if err := p.Db.WithContext(ctx). + Where("seq > ?", cursor). + Order("seq asc"). + Limit(pageSize). + Find(&records).Error; err != nil { + return fmt.Errorf("failed to query events: %w", err) + } + + if len(records) == 0 { + return nil + } + + for _, rec := range records { + evt, err := deserializeEvent(rec.Type, rec.Data) + if err != nil { + return fmt.Errorf("failed to deserialize event %d: %w", rec.Seq, err) + } + + if err := cb(evt); err != nil { + return err + } + + cursor = rec.Seq + } + + if len(records) < pageSize { + return nil + } + } +} + +func (p *DbPersister) TakeDownRepo(ctx context.Context, uid indigomodels.Uid) error { + return nil +} + +func (p *DbPersister) Flush(ctx context.Context) error { + return nil +} + +func (p *DbPersister) Shutdown(ctx context.Context) error { + return nil +} + +func (p *DbPersister) cleanupRoutine() { + ticker := time.NewTicker(time.Hour) + defer ticker.Stop() + + for range ticker.C { + cutoff := time.Now().Add(-p.Retention) + if err := p.Db.Where("created_at < ?", cutoff).Delete(&models.EventRecord{}).Error; err != nil { + continue + } + } +} + +func serializeEvent(e *events.XRPCStreamEvent) ([]byte, error) { + buf := new(bytes.Buffer) + cw := cbg.NewCborWriter(buf) + + switch { + case e.RepoCommit != nil: + if err := e.RepoCommit.MarshalCBOR(cw); err != nil { + return nil, err + } + case e.RepoSync != nil: + if err := e.RepoSync.MarshalCBOR(cw); err != nil { + return nil, err + } + case e.RepoIdentity != nil: + if err := e.RepoIdentity.MarshalCBOR(cw); err != nil { + return nil, err + } + case e.RepoAccount != nil: + if err := e.RepoAccount.MarshalCBOR(cw); err != nil { + return nil, err + } + default: + return nil, fmt.Errorf("unknown event type") + } + + return buf.Bytes(), nil +} + +func deserializeEvent(evtType string, data []byte) (*events.XRPCStreamEvent, error) { + r := bytes.NewReader(data) + cr := cbg.NewCborReader(r) + + switch evtType { + case "commit": + evt := &atproto.SyncSubscribeRepos_Commit{} + if err := evt.UnmarshalCBOR(cr); err != nil { + return nil, err + } + return &events.XRPCStreamEvent{RepoCommit: evt}, nil + case "sync": + evt := &atproto.SyncSubscribeRepos_Sync{} + if err := evt.UnmarshalCBOR(cr); err != nil { + return nil, err + } + return &events.XRPCStreamEvent{RepoSync: evt}, nil + case "identity": + evt := &atproto.SyncSubscribeRepos_Identity{} + if err := evt.UnmarshalCBOR(cr); err != nil { + return nil, err + } + return &events.XRPCStreamEvent{RepoIdentity: evt}, nil + case "account": + evt := &atproto.SyncSubscribeRepos_Account{} + if err := evt.UnmarshalCBOR(cr); err != nil { + return nil, err + } + return &events.XRPCStreamEvent{RepoAccount: evt}, nil + default: + return nil, fmt.Errorf("unknown event type: %s", evtType) + } +} diff --git a/server/server.go b/server/server.go --- a/server/server.go +++ b/server/server.go @@ -405,6 +405,11 @@ } else { nonceSecret = maybeSecret } + evtPersister, err := NewDbPersister(gdb, 72*time.Hour) + if err != nil { + return nil, fmt.Errorf("failed to create event persister: %w", err) + } + s := &Server{ http: h, httpd: httpd, @@ -429,7 +434,7 @@ SessionCookieKey: args.SessionCookieKey, BlockstoreVariant: args.BlockstoreVariant, FallbackProxy: args.FallbackProxy, }, - evtman: events.NewEventManager(events.NewMemPersister()), + evtman: events.NewEventManager(evtPersister), passport: identity.NewPassport(h, identity.NewMemCache(10_000)), dbName: args.DbName, -- tangled.sh