From a254492a87dcce0891cbc6a14e139885cfaae117 Mon Sep 17 00:00:00 2001 From: Aly Raffauf Date: Fri, 3 Jul 2026 14:30:13 -0400 Subject: [PATCH] firehose: add cursor + resume logic; store: serialize writes, free up reads --- internal/firehose/commit.go | 4 ++++ internal/firehose/stream.go | 12 +++++++++++- internal/store/cursor.go | 34 ++++++++++++++++++++++++++++++++++ internal/store/links.go | 4 ++-- internal/store/query.go | 10 +++++----- internal/store/store.go | 28 +++++++++++++++++++++------- 6 files changed, 77 insertions(+), 15 deletions(-) create mode 100644 internal/store/cursor.go diff --git a/internal/firehose/commit.go b/internal/firehose/commit.go index 44d14da..4bd0009 100644 --- a/internal/firehose/commit.go +++ b/internal/firehose/commit.go @@ -14,6 +14,10 @@ import ( func (c *Consumer) HandleCommit(ctx context.Context, event *atproto.SyncSubscribeRepos_Commit) error { // fmt.Println("repo:", event.Repo, "commit:", event.Rev) + // + if err := c.Store.SaveCursor(ctx, event.Seq); err != nil { + fmt.Println("could not save cursor:", err) + } if event.TooBig { go func() { diff --git a/internal/firehose/stream.go b/internal/firehose/stream.go index 60f1a75..8a4cfb4 100644 --- a/internal/firehose/stream.go +++ b/internal/firehose/stream.go @@ -2,6 +2,7 @@ package firehose import ( "context" + "fmt" "log/slog" "net/http" "time" @@ -21,7 +22,16 @@ func (c *Consumer) Run(ctx context.Context, relayURL string, logger *slog.Logger backoff := minBackoff for { - conn, _, err := websocket.DefaultDialer.DialContext(ctx, relayURL, http.Header{}) + dialURL := relayURL + + if cursor, err := c.Store.GetCursor(ctx); err != nil { + logger.Warn("could not load cursor, starting from live tip", "err", err) + } else if cursor > 0 { + dialURL = fmt.Sprintf("%s?cursor=%d", relayURL, cursor) + } + + + conn, _, err := websocket.DefaultDialer.DialContext(ctx, dialURL, http.Header{}) if err != nil { logger.Warn("dial failed", "err", err, "retry in", backoff) } else { diff --git a/internal/store/cursor.go b/internal/store/cursor.go new file mode 100644 index 0000000..1dc6a7d --- /dev/null +++ b/internal/store/cursor.go @@ -0,0 +1,34 @@ +package store + +import ( + "context" + "database/sql" + "errors" + "fmt" +) + +func (s *Store) SaveCursor(ctx context.Context, seq int64) error { + _, err := s.writeDB.ExecContext(ctx, + `INSERT INTO cursor (id, seq) VALUES (0, ?) ON CONFLICT (id) DO UPDATE SET seq = excluded.seq`, seq, + ) + + if err != nil { + return fmt.Errorf("save cursor: %w", err) + } + return nil +} + + +func (s *Store) GetCursor(ctx context.Context) (int64, error) { + var seq int64 + err := s.readDB.QueryRowContext(ctx, `SELECT seq FROM cursor WHERE id = 0`).Scan(&seq) + if errors.Is(err, sql.ErrNoRows) { + return 0, nil + } + + if err != nil { + return 0, fmt.Errorf("get cursor: %w", err) + } + + return seq, nil +} diff --git a/internal/store/links.go b/internal/store/links.go index 134119a..f1f41f3 100644 --- a/internal/store/links.go +++ b/internal/store/links.go @@ -21,7 +21,7 @@ func deleteLinks(ctx context.Context, tx *sql.Tx, actorDid, collection, recordKe } func (s *Store) DeleteLinks(ctx context.Context, actorDid, collection, recordKey string) error { - tx, err := s.db.BeginTx(ctx, nil) + tx, err := s.writeDB.BeginTx(ctx, nil) if err != nil { return fmt.Errorf("begin transaction: %w", err) } @@ -35,7 +35,7 @@ func (s *Store) DeleteLinks(ctx context.Context, actorDid, collection, recordKey } func (s *Store) SaveLinks(ctx context.Context, actorDid, collection, recordKey string, links []backlink.Link) error { - tx, err := s.db.BeginTx(ctx, nil) + tx, err := s.writeDB.BeginTx(ctx, nil) if err != nil { return fmt.Errorf("begin transaction: %w", err) } diff --git a/internal/store/query.go b/internal/store/query.go index f4f5524..90ae5d2 100644 --- a/internal/store/query.go +++ b/internal/store/query.go @@ -16,7 +16,7 @@ type Record struct { func (s *Store) CountBacklinks(ctx context.Context, target, collection, fieldPath string) (uint64, error) { var total uint64 - err := s.db.QueryRowContext(ctx, + err := s.readDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM links WHERE target = ? AND collection = ? AND field_path = ?`, target, collection, fieldPath, ).Scan(&total) @@ -28,7 +28,7 @@ func (s *Store) CountBacklinks(ctx context.Context, target, collection, fieldPat } func (s *Store) DistinctBacklinkDids(ctx context.Context, target, collection, fieldPath string, after string, limit uint64) (total uint64, dids []string, err error) { - err = s.db.QueryRowContext(ctx, + err = s.readDB.QueryRowContext(ctx, `SELECT COUNT(DISTINCT actor_did) FROM links WHERE target = ? AND collection = ? AND field_path = ?`, target, collection, fieldPath, ).Scan(&total) @@ -36,7 +36,7 @@ func (s *Store) DistinctBacklinkDids(ctx context.Context, target, collection, fi return 0, nil, fmt.Errorf("count distinct dids: %w", err) } - rows, err := s.db.QueryContext(ctx, + rows, err := s.readDB.QueryContext(ctx, `SELECT DISTINCT actor_did FROM links WHERE target = ? AND collection = ? AND field_path = ? AND actor_did > ? ORDER BY actor_did LIMIT ?`, @@ -73,7 +73,7 @@ func (s *Store) ListBacklinks(ctx context.Context, target, collection, fieldPath } where += `AND actor_did IN (` + strings.Join(placeholders, ", ") + `)` - err = s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM links WHERE `+where, args...).Scan(&total) + err = s.readDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM links WHERE `+where, args...).Scan(&total) if err != nil { return 0, nil, fmt.Errorf("count backlinks: %w", err) } @@ -103,7 +103,7 @@ func (s *Store) ListBacklinks(ctx context.Context, target, collection, fieldPath } listArgs = append(listArgs, limit) - rows, err := s.db.QueryContext(ctx, query, listArgs...) + rows, err := s.readDB.QueryContext(ctx, query, listArgs...) if err != nil { return 0, nil, fmt.Errorf("query backlinks: %w", err) } diff --git a/internal/store/store.go b/internal/store/store.go index 7f8c220..d672a7c 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -8,7 +8,8 @@ import ( ) type Store struct { - db *sql.DB + writeDB *sql.DB + readDB *sql.DB } const schema = ` @@ -25,24 +26,37 @@ CREATE TABLE IF NOT EXISTS links ( ); CREATE INDEX IF NOT EXISTS idx_links_target ON links(target); CREATE INDEX IF NOT EXISTS idx_links_source ON links(actor_did, collection, record_key); +CREATE TABLE IF NOT EXISTS cursor ( + id INTEGER PRIMARY KEY CHECK (id = 0), + seq INTEGER NOT NULL +); ` func Open(path string) (*Store, error) { - db, err := sql.Open("sqlite", path) + writeDB, err := sql.Open("sqlite", path) if err != nil { - return nil, fmt.Errorf("open database: %w", err) + return nil, fmt.Errorf("open write database: %w", err) } - if _, err := db.Exec("PRAGMA journal_mode=WAL;"); err != nil { + if _, err := writeDB.Exec("PRAGMA journal_mode=WAL;"); err != nil { return nil, fmt.Errorf("set journal mode: %w", err) } - if _, err := db.Exec("PRAGMA busy_timeout=5000;"); err != nil { + if _, err := writeDB.Exec("PRAGMA busy_timeout=5000;"); err != nil { return nil, fmt.Errorf("set busy timeout: %w", err) } + writeDB.SetMaxOpenConns(1) - if _, err := db.Exec(schema); err != nil { + if _, err := writeDB.Exec(schema); err != nil { return nil, fmt.Errorf("create schema: %w", err) } - return &Store{db: db}, nil + readDB, err := sql.Open("sqlite", path) + if err != nil { + return nil, fmt.Errorf("open read database: %w", err) + } + if _, err := readDB.Exec("PRAGMA busy_timeout=5000;"); err != nil { + return nil, fmt.Errorf("set read busy timeout: %w", err) + } + + return &Store{writeDB: writeDB, readDB: readDB}, nil } -- 2.51.2