Advance the search read pointer from the reindex instead of a chain (PP-4908)

src/palace/manager/celery/tasks/search.py

@@ -5,18 +5,20 @@
 from datetime import timedelta
 from typing import Any
 
-from celery import chain, shared_task
+from celery import shared_task
 from celery.exceptions import Ignore, Retry
 from opensearchpy import OpenSearchException
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from palace.util.exceptions import BasePalaceException
-from palace.util.log import elapsed_time_logging
+from palace.util.log import LoggerType, elapsed_time_logging
 
 from palace.manager.celery.task import Task
 from palace.manager.celery.utils import signature_with
 from palace.manager.search.external_search import ExternalSearchIndex
+from palace.manager.search.revision import SearchSchemaRevision
+from palace.manager.search.service import SearchService
 from palace.manager.service.celery.celery import QueueNames
 from palace.manager.service.redis.models.lock import LockNotAcquired, TaskLock
 from palace.manager.service.redis.models.search import WaitingForIndexing
@@ -44,93 +46,188 @@ def get_work_search_documents(
 
 
 def add_documents_to_index(
-    task: Task, index: ExternalSearchIndex, documents: Sequence[dict[str, Any]]
+    log: LoggerType, index: ExternalSearchIndex, documents: Sequence[dict[str, Any]]
 ) -> None:
-    try:
-        with elapsed_time_logging(
-            log_method=task.log.info,
-            message_prefix="Works added to index",
-            skip_start=True,
-        ):
-            failed_documents = index.add_documents(documents=documents)
-        if failed_documents:
-            raise FailedToIndex(f"Failed to index {len(failed_documents)} works.")
-    except (FailedToIndex, OpenSearchException) as e:
-        wait_time = exponential_backoff(task.request.retries)
-        task.log.error(f"{e}. Retrying in {wait_time} seconds.")
-        raise task.retry(countdown=wait_time)
+    """
+    Submit a batch of documents to the search index.
+
+    :raises FailedToIndex: If the index rejected some of the documents.
+    :raises OpenSearchException: If the index rejected the request.
+    """
+    with elapsed_time_logging(
+        log_method=log.info,
+        message_prefix="Works added to index",
+        skip_start=True,
+    ):
+        failed_documents = index.add_documents(documents=documents)
+    if failed_documents:
+        raise FailedToIndex(f"Failed to index {len(failed_documents)} works.")
 
 
 class FailedToIndex(BasePalaceException): ...
 
 
-@shared_task(queue=QueueNames.default, bind=True, max_retries=4)
-def search_reindex(task: Task, offset: int = 0, batch_size: int = 500) -> None:
+def set_read_pointer(
+    log: LoggerType, service: SearchService, revision: SearchSchemaRevision
+) -> None:
+    """
+    Publish a revision's index for reads.
+
+    :param log: The logger of the task doing the update.
+    :param service: The search service whose read pointer is being moved.
+    :param revision: The revision whose index should serve reads.
+    :raises OpenSearchException: If OpenSearch rejects the alias update.
+    """
+    service.read_pointer_set(revision)
+    log.info(
+        f"Updated read pointer ({service.base_revision_name} v{revision.version})."
+    )
+
+
+def resolve_target_index(service: SearchService) -> str | None:
+    """
+    Identify the index a reindex starting now is going to fill.
+
+    :param service: The search service to read the write pointer from.
+    :return: The index writes are currently going to, or None if there is no write pointer.
+    :raises OpenSearchException: If the write pointer cannot be read.
+    """
+    write_pointer = service.write_pointer()
+    return write_pointer.index if write_pointer is not None else None
+
+
+def advance_read_pointer(
+    log: LoggerType,
+    service: SearchService,
+    revision: SearchSchemaRevision,
+    target_index: str | None,
+) -> None:
+    """
+    Point reads at the index a completed reindex just filled, if reads are still being
+    served from an older one.
+
+    This is what makes a schema migration finish on its own: a new revision is created
+    with the write pointer aimed at it, and whichever reindex first completes a full
+    pass over that index publishes it. No caller has to remember to advance the pointer,
+    and a run that dies partway through simply leaves the pointer where it was.
+
+    :param log: The logger of the task that completed the reindex.
+    :param service: The search service whose pointers are being read and moved.
+    :param revision: The latest revision, whose index is the one worth publishing.
+    :param target_index: The index the completed run was filling.
+    :raises OpenSearchException: If the pointers cannot be read, or the read pointer cannot
+        be updated.
+    """
+    latest_index = revision.name_for_index(service.base_revision_name)
+
+    read_pointer = service.read_pointer()
+    if read_pointer is not None and read_pointer.version >= revision.version:
+        return
+
+    if target_index is None:
+        # Either the run started partway through, or there was no write pointer for it to
+        # record when it started.
+        log.warning(
+            "Not advancing read pointer: this run did not record an index to fill."
+        )
+        return
+
+    if target_index != latest_index:
+        log.warning(
+            f"Not advancing read pointer: this reindex filled {target_index}, "
+            f"but the latest revision is {latest_index}."
+        )
+        return
+
+    write_pointer = service.write_pointer()
+    if write_pointer is None or write_pointer.index != target_index:
+        # The write pointer moved partway through, so our documents are split across the
+        # old and new indexes and neither one received a complete pass.
+        log.warning(
+            f"Not advancing read pointer: the write pointer moved to "
+            f"{write_pointer.index if write_pointer is not None else None} "
+            f"while this reindex was filling {target_index}."
+        )
+        return
+
+    set_read_pointer(log, service, revision)
+
+
+@shared_task(
+    queue=QueueNames.default, bind=True, max_retries=4, throws=(LockNotAcquired,)
+)
+def search_reindex(
+    task: Task, offset: int = 0, batch_size: int = 500, target_index: str | None = None
+) -> None:
     """
     Submit all works that are presentation ready to the search index.
 
     This is done in batches, with the batch size determined by the batch_size parameter. This
-    task will do a batch, then requeue itself until all works have been indexed.
+    task will do a batch, then requeue itself until all works have been indexed. Once every
+    work has been indexed, the read pointer is advanced to the index we filled if it is still
+    behind. See advance_read_pointer.
+
+    :param target_index: The index this run is filling, carried across requeues so the final
+        batch can tell whether the write pointer moved partway through. Resolved from the write
+        pointer when a run starts from the beginning; callers should not pass it. A run started
+        at a non-zero offset skips the works before it, so it leaves this unset and never
+        advances the read pointer.
     """
     index = task.services.search.index()
     task_lock = TaskLock(task, lock_name="search_reindex")
 
     with task_lock.lock(release_on_exit=False, ignored_exceptions=(Retry, Ignore)):
-        task.log.info(
-            f"Running search reindex at offset {offset} with batch size {batch_size}."
-        )
+        try:
+            if target_index is None and offset == 0:
+                target_index = resolve_target_index(task.services.search.service())
 
-        with (
-            task.session() as session,
-            elapsed_time_logging(
-                log_method=task.log.info,
-                message_prefix="Works queried from database",
-                skip_start=True,
-            ),
-        ):
-            documents = get_work_search_documents(session, batch_size, offset)
-
-        add_documents_to_index(task, index, documents)
-
-        if len(documents) == batch_size:
-            # This task is complete, but there are more works waiting to be indexed. Requeue ourselves
-            # to process the next batch. We add a random delay to avoid hammering the search service
-            # when this task is running in parallel on multiple workers.
-            delay = random.uniform(5, 15)
-            raise task.replace(
-                signature_with(task, offset=offset + batch_size).set(countdown=delay)
+            task.log.info(
+                f"Running search reindex at offset {offset} with batch size {batch_size}."
             )
 
+            with (
+                task.session() as session,
+                elapsed_time_logging(
+                    log_method=task.log.info,
+                    message_prefix="Works queried from database",
+                    skip_start=True,
+                ),
+            ):
+                documents = get_work_search_documents(session, batch_size, offset)
+
+            add_documents_to_index(task.log, index, documents)
+
+            if len(documents) == batch_size:
+                # This task is complete, but there are more works waiting to be indexed. Requeue ourselves
+                # to process the next batch. We add a random delay to avoid hammering the search service
+                # when this task is running in parallel on multiple workers.
+                delay = random.uniform(5, 15)
+                raise task.replace(
+                    signature_with(
+                        task, offset=offset + batch_size, target_index=target_index
+                    ).set(countdown=delay)
+                )
+
+            advance_read_pointer(
+                task.log,
+                task.services.search.service(),
+                task.services.search.revision_directory().highest(),
+                target_index,
+            )
+        except (FailedToIndex, OpenSearchException) as e:
+            # A full pass takes days on a large collection, so a transient search failure
+            # retries this batch rather than discarding the run's work. The retry happens
+            # here, inside the lock, so the run keeps the lock while it waits.
+            wait_time = exponential_backoff(task.request.retries)
+            task.log.exception(
+                f"Search reindex failed ({e}). Retrying in {wait_time} seconds."
+            )
+            raise task.retry(countdown=wait_time)
+
     task.log.info("Finished search reindex.")
     task_lock.release()
 
 
-@shared_task(queue=QueueNames.default, bind=True, max_retries=4)
-def update_read_pointer(task: Task) -> None:
-    """
-    Update the read pointer to the latest revision.
-
-    This is used to indicate that the search index has been updated to a specific version. We
-    chain this task with search_reindex when doing a migration to ensure that the read pointer is
-    updated after all works have been indexed. See get_migrate_search_chain.
-    """
-    task.log.info("Updating read pointer.")
-    service = task.services.search.service()
-    revision_directory = task.services.search.revision_directory()
-    revision = revision_directory.highest()
-    try:
-        service.read_pointer_set(revision)
-    except OpenSearchException as e:
-        wait_time = exponential_backoff(task.request.retries)
-        task.log.error(
-            f"Failed to update read pointer: {e}. Retrying in {wait_time} seconds."
-        )
-        raise task.retry(countdown=wait_time)
-    task.log.info(
-        f"Updated read pointer ({service.base_revision_name} v{revision.version})."
-    )
-
-
 @shared_task(queue=QueueNames.default, bind=True, throws=(LockNotAcquired,))
 def search_indexing(task: Task, batch_size: int = 500) -> None:
     redis_client = task.services.redis.client()
@@ -175,11 +272,9 @@ def index_works(task: Task, works: Sequence[int]) -> None:
     ):
         documents = Work.to_search_documents(session, works)
 
-    add_documents_to_index(task, index, documents)
-
-
-def get_migrate_search_chain() -> chain:
-    """
-    Get the chain of tasks to run when migrating the search index to a new schema.
-    """
-    return chain(search_reindex.si(), update_read_pointer.si())
+    try:
+        add_documents_to_index(task.log, index, documents)
+    except (FailedToIndex, OpenSearchException) as e:
+        wait_time = exponential_backoff(task.request.retries)
+        task.log.exception(f"{e}. Retrying in {wait_time} seconds.")
+        raise task.retry(countdown=wait_time)

src/palace/manager/scripts/initialization.py

@@ -12,7 +12,7 @@
 
 from palace.util.log import LoggerMixin
 
-from palace.manager.celery.tasks.search import get_migrate_search_chain
+from palace.manager.celery.tasks.search import search_reindex
 from palace.manager.scripts.startup import run_startup_tasks as _run_startup_tasks
 from palace.manager.search.revision import SearchSchemaRevision
 from palace.manager.search.service import SearchService
@@ -126,10 +126,11 @@ def migrate_search(
     ) -> None:
         # The revision is not the most recent. We need to create a new index.
         # and start reindexing our data into it asynchronously. When the reindex
-        # is complete, we will switch the read pointer to the new index.
+        # is complete, we will switch the read pointer to the new index. The normal
+        # search_reindex task handles this transition.
         cls.logger().info(f"Creating a new index for revision (v{revision.version}).")
         cls.create_search_index(service, revision)
-        task = get_migrate_search_chain().apply_async()
+        task = search_reindex.apply_async()
         cls.logger().info(
             f"Task queued to index data into new search index (Task ID: {task.id})."
         )

src/palace/manager/scripts/search.py

@@ -2,7 +2,7 @@
 
 from sqlalchemy.orm import Session
 
-from palace.manager.celery.tasks.search import get_migrate_search_chain, search_reindex
+from palace.manager.celery.tasks.search import search_reindex
 from palace.manager.scripts.base import Script
 from palace.manager.search.external_search import ExternalSearchIndex
 from palace.manager.service.container import Services
@@ -23,7 +23,6 @@ def __init__(
         args = self.parse_command_line(self._db, cmd_args=cmd_args)
         self.blocking: bool = args.blocking
         self.delete: bool = args.delete
-        self.migration: bool = args.migration
 
     @classmethod
     def arg_parser(cls, _db: Session) -> argparse.ArgumentParser:
@@ -42,12 +41,6 @@ def arg_parser(cls, _db: Session) -> argparse.ArgumentParser:
             action="store_true",
             help="Delete the search index before rebuilding.",
         )
-        parser.add_argument(
-            "-m",
-            "--migration",
-            action="store_true",
-            help="Treat as a migration and update the read pointer after the rebuild is complete.",
-        )
         return parser
 
     def do_run(self) -> None:
@@ -58,10 +51,7 @@ def do_run(self) -> None:
 
         self.log.info("Rebuilding search index.")
 
-        if self.migration:
-            rebuild_task = get_migrate_search_chain()
-        else:
-            rebuild_task = search_reindex.s()
+        rebuild_task = search_reindex.s()
 
         if self.blocking:
             rebuild_task()