diff --git a/pkg/atproto/firehose.go b/pkg/atproto/firehose.go index c7d70546c..2c9c92d25 100644 --- a/pkg/atproto/firehose.go +++ b/pkg/atproto/firehose.go @@ -7,6 +7,7 @@ import ( "net/http" "net/url" "runtime" + "strconv" "strings" "sync/atomic" "time" @@ -156,6 +157,27 @@ func (atsync *ATProtoSynchronizer) StartFirehose(ctx context.Context) error { // is cancelled. func (atsync *ATProtoSynchronizer) consumeRelay(ctx context.Context, relay string) { ctx = log.WithLogValues(ctx, "relay", relay) + + cursor := atsync.newRelayCursor(ctx, relay) + // Persist progress on a timer and once more on the way out, so a restart + // resumes near where we left off rather than re-tailing from live. + flushDone := make(chan struct{}) + go func() { + defer close(flushDone) + ticker := time.NewTicker(cursorFlushInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + cursor.flush(ctx) + return + case <-ticker.C: + cursor.flush(ctx) + } + } + }() + defer func() { <-flushDone }() + const ( minBackoff = time.Second maxBackoff = 30 * time.Second @@ -166,7 +188,7 @@ func (atsync *ATProtoSynchronizer) consumeRelay(ctx context.Context, relay strin return } start := time.Now() - err := atsync.connectRelay(ctx, relay) + err := atsync.connectRelay(ctx, relay, cursor) if ctx.Err() != nil { return } @@ -198,12 +220,17 @@ func (atsync *ATProtoSynchronizer) consumeRelay(ctx context.Context, relay strin // the per-connection one) so an in-flight commit keeps indexing across a // reconnect — important because dedup has already claimed it, so no other relay // will re-deliver it to us. -func (atsync *ATProtoSynchronizer) connectRelay(ctx context.Context, relay string) error { +func (atsync *ATProtoSynchronizer) connectRelay(ctx context.Context, relay string, cursor *relayCursor) error { u, err := url.Parse(relay) if err != nil { return fmt.Errorf("invalid relay URI %q: %w", relay, err) } u.Path = "xrpc/com.atproto.sync.subscribeRepos" + if seq, ok := cursor.param(); ok { + q := u.Query() + q.Set("cursor", strconv.FormatInt(seq, 10)) + u.RawQuery = q.Encode() + } con, _, err := websocket.DefaultDialer.Dial(u.String(), http.Header{ "User-Agent": []string{aqhttp.UserAgent}, @@ -222,6 +249,7 @@ func (atsync *ATProtoSynchronizer) connectRelay(ctx context.Context, relay strin rsc := &events.RepoStreamCallbacks{ RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error { atsync.markSeen() + cursor.observe(evt.Seq) if atsync.commitDedup.seen(evt.Commit.String()) { spmetrics.FirehoseEventsDedupedTotal.WithLabelValues("commit").Inc() return nil @@ -231,6 +259,7 @@ func (atsync *ATProtoSynchronizer) connectRelay(ctx context.Context, relay strin }, RepoIdentity: func(evt *comatproto.SyncSubscribeRepos_Identity) error { atsync.markSeen() + cursor.observe(evt.Seq) if atsync.identityDedup.seen(identityDedupKey(evt)) { spmetrics.FirehoseEventsDedupedTotal.WithLabelValues("identity").Inc() return nil diff --git a/pkg/atproto/firehose_cursor.go b/pkg/atproto/firehose_cursor.go new file mode 100644 index 000000000..a79a4d694 --- /dev/null +++ b/pkg/atproto/firehose_cursor.go @@ -0,0 +1,85 @@ +package atproto + +import ( + "context" + "sync/atomic" + "time" + + "stream.place/streamplace/pkg/log" + "stream.place/streamplace/pkg/model" +) + +// cursorFlushInterval bounds how often a relay's progress is written to the +// index DB. The firehose is high-volume, so we persist on a timer rather than +// per-event; between flushes the latest seq lives in memory and already covers +// in-process reconnects. +const cursorFlushInterval = 5 * time.Second + +// relayCursor tracks how far we've consumed one relay's firehose so we can +// resume after a disconnect or restart instead of re-tailing from live (which +// would leave a gap). It keeps the high-water sequence number in memory, +// updated on every frame, and persists it periodically and once on shutdown. +// +// Because the parallel scheduler may surface frames slightly out of sequence +// order, the persisted cursor is the highest seq observed — on an unclean crash +// a handful of in-flight frames just below it can be skipped on resume. That is +// safe here: downstream handlers are idempotent, and with several relays plus a +// cold deduper after restart, those commits get re-delivered and re-indexed. +type relayCursor struct { + host string + model model.Model + + latest atomic.Int64 // highest seq seen; 0 = nothing yet (tail from live) + flushed int64 // last persisted value; only the flush loop touches it +} + +func (atsync *ATProtoSynchronizer) newRelayCursor(ctx context.Context, host string) *relayCursor { + rc := &relayCursor{host: host, model: atsync.Model} + stored, err := atsync.Model.GetRelayCursor(host) + if err != nil { + log.Error(ctx, "failed to load relay cursor; tailing from live", "err", err) + return rc + } + if stored != nil { + rc.latest.Store(stored.Cursor) + rc.flushed = stored.Cursor + log.Log(ctx, "resuming relay from stored cursor", "cursor", stored.Cursor) + } + return rc +} + +// observe advances the high-water mark. Safe for concurrent callers (the +// scheduler runs several event workers). +func (rc *relayCursor) observe(seq int64) { + for { + cur := rc.latest.Load() + if seq <= cur { + return + } + if rc.latest.CompareAndSwap(cur, seq) { + return + } + } +} + +// param returns the cursor to dial with and whether to send one at all. With no +// progress yet we send none, so a fresh external relay tails from live instead +// of backfilling its entire history. +func (rc *relayCursor) param() (int64, bool) { + v := rc.latest.Load() + return v, v > 0 +} + +// flush persists the high-water mark if it has advanced since the last write. +// Only ever called from the single flush goroutine, so flushed is unsynchronized. +func (rc *relayCursor) flush(ctx context.Context) { + v := rc.latest.Load() + if v == rc.flushed { + return + } + if err := rc.model.UpsertRelayCursor(rc.host, v); err != nil { + log.Error(ctx, "failed to persist relay cursor", "err", err, "cursor", v) + return + } + rc.flushed = v +} diff --git a/pkg/atproto/firehose_cursor_test.go b/pkg/atproto/firehose_cursor_test.go new file mode 100644 index 000000000..94987fce3 --- /dev/null +++ b/pkg/atproto/firehose_cursor_test.go @@ -0,0 +1,59 @@ +package atproto + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "stream.place/streamplace/pkg/model" +) + +func TestRelayCursorResume(t *testing.T) { + mod, err := model.MakeDB(":memory:") + require.NoError(t, err) + atsync := &ATProtoSynchronizer{Model: mod} + ctx := context.Background() + const host = "wss://relay.example" + + // A fresh relay has no stored cursor, so we dial with none and tail live + // rather than backfilling the relay's whole history. + rc := atsync.newRelayCursor(ctx, host) + if _, ok := rc.param(); ok { + t.Fatal("fresh relay should not send a cursor") + } + + // observe tracks the high-water mark and never regresses on out-of-order + // frames (the parallel scheduler can surface them out of sequence). + rc.observe(100) + rc.observe(50) + rc.observe(120) + if v, ok := rc.param(); !ok || v != 120 { + t.Fatalf("expected cursor 120, got %d (ok=%v)", v, ok) + } + + // flush persists the high-water mark to the index DB. + rc.flush(ctx) + stored, err := mod.GetRelayCursor(host) + require.NoError(t, err) + require.NotNil(t, stored) + require.Equal(t, int64(120), stored.Cursor) + + // A second flush with no advance is a no-op, and the stored value is stable. + rc.observe(120) + rc.flush(ctx) + stored, err = mod.GetRelayCursor(host) + require.NoError(t, err) + require.Equal(t, int64(120), stored.Cursor) + + // A new cursor (as if the process restarted) resumes from the stored value. + resumed := atsync.newRelayCursor(ctx, host) + v, ok := resumed.param() + require.True(t, ok) + require.Equal(t, int64(120), v) + + // Cursors are independent per relay. + other := atsync.newRelayCursor(ctx, "wss://other.example") + if _, ok := other.param(); ok { + t.Fatal("a different relay must not inherit another relay's cursor") + } +} diff --git a/pkg/model/model.go b/pkg/model/model.go index 0ae5e4e07..058b7b5ce 100644 --- a/pkg/model/model.go +++ b/pkg/model/model.go @@ -102,6 +102,9 @@ type Model interface { GetLabeler(did string) (*Labeler, error) UpdateLabelerCursor(did string, cursor int64) error + GetRelayCursor(host string) (*RelayCursor, error) + UpsertRelayCursor(host string, cursor int64) error + CreateLabel(label *Label) error GetActiveLabels(uri string) ([]*comatproto.LabelDefs_Label, error) @@ -243,6 +246,7 @@ func MakeDB(dbURL string) (Model, error) { PinnedRecord{}, ServerSettings{}, Labeler{}, + RelayCursor{}, Label{}, BroadcastOrigin{}, MetadataConfiguration{}, diff --git a/pkg/model/relay_cursor.go b/pkg/model/relay_cursor.go new file mode 100644 index 000000000..aaa1db1c8 --- /dev/null +++ b/pkg/model/relay_cursor.go @@ -0,0 +1,40 @@ +package model + +import ( + "errors" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// RelayCursor remembers how far we have consumed each relay's firehose, keyed by +// the relay's websocket URL. On reconnect or restart we resume from the stored +// sequence number instead of re-tailing from live (which would leave a gap) or +// replaying from the beginning. Cursors are per-relay because each relay +// assigns its own sequence numbers. +type RelayCursor struct { + Host string `gorm:"primaryKey;column:host"` + Cursor int64 `gorm:"column:cursor"` +} + +// GetRelayCursor returns the stored cursor for a relay, or nil if we have never +// recorded one (i.e. this is a fresh subscription). +func (m *DBModel) GetRelayCursor(host string) (*RelayCursor, error) { + var rc RelayCursor + err := m.DB.Where("host = ?", host).First(&rc).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &rc, nil +} + +// UpsertRelayCursor stores the latest consumed sequence number for a relay. +func (m *DBModel) UpsertRelayCursor(host string, cursor int64) error { + return m.DB.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "host"}}, + DoUpdates: clause.AssignmentColumns([]string{"cursor"}), + }).Create(&RelayCursor{Host: host, Cursor: cursor}).Error +}