diff --git a/spindle/embedtap.go b/spindle/embedtap.go index 1415ed79..e315c0a2 100644 --- a/spindle/embedtap.go +++ b/spindle/embedtap.go @@ -10,6 +10,7 @@ import ( "net" "net/http" "strings" + "sync/atomic" "time" "github.com/bluesky-social/indigo/service/tap" @@ -46,6 +47,7 @@ func assertLoopbackBind(bind string) error { type embeddedTap struct { tap *tap.Tap logger *slog.Logger + closed atomic.Bool } func startEmbeddedTap(ctx context.Context, cfg *config.Config, logger *slog.Logger) (*embeddedTap, error) { @@ -116,6 +118,9 @@ func (e *embeddedTap) Shutdown() { if e == nil || e.tap == nil { return } + if e.closed.Swap(true) { + return + } shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() if err := e.tap.Server.Shutdown(shutdownCtx); err != nil { diff --git a/spindle/ingester.go b/spindle/ingester.go index 2137b85d..6544a9b2 100644 --- a/spindle/ingester.go +++ b/spindle/ingester.go @@ -8,6 +8,7 @@ import ( "tangled.org/core/api/tangled" "tangled.org/core/spindle/db" + "tangled.org/core/tapc" "github.com/bluesky-social/indigo/atproto/syntax" "github.com/bluesky-social/jetstream/pkg/models" @@ -25,6 +26,10 @@ func (s *Spindle) ingest() Ingester { switch e.Commit.Collection { case tangled.SpindleMemberNSID: err = s.ingestMember(ctx, e) + case tangled.RepoNSID, tangled.RepoCollaboratorNSID: + if evt, ok := jetstreamToTapEvent(e); ok { + err = s.tap.processEvent(ctx, evt) + } } if err != nil { @@ -40,6 +45,37 @@ func (s *Spindle) ingest() Ingester { } } +func jetstreamToTapEvent(e *models.Event) (tapc.Event, bool) { + if e.Commit == nil { + return tapc.Event{}, false + } + did, err := syntax.ParseDID(e.Did) + if err != nil { + return tapc.Event{}, false + } + var action tapc.RecordAction + switch e.Commit.Operation { + case models.CommitOperationCreate: + action = tapc.RecordCreateAction + case models.CommitOperationUpdate: + action = tapc.RecordUpdateAction + case models.CommitOperationDelete: + action = tapc.RecordDeleteAction + default: + return tapc.Event{}, false + } + return tapc.Event{ + Type: tapc.EvtRecord, + Record: &tapc.RecordEventData{ + Did: did, + Rkey: syntax.RecordKey(e.Commit.RKey), + Collection: syntax.NSID(e.Commit.Collection), + Action: action, + Record: e.Commit.Record, + }, + }, true +} + func (s *Spindle) ingestMember(_ context.Context, e *models.Event) error { var err error did := e.Did diff --git a/spindle/server.go b/spindle/server.go index 487bb66b..b0eb9084 100644 --- a/spindle/server.go +++ b/spindle/server.go @@ -112,6 +112,8 @@ func New(ctx context.Context, cfg *config.Config, engines map[string]models.Engi collections := []string{ tangled.SpindleMemberNSID, + tangled.RepoNSID, + tangled.RepoCollaboratorNSID, } jc, err := jetstream.NewJetstreamClient(cfg.Server.JetstreamEndpoint, "spindle", collections, nil, log.SubLogger(logger, "jetstream"), d, true, true) if err != nil { @@ -128,6 +130,16 @@ func New(ctx context.Context, cfg *config.Config, engines map[string]models.Engi jc.AddDid(d) } + knownRepos, err := d.AllRepos() + if err != nil { + return nil, fmt.Errorf("failed to get known repos: %w", err) + } + for _, r := range knownRepos { + if r.Owner != "" { + jc.AddDid(r.Owner.String()) + } + } + resolver := idresolver.DefaultResolver(cfg.Server.PlcUrl) spindle := &Spindle{ @@ -252,13 +264,23 @@ func (s *Spindle) Start(ctx context.Context) error { defer stopper.Stop() } + tapCtx, tapCancel := context.WithCancel(ctx) + if s.cfg.Server.Tap.Embed { - emb, err := startEmbeddedTap(ctx, s.cfg, log.SubLogger(s.l, "embedtap")) + emb, err := startEmbeddedTap(tapCtx, s.cfg, log.SubLogger(s.l, "embedtap")) if err != nil { + tapCancel() return fmt.Errorf("starting embedded tap: %w", err) } s.embedTap = emb - defer s.embedTap.Shutdown() + defer func() { + tapCancel() + s.embedTap.Shutdown() + }() + + go s.watchTapDrain(tapCtx, tapCancel) + } else { + defer tapCancel() } go func() { @@ -267,7 +289,7 @@ func (s *Spindle) Start(ctx context.Context) error { }() s.l.Info("starting tap client", "url", s.cfg.Server.Tap.Url) - s.tap.Start(ctx) + s.tap.Start(tapCtx) s.l.Info("starting spindle server", "address", s.cfg.Server.ListenAddr) return http.ListenAndServe(s.cfg.Server.ListenAddr, s.Router()) diff --git a/spindle/tap_drain.go b/spindle/tap_drain.go new file mode 100644 index 00000000..71992dfa --- /dev/null +++ b/spindle/tap_drain.go @@ -0,0 +1,152 @@ +package spindle + +import ( + "context" + "database/sql" + "fmt" + "net/url" + "sync/atomic" + "time" + + comatproto "github.com/bluesky-social/indigo/api/atproto" + "github.com/bluesky-social/indigo/events" + "github.com/bluesky-social/indigo/events/schedulers/sequential" + "github.com/gorilla/websocket" + _ "github.com/mattn/go-sqlite3" +) + +const ( + tapDrainPollInterval = 3 * time.Second + tapDrainStableChecks = 2 + tapEmptyGraceChecks = 10 +) + +func (s *Spindle) watchTapDrain(ctx context.Context, stop context.CancelFunc) { + headSeq, err := relayHeadSeq(ctx, s.cfg.Server.Tap.RelayUrl) + if err != nil { + s.l.Warn("tap drain watcher: relay head checking failed, falling back to resync-drain only", "err", err) + headSeq = 0 + } else { + s.l.Info("tap drain watcher: relay head seq at startup", "head", headSeq, "relay", s.cfg.Server.Tap.RelayUrl) + } + + conn, err := sql.Open("sqlite3", s.cfg.Server.Tap.DBPath) + if err != nil { + s.l.Warn("tap drain watcher: opening tap db failed", "err", err) + return + } + defer conn.Close() + + ticker := time.NewTicker(tapDrainPollInterval) + defer ticker.Stop() + + sawWork := false + readyStreak := 0 + emptyStreak := 0 + queryFailed := false + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + var total, busy int + if err := conn.QueryRowContext(ctx, ` + select count(*), + coalesce(sum(case when state in ('pending','resyncing','desynchronized') then 1 else 0 end), 0) + from repos`).Scan(&total, &busy); err != nil { + if !queryFailed { + s.l.Warn("tap drain watcher: repos query failed", "err", err) + queryFailed = true + } + continue + } + queryFailed = false + + var cursor int64 + if headSeq > 0 { + if err := conn.QueryRowContext(ctx, + `select cursor from firehose_cursors where url = ?`, + s.cfg.Server.Tap.RelayUrl, + ).Scan(&cursor); err != nil { + cursor = 0 + } + } + + if total > 0 { + sawWork = true + emptyStreak = 0 + } else { + emptyStreak++ + } + + caughtUp := headSeq <= 0 || cursor >= headSeq + drained := sawWork && busy == 0 + + if caughtUp && drained { + readyStreak++ + } else { + readyStreak = 0 + } + + if readyStreak >= tapDrainStableChecks { + s.l.Info("tap caught up and backfill drained, shutting down embedded tap!", "tracked", total, "cursor", cursor, "head", headSeq) + stop() + s.embedTap.Shutdown() + return + } + if !sawWork && emptyStreak >= tapEmptyGraceChecks { + s.l.Info("tap has nothing to backfill, shutting down embedded tap!") + stop() + s.embedTap.Shutdown() + return + } + } + } +} + +func relayHeadSeq(ctx context.Context, relayURL string) (int64, error) { + u, err := url.Parse(relayURL) + if err != nil { + return 0, err + } + switch u.Scheme { + case "http": + u.Scheme = "ws" + case "https": + u.Scheme = "wss" + } + u.Path = "xrpc/com.atproto.sync.subscribeRepos" + + dialCtx, cancelDial := context.WithTimeout(ctx, 15*time.Second) + defer cancelDial() + + conn, _, err := websocket.DefaultDialer.DialContext(dialCtx, u.String(), nil) + if err != nil { + return 0, fmt.Errorf("dial relay: %w", err) + } + defer conn.Close() + + streamCtx, cancelStream := context.WithCancel(dialCtx) + defer cancelStream() + + var seq atomic.Int64 + capture := func(v int64) error { + seq.Store(v) + cancelStream() + return nil + } + rsc := &events.RepoStreamCallbacks{ + RepoCommit: func(e *comatproto.SyncSubscribeRepos_Commit) error { return capture(e.Seq) }, + RepoSync: func(e *comatproto.SyncSubscribeRepos_Sync) error { return capture(e.Seq) }, + RepoIdentity: func(e *comatproto.SyncSubscribeRepos_Identity) error { return capture(e.Seq) }, + RepoAccount: func(e *comatproto.SyncSubscribeRepos_Account) error { return capture(e.Seq) }, + } + sched := sequential.NewScheduler("spindle-head-probe", rsc.EventHandler) + _ = events.HandleRepoStream(streamCtx, conn, sched, nil) + + if h := seq.Load(); h > 0 { + return h, nil + } + return 0, fmt.Errorf("no head seq received from relay") +} diff --git a/spindle/tapclient.go b/spindle/tapclient.go index 23f6cfc0..7299fbde 100644 --- a/spindle/tapclient.go +++ b/spindle/tapclient.go @@ -53,12 +53,12 @@ func (t *Tap) AddOwnerDIDs(ctx context.Context, dids []syntax.DID) error { return t.tap.AddRepos(ctx, dids) } -func (t *Tap) Start(ctx context.Context) { - go t.tap.Connect(ctx, &tapc.SimpleIndexer{ +func (t *Tap) Start(connCtx context.Context) { + go t.tap.Connect(connCtx, &tapc.SimpleIndexer{ EventHandler: t.processEvent, ConnectHandler: t.onConnect, }) - go t.purgePendingCollabsLoop(ctx) + go t.purgePendingCollabsLoop(t.spindle.rootCtx) } func (t *Tap) onConnect(ctx context.Context) { @@ -146,9 +146,12 @@ func (t *Tap) processRepo(ctx context.Context, evt *tapc.RecordEventData) error l.Info("collapsed rename leftovers", "owner", ownerDid, "repo_did", repoDid, "removed", removed) } - if err := t.tap.AddRepos(ctx, []syntax.DID{ownerDid}); err != nil { - l.Warn("tap AddRepos rejected", "did", ownerDid, "err", err) + if e := t.spindle.embedTap; e == nil || !e.closed.Load() { + if err := t.tap.AddRepos(ctx, []syntax.DID{ownerDid}); err != nil { + l.Warn("tap AddRepos rejected", "did", ownerDid, "err", err) + } } + t.spindle.jc.AddDid(ownerDid.String()) t.drainPendingCollabs(ctx, repoDid)