From b75ab56684cccb09f6911f728a2a3d6f48fb4b1c Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Sat, 7 Nov 2026 03:59:31 +0000 Subject: [PATCH] feat(cloud_scheduler): add sample that prevents duplicate work from job requests Add cloudscheduler_deduplicate_job_requests, a Cloud Run function that does a Cloud Scheduler job's work once per scheduled run, even when Cloud Scheduler delivers the same run more than once. The function identifies a run by the job name and the X-CloudScheduler-ScheduleTime header, which stays the same across retries. It claims the run in a Firestore transaction with a lease: - A completed run is skipped. - An active claim returns 409. - An expired claim from a crashed attempt can be claimed again. If the work fails, the function releases the claim so that Cloud Scheduler's retry can process the run. This rewrites the sample proposed in #14702 and keeps its region tag. Unit tests use an in-memory fake Firestore. Integration tests run against a real Firestore database and are skipped unless RUN_INTEGRATION_TESTS is set. Co-authored-by: Camie Kim <12733285+camiekim@users.noreply.github.com> --- .../deduplicate_job_requests/.python-version | 1 + .../deduplicate_job_requests/__init__.py | 13 ++ .../deduplicate_job_requests/main.py | 90 ++++++++++ .../noxfile_config.py | 42 +++++ .../requirements-test.txt | 1 + .../deduplicate_job_requests/requirements.txt | 2 + .../tests/__init__.py | 13 ++ .../tests/conftest.py | 22 +++ .../tests/integration_test.py | 95 ++++++++++ .../tests/main_test.py | 170 ++++++++++++++++++ 10 files changed, 449 insertions(+) create mode 100644 cloud_scheduler/deduplicate_job_requests/.python-version create mode 100644 cloud_scheduler/deduplicate_job_requests/__init__.py create mode 100644 cloud_scheduler/deduplicate_job_requests/main.py create mode 100644 cloud_scheduler/deduplicate_job_requests/noxfile_config.py create mode 100644 cloud_scheduler/deduplicate_job_requests/requirements-test.txt create mode 100644 cloud_scheduler/deduplicate_job_requests/requirements.txt create mode 100644 cloud_scheduler/deduplicate_job_requests/tests/__init__.py create mode 100644 cloud_scheduler/deduplicate_job_requests/tests/conftest.py create mode 100644 cloud_scheduler/deduplicate_job_requests/tests/integration_test.py create mode 100644 cloud_scheduler/deduplicate_job_requests/tests/main_test.py diff --git a/cloud_scheduler/deduplicate_job_requests/.python-version b/cloud_scheduler/deduplicate_job_requests/.python-version new file mode 100644 index 00000000000..24ee5b1be99 --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/.python-version @@ -0,0 +1 @@ +3.13 diff --git a/cloud_scheduler/deduplicate_job_requests/__init__.py b/cloud_scheduler/deduplicate_job_requests/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/cloud_scheduler/deduplicate_job_requests/main.py b/cloud_scheduler/deduplicate_job_requests/main.py new file mode 100644 index 00000000000..f5832dc7f5f --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/main.py @@ -0,0 +1,90 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START cloudscheduler_deduplicate_job_requests] +import datetime +import hashlib + +import flask +import functions_framework +from google.cloud import firestore + +# Match the job's attempt deadline, so a crashed attempt's claim expires. +LEASE = datetime.timedelta(minutes=5) +# Keep records past the job's retry window. A TTL policy on expire_at deletes them. +RETENTION = datetime.timedelta(days=1) + +db = firestore.Client() + + +def run_job_logic() -> None: + """Does the scheduled work. Replace the body with your own task.""" + + +@firestore.transactional +def claim_run( + transaction: firestore.Transaction, + run_ref: firestore.DocumentReference, + now: datetime.datetime, +) -> str: + """Claims the run unless it's done or another attempt holds the lease.""" + snapshot = run_ref.get(transaction=transaction) + if snapshot.exists: + run = snapshot.to_dict() + if run["status"] == "completed": + return "completed" + if run["lease_expires"] > now: + return "in_progress" + transaction.set( + run_ref, + { + "status": "in_progress", + "lease_expires": now + LEASE, + "expire_at": now + RETENTION, + }, + ) + return "claimed" + + +@functions_framework.http +def handle_scheduled_job(request: flask.Request) -> tuple[str, int]: + """Runs a job's work once per scheduled run, even if the run is delivered twice.""" + job_name = request.headers.get("X-CloudScheduler-JobName") + schedule_time = request.headers.get("X-CloudScheduler-ScheduleTime") + if not job_name or not schedule_time: + return "error: missing Cloud Scheduler headers", 400 + + # Retries keep the same schedule time. A forced run uses the next + # scheduled time, so that scheduled run is then skipped. + run_id = hashlib.sha256(f"{job_name}/{schedule_time}".encode()).hexdigest() + run_ref = db.collection("scheduler_runs").document(run_id) + now = datetime.datetime.now(datetime.timezone.utc) + + status = claim_run(db.transaction(), run_ref, now) + if status == "completed": + return "Duplicate run skipped", 200 + if status == "in_progress": + # A non-2xx response keeps retries going in case that attempt fails. + return "Run in progress", 409 + + try: + run_job_logic() + except Exception: + run_ref.delete() # Lets a retry claim the run. + raise + run_ref.update({"status": "completed"}) + return "Run processed", 200 + + +# [END cloudscheduler_deduplicate_job_requests] diff --git a/cloud_scheduler/deduplicate_job_requests/noxfile_config.py b/cloud_scheduler/deduplicate_job_requests/noxfile_config.py new file mode 100644 index 00000000000..95c2d9ac875 --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/noxfile_config.py @@ -0,0 +1,42 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Default TEST_CONFIG_OVERRIDE for python repos. + +# You can copy this file into your directory, then it will be imported from +# the noxfile.py. + +# The source of truth: +# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py + +TEST_CONFIG_OVERRIDE = { + # You can opt out from the test for specific Python versions. + "ignored_versions": ["2.7", "3.7", "3.8", "3.9", "3.10", "3.11", "3.12"], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": True, + # An envvar key for determining the project id to use. Change it + # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a + # build specific Cloud project. You can also use your own string + # to use your own Cloud project. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT', + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. Don't put any + # secrets here. These values will override predefined values. + "envs": {}, +} diff --git a/cloud_scheduler/deduplicate_job_requests/requirements-test.txt b/cloud_scheduler/deduplicate_job_requests/requirements-test.txt new file mode 100644 index 00000000000..b3b2b3ca4cc --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/requirements-test.txt @@ -0,0 +1 @@ +pytest==9.0.3 diff --git a/cloud_scheduler/deduplicate_job_requests/requirements.txt b/cloud_scheduler/deduplicate_job_requests/requirements.txt new file mode 100644 index 00000000000..6d4f99e4b72 --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/requirements.txt @@ -0,0 +1,2 @@ +functions-framework==3.10.2 +google-cloud-firestore==2.34.1 diff --git a/cloud_scheduler/deduplicate_job_requests/tests/__init__.py b/cloud_scheduler/deduplicate_job_requests/tests/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/cloud_scheduler/deduplicate_job_requests/tests/conftest.py b/cloud_scheduler/deduplicate_job_requests/tests/conftest.py new file mode 100644 index 00000000000..b68ff8d1eb5 --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/tests/conftest.py @@ -0,0 +1,22 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from unittest import mock + +# main.py creates a Firestore client when it's imported. Without credentials +# that fails, so unit-test runs replace the client class. The unit tests then +# swap in a fake database, and the integration tests use a real one. +if os.environ.get("RUN_INTEGRATION_TESTS") != "true": + mock.patch("google.cloud.firestore.Client").start() diff --git a/cloud_scheduler/deduplicate_job_requests/tests/integration_test.py b/cloud_scheduler/deduplicate_job_requests/tests/integration_test.py new file mode 100644 index 00000000000..f22b98292cd --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/tests/integration_test.py @@ -0,0 +1,95 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests that run the handler against a real Firestore database. + +These tests are skipped unless RUN_INTEGRATION_TESTS is "true". They need a +Firestore database in Native mode; set FIRESTORE_DATABASE to use a named +database instead of "(default)". +""" + +import datetime +import hashlib +import os +from typing import Iterator +from unittest import mock +import uuid + +import flask +from google.cloud import firestore +import pytest + +from .. import main + +pytestmark = pytest.mark.skipif( + os.environ.get("RUN_INTEGRATION_TESTS") != "true", + reason="Set RUN_INTEGRATION_TESTS=true to run integration tests.", +) + +app = flask.Flask(__name__) + + +def run_ref(headers: dict[str, str]) -> firestore.DocumentReference: + """Returns the record that the handler keeps for a run.""" + key = ( + f"{headers['X-CloudScheduler-JobName']}/" + f"{headers['X-CloudScheduler-ScheduleTime']}" + ) + run_id = hashlib.sha256(key.encode()).hexdigest() + return main.db.collection("scheduler_runs").document(run_id) + + +@pytest.fixture +def headers(monkeypatch: pytest.MonkeyPatch) -> Iterator[dict[str, str]]: + """Uses a real database and a unique run, and deletes the run's record.""" + database = firestore.Client( + database=os.environ.get("FIRESTORE_DATABASE", "(default)") + ) + monkeypatch.setattr(main, "db", database) + run_headers = { + "X-CloudScheduler-JobName": f"test-job-{uuid.uuid4().hex[:12]}", + "X-CloudScheduler-ScheduleTime": "2026-11-06T19:29:01.336893-08:00", + } + yield run_headers + run_ref(run_headers).delete() + + +def call_function(headers: dict[str, str]) -> tuple[str, int]: + with app.test_request_context("/", method="POST", headers=headers): + return main.handle_scheduled_job(flask.request) + + +def test_duplicate_delivery_is_skipped(headers: dict[str, str]) -> None: + assert call_function(headers) == ("Run processed", 200) + assert call_function(headers) == ("Duplicate run skipped", 200) + assert run_ref(headers).get().to_dict()["status"] == "completed" + + +def test_failed_work_releases_the_claim(headers: dict[str, str]) -> None: + with mock.patch.object(main, "run_job_logic", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError): + call_function(headers) + + assert not run_ref(headers).get().exists + assert call_function(headers) == ("Run processed", 200) + + +def test_active_and_expired_leases(headers: dict[str, str]) -> None: + now = datetime.datetime.now(datetime.timezone.utc) + run_ref(headers).set({"status": "in_progress", "lease_expires": now + main.LEASE}) + assert call_function(headers) == ("Run in progress", 409) + + expired = now - datetime.timedelta(seconds=1) + run_ref(headers).set({"status": "in_progress", "lease_expires": expired}) + assert call_function(headers) == ("Run processed", 200) diff --git a/cloud_scheduler/deduplicate_job_requests/tests/main_test.py b/cloud_scheduler/deduplicate_job_requests/tests/main_test.py new file mode 100644 index 00000000000..23b12feae9e --- /dev/null +++ b/cloud_scheduler/deduplicate_job_requests/tests/main_test.py @@ -0,0 +1,170 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +from typing import Any +from unittest import mock + +import flask +import pytest + +from .. import main + +app = flask.Flask(__name__) + +HEADERS = { + "X-CloudScheduler-JobName": "test-job", + "X-CloudScheduler-ScheduleTime": "2026-11-06T19:29:01.336893-08:00", +} + + +class FakeSnapshot: + def __init__(self, data: dict[str, Any] | None) -> None: + self._data = data + self.exists = data is not None + + def to_dict(self) -> dict[str, Any]: + return dict(self._data or {}) + + +class FakeDocument: + """A Firestore document reference backed by a dictionary.""" + + def __init__(self, store: dict[str, dict[str, Any]], doc_id: str) -> None: + self.store = store + self.id = doc_id + + def get(self, transaction: object = None) -> FakeSnapshot: + return FakeSnapshot(self.store.get(self.id)) + + def update(self, fields: dict[str, Any]) -> None: + self.store[self.id].update(fields) + + def delete(self) -> None: + self.store.pop(self.id, None) + + +class FakeDatabase: + def __init__(self) -> None: + self.store: dict[str, dict[str, Any]] = {} + + def collection(self, name: str) -> mock.MagicMock: + collection = mock.MagicMock() + collection.document.side_effect = lambda doc_id: FakeDocument( + self.store, doc_id + ) + return collection + + def transaction(self) -> mock.MagicMock: + # firestore.transactional begins, commits, and retries the transaction. + transaction = mock.MagicMock(_max_attempts=1, _id=b"test", _read_only=False) + transaction.set.side_effect = lambda ref, data: self.store.__setitem__( + ref.id, dict(data) + ) + return transaction + + +@pytest.fixture(autouse=True) +def fake_db(monkeypatch: pytest.MonkeyPatch) -> FakeDatabase: + database = FakeDatabase() + monkeypatch.setattr(main, "db", database) + return database + + +def call_function(headers: dict[str, str] | None = None) -> tuple[str, int]: + with app.test_request_context("/", method="POST", headers=headers): + return main.handle_scheduled_job(flask.request) + + +def only_record(fake_db: FakeDatabase) -> dict[str, Any]: + assert len(fake_db.store) == 1 + return next(iter(fake_db.store.values())) + + +def test_first_delivery_runs_the_work(fake_db: FakeDatabase) -> None: + with mock.patch.object(main, "run_job_logic") as run_job_logic: + body, status = call_function(HEADERS) + + assert (body, status) == ("Run processed", 200) + run_job_logic.assert_called_once() + record = only_record(fake_db) + assert record["status"] == "completed" + assert record["expire_at"] > record["lease_expires"] + + +def test_duplicate_after_completion_is_skipped(fake_db: FakeDatabase) -> None: + call_function(HEADERS) + + with mock.patch.object(main, "run_job_logic") as run_job_logic: + body, status = call_function(HEADERS) + + assert (body, status) == ("Duplicate run skipped", 200) + run_job_logic.assert_not_called() + + +def test_attempt_with_active_lease_returns_409(fake_db: FakeDatabase) -> None: + with mock.patch.object(main, "run_job_logic", side_effect=lambda: None): + # Simulate another attempt that holds the lease. + future = datetime.datetime.now(datetime.timezone.utc) + main.LEASE + call_function(HEADERS) + record = only_record(fake_db) + record.update(status="in_progress", lease_expires=future) + + body, status = call_function(HEADERS) + + assert (body, status) == ("Run in progress", 409) + + +def test_expired_lease_from_crashed_attempt_is_reclaimed( + fake_db: FakeDatabase, +) -> None: + call_function(HEADERS) + past = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=1) + only_record(fake_db).update(status="in_progress", lease_expires=past) + + with mock.patch.object(main, "run_job_logic") as run_job_logic: + body, status = call_function(HEADERS) + + assert (body, status) == ("Run processed", 200) + run_job_logic.assert_called_once() + assert only_record(fake_db)["status"] == "completed" + + +def test_failed_work_releases_the_claim(fake_db: FakeDatabase) -> None: + with mock.patch.object(main, "run_job_logic", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError): + call_function(HEADERS) + + assert fake_db.store == {} + body, status = call_function(HEADERS) + assert (body, status) == ("Run processed", 200) + + +def test_different_schedule_times_are_different_runs(fake_db: FakeDatabase) -> None: + next_run = dict( + HEADERS, **{"X-CloudScheduler-ScheduleTime": "2026-11-06T19:30:01-08:00"} + ) + + call_function(HEADERS) + body, status = call_function(next_run) + + assert (body, status) == ("Run processed", 200) + assert len(fake_db.store) == 2 + + +def test_missing_scheduler_headers_are_rejected() -> None: + body, status = call_function() + + assert status == 400 + assert body.startswith("error:")