nandi/glean · 7d7fc0d

modified internal/atproto/auth.go +1 -1 Viewed

@@ -149,7 +149,7 @@ func pruneProfileCache() {
149	149	for _, k := range toDelete {
150	150	profileCache.Delete(k)
151	151	}
152		- profileCacheSize.Store(int64(len(toDelete)))
	152	+ profileCacheSize.Add(-int64(len(toDelete)))
153	153	}
154	154
155	155	func ResolveProfile(ctx context.Context, did string) Profile {

  	for _, k := range toDelete {
  		profileCache.Delete(k)
  	}
--	profileCacheSize.Store(int64(len(toDelete)))
++	profileCacheSize.Add(-int64(len(toDelete)))
  }
  func ResolveProfile(ctx context.Context, did string) Profile {

modified internal/atproto/jetstream.go +63 -7 Viewed

@@ -6,6 +6,7 @@ import (
6	6	"fmt"
7	7	"log/slog"
8	8	"strings"
	9	+ "sync"
9	10	"time"
10	11
11	12	jsc "github.com/bluesky-social/jetstream/pkg/client"
@@ -35,14 +36,12 @@ type EventHandler func(ctx context.Context, event *Event) error
35	36	type jetstreamScheduler struct {
36	37	handler EventHandler
37	38	logger *slog.Logger
38		- cursorStore CursorStore
	39	+ cursorStore *BatchCursorStore
39	40	}
40	41
41	42	func (s jetstreamScheduler) AddWork(ctx context.Context, _ string, evt models.Event) error {
42	43	if evt.TimeUS > 0 {
43		- if err := s.cursorStore.SaveCursor(ctx, evt.TimeUS); err != nil {
44		- s.logger.Warn("failed to save cursor", "error", err)
45		- }
	44	+ s.cursorStore.Record(ctx, evt.TimeUS)
46	45	}
47	46
48	47	if evt.Kind != models.EventKindCommit \|\| evt.Commit == nil {
@@ -81,15 +80,17 @@ type JetstreamConsumer struct {
81	80	client *jsc.Client
82	81	logger *slog.Logger
83	82	sched *jetstreamScheduler
84		- cursorStore CursorStore
	83	+ cursorStore *BatchCursorStore
85	84	rewind time.Duration
86	85	}
87	86
88	87	func NewJetstreamConsumer(jetstreamURL string, handler EventHandler, logger slog.Logger, cursorStore CursorStore) JetstreamConsumer {
	88	+ batchCursor := NewBatchCursorStore(cursorStore, 5*time.Second, logger)
	89	+
89	90	sched := &jetstreamScheduler{
90	91	handler: handler,
91	92	logger: logger,
92		- cursorStore: cursorStore,
	93	+ cursorStore: batchCursor,
93	94	}
94	95
95	96	wsURL := jetstreamURL
@@ -130,12 +131,14 @@ func NewJetstreamConsumer(jetstreamURL string, handler EventHandler, logger *slo
130	131	client: c,
131	132	logger: logger,
132	133	sched: sched,
133		- cursorStore: cursorStore,
	134	+ cursorStore: batchCursor,
134	135	rewind: rewind,
135	136	}
136	137	}
137	138
138	139	func (jc *JetstreamConsumer) Start(ctx context.Context) error {
	140	+ go jc.cursorStore.Run(ctx)
	141	+
139	142	for {
140	143	var cursorPtr *int64
141	144	if jc.cursorStore != nil {
@@ -165,3 +168,56 @@ func (jc *JetstreamConsumer) Start(ctx context.Context) error {
165	168	}
166	169	}
167	170	}
	171	+
	172	+type BatchCursorStore struct {
	173	+ store CursorStore
	174	+ mu sync.Mutex
	175	+ cursor int64
	176	+ dirty bool
	177	+ flush time.Duration
	178	+ logger *slog.Logger
	179	+}
	180	+
	181	+func NewBatchCursorStore(store CursorStore, flushInterval time.Duration, logger slog.Logger) BatchCursorStore {
	182	+ return &BatchCursorStore{store: store, flush: flushInterval, logger: logger}
	183	+}
	184	+
	185	+func (b BatchCursorStore) LoadCursor(ctx context.Context) (int64, error) {
	186	+ return b.store.LoadCursor(ctx)
	187	+}
	188	+
	189	+func (b *BatchCursorStore) Record(_ context.Context, cursor int64) {
	190	+ b.mu.Lock()
	191	+ b.cursor = cursor
	192	+ b.dirty = true
	193	+ b.mu.Unlock()
	194	+}
	195	+
	196	+func (b *BatchCursorStore) Run(ctx context.Context) {
	197	+ ticker := time.NewTicker(b.flush)
	198	+ defer ticker.Stop()
	199	+ for {
	200	+ select {
	201	+ case <-ctx.Done():
	202	+ b.flushNow(context.Background())
	203	+ return
	204	+ case <-ticker.C:
	205	+ b.flushNow(ctx)
	206	+ }
	207	+ }
	208	+}
	209	+
	210	+func (b *BatchCursorStore) flushNow(ctx context.Context) {
	211	+ b.mu.Lock()
	212	+ if !b.dirty {
	213	+ b.mu.Unlock()
	214	+ return
	215	+ }
	216	+ cursor := b.cursor
	217	+ b.dirty = false
	218	+ b.mu.Unlock()
	219	+
	220	+ if err := b.store.SaveCursor(ctx, cursor); err != nil {
	221	+ b.logger.Warn("failed to save cursor", "error", err)
	222	+ }
	223	+}

  	"fmt"
  	"log/slog"
  	"strings"
++	"sync"
  	"time"
  	jsc "github.com/bluesky-social/jetstream/pkg/client"
  type jetstreamScheduler struct {
  	handler     EventHandler
  	logger      *slog.Logger
--	cursorStore CursorStore
++	cursorStore *BatchCursorStore
  }
  func (s *jetstreamScheduler) AddWork(ctx context.Context, _ string, evt *models.Event) error {
  	if evt.TimeUS > 0 {
--		if err := s.cursorStore.SaveCursor(ctx, evt.TimeUS); err != nil {
++		s.cursorStore.Record(ctx, evt.TimeUS)
--			s.logger.Warn("failed to save cursor", "error", err)
--		}
  	}
  	if evt.Kind != models.EventKindCommit || evt.Commit == nil {
  	client      *jsc.Client
  	logger      *slog.Logger
  	sched       *jetstreamScheduler
--	cursorStore CursorStore
++	cursorStore *BatchCursorStore
  	rewind      time.Duration
  }
  func NewJetstreamConsumer(jetstreamURL string, handler EventHandler, logger *slog.Logger, cursorStore CursorStore) *JetstreamConsumer {
++	batchCursor := NewBatchCursorStore(cursorStore, 5*time.Second, logger)
++
  	sched := &jetstreamScheduler{
  		handler:     handler,
  		logger:      logger,
--		cursorStore: cursorStore,
++		cursorStore: batchCursor,
  	}
  	wsURL := jetstreamURL
  		client:      c,
  		logger:      logger,
  		sched:       sched,
--		cursorStore: cursorStore,
++		cursorStore: batchCursor,
  		rewind:      rewind,
  	}
  }
  func (jc *JetstreamConsumer) Start(ctx context.Context) error {
++	go jc.cursorStore.Run(ctx)
++
  	for {
  		var cursorPtr *int64
  		if jc.cursorStore != nil {
  		}
  	}
  }
++
++type BatchCursorStore struct {
++	store  CursorStore
++	mu     sync.Mutex
++	cursor int64
++	dirty  bool
++	flush  time.Duration
++	logger *slog.Logger
++}
++
++func NewBatchCursorStore(store CursorStore, flushInterval time.Duration, logger *slog.Logger) *BatchCursorStore {
++	return &BatchCursorStore{store: store, flush: flushInterval, logger: logger}
++}
++
++func (b *BatchCursorStore) LoadCursor(ctx context.Context) (*int64, error) {
++	return b.store.LoadCursor(ctx)
++}
++
++func (b *BatchCursorStore) Record(_ context.Context, cursor int64) {
++	b.mu.Lock()
++	b.cursor = cursor
++	b.dirty = true
++	b.mu.Unlock()
++}
++
++func (b *BatchCursorStore) Run(ctx context.Context) {
++	ticker := time.NewTicker(b.flush)
++	defer ticker.Stop()
++	for {
++		select {
++		case <-ctx.Done():
++			b.flushNow(context.Background())
++			return
++		case <-ticker.C:
++			b.flushNow(ctx)
++		}
++	}
++}
++
++func (b *BatchCursorStore) flushNow(ctx context.Context) {
++	b.mu.Lock()
++	if !b.dirty {
++		b.mu.Unlock()
++		return
++	}
++	cursor := b.cursor
++	b.dirty = false
++	b.mu.Unlock()
++
++	if err := b.store.SaveCursor(ctx, cursor); err != nil {
++		b.logger.Warn("failed to save cursor", "error", err)
++	}
++}

modified internal/cluster/article.go +89 -83 Viewed

@@ -52,50 +52,48 @@ func (e *Engine) ComputeArticleEmbeddings(ctx context.Context) error {
52	52	return fmt.Errorf("clean stale article embeddings: %w", err)
53	53	}
54	54
55		- rows, err := conn.QueryContext(ctx, `
56		- SELECT a.id, COALESCE(a.title, '') \|\| ' ' \|\| COALESCE(a.summary, '') \|\| ' ' \|\| COALESCE(a.content, '')
57		- FROM articles.articles a
58		- WHERE (COALESCE(a.title, '') != '' OR COALESCE(a.summary, '') != '' OR COALESCE(a.content, '') != '')
59		- AND a.id NOT IN (SELECT article_id FROM recs.article_embeddings)
60		- ORDER BY a.id
61		- `)
62		- if err != nil {
63		- return err
64		- }
65		-
66		- type article struct {
67		- id int64
68		- text string
69		- }
70		- var batch []article
71		- for rows.Next() {
72		- var a article
73		- if err := rows.Scan(&a.id, &a.text); err != nil {
74		- rows.Close()
	55	+ totalComputed := 0
	56	+ for {
	57	+ rows, err := conn.QueryContext(ctx, `
	58	+ SELECT a.id, COALESCE(a.title, '') \|\| ' ' \|\| COALESCE(a.summary, '') \|\| ' ' \|\| COALESCE(a.content, '')
	59	+ FROM articles.articles a
	60	+ WHERE (COALESCE(a.title, '') != '' OR COALESCE(a.summary, '') != '' OR COALESCE(a.content, '') != '')
	61	+ AND a.id NOT IN (SELECT article_id FROM recs.article_embeddings)
	62	+ ORDER BY a.id
	63	+ LIMIT ?
	64	+ `, embedBatchSize)
	65	+ if err != nil {
75	66	return err
76	67	}
77		- a.text = truncateForEmbed(a.text)
78		- batch = append(batch, a)
79		- }
80		- rows.Close()
81	68
82		- if len(batch) == 0 {
83		- e.logger.Info("article embeddings up to date")
84		- return nil
85		- }
	69	+ type article struct {
	70	+ id int64
	71	+ text string
	72	+ }
	73	+ var batch []article
	74	+ for rows.Next() {
	75	+ var a article
	76	+ if err := rows.Scan(&a.id, &a.text); err != nil {
	77	+ rows.Close()
	78	+ return err
	79	+ }
	80	+ a.text = truncateForEmbed(a.text)
	81	+ batch = append(batch, a)
	82	+ }
	83	+ rows.Close()
86	84
87		- for i := 0; i < len(batch); i += embedBatchSize {
88		- end := min(i+embedBatchSize, len(batch))
89		- sub := batch[i:end]
	85	+ if len(batch) == 0 {
	86	+ break
	87	+ }
90	88
91		- texts := make([]string, len(sub))
92		- for j, a := range sub {
	89	+ texts := make([]string, len(batch))
	90	+ for j, a := range batch {
93	91	texts[j] = a.text
94	92	}
95	93
96	94	embeddings, err := e.embedder.Embed(ctx, texts, "Represent this news article for retrieving topically similar articles. Focus on the subjects, themes, and key entities discussed.")
97	95	if err != nil {
98		- return fmt.Errorf("embed batch %d: %w", i/embedBatchSize, err)
	96	+ return fmt.Errorf("embed batch starting at total %d: %w", totalComputed, err)
99	97	}
100	98
101	99	tx, err := conn.BeginTx(ctx, nil)
@@ -111,7 +109,7 @@ func (e *Engine) ComputeArticleEmbeddings(ctx context.Context) error {
111	109	}
112	110	if _, err := tx.ExecContext(ctx,
113	111	`INSERT OR IGNORE INTO recs.article_embeddings(article_id, embedding) VALUES (?, ?)`,
114		- sub[j].id, blob,
	112	+ batch[j].id, blob,
115	113	); err != nil {
116	114	return fmt.Errorf("insert embedding: %w", err)
117	115	}
@@ -121,13 +119,18 @@ func (e *Engine) ComputeArticleEmbeddings(ctx context.Context) error {
121	119	return err
122	120	}
123	121
	122	+ totalComputed += len(batch)
124	123	e.logger.Info("article embeddings batch computed",
125		- slog.Int("batch", i/embedBatchSize),
126		- slog.Int("count", len(sub)),
	124	+ slog.Int("batch_total", totalComputed),
	125	+ slog.Int("count", len(batch)),
127	126	)
128	127	}
129	128
130		- e.logger.Info("article embeddings computed", slog.Int("total", len(batch)))
	129	+ if totalComputed == 0 {
	130	+ e.logger.Info("article embeddings up to date")
	131	+ } else {
	132	+ e.logger.Info("article embeddings computed", slog.Int("total", totalComputed))
	133	+ }
131	134	return nil
132	135	}
133	136
@@ -268,56 +271,54 @@ func (e *Engine) ComputeFeedEmbeddings(ctx context.Context) error {
268	271	return fmt.Errorf("clean stale feed embeddings: %w", err)
269	272	}
270	273
271		- rows, err := conn.QueryContext(ctx, `
272		- SELECT f.feed_url, COALESCE(f.title, '') \|\| ' ' \|\| COALESCE(f.description, '')
273		- FROM articles.feeds f
274		- WHERE (COALESCE(f.title, '') != '' OR COALESCE(f.description, '') != '')
275		- AND (
276		- f.feed_url NOT IN (SELECT feed_url FROM recs.feed_embedding_meta)
277		- OR EXISTS (
278		- SELECT 1 FROM recs.feed_embedding_meta fm
279		- WHERE fm.feed_url = f.feed_url
280		- AND fm.source_text != COALESCE(f.title, '') \|\| ' ' \|\| COALESCE(f.description, '')
	274	+ totalComputed := 0
	275	+ for {
	276	+ rows, err := conn.QueryContext(ctx, `
	277	+ SELECT f.feed_url, COALESCE(f.title, '') \|\| ' ' \|\| COALESCE(f.description, '')
	278	+ FROM articles.feeds f
	279	+ WHERE (COALESCE(f.title, '') != '' OR COALESCE(f.description, '') != '')
	280	+ AND (
	281	+ f.feed_url NOT IN (SELECT feed_url FROM recs.feed_embedding_meta)
	282	+ OR EXISTS (
	283	+ SELECT 1 FROM recs.feed_embedding_meta fm
	284	+ WHERE fm.feed_url = f.feed_url
	285	+ AND fm.source_text != COALESCE(f.title, '') \|\| ' ' \|\| COALESCE(f.description, '')
	286	+ )
281	287	)
282		- )
283		- ORDER BY f.feed_url
284		- `)
285		- if err != nil {
286		- return err
287		- }
288		-
289		- type feed struct {
290		- url string
291		- text string
292		- }
293		- var batch []feed
294		- for rows.Next() {
295		- var f feed
296		- if err := rows.Scan(&f.url, &f.text); err != nil {
297		- rows.Close()
	288	+ ORDER BY f.feed_url
	289	+ LIMIT ?
	290	+ `, embedBatchSize)
	291	+ if err != nil {
298	292	return err
299	293	}
300		- batch = append(batch, f)
301		- }
302		- rows.Close()
303	294
304		- if len(batch) == 0 {
305		- e.logger.Info("feed embeddings up to date")
306		- return nil
307		- }
	295	+ type feed struct {
	296	+ url string
	297	+ text string
	298	+ }
	299	+ var batch []feed
	300	+ for rows.Next() {
	301	+ var f feed
	302	+ if err := rows.Scan(&f.url, &f.text); err != nil {
	303	+ rows.Close()
	304	+ return err
	305	+ }
	306	+ batch = append(batch, f)
	307	+ }
	308	+ rows.Close()
308	309
309		- for i := 0; i < len(batch); i += embedBatchSize {
310		- end := min(i+embedBatchSize, len(batch))
311		- sub := batch[i:end]
	310	+ if len(batch) == 0 {
	311	+ break
	312	+ }
312	313
313		- texts := make([]string, len(sub))
314		- for j, f := range sub {
	314	+ texts := make([]string, len(batch))
	315	+ for j, f := range batch {
315	316	texts[j] = f.text
316	317	}
317	318
318	319	embeddings, err := e.embedder.Embed(ctx, texts, "Represent this RSS feed description for discovering feeds with similar editorial focus and topic coverage.")
319	320	if err != nil {
320		- return fmt.Errorf("embed feed batch %d: %w", i/embedBatchSize, err)
	321	+ return fmt.Errorf("embed feed batch starting at total %d: %w", totalComputed, err)
321	322	}
322	323
323	324	tx, err := conn.BeginTx(ctx, nil)
@@ -332,19 +333,19 @@ func (e *Engine) ComputeFeedEmbeddings(ctx context.Context) error {
332	333	return fmt.Errorf("serialize feed embedding: %w", err)
333	334	}
334	335	if _, err := tx.ExecContext(ctx,
335		- `DELETE FROM recs.feed_embeddings WHERE feed_url = ?`, sub[j].url,
	336	+ `DELETE FROM recs.feed_embeddings WHERE feed_url = ?`, batch[j].url,
336	337	); err != nil {
337	338	return fmt.Errorf("delete feed embedding: %w", err)
338	339	}
339	340	if _, err := tx.ExecContext(ctx,
340	341	`INSERT INTO recs.feed_embeddings(feed_url, embedding) VALUES (?, ?)`,
341		- sub[j].url, blob,
	342	+ batch[j].url, blob,
342	343	); err != nil {
343	344	return fmt.Errorf("insert feed embedding: %w", err)
344	345	}
345	346	if _, err := tx.ExecContext(ctx,
346	347	`INSERT OR REPLACE INTO recs.feed_embedding_meta(feed_url, source_text) VALUES (?, ?)`,
347		- sub[j].url, sub[j].text,
	348	+ batch[j].url, batch[j].text,
348	349	); err != nil {
349	350	return fmt.Errorf("insert feed embedding: %w", err)
350	351	}
@@ -354,13 +355,18 @@ func (e *Engine) ComputeFeedEmbeddings(ctx context.Context) error {
354	355	return err
355	356	}
356	357
	358	+ totalComputed += len(batch)
357	359	e.logger.Info("feed embeddings batch computed",
358		- slog.Int("batch", i/embedBatchSize),
359		- slog.Int("count", len(sub)),
	360	+ slog.Int("batch_total", totalComputed),
	361	+ slog.Int("count", len(batch)),
360	362	)
361	363	}
362	364
363		- e.logger.Info("feed embeddings computed", slog.Int("total", len(batch)))
	365	+ if totalComputed == 0 {
	366	+ e.logger.Info("feed embeddings up to date")
	367	+ } else {
	368	+ e.logger.Info("feed embeddings computed", slog.Int("total", totalComputed))
	369	+ }
364	370	return nil
365	371	}
366	372

  		return fmt.Errorf("clean stale article embeddings: %w", err)
  	}
--	rows, err := conn.QueryContext(ctx, `
++	totalComputed := 0
--		SELECT a.id, COALESCE(a.title, '') || ' ' || COALESCE(a.summary, '') || ' ' || COALESCE(a.content, '')
++	for {
--		FROM articles.articles a
++		rows, err := conn.QueryContext(ctx, `
--		WHERE (COALESCE(a.title, '') != '' OR COALESCE(a.summary, '') != '' OR COALESCE(a.content, '') != '')
++			SELECT a.id, COALESCE(a.title, '') || ' ' || COALESCE(a.summary, '') || ' ' || COALESCE(a.content, '')
--		AND a.id NOT IN (SELECT article_id FROM recs.article_embeddings)
++			FROM articles.articles a
--		ORDER BY a.id
++			WHERE (COALESCE(a.title, '') != '' OR COALESCE(a.summary, '') != '' OR COALESCE(a.content, '') != '')
--	`)
++			AND a.id NOT IN (SELECT article_id FROM recs.article_embeddings)
--	if err != nil {
++			ORDER BY a.id
--		return err
++			LIMIT ?
--	}
++		`, embedBatchSize)
--
++		if err != nil {
--	type article struct {
--		id   int64
--		text string
--	}
--	var batch []article
--	for rows.Next() {
--		var a article
--		if err := rows.Scan(&a.id, &a.text); err != nil {
--			rows.Close()
  			return err
  		}
--		a.text = truncateForEmbed(a.text)
--		batch = append(batch, a)
--	}
--	rows.Close()
--	if len(batch) == 0 {
++		type article struct {
--		e.logger.Info("article embeddings up to date")
++			id   int64
--		return nil
++			text string
--	}
++		}
++		var batch []article
++		for rows.Next() {
++			var a article
++			if err := rows.Scan(&a.id, &a.text); err != nil {
++				rows.Close()
++				return err
++			}
++			a.text = truncateForEmbed(a.text)
++			batch = append(batch, a)
++		}
++		rows.Close()
--	for i := 0; i < len(batch); i += embedBatchSize {
++		if len(batch) == 0 {
--		end := min(i+embedBatchSize, len(batch))
++			break
--		sub := batch[i:end]
++		}
--		texts := make([]string, len(sub))
++		texts := make([]string, len(batch))
--		for j, a := range sub {
++		for j, a := range batch {
  			texts[j] = a.text
  		}
  		embeddings, err := e.embedder.Embed(ctx, texts, "Represent this news article for retrieving topically similar articles. Focus on the subjects, themes, and key entities discussed.")
  		if err != nil {
--			return fmt.Errorf("embed batch %d: %w", i/embedBatchSize, err)
++			return fmt.Errorf("embed batch starting at total %d: %w", totalComputed, err)
  		}
  		tx, err := conn.BeginTx(ctx, nil)
  			}
  			if _, err := tx.ExecContext(ctx,
  				`INSERT OR IGNORE INTO recs.article_embeddings(article_id, embedding) VALUES (?, ?)`,
--				sub[j].id, blob,
++				batch[j].id, blob,
  			); err != nil {
  				return fmt.Errorf("insert embedding: %w", err)
  			}
  			return err
  		}
++		totalComputed += len(batch)
  		e.logger.Info("article embeddings batch computed",
--			slog.Int("batch", i/embedBatchSize),
++			slog.Int("batch_total", totalComputed),
--			slog.Int("count", len(sub)),
++			slog.Int("count", len(batch)),
  		)
  	}
--	e.logger.Info("article embeddings computed", slog.Int("total", len(batch)))
++	if totalComputed == 0 {
++		e.logger.Info("article embeddings up to date")
++	} else {
++		e.logger.Info("article embeddings computed", slog.Int("total", totalComputed))
++	}
  	return nil
  }
  		return fmt.Errorf("clean stale feed embeddings: %w", err)
  	}
--	rows, err := conn.QueryContext(ctx, `
++	totalComputed := 0
--		SELECT f.feed_url, COALESCE(f.title, '') || ' ' || COALESCE(f.description, '')
++	for {
--		FROM articles.feeds f
++		rows, err := conn.QueryContext(ctx, `
--		WHERE (COALESCE(f.title, '') != '' OR COALESCE(f.description, '') != '')
++			SELECT f.feed_url, COALESCE(f.title, '') || ' ' || COALESCE(f.description, '')
--		AND (
++			FROM articles.feeds f
--			f.feed_url NOT IN (SELECT feed_url FROM recs.feed_embedding_meta)
++			WHERE (COALESCE(f.title, '') != '' OR COALESCE(f.description, '') != '')
--			OR EXISTS (
++			AND (
--				SELECT 1 FROM recs.feed_embedding_meta fm
++				f.feed_url NOT IN (SELECT feed_url FROM recs.feed_embedding_meta)
--				WHERE fm.feed_url = f.feed_url
++				OR EXISTS (
--				AND fm.source_text != COALESCE(f.title, '') || ' ' || COALESCE(f.description, '')
++					SELECT 1 FROM recs.feed_embedding_meta fm
++					WHERE fm.feed_url = f.feed_url
++					AND fm.source_text != COALESCE(f.title, '') || ' ' || COALESCE(f.description, '')
++				)
  			)
--		)
++			ORDER BY f.feed_url
--		ORDER BY f.feed_url
++			LIMIT ?
--	`)
++		`, embedBatchSize)
--	if err != nil {
++		if err != nil {
--		return err
--	}
--
--	type feed struct {
--		url  string
--		text string
--	}
--	var batch []feed
--	for rows.Next() {
--		var f feed
--		if err := rows.Scan(&f.url, &f.text); err != nil {
--			rows.Close()
  			return err
  		}
--		batch = append(batch, f)
--	}
--	rows.Close()
--	if len(batch) == 0 {
++		type feed struct {
--		e.logger.Info("feed embeddings up to date")
++			url  string
--		return nil
++			text string
--	}
++		}
++		var batch []feed
++		for rows.Next() {
++			var f feed
++			if err := rows.Scan(&f.url, &f.text); err != nil {
++				rows.Close()
++				return err
++			}
++			batch = append(batch, f)
++		}
++		rows.Close()
--	for i := 0; i < len(batch); i += embedBatchSize {
++		if len(batch) == 0 {
--		end := min(i+embedBatchSize, len(batch))
++			break
--		sub := batch[i:end]
++		}
--		texts := make([]string, len(sub))
++		texts := make([]string, len(batch))
--		for j, f := range sub {
++		for j, f := range batch {
  			texts[j] = f.text
  		}
  		embeddings, err := e.embedder.Embed(ctx, texts, "Represent this RSS feed description for discovering feeds with similar editorial focus and topic coverage.")
  		if err != nil {
--			return fmt.Errorf("embed feed batch %d: %w", i/embedBatchSize, err)
++			return fmt.Errorf("embed feed batch starting at total %d: %w", totalComputed, err)
  		}
  		tx, err := conn.BeginTx(ctx, nil)
  				return fmt.Errorf("serialize feed embedding: %w", err)
  			}
  			if _, err := tx.ExecContext(ctx,
--				`DELETE FROM recs.feed_embeddings WHERE feed_url = ?`, sub[j].url,
++				`DELETE FROM recs.feed_embeddings WHERE feed_url = ?`, batch[j].url,
  			); err != nil {
  				return fmt.Errorf("delete feed embedding: %w", err)
  			}
  			if _, err := tx.ExecContext(ctx,
  				`INSERT INTO recs.feed_embeddings(feed_url, embedding) VALUES (?, ?)`,
--				sub[j].url, blob,
++				batch[j].url, blob,
  			); err != nil {
  				return fmt.Errorf("insert feed embedding: %w", err)
  			}
  			if _, err := tx.ExecContext(ctx,
  				`INSERT OR REPLACE INTO recs.feed_embedding_meta(feed_url, source_text) VALUES (?, ?)`,
--				sub[j].url, sub[j].text,
++				batch[j].url, batch[j].text,
  			); err != nil {
  				return fmt.Errorf("insert feed embedding: %w", err)
  			}
  			return err
  		}
++		totalComputed += len(batch)
  		e.logger.Info("feed embeddings batch computed",
--			slog.Int("batch", i/embedBatchSize),
++			slog.Int("batch_total", totalComputed),
--			slog.Int("count", len(sub)),
++			slog.Int("count", len(batch)),
  		)
  	}
--	e.logger.Info("feed embeddings computed", slog.Int("total", len(batch)))
++	if totalComputed == 0 {
++		e.logger.Info("feed embeddings up to date")
++	} else {
++		e.logger.Info("feed embeddings computed", slog.Int("total", totalComputed))
++	}
  	return nil
  }

modified internal/cluster/jaccard.go +3 -3 Viewed

@@ -155,7 +155,7 @@ func (e Engine) computeEmbeddingSimilarity(ctx context.Context, tx sql.Tx) err
155	155	}
156	156
157	157	const knnLimit = 50
158		- stmt, err := tx.PrepareContext(ctx, `
	158	+ knnStmt, err := tx.PrepareContext(ctx, `
159	159	SELECT feed_url, distance
160	160	FROM recs.feed_embeddings
161	161	WHERE embedding MATCH ? AND k = ?
@@ -164,7 +164,7 @@ func (e Engine) computeEmbeddingSimilarity(ctx context.Context, tx sql.Tx) err
164	164	if err != nil {
165	165	return err
166	166	}
167		- defer stmt.Close()
	167	+ defer knnStmt.Close()
168	168
169	169	updateStmt, err := tx.PrepareContext(ctx,
170	170	`UPDATE _feed_sim_staging SET jaccard = jaccard + ? WHERE feed_a = ? AND feed_b = ?`,
@@ -175,7 +175,7 @@ func (e Engine) computeEmbeddingSimilarity(ctx context.Context, tx sql.Tx) err
175	175	defer updateStmt.Close()
176	176
177	177	for _, f := range feeds {
178		- knnRows, err := stmt.QueryContext(ctx, f.vec, knnLimit)
	178	+ knnRows, err := knnStmt.QueryContext(ctx, f.vec, knnLimit)
179	179	if err != nil {
180	180	return err
181	181	}

  	}
  	const knnLimit = 50
--	stmt, err := tx.PrepareContext(ctx, `
++	knnStmt, err := tx.PrepareContext(ctx, `
  		SELECT feed_url, distance
  		FROM recs.feed_embeddings
  		WHERE embedding MATCH ? AND k = ?
  	if err != nil {
  		return err
  	}
--	defer stmt.Close()
++	defer knnStmt.Close()
  	updateStmt, err := tx.PrepareContext(ctx,
  		`UPDATE _feed_sim_staging SET jaccard = jaccard + ? WHERE feed_a = ? AND feed_b = ?`,
  	defer updateStmt.Close()
  	for _, f := range feeds {
--		knnRows, err := stmt.QueryContext(ctx, f.vec, knnLimit)
++		knnRows, err := knnStmt.QueryContext(ctx, f.vec, knnLimit)
  		if err != nil {
  			return err
  		}

modified internal/cluster/social.go +48 -36 Viewed

@@ -6,6 +6,7 @@ import (
6	6	)
7	7
8	8	const maxFollowDepth = 3
	9	+const maxReachablePerUser = 10000
9	10
10	11	func chunk[T any](s []T, size int) [][]T {
11	12	var chunks [][]T
@@ -33,43 +34,10 @@ func (e *Engine) ComputeFollowDistancesData(ctx context.Context, sources []strin
33	34	distances := make(map[pair]int)
34	35
35	36	for _, src := range sources {
36		- frontier := []string{src}
37		- reachable := map[string]int{src: 0}
38		-
39		- for depth := 0; depth < maxFollowDepth && len(frontier) > 0; depth++ {
40		- var nextLevel []string
41		- for _, batch := range chunk(frontier, 500) {
42		- ph := make([]string, len(batch))
43		- args := make([]any, len(batch))
44		- for i, did := range batch {
45		- ph[i] = "?"
46		- args[i] = did
47		- }
48		-
49		- rows, err := e.db.QueryContext(ctx,
50		- fmt.Sprintf(`SELECT target_did FROM main.follows WHERE user_did IN (%s) AND user_did != target_did`, joinPh(ph)),
51		- args...,
52		- )
53		- if err != nil {
54		- return nil, err
55		- }
56		-
57		- for rows.Next() {
58		- var dst string
59		- if err := rows.Scan(&dst); err != nil {
60		- rows.Close()
61		- return nil, err
62		- }
63		- if _, ok := reachable[dst]; !ok {
64		- reachable[dst] = depth + 1
65		- nextLevel = append(nextLevel, dst)
66		- }
67		- }
68		- rows.Close()
69		- }
70		- frontier = nextLevel
	37	+ reachable, err := e.bfsReachable(ctx, src)
	38	+ if err != nil {
	39	+ return nil, err
71	40	}
72		-
73	41	for other, d := range reachable {
74	42	if d > 0 {
75	43	distances[pair{src, other}] = d
@@ -84,6 +52,50 @@ func (e *Engine) ComputeFollowDistancesData(ctx context.Context, sources []strin
84	52	return result, nil
85	53	}
86	54
	55	+func (e *Engine) bfsReachable(ctx context.Context, src string) (map[string]int, error) {
	56	+ reachable := map[string]int{src: 0}
	57	+ frontier := []string{src}
	58	+
	59	+ for depth := 0; depth < maxFollowDepth && len(frontier) > 0; depth++ {
	60	+ var nextLevel []string
	61	+ for _, batch := range chunk(frontier, 500) {
	62	+ ph := make([]string, len(batch))
	63	+ args := make([]any, len(batch))
	64	+ for i, did := range batch {
	65	+ ph[i] = "?"
	66	+ args[i] = did
	67	+ }
	68	+
	69	+ rows, err := e.db.QueryContext(ctx,
	70	+ fmt.Sprintf(`SELECT target_did FROM main.follows WHERE user_did IN (%s) AND user_did != target_did`, joinPh(ph)),
	71	+ args...,
	72	+ )
	73	+ if err != nil {
	74	+ return reachable, err
	75	+ }
	76	+
	77	+ for rows.Next() {
	78	+ var dst string
	79	+ if err := rows.Scan(&dst); err != nil {
	80	+ rows.Close()
	81	+ return reachable, err
	82	+ }
	83	+ if _, ok := reachable[dst]; !ok {
	84	+ if len(reachable) >= maxReachablePerUser {
	85	+ rows.Close()
	86	+ return reachable, nil
	87	+ }
	88	+ reachable[dst] = depth + 1
	89	+ nextLevel = append(nextLevel, dst)
	90	+ }
	91	+ }
	92	+ rows.Close()
	93	+ }
	94	+ frontier = nextLevel
	95	+ }
	96	+ return reachable, nil
	97	+}
	98	+
87	99	func (e *Engine) WriteFollowDistances(ctx context.Context, distances []followDistance) error {
88	100	tx, err := e.db.BeginTx(ctx, nil)
89	101	if err != nil {

  )
  const maxFollowDepth = 3
++const maxReachablePerUser = 10000
  func chunk[T any](s []T, size int) [][]T {
  	var chunks [][]T
  	distances := make(map[pair]int)
  	for _, src := range sources {
--		frontier := []string{src}
++		reachable, err := e.bfsReachable(ctx, src)
--		reachable := map[string]int{src: 0}
++		if err != nil {
--
++			return nil, err
--		for depth := 0; depth < maxFollowDepth && len(frontier) > 0; depth++ {
--			var nextLevel []string
--			for _, batch := range chunk(frontier, 500) {
--				ph := make([]string, len(batch))
--				args := make([]any, len(batch))
--				for i, did := range batch {
--					ph[i] = "?"
--					args[i] = did
--				}
--
--				rows, err := e.db.QueryContext(ctx,
--					fmt.Sprintf(`SELECT target_did FROM main.follows WHERE user_did IN (%s) AND user_did != target_did`, joinPh(ph)),
--					args...,
--				)
--				if err != nil {
--					return nil, err
--				}
--
--				for rows.Next() {
--					var dst string
--					if err := rows.Scan(&dst); err != nil {
--						rows.Close()
--						return nil, err
--					}
--					if _, ok := reachable[dst]; !ok {
--						reachable[dst] = depth + 1
--						nextLevel = append(nextLevel, dst)
--					}
--				}
--				rows.Close()
--			}
--			frontier = nextLevel
  		}
--
  		for other, d := range reachable {
  			if d > 0 {
  				distances[pair{src, other}] = d
  	return result, nil
  }
++func (e *Engine) bfsReachable(ctx context.Context, src string) (map[string]int, error) {
++	reachable := map[string]int{src: 0}
++	frontier := []string{src}
++
++	for depth := 0; depth < maxFollowDepth && len(frontier) > 0; depth++ {
++		var nextLevel []string
++		for _, batch := range chunk(frontier, 500) {
++			ph := make([]string, len(batch))
++			args := make([]any, len(batch))
++			for i, did := range batch {
++				ph[i] = "?"
++				args[i] = did
++			}
++
++			rows, err := e.db.QueryContext(ctx,
++				fmt.Sprintf(`SELECT target_did FROM main.follows WHERE user_did IN (%s) AND user_did != target_did`, joinPh(ph)),
++				args...,
++			)
++			if err != nil {
++				return reachable, err
++			}
++
++			for rows.Next() {
++				var dst string
++				if err := rows.Scan(&dst); err != nil {
++					rows.Close()
++					return reachable, err
++				}
++				if _, ok := reachable[dst]; !ok {
++					if len(reachable) >= maxReachablePerUser {
++						rows.Close()
++						return reachable, nil
++					}
++					reachable[dst] = depth + 1
++					nextLevel = append(nextLevel, dst)
++				}
++			}
++			rows.Close()
++		}
++		frontier = nextLevel
++	}
++	return reachable, nil
++}
++
  func (e *Engine) WriteFollowDistances(ctx context.Context, distances []followDistance) error {
  	tx, err := e.db.BeginTx(ctx, nil)
  	if err != nil {

modified internal/db/db.go +3 -3 Viewed

@@ -46,7 +46,7 @@ func Open(basePath string) (*Store, error) {
46	46	`PRAGMA busy_timeout = 30000`,
47	47	`PRAGMA synchronous = NORMAL`,
48	48	`PRAGMA cache = shared`,
49		- `PRAGMA temp_store = MEMORY`,
	49	+ `PRAGMA temp_store = FILE`,
50	50	`PRAGMA mmap_size = 268435456`,
51	51	} {
52	52	if _, err := conn.Exec(p, nil); err != nil {
@@ -62,7 +62,7 @@ func Open(basePath string) (*Store, error) {
62	62	`PRAGMA articles.busy_timeout = 30000`,
63	63	`PRAGMA articles.synchronous = NORMAL`,
64	64	`PRAGMA articles.cache = shared`,
65		- `PRAGMA articles.temp_store = MEMORY`,
	65	+ `PRAGMA articles.temp_store = FILE`,
66	66	`PRAGMA articles.mmap_size = 268435456`,
67	67	} {
68	68	if _, err := conn.Exec(p, nil); err != nil {
@@ -78,7 +78,7 @@ func Open(basePath string) (*Store, error) {
78	78	`PRAGMA recs.busy_timeout = 30000`,
79	79	`PRAGMA recs.synchronous = NORMAL`,
80	80	`PRAGMA recs.cache = shared`,
81		- `PRAGMA recs.temp_store = MEMORY`,
	81	+ `PRAGMA recs.temp_store = FILE`,
82	82	`PRAGMA recs.mmap_size = 268435456`,
83	83	} {
84	84	if _, err := conn.Exec(p, nil); err != nil {

  				`PRAGMA busy_timeout = 30000`,
  				`PRAGMA synchronous = NORMAL`,
  				`PRAGMA cache = shared`,
--				`PRAGMA temp_store = MEMORY`,
++				`PRAGMA temp_store = FILE`,
  				`PRAGMA mmap_size = 268435456`,
  			} {
  				if _, err := conn.Exec(p, nil); err != nil {
  				`PRAGMA articles.busy_timeout = 30000`,
  				`PRAGMA articles.synchronous = NORMAL`,
  				`PRAGMA articles.cache = shared`,
--				`PRAGMA articles.temp_store = MEMORY`,
++				`PRAGMA articles.temp_store = FILE`,
  				`PRAGMA articles.mmap_size = 268435456`,
  			} {
  				if _, err := conn.Exec(p, nil); err != nil {
  				`PRAGMA recs.busy_timeout = 30000`,
  				`PRAGMA recs.synchronous = NORMAL`,
  				`PRAGMA recs.cache = shared`,
--				`PRAGMA recs.temp_store = MEMORY`,
++				`PRAGMA recs.temp_store = FILE`,
  				`PRAGMA recs.mmap_size = 268435456`,
  			} {
  				if _, err := conn.Exec(p, nil); err != nil {

refactor: reduce memory usage, sqlite temp tables to disk, cap allocsUnverified