fix: prevent deadlock on index/search access in sqlitevec indexer

Author: William Petit

internal/adapter/memory/task_runner.go

@@ -47,6 +47,7 @@ func (r *TaskRunner) GetTask(ctx context.Context, id model.TaskID) (model.Task,
 func (r *TaskRunner) Run(ctx context.Context) error {
 	r.runningMutex.Lock()
 	r.running = true
+	slog.Debug("TaskRunner: setting running=true and broadcasting")
 	r.runningCond.Broadcast()
 	r.runningMutex.Unlock()
 
@@ -132,9 +133,12 @@ func (r *TaskRunner) ScheduleTask(ctx context.Context, task model.Task) error {
 		}()
 
 		r.runningMutex.Lock()
+		slog.Debug("TaskRunner: waiting for running flag", slog.Bool("running", r.running))
 		for !r.running {
+			slog.Debug("TaskRunner: waiting on condition variable")
 			r.runningCond.Wait()
 		}
+		slog.Debug("TaskRunner: running flag is true, proceeding")
 		r.runningMutex.Unlock()
 
 		r.semaphore <- struct{}{}
@@ -210,6 +214,7 @@ func (r *TaskRunner) ScheduleTask(ctx context.Context, task model.Task) error {
 }
 
 func (r *TaskRunner) updateState(task model.Task, fn func(s *port.TaskState)) {
+	slog.Debug("TaskRunner: acquiring state mutex for task", slog.String("taskID", string(task.ID())))
 	r.stateMutex.Lock()
 	defer r.stateMutex.Unlock()
 

internal/adapter/pipeline/index.go

@@ -168,7 +168,7 @@ func (i *Index) Index(ctx context.Context, document model.Document, funcs ...por
 
 	ctx = slogx.WithAttrs(ctx, slog.String("documentID", string(document.ID())))
 
-	slog.DebugContext(ctx, "indexing document")
+	slog.DebugContext(ctx, "pipeline: indexing document", slog.Int("indexCount", count))
 
 	for index := range i.indexes {
 		go func(index *IdentifiedIndex) {
@@ -202,6 +202,7 @@ func (i *Index) Index(ctx context.Context, document model.Document, funcs ...por
 				defer opts.OnProgress(1)
 			}
 
+			slog.DebugContext(indexCtx, "pipeline: calling Index() on underlying index")
 			if err := index.Index().Index(indexCtx, document, indexOptions...); err != nil {
 				err = errors.WithStack(err)
 				slog.ErrorContext(indexCtx, "could not index document", slog.Any("error", err))

internal/adapter/sqlitevec/index.go

@@ -24,11 +24,15 @@ type Index struct {
 	getConn  func(ctx context.Context) (*sqlite3.Conn, error)
 	llm      llm.Client
 	model    string
-	lock     sync.Mutex
+	// rwLock allows concurrent Search operations while serializing Index/Delete
+	rwLock sync.RWMutex
 }
 
 // DeleteByID implements port.Index.
 func (i *Index) DeleteByID(ctx context.Context, ids ...model.SectionID) error {
+	i.rwLock.Lock()
+	defer i.rwLock.Unlock()
+
 	err := i.withRetry(ctx, func(ctx context.Context, conn *sqlite3.Conn) error {
 		stmt, _, err := conn.Prepare("DELETE FROM embeddings WHERE section_id IN ( SELECT value FROM json_each(?) );")
 		if err != nil {
@@ -61,6 +65,9 @@ func (i *Index) DeleteByID(ctx context.Context, ids ...model.SectionID) error {
 
 // DeleteBySource implements port.Index.
 func (i *Index) DeleteBySource(ctx context.Context, source *url.URL) error {
+	i.rwLock.Lock()
+	defer i.rwLock.Unlock()
+
 	err := i.withRetry(ctx, func(ctx context.Context, conn *sqlite3.Conn) error {
 		stmt, _, err := conn.Prepare("DELETE FROM embeddings WHERE source = ?;")
 		if err != nil {
@@ -86,25 +93,6 @@ func (i *Index) DeleteBySource(ctx context.Context, source *url.URL) error {
 	return nil
 }
 
-func (i *Index) deleteBySource(ctx context.Context, conn *sqlite3.Conn, source *url.URL) error {
-	stmt, _, err := conn.Prepare("DELETE FROM embeddings WHERE source = ?;")
-	if err != nil {
-		return errors.WithStack(err)
-	}
-
-	defer stmt.Close()
-
-	if err := stmt.BindText(1, source.String()); err != nil {
-		return errors.WithStack(err)
-	}
-
-	if err := stmt.Exec(); err != nil {
-		return errors.WithStack(err)
-	}
-
-	return nil
-}
-
 type indexableChunk struct {
 	Section  model.Section
 	Text     string
@@ -127,6 +115,9 @@ const (
 
 // Index implements port.Index.
 func (i *Index) Index(ctx context.Context, document model.Document, funcs ...port.IndexOptionFunc) error {
+	i.rwLock.Lock()
+	defer i.rwLock.Unlock()
+
 	opts := port.NewIndexOptions(funcs...)
 
 	var chunksToProcess []*indexableChunk
@@ -348,6 +339,9 @@ func (i *Index) insertCollection(ctx context.Context, conn *sqlite3.Conn, embedd
 
 // Search implements port.Index.
 func (i *Index) Search(ctx context.Context, query string, opts port.IndexSearchOptions) ([]*port.IndexSearchResult, error) {
+	i.rwLock.RLock()
+	defer i.rwLock.RUnlock()
+
 	var searchResults []*port.IndexSearchResult
 	err := i.withRetry(ctx, func(ctx context.Context, conn *sqlite3.Conn) error {
 		res, err := i.llm.Embeddings(ctx, []string{query})
@@ -478,9 +472,6 @@ func (i *Index) Search(ctx context.Context, query string, opts port.IndexSearchO
 }
 
 func (i *Index) withRetry(ctx context.Context, fn func(ctx context.Context, conn *sqlite3.Conn) error, codes ...sqlite3.ErrorCode) error {
-	i.lock.Lock()
-	defer i.lock.Unlock()
-
 	conn, err := i.getConn(ctx)
 	if err != nil {
 		return errors.WithStack(err)