From 86eab9c7c834f362ddf79ee51c0f54a97f5c57e8 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 13 Nov 2026 15:07:13 +0000 Subject: [PATCH 1/2] test(bigtable): add accelerator read/write correctness and concurrency tests --- .../accelerator/test_concurrency_async.py | 197 ++++++++++++++++++ .../accelerator/test_concurrency_autogen.py | 188 +++++++++++++++++ .../test_fuzz_correctness_async.py | 143 +++++++++++++ .../test_fuzz_correctness_autogen.py | 124 +++++++++++ 4 files changed, 652 insertions(+) create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_async.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_autogen.py diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py new file mode 100644 index 000000000000..40199786cfa0 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py @@ -0,0 +1,197 @@ +# 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 concurrent.futures +import functools +import multiprocessing +import os +import queue +import time +import uuid + +from google.cloud.bigtable.data._accelerator import _daemon +from google.cloud.bigtable.data._cross_sync import CrossSync + +from . import _harness + +if CrossSync.is_async: + from ._base_async import AcceleratorTestBaseAsync as AcceleratorTestBase +else: + from ._base_autogen import AcceleratorTestBase + +__CROSS_SYNC_OUTPUT__ = "tests.system.data.accelerator.test_concurrency_autogen" + + +def _run_worker(table, worker_id, ops_per_worker): + prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() + ops = _harness.RandomOps(worker_id, key_prefix=prefix) + expected_state = _harness.ExpectedState() + touched = set() + for _ in range(ops_per_worker): + key, mutation = ops.build_mutation(expected_state) + table.mutate_row(key, mutation) + touched.add(key) + for key in touched: + row = table.read_row(key) + _harness.assert_rows_equivalent( + f"worker-{worker_id}", + _harness.normalize_row(row), + "model", + expected_state.expected_cells(key), + ) + return sorted(touched) + + +def _spawn_worker( + project, + instance_id, + table_id, + app_profile_id, + base_worker_id, + workers_per_process, + ops_per_worker, + result_queue, +): + from google.cloud.bigtable.data import BigtableDataClient + + try: + override = os.environ.get(_harness.BIN_ENV_VAR) + if override: + _daemon._resolve_binary_path = lambda: override + with BigtableDataClient(project=project, use_accelerator=True) as client: + table_kwargs = {"app_profile_id": app_profile_id} if app_profile_id else {} + with client.get_table(instance_id, table_id, **table_kwargs) as table: + if table._accelerator_client is None: + result_queue.put((False, [], "accelerator fell back to native")) + return + with concurrent.futures.ThreadPoolExecutor( + max_workers=workers_per_process + ) as executor: + futures = [ + executor.submit( + _run_worker, table, base_worker_id + w, ops_per_worker + ) + for w in range(workers_per_process) + ] + touched = [] + for future in futures: + touched.extend(future.result()) + result_queue.put((True, touched, None)) + except BaseException as exc: # noqa: BLE001 - reported to the parent process + result_queue.put((False, [], repr(exc))) + + +@CrossSync.convert_class(sync_name="TestConcurrency") +class TestConcurrencyAsync(AcceleratorTestBase): + NUM_WORKERS = 16 + OPS_PER_WORKER = 50 + NUM_PROCESSES = 4 + WORKERS_PER_PROCESS = 8 + + @CrossSync.convert + async def _worker(self, table, worker_id): + prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() + ops = _harness.RandomOps(worker_id, key_prefix=prefix) + expected_state = _harness.ExpectedState() + touched = set() + for _ in range(self.OPS_PER_WORKER): + key, mutation = ops.build_mutation(expected_state) + await table.mutate_row(key, mutation) + touched.add(key) + for key in touched: + row = await table.read_row(key) + _harness.assert_rows_equivalent( + f"worker-{worker_id}", + _harness.normalize_row(row), + "model", + expected_state.expected_cells(key), + ) + return sorted(touched) + + @CrossSync.pytest + async def test_concurrent_workers_stay_consistent(self, accel_table, janitor): + executor = ( + concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_WORKERS) + if not CrossSync.is_async + else None + ) + try: + partials = [ + functools.partial(self._worker, accel_table, i) + for i in range(self.NUM_WORKERS) + ] + results = await CrossSync.gather_partials( + partials, return_exceptions=True, sync_executor=executor + ) + finally: + if executor is not None: + executor.shutdown(wait=True) + + for i, result in enumerate(results): + assert not isinstance(result, BaseException), ( + f"worker {i} failed: {result!r}" + ) + for key in result: + janitor.track(key) + + @CrossSync.pytest + async def test_multiprocess_concurrent_clients_stay_consistent( + self, instance_id, table_id, janitor + ): + project = self._project_id() + app_profile_id = os.getenv("BIGTABLE_TEST_APP_PROFILE") or os.getenv( + "CBT_APP_PROFILE_ID" + ) + ctx = multiprocessing.get_context("spawn") + result_queue = ctx.Queue() + procs = [] + for p in range(self.NUM_PROCESSES): + proc = ctx.Process( + target=_spawn_worker, + args=( + project, + instance_id, + table_id, + app_profile_id, + p * self.WORKERS_PER_PROCESS, + self.WORKERS_PER_PROCESS, + self.OPS_PER_WORKER, + result_queue, + ), + ) + proc.start() + procs.append(proc) + results = [] + deadline = time.monotonic() + 300.0 + try: + while len(results) < self.NUM_PROCESSES and time.monotonic() < deadline: + try: + results.append(result_queue.get(timeout=0.2)) + except queue.Empty: + if not any(p.is_alive() for p in procs) and result_queue.empty(): + break + finally: + for proc in procs: + proc.join(timeout=30) + if proc.is_alive(): + proc.kill() + proc.join() + + assert len(results) == self.NUM_PROCESSES, ( + f"expected {self.NUM_PROCESSES} results, got {len(results)}" + ) + for ok, touched, error in results: + assert ok, f"a worker process failed: {error}" + for key in touched: + janitor.track(key) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py new file mode 100644 index 000000000000..877b5e6b436a --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py @@ -0,0 +1,188 @@ +# 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. + + +# This file is automatically generated by CrossSync. Do not edit manually. + +import concurrent.futures +import functools +import multiprocessing +import os +import queue +import time +import uuid + +from google.cloud.bigtable.data._accelerator import _daemon +from google.cloud.bigtable.data._cross_sync import CrossSync + +from . import _harness +from ._base_autogen import AcceleratorTestBase + + +def _run_worker(table, worker_id, ops_per_worker): + prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() + ops = _harness.RandomOps(worker_id, key_prefix=prefix) + expected_state = _harness.ExpectedState() + touched = set() + for _ in range(ops_per_worker): + key, mutation = ops.build_mutation(expected_state) + table.mutate_row(key, mutation) + touched.add(key) + for key in touched: + row = table.read_row(key) + _harness.assert_rows_equivalent( + f"worker-{worker_id}", + _harness.normalize_row(row), + "model", + expected_state.expected_cells(key), + ) + return sorted(touched) + + +def _spawn_worker( + project, + instance_id, + table_id, + app_profile_id, + base_worker_id, + workers_per_process, + ops_per_worker, + result_queue, +): + from google.cloud.bigtable.data import BigtableDataClient + + try: + override = os.environ.get(_harness.BIN_ENV_VAR) + if override: + _daemon._resolve_binary_path = lambda: override + with BigtableDataClient(project=project, use_accelerator=True) as client: + table_kwargs = {"app_profile_id": app_profile_id} if app_profile_id else {} + with client.get_table(instance_id, table_id, **table_kwargs) as table: + if table._accelerator_client is None: + result_queue.put((False, [], "accelerator fell back to native")) + return + with concurrent.futures.ThreadPoolExecutor( + max_workers=workers_per_process + ) as executor: + futures = [ + executor.submit( + _run_worker, table, base_worker_id + w, ops_per_worker + ) + for w in range(workers_per_process) + ] + touched = [] + for future in futures: + touched.extend(future.result()) + result_queue.put((True, touched, None)) + except BaseException as exc: + result_queue.put((False, [], repr(exc))) + + +class TestConcurrency(AcceleratorTestBase): + NUM_WORKERS = 16 + OPS_PER_WORKER = 50 + NUM_PROCESSES = 4 + WORKERS_PER_PROCESS = 8 + + def _worker(self, table, worker_id): + prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() + ops = _harness.RandomOps(worker_id, key_prefix=prefix) + expected_state = _harness.ExpectedState() + touched = set() + for _ in range(self.OPS_PER_WORKER): + key, mutation = ops.build_mutation(expected_state) + table.mutate_row(key, mutation) + touched.add(key) + for key in touched: + row = table.read_row(key) + _harness.assert_rows_equivalent( + f"worker-{worker_id}", + _harness.normalize_row(row), + "model", + expected_state.expected_cells(key), + ) + return sorted(touched) + + def test_concurrent_workers_stay_consistent(self, accel_table, janitor): + executor = ( + concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_WORKERS) + if not CrossSync._Sync_Impl.is_async + else None + ) + try: + partials = [ + functools.partial(self._worker, accel_table, i) + for i in range(self.NUM_WORKERS) + ] + results = CrossSync._Sync_Impl.gather_partials( + partials, return_exceptions=True, sync_executor=executor + ) + finally: + if executor is not None: + executor.shutdown(wait=True) + for i, result in enumerate(results): + assert not isinstance(result, BaseException), ( + f"worker {i} failed: {result!r}" + ) + for key in result: + janitor.track(key) + + def test_multiprocess_concurrent_clients_stay_consistent( + self, instance_id, table_id, janitor + ): + project = self._project_id() + app_profile_id = os.getenv("BIGTABLE_TEST_APP_PROFILE") or os.getenv( + "CBT_APP_PROFILE_ID" + ) + ctx = multiprocessing.get_context("spawn") + result_queue = ctx.Queue() + procs = [] + for p in range(self.NUM_PROCESSES): + proc = ctx.Process( + target=_spawn_worker, + args=( + project, + instance_id, + table_id, + app_profile_id, + p * self.WORKERS_PER_PROCESS, + self.WORKERS_PER_PROCESS, + self.OPS_PER_WORKER, + result_queue, + ), + ) + proc.start() + procs.append(proc) + results = [] + deadline = time.monotonic() + 300.0 + try: + while len(results) < self.NUM_PROCESSES and time.monotonic() < deadline: + try: + results.append(result_queue.get(timeout=0.2)) + except queue.Empty: + if not any((p.is_alive() for p in procs)) and result_queue.empty(): + break + finally: + for proc in procs: + proc.join(timeout=30) + if proc.is_alive(): + proc.kill() + proc.join() + assert len(results) == self.NUM_PROCESSES, ( + f"expected {self.NUM_PROCESSES} results, got {len(results)}" + ) + for ok, touched, error in results: + assert ok, f"a worker process failed: {error}" + for key in touched: + janitor.track(key) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_async.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_async.py new file mode 100644 index 000000000000..a258703628b3 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_async.py @@ -0,0 +1,143 @@ +# 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. +"""Differential read/write correctness tests for the Cloud Bigtable accelerator.""" + +import uuid + +import pytest + +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn, SetCell +from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter + +from . import _harness + +if CrossSync.is_async: + from ._base_async import AcceleratorTestBaseAsync as AcceleratorTestBase +else: + from ._base_autogen import AcceleratorTestBase + +__CROSS_SYNC_OUTPUT__ = "tests.system.data.accelerator.test_fuzz_correctness_autogen" + + +@CrossSync.convert_class(sync_name="TestFuzzCorrectness") +class TestFuzzCorrectnessAsync(AcceleratorTestBase): + NUM_OPS = 500 + + @CrossSync.convert + async def _verify_key(self, accel_table, native_table, expected_state, key): + accel_row = await accel_table.read_row(key) + native_row = await native_table.read_row(key) + expected = expected_state.expected_cells(key) + accel_cells = _harness.normalize_row(accel_row) + native_cells = _harness.normalize_row(native_row) + _harness.assert_rows_equivalent( + "accelerator", accel_cells, "expected", expected + ) + _harness.assert_rows_equivalent( + "accelerator", accel_cells, "native", native_cells + ) + + @pytest.mark.parametrize("seed", [0, 1]) + @CrossSync.pytest + async def test_random_mutations_match_model_and_native( + self, accel_table, native_table, janitor, seed + ): + prefix = f"fuzz-{seed}-{uuid.uuid4().hex[:8]}-".encode() + ops = _harness.RandomOps(seed, key_prefix=prefix, include_aggregate=True) + expected_state = _harness.ExpectedState() + touched: set[bytes] = set() + + for i in range(self.NUM_OPS): + key, mutation = ops.build_mutation(expected_state) + await accel_table.mutate_row(key, mutation) + janitor.track(key) + touched.add(key) + if touched and i % 10 == 9: + vkey = ops._rand.choice(sorted(touched)) + await self._verify_key(accel_table, native_table, expected_state, vkey) + + for key in sorted(touched): + await self._verify_key(accel_table, native_table, expected_state, key) + + @CrossSync.pytest + async def test_multiple_versions_range_delete_and_row_filter( + self, accel_table, native_table, janitor + ): + key = janitor.track(f"versions-{uuid.uuid4().hex}".encode()) + family = _harness.TEST_FAMILY + qualifier = b"q0" + expected_state = _harness.ExpectedState() + + for ts_ms, value in [(1000, b"v1"), (2000, b"v2"), (3000, b"v3")]: + ts = ts_ms * _harness.MS + await accel_table.mutate_row( + key, expected_state.set_cell(key, family, qualifier, value, ts) + ) + await self._verify_key(accel_table, native_table, expected_state, key) + + # Verify row_filter parity (latest cell only). + row_filter = CellsRowLimitFilter(1) + accel_filtered = await accel_table.read_row(key, row_filter=row_filter) + native_filtered = await native_table.read_row(key, row_filter=row_filter) + _harness.assert_rows_equivalent( + "accel-filtered", + _harness.normalize_row(accel_filtered), + "native-filtered", + _harness.normalize_row(native_filtered), + ) + assert len(accel_filtered.cells) == 1 + assert accel_filtered.cells[0].value == b"v3" + + # Delete middle version [2000ms, 3000ms). + expected_state.delete_range_from_column( + key, family, qualifier, 2000 * _harness.MS, 3000 * _harness.MS + ) + await accel_table.mutate_row( + key, + DeleteRangeFromColumn( + family, qualifier, 2000 * _harness.MS, 3000 * _harness.MS + ), + ) + await self._verify_key(accel_table, native_table, expected_state, key) + + @CrossSync.pytest + async def test_shared_daemon_across_multiple_tables( + self, accel_client, instance_id, table_id, native_table, janitor + ): + async with ( + accel_client.get_table(instance_id, table_id) as table_a, + accel_client.get_table(instance_id, table_id) as table_b, + ): + self.assert_accelerator_active(table_a) + self.assert_accelerator_active(table_b) + assert _harness.daemon_pid(table_a) == _harness.daemon_pid(table_b) + + key_a = janitor.track(f"shared-a-{uuid.uuid4().hex}".encode()) + key_b = janitor.track(f"shared-b-{uuid.uuid4().hex}".encode()) + ts = 1000 * _harness.MS + + await table_a.mutate_row( + key_a, SetCell(_harness.TEST_FAMILY, b"qa", b"va", timestamp_micros=ts) + ) + await table_b.mutate_row( + key_b, SetCell(_harness.TEST_FAMILY, b"qb", b"vb", timestamp_micros=ts) + ) + + row_a = await table_b.read_row(key_a) + row_b = await table_a.read_row(key_b) + assert row_a is not None and row_a.cells[0].value == b"va" + assert row_b is not None and row_b.cells[0].value == b"vb" + assert (await native_table.read_row(key_a)).cells[0].value == b"va" + assert (await native_table.read_row(key_b)).cells[0].value == b"vb" diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_autogen.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_autogen.py new file mode 100644 index 000000000000..c8679d9b1cee --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_autogen.py @@ -0,0 +1,124 @@ +# 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. + +# This file is automatically generated by CrossSync. Do not edit manually. + +"""Differential read/write correctness tests for the Cloud Bigtable accelerator.""" + +import uuid + +import pytest + +from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn, SetCell +from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter + +from . import _harness +from ._base_autogen import AcceleratorTestBase + + +class TestFuzzCorrectness(AcceleratorTestBase): + NUM_OPS = 500 + + def _verify_key(self, accel_table, native_table, expected_state, key): + accel_row = accel_table.read_row(key) + native_row = native_table.read_row(key) + expected = expected_state.expected_cells(key) + accel_cells = _harness.normalize_row(accel_row) + native_cells = _harness.normalize_row(native_row) + _harness.assert_rows_equivalent( + "accelerator", accel_cells, "expected", expected + ) + _harness.assert_rows_equivalent( + "accelerator", accel_cells, "native", native_cells + ) + + @pytest.mark.parametrize("seed", [0, 1]) + def test_random_mutations_match_model_and_native( + self, accel_table, native_table, janitor, seed + ): + prefix = f"fuzz-{seed}-{uuid.uuid4().hex[:8]}-".encode() + ops = _harness.RandomOps(seed, key_prefix=prefix, include_aggregate=True) + expected_state = _harness.ExpectedState() + touched: set[bytes] = set() + for i in range(self.NUM_OPS): + key, mutation = ops.build_mutation(expected_state) + accel_table.mutate_row(key, mutation) + janitor.track(key) + touched.add(key) + if touched and i % 10 == 9: + vkey = ops._rand.choice(sorted(touched)) + self._verify_key(accel_table, native_table, expected_state, vkey) + for key in sorted(touched): + self._verify_key(accel_table, native_table, expected_state, key) + + def test_multiple_versions_range_delete_and_row_filter( + self, accel_table, native_table, janitor + ): + key = janitor.track(f"versions-{uuid.uuid4().hex}".encode()) + family = _harness.TEST_FAMILY + qualifier = b"q0" + expected_state = _harness.ExpectedState() + for ts_ms, value in [(1000, b"v1"), (2000, b"v2"), (3000, b"v3")]: + ts = ts_ms * _harness.MS + accel_table.mutate_row( + key, expected_state.set_cell(key, family, qualifier, value, ts) + ) + self._verify_key(accel_table, native_table, expected_state, key) + row_filter = CellsRowLimitFilter(1) + accel_filtered = accel_table.read_row(key, row_filter=row_filter) + native_filtered = native_table.read_row(key, row_filter=row_filter) + _harness.assert_rows_equivalent( + "accel-filtered", + _harness.normalize_row(accel_filtered), + "native-filtered", + _harness.normalize_row(native_filtered), + ) + assert len(accel_filtered.cells) == 1 + assert accel_filtered.cells[0].value == b"v3" + expected_state.delete_range_from_column( + key, family, qualifier, 2000 * _harness.MS, 3000 * _harness.MS + ) + accel_table.mutate_row( + key, + DeleteRangeFromColumn( + family, qualifier, 2000 * _harness.MS, 3000 * _harness.MS + ), + ) + self._verify_key(accel_table, native_table, expected_state, key) + + def test_shared_daemon_across_multiple_tables( + self, accel_client, instance_id, table_id, native_table, janitor + ): + with ( + accel_client.get_table(instance_id, table_id) as table_a, + accel_client.get_table(instance_id, table_id) as table_b, + ): + self.assert_accelerator_active(table_a) + self.assert_accelerator_active(table_b) + assert _harness.daemon_pid(table_a) == _harness.daemon_pid(table_b) + key_a = janitor.track(f"shared-a-{uuid.uuid4().hex}".encode()) + key_b = janitor.track(f"shared-b-{uuid.uuid4().hex}".encode()) + ts = 1000 * _harness.MS + table_a.mutate_row( + key_a, SetCell(_harness.TEST_FAMILY, b"qa", b"va", timestamp_micros=ts) + ) + table_b.mutate_row( + key_b, SetCell(_harness.TEST_FAMILY, b"qb", b"vb", timestamp_micros=ts) + ) + row_a = table_b.read_row(key_a) + row_b = table_a.read_row(key_b) + assert row_a is not None and row_a.cells[0].value == b"va" + assert row_b is not None and row_b.cells[0].value == b"vb" + assert native_table.read_row(key_a).cells[0].value == b"va" + assert native_table.read_row(key_b).cells[0].value == b"vb" From 281a4f238ef786313810b75fe7791263fcaab0d6 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 13 Nov 2026 21:30:15 +0000 Subject: [PATCH 2/2] test(bigtable): test multi-table concurrency and parallelize fuzz test Update the accelerator concurrency test to run 8 concurrent workers across two tables from the same client. Update the fuzz correctness test to run 8 parallel workers per seed. --- .../accelerator/test_concurrency_async.py | 237 +++++++----------- .../accelerator/test_concurrency_autogen.py | 229 +++++++---------- .../test_fuzz_correctness_async.py | 56 ++++- .../test_fuzz_correctness_autogen.py | 53 +++- 4 files changed, 271 insertions(+), 304 deletions(-) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py index 40199786cfa0..fd12888d0ad8 100644 --- a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py @@ -14,13 +14,10 @@ import concurrent.futures import functools -import multiprocessing -import os -import queue -import time import uuid -from google.cloud.bigtable.data._accelerator import _daemon +import pytest + from google.cloud.bigtable.data._cross_sync import CrossSync from . import _harness @@ -33,76 +30,49 @@ __CROSS_SYNC_OUTPUT__ = "tests.system.data.accelerator.test_concurrency_autogen" -def _run_worker(table, worker_id, ops_per_worker): - prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() - ops = _harness.RandomOps(worker_id, key_prefix=prefix) - expected_state = _harness.ExpectedState() - touched = set() - for _ in range(ops_per_worker): - key, mutation = ops.build_mutation(expected_state) - table.mutate_row(key, mutation) - touched.add(key) - for key in touched: - row = table.read_row(key) - _harness.assert_rows_equivalent( - f"worker-{worker_id}", - _harness.normalize_row(row), - "model", - expected_state.expected_cells(key), - ) - return sorted(touched) - - -def _spawn_worker( - project, - instance_id, - table_id, - app_profile_id, - base_worker_id, - workers_per_process, - ops_per_worker, - result_queue, -): - from google.cloud.bigtable.data import BigtableDataClient - - try: - override = os.environ.get(_harness.BIN_ENV_VAR) - if override: - _daemon._resolve_binary_path = lambda: override - with BigtableDataClient(project=project, use_accelerator=True) as client: - table_kwargs = {"app_profile_id": app_profile_id} if app_profile_id else {} - with client.get_table(instance_id, table_id, **table_kwargs) as table: - if table._accelerator_client is None: - result_queue.put((False, [], "accelerator fell back to native")) - return - with concurrent.futures.ThreadPoolExecutor( - max_workers=workers_per_process - ) as executor: - futures = [ - executor.submit( - _run_worker, table, base_worker_id + w, ops_per_worker - ) - for w in range(workers_per_process) - ] - touched = [] - for future in futures: - touched.extend(future.result()) - result_queue.put((True, touched, None)) - except BaseException as exc: # noqa: BLE001 - reported to the parent process - result_queue.put((False, [], repr(exc))) - - @CrossSync.convert_class(sync_name="TestConcurrency") class TestConcurrencyAsync(AcceleratorTestBase): - NUM_WORKERS = 16 + CONCURRENCY = 8 OPS_PER_WORKER = 50 - NUM_PROCESSES = 4 - WORKERS_PER_PROCESS = 8 + + @pytest.fixture(scope="session") + def second_table_id( + self, + admin_client, + project_id, + instance_id, + column_family_config, + column_split_config, + ): + from google.api_core import exceptions, retry + + second_id = f"test-table-2-{uuid.uuid4().hex}" + retry_fn = retry.Retry( + predicate=retry.if_exception_type(exceptions.FailedPrecondition) + ) + parent_path = f"projects/{project_id}/instances/{instance_id}" + admin_client.table_admin_client.create_table( + request={ + "parent": parent_path, + "table_id": second_id, + "table": {"column_families": column_family_config}, + "initial_splits": [{"key": key} for key in column_split_config], + }, + retry=retry_fn, + ) + try: + yield second_id + finally: + try: + admin_client.table_admin_client.delete_table( + name=f"{parent_path}/tables/{second_id}" + ) + except Exception as exc: + print(f"Failed to delete table {second_id}: {exc}") @CrossSync.convert - async def _worker(self, table, worker_id): - prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() - ops = _harness.RandomOps(worker_id, key_prefix=prefix) + async def _worker(self, table, worker_id, key_prefix): + ops = _harness.RandomOps(worker_id, key_prefix=key_prefix) expected_state = _harness.ExpectedState() touched = set() for _ in range(self.OPS_PER_WORKER): @@ -120,78 +90,59 @@ async def _worker(self, table, worker_id): return sorted(touched) @CrossSync.pytest - async def test_concurrent_workers_stay_consistent(self, accel_table, janitor): - executor = ( - concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_WORKERS) - if not CrossSync.is_async - else None - ) - try: - partials = [ - functools.partial(self._worker, accel_table, i) - for i in range(self.NUM_WORKERS) - ] - results = await CrossSync.gather_partials( - partials, return_exceptions=True, sync_executor=executor - ) - finally: - if executor is not None: - executor.shutdown(wait=True) - - for i, result in enumerate(results): - assert not isinstance(result, BaseException), ( - f"worker {i} failed: {result!r}" - ) - for key in result: - janitor.track(key) - - @CrossSync.pytest - async def test_multiprocess_concurrent_clients_stay_consistent( - self, instance_id, table_id, janitor + async def test_concurrent_multitable_workers_stay_consistent( + self, accel_client, instance_id, table_id, second_table_id, janitor ): - project = self._project_id() - app_profile_id = os.getenv("BIGTABLE_TEST_APP_PROFILE") or os.getenv( - "CBT_APP_PROFILE_ID" - ) - ctx = multiprocessing.get_context("spawn") - result_queue = ctx.Queue() - procs = [] - for p in range(self.NUM_PROCESSES): - proc = ctx.Process( - target=_spawn_worker, - args=( - project, - instance_id, - table_id, - app_profile_id, - p * self.WORKERS_PER_PROCESS, - self.WORKERS_PER_PROCESS, - self.OPS_PER_WORKER, - result_queue, - ), - ) - proc.start() - procs.append(proc) - results = [] - deadline = time.monotonic() + 300.0 - try: - while len(results) < self.NUM_PROCESSES and time.monotonic() < deadline: + async with accel_client.get_table(instance_id, second_table_id) as table_b: + self.assert_accelerator_active(table_b) + daemon_pid_b = _harness.daemon_pid(table_b) + + async with accel_client.get_table(instance_id, table_id) as table_a: + self.assert_accelerator_active(table_a) + assert _harness.daemon_pid(table_a) == daemon_pid_b + + run_id = uuid.uuid4().hex[:8] + tables = (table_a, table_b) + executor = ( + concurrent.futures.ThreadPoolExecutor(max_workers=self.CONCURRENCY) + if not CrossSync.is_async + else None + ) try: - results.append(result_queue.get(timeout=0.2)) - except queue.Empty: - if not any(p.is_alive() for p in procs) and result_queue.empty(): - break - finally: - for proc in procs: - proc.join(timeout=30) - if proc.is_alive(): - proc.kill() - proc.join() - - assert len(results) == self.NUM_PROCESSES, ( - f"expected {self.NUM_PROCESSES} results, got {len(results)}" - ) - for ok, touched, error in results: - assert ok, f"a worker process failed: {error}" - for key in touched: - janitor.track(key) + partials = [ + functools.partial( + self._worker, + tables[w % 2], + w, + f"conc-mt-{run_id}-p{w // 2}-".encode(), + ) + for w in range(self.CONCURRENCY) + ] + results = await CrossSync.gather_partials( + partials, return_exceptions=True, sync_executor=executor + ) + finally: + if executor is not None: + executor.shutdown(wait=True) + + for w, result in enumerate(results): + assert not isinstance(result, BaseException), ( + f"worker {w} failed: {result!r}" + ) + if w % 2 == 0: + for key in result: + janitor.track(key) + + # Closing table_a must leave table_b and the shared daemon active. + assert _harness.daemon_pid(table_b) == daemon_pid_b + verify_ops = _harness.RandomOps(999, key_prefix=b"conc-mt-after-close-") + verify_state = _harness.ExpectedState() + vkey, vmut = verify_ops.build_mutation(verify_state) + await table_b.mutate_row(vkey, vmut) + vrow = await table_b.read_row(vkey) + _harness.assert_rows_equivalent( + "table_b-after-close", + _harness.normalize_row(vrow), + "model", + verify_state.expected_cells(vkey), + ) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py index 877b5e6b436a..e99b7689aded 100644 --- a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py @@ -17,87 +17,57 @@ import concurrent.futures import functools -import multiprocessing -import os -import queue -import time import uuid -from google.cloud.bigtable.data._accelerator import _daemon +import pytest + from google.cloud.bigtable.data._cross_sync import CrossSync from . import _harness from ._base_autogen import AcceleratorTestBase -def _run_worker(table, worker_id, ops_per_worker): - prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() - ops = _harness.RandomOps(worker_id, key_prefix=prefix) - expected_state = _harness.ExpectedState() - touched = set() - for _ in range(ops_per_worker): - key, mutation = ops.build_mutation(expected_state) - table.mutate_row(key, mutation) - touched.add(key) - for key in touched: - row = table.read_row(key) - _harness.assert_rows_equivalent( - f"worker-{worker_id}", - _harness.normalize_row(row), - "model", - expected_state.expected_cells(key), - ) - return sorted(touched) - - -def _spawn_worker( - project, - instance_id, - table_id, - app_profile_id, - base_worker_id, - workers_per_process, - ops_per_worker, - result_queue, -): - from google.cloud.bigtable.data import BigtableDataClient - - try: - override = os.environ.get(_harness.BIN_ENV_VAR) - if override: - _daemon._resolve_binary_path = lambda: override - with BigtableDataClient(project=project, use_accelerator=True) as client: - table_kwargs = {"app_profile_id": app_profile_id} if app_profile_id else {} - with client.get_table(instance_id, table_id, **table_kwargs) as table: - if table._accelerator_client is None: - result_queue.put((False, [], "accelerator fell back to native")) - return - with concurrent.futures.ThreadPoolExecutor( - max_workers=workers_per_process - ) as executor: - futures = [ - executor.submit( - _run_worker, table, base_worker_id + w, ops_per_worker - ) - for w in range(workers_per_process) - ] - touched = [] - for future in futures: - touched.extend(future.result()) - result_queue.put((True, touched, None)) - except BaseException as exc: - result_queue.put((False, [], repr(exc))) - - class TestConcurrency(AcceleratorTestBase): - NUM_WORKERS = 16 + CONCURRENCY = 8 OPS_PER_WORKER = 50 - NUM_PROCESSES = 4 - WORKERS_PER_PROCESS = 8 - def _worker(self, table, worker_id): - prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() - ops = _harness.RandomOps(worker_id, key_prefix=prefix) + @pytest.fixture(scope="session") + def second_table_id( + self, + admin_client, + project_id, + instance_id, + column_family_config, + column_split_config, + ): + from google.api_core import exceptions, retry + + second_id = f"test-table-2-{uuid.uuid4().hex}" + retry_fn = retry.Retry( + predicate=retry.if_exception_type(exceptions.FailedPrecondition) + ) + parent_path = f"projects/{project_id}/instances/{instance_id}" + admin_client.table_admin_client.create_table( + request={ + "parent": parent_path, + "table_id": second_id, + "table": {"column_families": column_family_config}, + "initial_splits": [{"key": key} for key in column_split_config], + }, + retry=retry_fn, + ) + try: + yield second_id + finally: + try: + admin_client.table_admin_client.delete_table( + name=f"{parent_path}/tables/{second_id}" + ) + except Exception as exc: + print(f"Failed to delete table {second_id}: {exc}") + + def _worker(self, table, worker_id, key_prefix): + ops = _harness.RandomOps(worker_id, key_prefix=key_prefix) expected_state = _harness.ExpectedState() touched = set() for _ in range(self.OPS_PER_WORKER): @@ -114,75 +84,54 @@ def _worker(self, table, worker_id): ) return sorted(touched) - def test_concurrent_workers_stay_consistent(self, accel_table, janitor): - executor = ( - concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_WORKERS) - if not CrossSync._Sync_Impl.is_async - else None - ) - try: - partials = [ - functools.partial(self._worker, accel_table, i) - for i in range(self.NUM_WORKERS) - ] - results = CrossSync._Sync_Impl.gather_partials( - partials, return_exceptions=True, sync_executor=executor - ) - finally: - if executor is not None: - executor.shutdown(wait=True) - for i, result in enumerate(results): - assert not isinstance(result, BaseException), ( - f"worker {i} failed: {result!r}" - ) - for key in result: - janitor.track(key) - - def test_multiprocess_concurrent_clients_stay_consistent( - self, instance_id, table_id, janitor + def test_concurrent_multitable_workers_stay_consistent( + self, accel_client, instance_id, table_id, second_table_id, janitor ): - project = self._project_id() - app_profile_id = os.getenv("BIGTABLE_TEST_APP_PROFILE") or os.getenv( - "CBT_APP_PROFILE_ID" - ) - ctx = multiprocessing.get_context("spawn") - result_queue = ctx.Queue() - procs = [] - for p in range(self.NUM_PROCESSES): - proc = ctx.Process( - target=_spawn_worker, - args=( - project, - instance_id, - table_id, - app_profile_id, - p * self.WORKERS_PER_PROCESS, - self.WORKERS_PER_PROCESS, - self.OPS_PER_WORKER, - result_queue, - ), - ) - proc.start() - procs.append(proc) - results = [] - deadline = time.monotonic() + 300.0 - try: - while len(results) < self.NUM_PROCESSES and time.monotonic() < deadline: + with accel_client.get_table(instance_id, second_table_id) as table_b: + self.assert_accelerator_active(table_b) + daemon_pid_b = _harness.daemon_pid(table_b) + with accel_client.get_table(instance_id, table_id) as table_a: + self.assert_accelerator_active(table_a) + assert _harness.daemon_pid(table_a) == daemon_pid_b + run_id = uuid.uuid4().hex[:8] + tables = (table_a, table_b) + executor = ( + concurrent.futures.ThreadPoolExecutor(max_workers=self.CONCURRENCY) + if not CrossSync._Sync_Impl.is_async + else None + ) try: - results.append(result_queue.get(timeout=0.2)) - except queue.Empty: - if not any((p.is_alive() for p in procs)) and result_queue.empty(): - break - finally: - for proc in procs: - proc.join(timeout=30) - if proc.is_alive(): - proc.kill() - proc.join() - assert len(results) == self.NUM_PROCESSES, ( - f"expected {self.NUM_PROCESSES} results, got {len(results)}" - ) - for ok, touched, error in results: - assert ok, f"a worker process failed: {error}" - for key in touched: - janitor.track(key) + partials = [ + functools.partial( + self._worker, + tables[w % 2], + w, + f"conc-mt-{run_id}-p{w // 2}-".encode(), + ) + for w in range(self.CONCURRENCY) + ] + results = CrossSync._Sync_Impl.gather_partials( + partials, return_exceptions=True, sync_executor=executor + ) + finally: + if executor is not None: + executor.shutdown(wait=True) + for w, result in enumerate(results): + assert not isinstance(result, BaseException), ( + f"worker {w} failed: {result!r}" + ) + if w % 2 == 0: + for key in result: + janitor.track(key) + assert _harness.daemon_pid(table_b) == daemon_pid_b + verify_ops = _harness.RandomOps(999, key_prefix=b"conc-mt-after-close-") + verify_state = _harness.ExpectedState() + vkey, vmut = verify_ops.build_mutation(verify_state) + table_b.mutate_row(vkey, vmut) + vrow = table_b.read_row(vkey) + _harness.assert_rows_equivalent( + "table_b-after-close", + _harness.normalize_row(vrow), + "model", + verify_state.expected_cells(vkey), + ) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_async.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_async.py index a258703628b3..18593dbab26d 100644 --- a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_async.py @@ -13,6 +13,8 @@ # limitations under the License. """Differential read/write correctness tests for the Cloud Bigtable accelerator.""" +import concurrent.futures +import functools import uuid import pytest @@ -33,7 +35,8 @@ @CrossSync.convert_class(sync_name="TestFuzzCorrectness") class TestFuzzCorrectnessAsync(AcceleratorTestBase): - NUM_OPS = 500 + CONCURRENCY = 8 + OPS_PER_WORKER = 64 @CrossSync.convert async def _verify_key(self, accel_table, native_table, expected_state, key): @@ -49,27 +52,58 @@ async def _verify_key(self, accel_table, native_table, expected_state, key): "accelerator", accel_cells, "native", native_cells ) - @pytest.mark.parametrize("seed", [0, 1]) - @CrossSync.pytest - async def test_random_mutations_match_model_and_native( - self, accel_table, native_table, janitor, seed - ): - prefix = f"fuzz-{seed}-{uuid.uuid4().hex[:8]}-".encode() - ops = _harness.RandomOps(seed, key_prefix=prefix, include_aggregate=True) + @CrossSync.convert + async def _fuzz_worker(self, accel_table, native_table, seed, worker_id, run_id): + prefix = f"fuzz-{seed}-{run_id}-w{worker_id}-".encode() + ops = _harness.RandomOps( + seed * 100 + worker_id, key_prefix=prefix, include_aggregate=True + ) expected_state = _harness.ExpectedState() touched: set[bytes] = set() - for i in range(self.NUM_OPS): + for i in range(self.OPS_PER_WORKER): key, mutation = ops.build_mutation(expected_state) await accel_table.mutate_row(key, mutation) - janitor.track(key) touched.add(key) - if touched and i % 10 == 9: + if touched and i % 8 == 7: vkey = ops._rand.choice(sorted(touched)) await self._verify_key(accel_table, native_table, expected_state, vkey) for key in sorted(touched): await self._verify_key(accel_table, native_table, expected_state, key) + return sorted(touched) + + @pytest.mark.parametrize("seed", [0, 1]) + @CrossSync.pytest + async def test_random_mutations_match_model_and_native( + self, accel_table, native_table, janitor, seed + ): + run_id = uuid.uuid4().hex[:8] + executor = ( + concurrent.futures.ThreadPoolExecutor(max_workers=self.CONCURRENCY) + if not CrossSync.is_async + else None + ) + try: + partials = [ + functools.partial( + self._fuzz_worker, accel_table, native_table, seed, w, run_id + ) + for w in range(self.CONCURRENCY) + ] + results = await CrossSync.gather_partials( + partials, return_exceptions=True, sync_executor=executor + ) + finally: + if executor is not None: + executor.shutdown(wait=True) + + for w, result in enumerate(results): + assert not isinstance(result, BaseException), ( + f"fuzz worker {w} failed: {result!r}" + ) + for key in result: + janitor.track(key) @CrossSync.pytest async def test_multiple_versions_range_delete_and_row_filter( diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_autogen.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_autogen.py index c8679d9b1cee..cbc8b87c2c59 100644 --- a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_fuzz_correctness_autogen.py @@ -16,10 +16,13 @@ """Differential read/write correctness tests for the Cloud Bigtable accelerator.""" +import concurrent.futures +import functools import uuid import pytest +from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn, SetCell from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter @@ -28,7 +31,8 @@ class TestFuzzCorrectness(AcceleratorTestBase): - NUM_OPS = 500 + CONCURRENCY = 8 + OPS_PER_WORKER = 64 def _verify_key(self, accel_table, native_table, expected_state, key): accel_row = accel_table.read_row(key) @@ -43,24 +47,53 @@ def _verify_key(self, accel_table, native_table, expected_state, key): "accelerator", accel_cells, "native", native_cells ) - @pytest.mark.parametrize("seed", [0, 1]) - def test_random_mutations_match_model_and_native( - self, accel_table, native_table, janitor, seed - ): - prefix = f"fuzz-{seed}-{uuid.uuid4().hex[:8]}-".encode() - ops = _harness.RandomOps(seed, key_prefix=prefix, include_aggregate=True) + def _fuzz_worker(self, accel_table, native_table, seed, worker_id, run_id): + prefix = f"fuzz-{seed}-{run_id}-w{worker_id}-".encode() + ops = _harness.RandomOps( + seed * 100 + worker_id, key_prefix=prefix, include_aggregate=True + ) expected_state = _harness.ExpectedState() touched: set[bytes] = set() - for i in range(self.NUM_OPS): + for i in range(self.OPS_PER_WORKER): key, mutation = ops.build_mutation(expected_state) accel_table.mutate_row(key, mutation) - janitor.track(key) touched.add(key) - if touched and i % 10 == 9: + if touched and i % 8 == 7: vkey = ops._rand.choice(sorted(touched)) self._verify_key(accel_table, native_table, expected_state, vkey) for key in sorted(touched): self._verify_key(accel_table, native_table, expected_state, key) + return sorted(touched) + + @pytest.mark.parametrize("seed", [0, 1]) + def test_random_mutations_match_model_and_native( + self, accel_table, native_table, janitor, seed + ): + run_id = uuid.uuid4().hex[:8] + executor = ( + concurrent.futures.ThreadPoolExecutor(max_workers=self.CONCURRENCY) + if not CrossSync._Sync_Impl.is_async + else None + ) + try: + partials = [ + functools.partial( + self._fuzz_worker, accel_table, native_table, seed, w, run_id + ) + for w in range(self.CONCURRENCY) + ] + results = CrossSync._Sync_Impl.gather_partials( + partials, return_exceptions=True, sync_executor=executor + ) + finally: + if executor is not None: + executor.shutdown(wait=True) + for w, result in enumerate(results): + assert not isinstance(result, BaseException), ( + f"fuzz worker {w} failed: {result!r}" + ) + for key in result: + janitor.track(key) def test_multiple_versions_range_delete_and_row_filter( self, accel_table, native_table, janitor