# This file is part of pipe_base.
#
# Developed for the LSST Data Management System.
# This product includes software developed by the LSST Project
# (http://www.lsst.org).
# See the COPYRIGHT file at the top-level directory of this distribution
# for details of code ownership.
#
# This software is dual licensed under the GNU General Public License and also
# under a 3-clause BSD license. Recipients may choose which of these licenses
# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
# respectively.  If you choose the GPL option then the following text applies
# (but note that there is still no warranty even if you opt for BSD instead):
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import annotations

__all__ = (
    "FatalWorkerError",
    "IngesterCommunicator",
    "ScannerCommunicator",
    "SupervisorCommunicator",
)

import cProfile
import dataclasses
import enum
import logging
import os
import signal
import time
import uuid
from collections.abc import Iterable, Iterator
from contextlib import ExitStack
from traceback import format_exception
from types import TracebackType
from typing import Literal, Self, overload

from lsst.utils.logging import LsstLogAdapter

from .._provenance import ProvenanceQuantumScanData
from ._config import AggregatorConfig
from ._progress import ProgressManager, make_worker_log
from ._structs import IngestRequest, ScanReport
from ._workers import Event, Queue, Worker, WorkerFactory

_TINY_TIMEOUT = 0.01


class FatalWorkerError(BaseException):
    """An exception raised by communicators when one worker (including the
    supervisor) has caught an exception in order to signal the others to shut
    down.
    """


class _WorkerCommunicationError(Exception):
    """An exception raised by communicators when a worker has died unexpectedly
    or become unresponsive.
    """


class _Sentinel(enum.Enum):
    """Sentinel values used to indicate sequence points or worker shutdown
    conditions.
    """

    NO_MORE_SCAN_REQUESTS = enum.auto()
    """Sentinel sent from the supervisor to scanners to indicate that there are
    no more quanta left to be scanned.
    """

    NO_MORE_INGEST_REQUESTS = enum.auto()
    """Sentinel sent from scanners to the ingester to indicate that there are
    will be no more ingest requests from a particular worker.
    """

    NO_MORE_WRITE_REQUESTS = enum.auto()
    """Sentinel sent from scanners and the supervisor to the writer to
    indicate that there are will be no more write requests from a particular
    worker.
    """

    WRITE_REPORT = enum.auto()
    """Sentinel sent from the writer to the supervisor to report that a
    quantum's provenance was written.
    """


@dataclasses.dataclass
class _WorkerErrorMessage:
    """An internal worker used to pass information about an error that occurred
    on a worker back to the supervisor.

    As a rule, these are unexpected, unrecoverable exceptions.
    """

    worker: str
    """Name of the originating worker."""

    traceback: str
    """A logged exception traceback.

    Note that this is not a `BaseException` subclass that can actually be
    re-raised on the supervisor; it's just something we can log to make the
    right traceback appear on the screen.  If something silences that printing
    in favor of its own exception management (pytest!) this information
    disappears.
    """


@dataclasses.dataclass
class _ScanRequest:
    """An internal struct passed from the supervisor to the scanners to request
    a quantum be scanned.
    """

    quantum_id: uuid.UUID
    """ID of the quantum to be scanned."""


@dataclasses.dataclass
class _IngestReport:
    """An internal struct passed from the ingester to the supervisor to report
    a completed ingest batch.
    """

    n_producers: int
    """Number of producing quanta whose datasets were ingested.

    We use quanta rather than datasets as the count here because the supervisor
    knows the total number of quanta in advance but not the total number of
    datasets to be ingested, so it's a lot easier to attach a denominator
    and/or progress bar to this number.
    """


@dataclasses.dataclass
class _WorkerDone:
    """An internal struct passed from a worker to the supervisor when it has
    successfully completed all work.
    """

    name: str
    """Name of the worker reporting completion."""


@dataclasses.dataclass
class _ProgressLog:
    """A high-level log message sent from a worker to the supervisor.

    These are messages that should appear to come from the main
    'aggregate-graph' logger, not a worker-specific one.
    """

    message: str
    """Log message."""

    level: int
    """Log level."""


@dataclasses.dataclass
class _CompressionDictionary:
    """An internal struct used to send the compression dictionary from the
    writer to the scanners.
    """

    data: bytes
    """The `bytes` representation of a `zstandard.ZstdCompressionDict`.
    """


type Report = (
    ScanReport
    | _IngestReport
    | _WorkerErrorMessage
    | _ProgressLog
    | _WorkerDone
    | Literal[_Sentinel.WRITE_REPORT]
)


def _disable_resources_parallelism() -> None:
    os.environ["LSST_RESOURCES_NUM_WORKERS"] = "1"
    os.environ.pop("LSST_RESOURCES_EXECUTOR", None)
    os.environ["LSST_S3_USE_THREADS"] = "False"


class SupervisorCommunicator:
    """A helper object that lets the supervisor direct the other workers.

    Parameters
    ----------
    log : `lsst.utils.logging.LsstLogAdapter`
        LSST-customized logger.
    n_scanners : `int`
        Number of scanner workers.
    worker_factory : `WorkerFactory`
        Abstraction over threading vs. multiprocessing.
    config : `AggregatorConfig`
        Configuration for the aggregator.
    """

    def __init__(
        self,
        log: LsstLogAdapter,
        n_scanners: int,
        worker_factory: WorkerFactory,
        config: AggregatorConfig,
    ) -> None:
        self.config = config
        self.progress = ProgressManager(log, config)
        self.n_scanners = n_scanners
        # The supervisor sends scan requests to scanners on this queue.
        # When complete, the supervisor sends n_scanners sentinals and each
        # scanner is careful to only take one before it starts its shutdown.
        self._scan_requests: Queue[_ScanRequest | Literal[_Sentinel.NO_MORE_SCAN_REQUESTS]] = (
            worker_factory.make_queue()
        )
        # The scanners send ingest requests to the ingester on this queue. Each
        # scanner sends one sentinal when it is done, and the ingester is
        # careful to wait for n_scanners sentinals to arrive before it starts
        # its shutdown.
        self._ingest_requests: Queue[IngestRequest | Literal[_Sentinel.NO_MORE_INGEST_REQUESTS]] = (
            worker_factory.make_queue()
        )
        # The scanners send write requests to the writer on this queue (which
        # will be `None` if we're not writing).  The supervisor also sends
        # write requests for blocked quanta (which we don't scan).  Each
        # scanner and the supervisor send one sentinal when done, and the
        # writer waits for (n_scanners + 1) sentinals to arrive before it
        # starts its shutdown.
        self._write_requests: (
            Queue[ProvenanceQuantumScanData | Literal[_Sentinel.NO_MORE_WRITE_REQUESTS]] | None
        ) = worker_factory.make_queue() if config.is_writing_provenance else None
        # All other workers use this queue to send many different kinds of
        # reports the supervisor.  The supervisor waits for a _DONE sentinal
        # from each worker before it finishes its shutdown.
        self._reports: Queue[Report] = worker_factory.make_queue()
        # The writer sends the compression dictionary to the scanners on this
        # queue.  It puts n_scanners copies on the queue, and each scanner only
        # takes one.  The compression_dict queue has no sentinal because it is
        # only used at most once; the supervisor takes responsibility for
        # clearing it out shutting down.
        self._compression_dict: Queue[_CompressionDictionary] = worker_factory.make_queue()
        # The supervisor sets this event when it receives an interrupt request
        # from an exception in the main process (usually KeyboardInterrupt).
        # Worker communicators check this in their polling loops and raise
        # FatalWorkerError when they see it set.
        self._cancel_event: Event = worker_factory.make_event()
        # Track what state we are in closing down, so we can start at the right
        # point if we're interrupted and __exit__ needs to clean up.  Note that
        # we can't rely on a non-exception __exit__ to do any shutdown work
        # that might be slow, since a KeyboardInterrupt that occurs when
        # __exit__ is already running can't be caught inside __exit__.
        self._sent_no_more_scan_requests = False
        self._sent_no_more_write_requests = False
        self._n_scanners_done = 0
        self.workers: dict[str, Worker] = {}

    def _wait_for_workers_to_finish(self, already_failing: bool = False) -> None:
        # Orderly shutdown, including exceptions: let workers clear out the
        # queues they're responsible for reading from.
        if not self._sent_no_more_scan_requests:
            for _ in range(self.n_scanners):
                self._scan_requests.put(_Sentinel.NO_MORE_SCAN_REQUESTS)
            self._sent_no_more_scan_requests = True
        if not self._sent_no_more_write_requests and self._write_requests is not None:
            self._write_requests.put(_Sentinel.NO_MORE_WRITE_REQUESTS)
            self._sent_no_more_write_requests = True
        while not all(w.successful for w in self.workers.values()):
            match self._handle_progress_reports(
                self._get_report(block=True), already_failing=already_failing
            ):
                case None | ScanReport():
                    pass
                case _WorkerDone(name=worker_name):
                    self.workers[worker_name].successful = True
                    if worker_name == IngesterCommunicator.get_worker_name():
                        self.progress.quantum_ingests.close()
                    elif worker_name == WriterCommunicator.get_worker_name():
                        self.progress.writes.close()
                    else:
                        self._n_scanners_done += 1
                        if self._n_scanners_done == self.n_scanners:
                            self.progress.scans.close()
                case unexpected:
                    raise AssertionError(f"Unexpected message {unexpected!r} to supervisor.")
            self.log.verbose(
                "Waiting for workers [%s] to report successful completion.",
                ", ".join(w.name for w in self.workers.values() if not w.successful),
            )
        self.log.verbose("Checking that all queues are empty.")
        if self._scan_requests.clear():
            self.progress.log.warning("Scan request queue was not empty at shutdown.")
            self._scan_requests.kill()
        if self._ingest_requests.clear():
            self.progress.log.warning("Ingest request queue was not empty at shutdown.")
            self._ingest_requests.kill()
        if self._write_requests is not None and self._write_requests.clear():
            self.progress.log.warning("Write request queue was not empty at shutdown.")
            self._write_requests.kill()
        if self._reports.clear():
            self.progress.log.warning("Reports queue was not empty at shutdown.")
            self._reports.kill()
        if self._compression_dict.clear():
            self.progress.log.warning("Compression dictionary queue was not empty at shutdown.")
            self._compression_dict.kill()
        for worker in self.workers.values():
            self.log.verbose("Waiting for %s to shut down.", worker.name)
            worker.join()

    def _terminate(self) -> None:
        # Disorderly shutdown: we cannot assume any of the
        # multiprocessing.Queue object work, and in fact they may hang
        # if we try to do anything with them.
        self._scan_requests.kill()
        self._ingest_requests.kill()
        if self._write_requests is not None:
            self._write_requests.kill()
        self._compression_dict.kill()
        self._reports.kill()
        for name, worker in self.workers.items():
            if worker.is_alive():
                self.progress.log.critical("Terminating worker %r.", name)
                worker.kill()

    def __enter__(self) -> Self:
        _disable_resources_parallelism()
        self.progress.__enter__()
        # We make the low-level logger in __enter__ instead of __init__ only
        # because that's the pattern used by true workers (where it matters).
        self.log = make_worker_log("supervisor", self.config)
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        if exc_type is not None:
            self._cancel_event.set()
            if exc_type is _WorkerCommunicationError:
                self.progress.log.critical("Worker '%s' was terminated before it could finish.", exc_value)
                self._terminate()
                return None
            if exc_type is not FatalWorkerError:
                self.progress.log.critical("Caught %s; attempting to shut down cleanly.", exc_type)
        try:
            self._wait_for_workers_to_finish(already_failing=exc_type is not None)
        except _WorkerCommunicationError as err:
            self.progress.log.critical(
                "Worker '%s' was terminated before it could finish (after scanning).", err
            )
            self._terminate()
            raise
        self.progress.__exit__(exc_type, exc_value, traceback)
        return None

    def request_scan(self, quantum_id: uuid.UUID) -> None:
        """Send a request to the scanners to scan the given quantum.

        Parameters
        ----------
        quantum_id : `uuid.UUID`
            ID of the quantum to scan.
        """
        self._scan_requests.put(_ScanRequest(quantum_id))

    def request_write(self, request: ProvenanceQuantumScanData) -> None:
        """Send a request to the writer to write provenance for the given scan.

        Parameters
        ----------
        request : `ProvenanceQuantumScanData`
            Information from scanning a quantum (or knowing you don't have to,
            in the case of blocked quanta).
        """
        assert self._write_requests is not None, "Writer should not be used if writing is disabled."
        self._write_requests.put(request)

    def poll(self) -> Iterator[ScanReport]:
        """Poll for reports from workers while sending scan requests.

        Yields
        ------
        scan_report : `ScanReport`
            A report from a scanner that a quantum was scanned.

        Notes
        -----
        This iterator blocks until the first scan report is received, and then
        it continues until the report queue is empty.
        """
        block = True
        while report := self._get_report(block=block):
            match self._handle_progress_reports(report):
                case ScanReport() as scan_report:
                    block = False
                    yield scan_report
                case None:
                    pass
                case unexpected:
                    raise AssertionError(f"Unexpected message {unexpected!r} to supervisor.")

    @overload
    def _get_report(self, block: Literal[True]) -> Report: ...

    @overload
    def _get_report(self, block: bool) -> Report | None: ...

    def _get_report(self, block: bool) -> Report | None:
        """Get a report from the reports queue, with timeout guards on
        blocking requests.

        This method may *return* WorkerCommunicatorError (rather than raise it)
        when a serious error occurred communicating with a subprocess.  This
        is to avoid raising an exception in an __exit__ method (which calls
        method).
        """
        report = self._reports.get(block=block, timeout=self.config.worker_check_timeout)
        while report is None and block:
            # We hit the timeout; make sure all of the workers
            # that should be alive actually are.
            for name, worker in self.workers.items():
                if not worker.successful and not worker.is_alive():
                    # Delete this worker from the list of workers so we don't
                    # hit this condition again when we try to handle the
                    # exception we raise.
                    raise _WorkerCommunicationError(name)
            # If nothing is dead and we didn't hit the hang timeout, keep
            # trying.
            report = self._reports.get(block=block, timeout=self.config.worker_check_timeout)
        return report

    def _handle_progress_reports(
        self, report: Report, already_failing: bool = False
    ) -> ScanReport | _WorkerDone | None:
        """Handle reports to the supervisor that can appear at any time, and
        are typically just updates to the progress we've made.

        This includes:

        - exceptions from workers (which raise `FatalWorkerError` here to
          trigger ``__exit__``);
        - ingest reports;
        - write reports;
        - progress logs.

        If one of these is handled, `None` is returned; otherwise the original
        report is returned.
        """
        match report:
            case _WorkerErrorMessage(traceback=traceback, worker=worker):
                self.progress.log.fatal("Exception raised on %s: \n%s", worker, traceback)
                if not already_failing:
                    raise FatalWorkerError()
            case _IngestReport(n_producers=n_producers):
                self.progress.quantum_ingests.update(n_producers)
            case _Sentinel.WRITE_REPORT:
                self.progress.writes.update(1)
            case _ProgressLog(message=message, level=level):
                self.progress.log.log(level, "%s [after %0.1fs]", message, self.progress.elapsed_time)
            case _:
                return report
        return None


class WorkerCommunicator:
    """A base class for non-supervisor worker communicators.

    Parameters
    ----------
    supervisor : `SupervisorCommunicator`
        Communicator for the supervisor to grab queues and information from.
    name : `str`
        Human-readable name for this worker.

    Notes
    -----
    Each worker communicator is constructed in the main process and entered as
    a context manager *only* on the actual worker process, so attributes that
    cannot be pickled are constructed in ``__enter__`` instead of ``__init__``.

    Worker communicators provide access to an `AggregatorConfig` and a logger
    to their workers.  As context managers, they handle exceptions and ensure
    clean shutdowns, and since most workers need to use a lot of other context
    managers (for file reading and writing, mostly), they provide an `enter`
    method to keep every worker from also having to be a context manager just
    to hold a context manager instance attribute.

    Worker communicators can also be configured to record and dump profiling
    information.
    """

    def __init__(self, supervisor: SupervisorCommunicator, name: str):
        self.name = name
        self.config = supervisor.config
        self._reports = supervisor._reports
        self._cancel_event = supervisor._cancel_event

    def __enter__(self) -> Self:
        _disable_resources_parallelism()
        self.log = make_worker_log(self.name, self.config)
        self.log.verbose("%s has PID %s (parent is %s).", self.name, os.getpid(), os.getppid())
        self._exit_stack = ExitStack().__enter__()
        if self.config.n_processes > 1:
            # Multiprocessing: ignore interrupts so we can shut down cleanly.
            signal.signal(signal.SIGINT, signal.SIG_IGN)
            if self.config.worker_profile_dir is not None:
                # We use time.time because we're interested in wall-clock time,
                # not just CPU effort, since this is I/O-bound work.
                self._profiler = cProfile.Profile(timer=time.time)
                self._profiler.enable()
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> bool | None:
        if self.config.worker_profile_dir is not None and self.config.n_processes > 1:
            self._profiler.disable()
            os.makedirs(self.config.worker_profile_dir, exist_ok=True)
            self._profiler.dump_stats(os.path.join(self.config.worker_profile_dir, f"{self.name}.profile"))
        if exc_value is not None:
            assert exc_type is not None, "Should be guaranteed by Python, but MyPy doesn't know that."
            if exc_type is not FatalWorkerError:
                self.log.warning("Error raised on this worker.", exc_info=(exc_type, exc_value, traceback))
                assert exc_type is not None and traceback is not None
                self._reports.put(
                    _WorkerErrorMessage(
                        self.name,
                        "".join(format_exception(exc_type, exc_value, traceback)),
                    )
                )
                self.log.debug("Error message sent to supervisor.")
            else:
                self.log.warning("Shutting down due to exception raised on another worker.")
        self._exit_stack.__exit__(exc_type, exc_value, traceback)
        return True

    @property
    def exit_stack(self) -> ExitStack:
        """A `contextlib.ExitStack` tied to the communicator."""
        return self._exit_stack

    def log_progress(self, level: int, message: str) -> None:
        """Send a high-level log message to the supervisor.

        Parameters
        ----------
        level : `int`
            Log level.  Should be ``VERBOSE`` or higher.
        message : `str`
            Log message.
        """
        self._reports.put(_ProgressLog(message=message, level=level))

    def check_for_cancel(self) -> None:
        """Check for a cancel signal from the supervisor and raise
        `FatalWorkerError` if it is present.
        """
        if self._cancel_event.is_set():
            raise FatalWorkerError()


class ScannerCommunicator(WorkerCommunicator):
    """A communicator for scanner workers.

    Parameters
    ----------
    supervisor : `SupervisorCommunicator`
        Communicator for the supervisor to grab queues and information from.
    scanner_id : `int`
        Integer ID for this canner.
    """

    def __init__(self, supervisor: SupervisorCommunicator, scanner_id: int):
        super().__init__(supervisor, self.get_worker_name(scanner_id))
        self.scanner_id = scanner_id
        self._scan_requests = supervisor._scan_requests
        self._ingest_requests = supervisor._ingest_requests
        self._write_requests = supervisor._write_requests
        self._compression_dict = supervisor._compression_dict
        self._got_no_more_scan_requests: bool = False
        self._sent_no_more_ingest_requests: bool = False

    @staticmethod
    def get_worker_name(scanner_id: int) -> str:
        return f"scanner-{scanner_id:03d}"

    def report_scan(self, msg: ScanReport) -> None:
        """Report a completed scan to the supervisor.

        Parameters
        ----------
        msg : `ScanReport`
            Report to send.
        """
        self._reports.put(msg)

    def request_ingest(self, request: IngestRequest) -> None:
        """Ask the ingester to ingest a quantum's outputs.

        Parameters
        ----------
        request : `IngestRequest`
            Description of the datasets to ingest.

        Notes
        -----
        If this request has no datasets, this automatically reports the ingest
        as complete to the supervisor instead of sending it to the ingester.
        """
        if request:
            self._ingest_requests.put(request)
        else:
            self._reports.put(_IngestReport(1))

    def request_write(self, request: ProvenanceQuantumScanData) -> None:
        """Ask the writer to write provenance for a quantum.

        Parameters
        ----------
        request : `ProvenanceQuantumScanData`
            Result of scanning a quantum.
        """
        assert self._write_requests is not None, "Writer should not be used if writing is disabled."
        self._write_requests.put(request)

    def get_compression_dict(self) -> bytes | None:
        """Attempt to get the compression dict from the writer.

        Returns
        -------
        data : `bytes` or `None`
            The `bytes` representation of the compression dictionary, or `None`
            if the compression dictionary is not yet available.

        Notes
        -----
        A scanner should only call this method before it actually has the
        compression dict.
        """
        if (cdict := self._compression_dict.get()) is not None:
            return cdict.data
        return None

    def poll(self) -> Iterator[uuid.UUID]:
        """Poll for scan requests to process.

        Yields
        ------
        quantum_id : `uuid.UUID`
            ID of a new quantum to scan.

        Notes
        -----
        This iterator ends when the supervisor reports that it is done
        traversing the graph.
        """
        while True:
            self.check_for_cancel()
            scan_request = self._scan_requests.get(block=True, timeout=self.config.worker_sleep)
            if scan_request is _Sentinel.NO_MORE_SCAN_REQUESTS:
                self._got_no_more_scan_requests = True
                return
            if scan_request is not None:
                yield scan_request.quantum_id

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> bool | None:
        result = super().__exit__(exc_type, exc_value, traceback)
        self._ingest_requests.put(_Sentinel.NO_MORE_INGEST_REQUESTS)
        if self._write_requests is not None:
            self._write_requests.put(_Sentinel.NO_MORE_WRITE_REQUESTS)
        while not self._got_no_more_scan_requests:
            if (
                not self._got_no_more_scan_requests
                and self._scan_requests.get(block=True) is _Sentinel.NO_MORE_SCAN_REQUESTS
            ):
                self._got_no_more_scan_requests = True
        # We let the writer clear out the compression dict queue.
        self.log.verbose("Sending completion message.")
        self._reports.put(_WorkerDone(self.name))
        return result


class IngesterCommunicator(WorkerCommunicator):
    """A communicator for the ingester worker.

    Parameters
    ----------
    supervisor : `SupervisorCommunicator`
        Communicator for the supervisor to grab queues and information from.
    """

    def __init__(self, supervisor: SupervisorCommunicator):
        super().__init__(supervisor, self.get_worker_name())
        self.n_scanners = supervisor.n_scanners
        self._ingest_requests = supervisor._ingest_requests
        self._n_requesters_done = 0

    @staticmethod
    def get_worker_name() -> str:
        return "ingester"

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> bool | None:
        result = super().__exit__(exc_type, exc_value, traceback)
        while self._n_requesters_done != self.n_scanners:
            self.log.debug(
                "Waiting for %d requesters to be done (currently %d).",
                self.n_scanners,
                self._n_requesters_done,
            )
            if self._ingest_requests.get(block=True) is _Sentinel.NO_MORE_INGEST_REQUESTS:
                self._n_requesters_done += 1
        self.log.verbose("Sending completion message.")
        self._reports.put(_WorkerDone(self.name))
        return result

    def report_ingest(self, n_producers: int) -> None:
        """Report to the supervisor that an ingest batch was completed.

        Parameters
        ----------
        n_producers : `int`
            Number of producing quanta whose datasets were ingested.
        """
        self._reports.put(_IngestReport(n_producers))

    def poll(self) -> Iterator[IngestRequest]:
        """Poll for ingest requests from the scanner workers.

        Yields
        ------
        request : `IngestRequest`
            A request to ingest datasets produced by a single quantum.

        Notes
        -----
        This iterator ends when all scanners indicate that they are done making
        ingest requests.
        """
        while True:
            self.check_for_cancel()
            ingest_request = self._ingest_requests.get(block=True, timeout=_TINY_TIMEOUT)
            if ingest_request is _Sentinel.NO_MORE_INGEST_REQUESTS:
                self._n_requesters_done += 1
                if self._n_requesters_done == self.n_scanners:
                    return
                else:
                    continue
            if ingest_request is not None:
                yield ingest_request


class WriterCommunicator(WorkerCommunicator):
    """A communicator for the writer worker.

    Parameters
    ----------
    supervisor : `SupervisorCommunicator`
        Communicator for the supervisor to grab queues and information from.
    """

    def __init__(self, supervisor: SupervisorCommunicator):
        assert supervisor._write_requests is not None
        super().__init__(supervisor, self.get_worker_name())
        self.n_scanners = supervisor.n_scanners
        self._write_requests = supervisor._write_requests
        self._compression_dict = supervisor._compression_dict
        self._n_requesters = supervisor.n_scanners + 1
        self._n_requesters_done = 0
        self._sent_compression_dict = False

    @staticmethod
    def get_worker_name() -> str:
        return "writer"

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> bool | None:
        result = super().__exit__(exc_type, exc_value, traceback)
        if exc_type is None:
            self.log_progress(logging.INFO, "Provenance quantum graph written successfully.")
        while self._n_requesters_done != self._n_requesters:
            self.log.debug(
                "Waiting for %d requesters to be done (currently %d).",
                self._n_requesters,
                self._n_requesters_done,
            )
            if self._write_requests.get(block=True) is _Sentinel.NO_MORE_WRITE_REQUESTS:
                self._n_requesters_done += 1
        if self._compression_dict.clear():
            self.log.verbose("Cleared out compression dictionary queue.")
        else:
            self.log.verbose("Compression dictionary queue was already empty.")
        self.log.verbose("Sending completion message.")
        self._reports.put(_WorkerDone(self.name))
        return result

    def poll(self) -> Iterator[ProvenanceQuantumScanData]:
        """Poll for writer requests from the scanner workers and supervisor.

        Yields
        ------
        request : `ProvenanceQuantumScanData`
            The result of a quantum scan.

        Notes
        -----
        This iterator ends when all scanners and the supervisor indicate that
        they are done making write requests.
        """
        while True:
            self.check_for_cancel()
            write_request = self._write_requests.get(block=True, timeout=_TINY_TIMEOUT)
            if write_request is _Sentinel.NO_MORE_WRITE_REQUESTS:
                self._n_requesters_done += 1
                if self._n_requesters_done == self._n_requesters:
                    return
                else:
                    continue
            if write_request is not None:
                yield write_request

    def send_compression_dict(self, cdict_data: bytes) -> None:
        """Send the compression dictionary to the scanners.

        Parameters
        ----------
        cdict_data : `bytes`
            The `bytes` representation of the compression dictionary.
        """
        self.log.debug("Sending compression dictionary.")
        for _ in range(self.n_scanners):
            self._compression_dict.put(_CompressionDictionary(cdict_data))
        self._sent_compression_dict = True

    def report_write(self) -> None:
        """Report to the supervisor that provenance for a quantum was written
        to the graph.
        """
        self._reports.put(_Sentinel.WRITE_REPORT)

    def periodically_check_for_cancel[T](self, iterable: Iterable[T], n: int = 100) -> Iterator[T]:
        """Iterate while checking for a cancellation signal every ``n``
        iterations.

        Parameters
        ----------
        iterable : `~collections.abc.Iterable`
            Object to iterate over.
        n : `int`
            Check for cancellation every ``n`` iterations.

        Returns
        -------
        iterator : `~collections.abc.Iterator`
            Iterator.
        """
        i = 0
        for entry in iterable:
            yield entry
            i += 1
            if i % n == 0:
                self.check_for_cancel()
