Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 81 additions & 1 deletion clients/python/src/taskbroker_client/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
CompressionType,
)
from taskbroker_client.metrics import MetricsBackend
from taskbroker_client.retry import Retry
from taskbroker_client.retry import LastAction, Retry
from taskbroker_client.router import TaskRouter
from taskbroker_client.sdk import start_span
from taskbroker_client.task import ExternalTask, P, R, Task
Expand Down Expand Up @@ -47,13 +47,15 @@ def __init__(
processing_deadline_duration: int = DEFAULT_PROCESSING_DEADLINE,
app_feature: str | None = None,
context_hooks: list[ContextHook] | None = None,
is_raw_mode: bool = False,
):
self.name = name
self.application = application
self.router = router
self.default_retry = retry
self.default_expires = expires # seconds
self.default_processing_deadline_duration = processing_deadline_duration # seconds
self.is_raw_mode = is_raw_mode
self.app_feature = app_feature or name
self.context_hooks: list[ContextHook] = context_hooks or []
self._registered_tasks: dict[str, Task[Any, Any]] = {}
Expand Down Expand Up @@ -82,6 +84,66 @@ def topic(self) -> str:
"""The topic that a namespace is routed to."""
return self.router.route_namespace(self.name)

def _validate_raw_mode(
self,
*,
name: str,
retry: Retry | None,
expires: int | datetime.timedelta | None,
processing_deadline_duration: int | datetime.timedelta | None,
at_most_once: bool,
compression_type: CompressionType,
) -> None:
"""
Reject task options that taskbroker ignores in raw mode.

Nothing produces a raw-mode task: taskbroker consumes the topic and builds the
activation itself. Options that a producer would normally embed in the
activation therefore never reach the broker. Some of them it takes from the
topic's `raw:` config instead, and the rest it hardcodes. Either way the
decorator value is inert, so reject it here instead of dropping it silently at
runtime.
"""
if self._registered_tasks:
registered = ", ".join(sorted(self._registered_tasks))
raise ValueError(
f"Raw-mode namespace {self.name!r} already has a registered task "
f"({registered}). taskbroker only spawns the single taskname configured "
f"for the raw topic, so {name!r} would never run. Give it its own namespace."
)

from_topic_config = []
if expires is not None or self.default_expires is not None:
from_topic_config.append("expires")
if (
processing_deadline_duration is not None
or self.default_processing_deadline_duration != DEFAULT_PROCESSING_DEADLINE
):
from_topic_config.append("processing_deadline_duration")
if compression_type != CompressionType.PLAINTEXT:
from_topic_config.append("compression_type")
if from_topic_config:
raise ValueError(
f"Task {name!r} in raw-mode namespace {self.name!r} sets "
f"{', '.join(from_topic_config)} (on the task or as a namespace default). "
"In raw mode taskbroker takes these from the topic's `raw:` block in its "
"own config, not from the decorator, so the value here has no effect and "
"will silently drift from what is deployed."
)

if at_most_once:
raise ValueError(
f"Task {name!r} in raw-mode namespace {self.name!r} sets at_most_once. "
"taskbroker hardcodes at_most_once=false for raw activations, so the task "
"would still be retried past its processing deadline."
)
if retry is not None and retry._times_exceeded == LastAction.Deadletter:
raise ValueError(
Comment on lines +140 to +141

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: The deadletter validation in _validate_raw_mode doesn't check the namespace's default_retry policy when a task is registered with at_most_once=True, allowing an invalid configuration.
Severity: MEDIUM

Suggested Fix

Update _validate_raw_mode to check both the passed retry parameter and the namespace's self.default_retry policy for LastAction.Deadletter. This ensures any deadletter configuration is caught for raw mode tasks, regardless of the at_most_once setting.

Prompt for AI Agent
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.

Location: clients/python/src/taskbroker_client/registry.py#L140-L141

Potential issue: When registering a raw mode task with `at_most_once=True` in a
namespace that has a `default_retry` policy configured with `LastAction.Deadletter`, the
validation logic fails to detect this conflict. The `_validate_raw_mode` function only
checks the `retry` policy passed to it, which is `None` in this scenario because
`at_most_once` prevents the namespace default from being applied to the task. However,
the deadlettering configuration at the namespace level is still incompatible with raw
mode and should be flagged. This leads to an incomplete validation, masking a
misconfiguration.

Did we get this right? 👍 / 👎 to inform future reviews.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think that matters. If at_most_once is set, an error will be raised and developers can iterate their way to correctness.

f"Task {name!r} in raw-mode namespace {self.name!r} sets "
"retry times_exceeded=LastAction.Deadletter. taskbroker hardcodes Discard "
"for raw activations, so exhausted tasks would be dropped, not deadlettered."
)

def register(
self,
*,
Expand Down Expand Up @@ -135,6 +197,15 @@ def wrapped(func: Callable[P, R]) -> Task[P, R]:
task_retry = retry
if not at_most_once:
task_retry = retry or self.default_retry
if self.is_raw_mode:
self._validate_raw_mode(
name=name,
retry=task_retry,
expires=expires,
processing_deadline_duration=processing_deadline_duration,
at_most_once=at_most_once,
compression_type=compression_type,
)
task = Task(
name=name,
func=func,
Expand Down Expand Up @@ -355,6 +426,7 @@ def create_namespace(
processing_deadline_duration: int = DEFAULT_PROCESSING_DEADLINE,
app_feature: str | None = None,
internal: bool = False,
is_raw_mode: bool = False,
) -> TaskNamespace:
"""
Create a task namespace.
Expand All @@ -363,6 +435,13 @@ def create_namespace(
infrastructure to be scaled based on a region's requirements.

Namespaces can define default behavior for tasks defined within a namespace.

Set `is_raw_mode` when the namespace backs a topic that taskbroker consumes in
"raw mode", where taskbroker builds activations from raw Kafka messages itself
rather than a producer sending them. A raw topic maps 1:1 onto a namespace and
onto a single task, both pinned in taskbroker's own config. The flag makes that
contract explicit and rejects task options taskbroker would ignore; see
`TaskNamespace._validate_raw_mode`.
"""
if name == INTERNAL_NAMESPACE and not internal:
raise ValueError(f"{INTERNAL_NAMESPACE!r} is reserved for internal taskbroker tasks.")
Expand All @@ -379,6 +458,7 @@ def create_namespace(
processing_deadline_duration=processing_deadline_duration,
app_feature=app_feature,
context_hooks=self._context_hooks,
is_raw_mode=is_raw_mode,
)
self._namespaces[name] = namespace

Expand Down
117 changes: 117 additions & 0 deletions clients/python/tests/test_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -398,3 +398,120 @@ def test_registry_create_namespace_duplicate() -> None:
registry.create_namespace(name="tests")
with pytest.raises(ValueError, match="tests already exists"):
registry.create_namespace(name="tests")


def raw_namespace(retry: Retry | None = None, **kwargs: Any) -> TaskNamespace:
return TaskNamespace(
name="tests.raw",
application="acme",
producer_factory=producer_factory,
router=DefaultRouter(),
metrics=NoOpMetricsBackend(),
retry=retry,
is_raw_mode=True,
**kwargs,
)


def test_raw_mode_namespace_allows_plain_task() -> None:
namespace = raw_namespace()

@namespace.register(name="tests.raw.consume", retry=Retry(times=3, delay=5))
def consume(message_bytes: bytes) -> None:
raise NotImplementedError

assert namespace.is_raw_mode
assert namespace.contains("tests.raw.consume")


def test_raw_mode_namespace_rejects_second_task() -> None:
namespace = raw_namespace()

@namespace.register(name="tests.raw.consume")
def consume(message_bytes: bytes) -> None:
raise NotImplementedError

with pytest.raises(ValueError, match="would never run"):

@namespace.register(name="tests.raw.other")
def other(message_bytes: bytes) -> None:
raise NotImplementedError


@pytest.mark.parametrize(
"kwargs",
[
pytest.param({"expires": 60}, id="expires"),
pytest.param({"processing_deadline_duration": 90}, id="processing_deadline_duration"),
pytest.param({"compression_type": CompressionType.ZSTD}, id="compression_type"),
],
)
def test_raw_mode_namespace_rejects_topic_config_params(kwargs: dict[str, Any]) -> None:
namespace = raw_namespace()

with pytest.raises(ValueError, match="takes these from the topic's `raw:` block"):

@namespace.register(name="tests.raw.consume", **kwargs)
def consume(message_bytes: bytes) -> None:
raise NotImplementedError


@pytest.mark.parametrize(
"namespace_kwargs",
[
pytest.param({"expires": 60}, id="expires"),
pytest.param({"processing_deadline_duration": 90}, id="processing_deadline_duration"),
],
)
def test_raw_mode_namespace_rejects_topic_config_defaults(namespace_kwargs: dict[str, Any]) -> None:
namespace = raw_namespace(**namespace_kwargs)

with pytest.raises(ValueError, match="takes these from the topic's `raw:` block"):

@namespace.register(name="tests.raw.consume")
def consume(message_bytes: bytes) -> None:
raise NotImplementedError


def test_raw_mode_namespace_rejects_at_most_once() -> None:
namespace = raw_namespace()

with pytest.raises(ValueError, match="hardcodes at_most_once=false"):

@namespace.register(name="tests.raw.consume", at_most_once=True)
def consume(message_bytes: bytes) -> None:
raise NotImplementedError


def test_raw_mode_namespace_rejects_deadletter() -> None:
namespace = raw_namespace()

with pytest.raises(ValueError, match="hardcodes Discard"):

@namespace.register(
name="tests.raw.consume",
retry=Retry(times=3, times_exceeded=LastAction.Deadletter),
)
def consume(message_bytes: bytes) -> None:
raise NotImplementedError


def test_raw_mode_namespace_rejects_deadletter_from_namespace_default() -> None:
namespace = raw_namespace(retry=Retry(times=3, times_exceeded=LastAction.Deadletter))

with pytest.raises(ValueError, match="hardcodes Discard"):

@namespace.register(name="tests.raw.consume")
def consume(message_bytes: bytes) -> None:
raise NotImplementedError


def test_registry_create_namespace_is_raw_mode() -> None:
registry = TaskRegistry(
application="acme",
producer_factory=producer_factory,
router=DefaultRouter(),
metrics=NoOpMetricsBackend(),
)
assert not registry.create_namespace("tests").is_raw_mode
assert registry.create_namespace("tests.raw", is_raw_mode=True).is_raw_mode
Loading