diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_lifecycle_async.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_lifecycle_async.py new file mode 100644 index 000000000000..0f920b722f1a --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_lifecycle_async.py @@ -0,0 +1,151 @@ +# 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 multiprocessing +import os +import queue +import shutil +import signal +import time +import uuid + +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data.mutations import DeleteAllFromRow, SetCell + +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_lifecycle_autogen" + + +def _daemon_owner_process(project, instance_id, table_id, result_queue): + from google.cloud.bigtable.data import BigtableDataClient + + with BigtableDataClient(project=project, use_accelerator=True) as client: + with client.get_table(instance_id, table_id) as table: + key = f"crash-{uuid.uuid4().hex}".encode() + ts = 1000 * _harness.MS + table.mutate_row( + key, SetCell(_harness.TEST_FAMILY, b"q", b"v", timestamp_micros=ts) + ) + table.mutate_row(key, DeleteAllFromRow()) + result_queue.put( + (_harness.daemon_pid(table), table._accelerator_daemon._tempdir) + ) + while True: + time.sleep(3600) + + +@CrossSync.convert_class(sync_name="TestLifecycle") +class TestLifecycleAsync(AcceleratorTestBase): + @CrossSync.convert + async def _one_lifecycle(self, instance_id, table_id): + async with self._make_client(use_accelerator=True) as client: + async with client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + key = f"life-{uuid.uuid4().hex}".encode() + await table.mutate_row( + key, + SetCell( + _harness.TEST_FAMILY, + b"q", + b"v", + timestamp_micros=1000 * _harness.MS, + ), + ) + assert (await table.read_row(key)) is not None + await table.mutate_row(key, DeleteAllFromRow()) + + @CrossSync.pytest + async def test_close_reaps_daemon_and_tempdir(self, instance_id, table_id): + async with self._make_client(use_accelerator=True) as client: + async with client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + pid = _harness.daemon_pid(table) + tempdir = table._accelerator_daemon._tempdir + assert pid is not None and _harness.pid_alive(pid) + assert tempdir is not None and os.path.exists(tempdir) + assert not _harness.pid_alive(pid), "daemon still alive after close" + assert not os.path.exists(tempdir), "UDS tempdir not cleaned up" + + @CrossSync.pytest + async def test_no_leaks_over_repeated_lifecycles(self, instance_id, table_id): + introspector = _harness.ProcessIntrospector() + await self._one_lifecycle(instance_id, table_id) + before = introspector.snapshot() + for _ in range(8): + await self._one_lifecycle(instance_id, table_id) + introspector.assert_no_leaks(before, label="8x accelerator lifecycle") + + @CrossSync.pytest + async def test_close_reclaims_resources_after_daemon_crash( + self, instance_id, table_id + ): + introspector = _harness.ProcessIntrospector() + before = introspector.snapshot() + async with self._make_client(use_accelerator=True) as client: + async with client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + pid = _harness.daemon_pid(table) + tempdir = table._accelerator_daemon._tempdir + assert pid is not None and _harness.pid_alive(pid) + assert tempdir is not None and os.path.exists(tempdir) + os.kill(pid, signal.SIGKILL) + deadline = time.monotonic() + 10.0 + while _harness.pid_alive(pid) and time.monotonic() < deadline: + time.sleep(0.05) + assert not _harness.pid_alive(pid), "daemon did not die after SIGKILL" + assert not os.path.exists(tempdir), "close() did not reclaim tempdir" + introspector.assert_no_leaks(before, label="daemon-crash teardown") + + @CrossSync.pytest + async def test_no_daemon_leak_when_owner_process_is_killed( + self, instance_id, table_id + ): + project = self._project_id() + ctx = multiprocessing.get_context("fork") + result_queue = ctx.Queue() + owner = ctx.Process( + target=_daemon_owner_process, + args=(project, instance_id, table_id, result_queue), + ) + owner.start() + daemon_pid = tempdir = None + try: + deadline = time.monotonic() + 60.0 + while time.monotonic() < deadline: + try: + daemon_pid, tempdir = result_queue.get(timeout=0.2) + break + except queue.Empty: + assert owner.is_alive(), "owner exited before starting daemon" + assert daemon_pid is not None, "accelerator fell back to native" + assert _harness.pid_alive(daemon_pid), "daemon not alive after owner start" + owner.kill() + owner.join(timeout=10) + assert not owner.is_alive(), "owner process not reaped after SIGKILL" + deadline = time.monotonic() + 30.0 + while _harness.pid_alive(daemon_pid) and time.monotonic() < deadline: + time.sleep(0.1) + assert not _harness.pid_alive(daemon_pid), f"orphaned daemon {daemon_pid}" + finally: + if owner.is_alive(): + owner.kill() + owner.join() + if tempdir and os.path.isdir(tempdir): + shutil.rmtree(tempdir, ignore_errors=True) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_lifecycle_autogen.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_lifecycle_autogen.py new file mode 100644 index 000000000000..a713edd50492 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_lifecycle_autogen.py @@ -0,0 +1,119 @@ +# 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 multiprocessing +import os +import queue +import shutil +import signal +import time +import uuid +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data.mutations import DeleteAllFromRow, SetCell +from . import _harness +from ._base_autogen import AcceleratorTestBase + +def _daemon_owner_process(project, instance_id, table_id, result_queue): + from google.cloud.bigtable.data import BigtableDataClient + with BigtableDataClient(project=project, use_accelerator=True) as client: + with client.get_table(instance_id, table_id) as table: + key = f'crash-{uuid.uuid4().hex}'.encode() + ts = 1000 * _harness.MS + table.mutate_row(key, SetCell(_harness.TEST_FAMILY, b'q', b'v', timestamp_micros=ts)) + table.mutate_row(key, DeleteAllFromRow()) + result_queue.put((_harness.daemon_pid(table), table._accelerator_daemon._tempdir)) + while True: + time.sleep(3600) + +class TestLifecycle(AcceleratorTestBase): + + def _one_lifecycle(self, instance_id, table_id): + with self._make_client(use_accelerator=True) as client: + with client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + key = f'life-{uuid.uuid4().hex}'.encode() + table.mutate_row(key, SetCell(_harness.TEST_FAMILY, b'q', b'v', timestamp_micros=1000 * _harness.MS)) + assert table.read_row(key) is not None + table.mutate_row(key, DeleteAllFromRow()) + + def test_close_reaps_daemon_and_tempdir(self, instance_id, table_id): + with self._make_client(use_accelerator=True) as client: + with client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + pid = _harness.daemon_pid(table) + tempdir = table._accelerator_daemon._tempdir + assert pid is not None and _harness.pid_alive(pid) + assert tempdir is not None and os.path.exists(tempdir) + assert not _harness.pid_alive(pid), 'daemon still alive after close' + assert not os.path.exists(tempdir), 'UDS tempdir not cleaned up' + + def test_no_leaks_over_repeated_lifecycles(self, instance_id, table_id): + introspector = _harness.ProcessIntrospector() + self._one_lifecycle(instance_id, table_id) + before = introspector.snapshot() + for _ in range(8): + self._one_lifecycle(instance_id, table_id) + introspector.assert_no_leaks(before, label='8x accelerator lifecycle') + + def test_close_reclaims_resources_after_daemon_crash(self, instance_id, table_id): + introspector = _harness.ProcessIntrospector() + before = introspector.snapshot() + with self._make_client(use_accelerator=True) as client: + with client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + pid = _harness.daemon_pid(table) + tempdir = table._accelerator_daemon._tempdir + assert pid is not None and _harness.pid_alive(pid) + assert tempdir is not None and os.path.exists(tempdir) + os.kill(pid, signal.SIGKILL) + deadline = time.monotonic() + 10.0 + while _harness.pid_alive(pid) and time.monotonic() < deadline: + time.sleep(0.05) + assert not _harness.pid_alive(pid), 'daemon did not die after SIGKILL' + assert not os.path.exists(tempdir), 'close() did not reclaim tempdir' + introspector.assert_no_leaks(before, label='daemon-crash teardown') + + def test_no_daemon_leak_when_owner_process_is_killed(self, instance_id, table_id): + project = self._project_id() + ctx = multiprocessing.get_context('fork') + result_queue = ctx.Queue() + owner = ctx.Process(target=_daemon_owner_process, args=(project, instance_id, table_id, result_queue)) + owner.start() + daemon_pid = tempdir = None + try: + deadline = time.monotonic() + 60.0 + while time.monotonic() < deadline: + try: + daemon_pid, tempdir = result_queue.get(timeout=0.2) + break + except queue.Empty: + assert owner.is_alive(), 'owner exited before starting daemon' + assert daemon_pid is not None, 'accelerator fell back to native' + assert _harness.pid_alive(daemon_pid), 'daemon not alive after owner start' + owner.kill() + owner.join(timeout=10) + assert not owner.is_alive(), 'owner process not reaped after SIGKILL' + deadline = time.monotonic() + 30.0 + while _harness.pid_alive(daemon_pid) and time.monotonic() < deadline: + time.sleep(0.1) + assert not _harness.pid_alive(daemon_pid), f'orphaned daemon {daemon_pid}' + finally: + if owner.is_alive(): + owner.kill() + owner.join() + if tempdir and os.path.isdir(tempdir): + shutil.rmtree(tempdir, ignore_errors=True)