Skip to content

fix(bigtable): close accelerator daemon and unregister instance on Table GC #19023

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mutianf wants to merge 1 commit into
base: bigtable-accel
Choose a base branch
from
Open

fix(bigtable): close accelerator daemon and unregister instance on Table GC #19023

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
66 changes: 66 additions & 0 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import sys
import tempfile
import time
import weakref
from typing import Any, Mapping, Sequence

import grpc
Expand Down Expand Up @@ -205,6 +206,7 @@ def __init__(
self._log_path: str | None = None
self._proc: subprocess.Popen[bytes] | None = None
self._auth_secret: str | None = None
self._finalizer: weakref.finalize | None = None

def __enter__(self) -> "AcceleratorDaemon":
"""Start the daemon on ``with`` entry and return it."""
Expand Down Expand Up @@ -359,8 +361,19 @@ def start(self) -> None:
self._cleanup_tempdir()
raise

# Close stdin and reap the daemon before Popen.__del__ runs if this
# instance is garbage-collected without close(). Otherwise Popen.__del__
# appends the live process to subprocess._active, keeping proc.stdin open.
self._finalizer = weakref.finalize(
self, AcceleratorDaemon._finalize_unclosed, self._proc, self._tempdir
)

def close(self) -> None:
"""Tear down the daemon and clean up the UDS tempdir."""
finalizer = self._finalizer
self._finalizer = None
if finalizer is not None:
finalizer.detach()
proc = self._proc
if proc is None:
return
Expand Down Expand Up @@ -540,3 +553,56 @@ def _cleanup_tempdir(self) -> None:
self._uds_path = None
self._log_path = None
self._auth_secret = None

@staticmethod
def _wait_for_proc_exit(
proc: subprocess.Popen[bytes] | None, timeout: float
) -> bool:
"""Wait up to ``timeout`` seconds for ``proc`` to exit."""
if proc is None:
return True
try:
proc.wait(timeout=timeout)
return True
except subprocess.TimeoutExpired:
return False
except Exception:
return True

@staticmethod
def _finalize_unclosed(
proc: subprocess.Popen[bytes] | None, tempdir: str | None
) -> None:
"""Stop the daemon subprocess and remove its tempdir on garbage collection.

Holding ``proc`` in ``weakref.finalize`` ensures stdin is closed and the
child exits before ``subprocess.Popen.__del__`` can resurrect it into
``subprocess._active``.
"""
try:
if proc is not None and proc.poll() is None:
if proc.stdin is not None:
try:
proc.stdin.close()
except Exception:
pass
if not AcceleratorDaemon._wait_for_proc_exit(
proc, _STDIN_GRACE_SECONDS
):
try:
proc.terminate()
except Exception:
pass
if not AcceleratorDaemon._wait_for_proc_exit(
proc, _SIGTERM_GRACE_SECONDS
):
try:
proc.kill()
except Exception:
pass
AcceleratorDaemon._wait_for_proc_exit(
proc, _SIGKILL_GRACE_SECONDS
)
finally:
if tempdir is not None and os.path.isdir(tempdir):
shutil.rmtree(tempdir, ignore_errors=True)
28 changes: 28 additions & 0 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1420,6 +1420,15 @@ def __init__(
raise RuntimeError(
f"{self.__class__.__name__} must be created within an async event loop context."
) from e
self._instance_finalizer = weakref.finalize(
self,
self._finalize_instance_registration,
self.client,
self.instance_id,
self.app_profile_id,
id(self),
self._register_instance_future,
)
Comment on lines +1423 to +1431

@gemini-code-assist gemini-code-assist Bot Nov 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

critical

Passing self._register_instance_future directly to weakref.finalize creates a reference cycle that prevents the TableAsync instance from ever being garbage collected.

Why this happens:

  1. weakref.finalize objects are stored in a global registry (weakref.finalize._registry) to keep them alive.
  2. This creates a strong reference path: global registry -> finalize object -> args -> self._register_instance_future.
  3. Since self._register_instance_future is a Task/Future wrapping a coroutine of self, it holds a strong reference to self.
  4. Therefore, self is reachable from a global root, meaning its reference count will never drop to 0, and the cyclic garbage collector will never collect it.

Solution:

Pass a weak reference to the future (weakref.ref(self._register_instance_future)) instead of the future itself. This breaks the reference cycle and allows the instance to be garbage collected as expected.

Suggested change
self._instance_finalizer = weakref.finalize(
self,
self._finalize_instance_registration,
self.client,
self.instance_id,
self.app_profile_id,
id(self),
self._register_instance_future,
)
self._instance_finalizer = weakref.finalize(
self,
self._finalize_instance_registration,
self.client,
self.instance_id,
self.app_profile_id,
id(self),
weakref.ref(self._register_instance_future),
)


# Optional in-process accelerator daemon, scoped to this data target's
# (project, instance_id, app_profile_id) tuple. Enabled by default;
Expand Down Expand Up @@ -1527,6 +1536,7 @@ def _start_accelerator(self) -> None:
server.uds_path, server.auth_secret
)
self._accelerator_daemon = server
self._daemon_finalizer = weakref.finalize(self, server.close)
except BaseException:
server.close()
raise
Expand Down Expand Up @@ -2493,18 +2503,36 @@ async def close(self):
finally:
self._accelerator_client = None
if self._accelerator_daemon is not None:
daemon_finalizer = getattr(self, "_daemon_finalizer", None)
if daemon_finalizer is not None:
daemon_finalizer.detach()
try:
self._accelerator_daemon.close()
except Exception as exc:
_LOGGER.warning("Failed to close accelerator daemon: %s", exc)
finally:
self._accelerator_daemon = None
instance_finalizer = getattr(self, "_instance_finalizer", None)
if instance_finalizer is not None:
instance_finalizer.detach()
if self._register_instance_future:
self._register_instance_future.cancel()
self.client._remove_instance_registration(
self.instance_id, self.app_profile_id, id(self)
)

@staticmethod
def _finalize_instance_registration(
client: Any,
instance_id: str,
app_profile_id: str | None,
owner_id: int,
register_future: Any,
) -> None:
if register_future:
register_future.cancel()
client._remove_instance_registration(instance_id, app_profile_id, owner_id)
Comment on lines +2524 to +2534

@gemini-code-assist gemini-code-assist Bot Nov 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

critical

Update _finalize_instance_registration to resolve the weak reference to the registration future before attempting to cancel it.

Suggested change
@staticmethod
def _finalize_instance_registration(
client: Any,
instance_id: str,
app_profile_id: str | None,
owner_id: int,
register_future: Any,
) -> None:
if register_future:
register_future.cancel()
client._remove_instance_registration(instance_id, app_profile_id, owner_id)
@staticmethod
def _finalize_instance_registration(
client: Any,
instance_id: str,
app_profile_id: str | None,
owner_id: int,
register_future_ref: Any,
) -> None:
register_future = register_future_ref() if register_future_ref else None
if register_future:
register_future.cancel()
client._remove_instance_registration(instance_id, app_profile_id, owner_id)


@CrossSync.convert(sync_name="__enter__")
async def __aenter__(self):
"""
Expand Down
33 changes: 30 additions & 3 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,12 +261,11 @@ def __init__(
if (
project is None
and credentials is None
and client_options is not None
and (client_options is not None)
and client_options.credentials_file
):
project = os.getenv(
environment_vars.PROJECT,
os.getenv(environment_vars.LEGACY_PROJECT),
environment_vars.PROJECT, os.getenv(environment_vars.LEGACY_PROJECT)
)
if project is None:
_, project = google.auth.load_credentials_from_file(
Expand Down Expand Up @@ -1122,6 +1121,15 @@ def __init__(
raise RuntimeError(
f"{self.__class__.__name__} must be created within an async event loop context."
) from e
self._instance_finalizer = weakref.finalize(
self,
self._finalize_instance_registration,
self.client,
self.instance_id,
self.app_profile_id,
id(self),
self._register_instance_future,
)
Comment on lines +1124 to +1132

@gemini-code-assist gemini-code-assist Bot Nov 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

critical

Passing self._register_instance_future directly to weakref.finalize creates a reference cycle that prevents the Table instance from ever being garbage collected.

Why this happens:

  1. weakref.finalize objects are stored in a global registry (weakref.finalize._registry) to keep them alive.
  2. This creates a strong reference path: global registry -> finalize object -> args -> self._register_instance_future.
  3. Since self._register_instance_future is a Task/Future wrapping a coroutine of self, it holds a strong reference to self.
  4. Therefore, self is reachable from a global root, meaning its reference count will never drop to 0, and the cyclic garbage collector will never collect it.

Solution:

Pass a weak reference to the future (weakref.ref(self._register_instance_future)) instead of the future itself. This breaks the reference cycle and allows the instance to be garbage collected as expected.

Suggested change
self._instance_finalizer = weakref.finalize(
self,
self._finalize_instance_registration,
self.client,
self.instance_id,
self.app_profile_id,
id(self),
self._register_instance_future,
)
self._instance_finalizer = weakref.finalize(
self,
self._finalize_instance_registration,
self.client,
self.instance_id,
self.app_profile_id,
id(self),
weakref.ref(self._register_instance_future),
)

self._accelerator_daemon: AcceleratorDaemon | None = None
self._accelerator_client: AcceleratorClientType | None = None
self._accelerator_health: AcceleratorHealthMonitorType | None = None
Expand Down Expand Up @@ -1194,6 +1202,7 @@ def _start_accelerator(self) -> None:
server.uds_path, server.auth_secret
)
self._accelerator_daemon = server
self._daemon_finalizer = weakref.finalize(self, server.close)
except BaseException:
server.close()
raise
Expand Down Expand Up @@ -2046,18 +2055,36 @@ def close(self):
finally:
self._accelerator_client = None
if self._accelerator_daemon is not None:
daemon_finalizer = getattr(self, "_daemon_finalizer", None)
if daemon_finalizer is not None:
daemon_finalizer.detach()
try:
self._accelerator_daemon.close()
except Exception as exc:
_LOGGER.warning("Failed to close accelerator daemon: %s", exc)
finally:
self._accelerator_daemon = None
instance_finalizer = getattr(self, "_instance_finalizer", None)
if instance_finalizer is not None:
instance_finalizer.detach()
if self._register_instance_future:
self._register_instance_future.cancel()
self.client._remove_instance_registration(
self.instance_id, self.app_profile_id, id(self)
)

@staticmethod
def _finalize_instance_registration(
client: Any,
instance_id: str,
app_profile_id: str | None,
owner_id: int,
register_future: Any,
) -> None:
if register_future:
register_future.cancel()
client._remove_instance_registration(instance_id, app_profile_id, owner_id)
Comment on lines +2076 to +2086

@gemini-code-assist gemini-code-assist Bot Nov 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

critical

Update _finalize_instance_registration to resolve the weak reference to the registration future before attempting to cancel it.

Suggested change
@staticmethod
def _finalize_instance_registration(
client: Any,
instance_id: str,
app_profile_id: str | None,
owner_id: int,
register_future: Any,
) -> None:
if register_future:
register_future.cancel()
client._remove_instance_registration(instance_id, app_profile_id, owner_id)
@staticmethod
def _finalize_instance_registration(
client: Any,
instance_id: str,
app_profile_id: str | None,
owner_id: int,
register_future_ref: Any,
) -> None:
register_future = register_future_ref() if register_future_ref else None
if register_future:
register_future.cancel()
client._remove_instance_registration(instance_id, app_profile_id, owner_id)


def __enter__(self):
"""Implement async context manager protocol

Expand Down
66 changes: 66 additions & 0 deletions packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -475,6 +475,72 @@ def test_close_clears_proc_reference(self, tmp_path):

assert d._proc is None

def test_gc_without_close_reaps_proc_and_removes_tempdir(
self, tmp_path, monkeypatch
):
"""Garbage-collecting an unclosed AcceleratorDaemon closes stdin, reaps
the subprocess before Popen.__del__ runs, and removes the tempdir."""
import gc
import subprocess
import sys
import warnings
import weakref

real_popen = subprocess.Popen

def fake_popen(argv, **kwargs):
return real_popen(
[sys.executable, "-c", "import sys; sys.stdin.read()"],
stdin=subprocess.PIPE,
stdout=kwargs.get("stdout"),
stderr=kwargs.get("stderr"),
)

monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen)
monkeypatch.setattr(
AcceleratorDaemon, "_wait_until_ready", lambda self, t: None
)

d = _make_daemon(tmp_path)
d.start()
tempdir = d._tempdir
assert tempdir is not None and os.path.isdir(tempdir)
proc_weak = weakref.ref(d._proc)

with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always", ResourceWarning)
del d
gc.collect()

res_warnings = [w for w in caught if issubclass(w.category, ResourceWarning)]
assert not res_warnings
assert proc_weak() is None
assert not os.path.exists(tempdir)

def test_close_detaches_finalizer(self, tmp_path, monkeypatch):
"""Calling close() detaches the weakref finalizer so it does not run twice."""
import subprocess

proc_mock = unittest.mock.Mock(spec=subprocess.Popen)
proc_mock.poll.return_value = None
proc_mock.stdin = unittest.mock.Mock()

monkeypatch.setattr(_daemon.subprocess, "Popen", lambda *a, **kw: proc_mock)
monkeypatch.setattr(
AcceleratorDaemon, "_wait_until_ready", lambda self, t: None
)

d = _make_daemon(tmp_path)
d.start()
finalizer = d._finalizer
assert finalizer is not None and finalizer.alive

with unittest.mock.patch.object(d, "_wait_for_exit", return_value=True):
d.close()

assert not finalizer.alive
assert d._finalizer is None


# ---------------------------------------------------------------------------
# _cleanup_tempdir: state clearing
Expand Down
73 changes: 73 additions & 0 deletions packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -987,6 +987,79 @@ def capture_daemon(cli_flags=(), **kwargs):

assert "--app-profile" not in captured["cli_flags"]

def test_gc_and_explicit_close_finalizers(self):
"""Dropping an unclosed Table/TableAsync stops health, closes the daemon,
and unregisters the instance; explicit close() detaches the finalizers."""
import gc

from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon
from google.cloud.bigtable.data._async.client import TableAsync
from google.cloud.bigtable.data._sync_autogen.client import Table

for table_cls, is_async, call_close in (
(TableAsync, True, False),
(Table, False, False),
(TableAsync, True, True),
(Table, False, True),
):
fake_daemon = mock.Mock(
spec=AcceleratorDaemon, uds_path="/tmp/s", auth_secret="sec"
)
fake_health = mock.Mock(close=mock.AsyncMock() if is_async else mock.Mock())
fake_client = mock.Mock(close=mock.AsyncMock() if is_async else mock.Mock())
fake_future = mock.Mock()

@gemini-code-assist gemini-code-assist Bot Nov 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

medium

Since mock.Mock objects do not support weak references by default, passing a standard mock to weakref.ref will raise a TypeError. Use a custom subclass of mock.Mock that defines __slots__ with __weakref__ to allow weak referencing in tests.

            class WeakableMock(mock.Mock):
                __slots__ = ("__weakref__",)
            fake_future = WeakableMock()

remove_reg = mock.Mock()
client = SimpleNamespace(
project="p",
_emulator_host=None,
_accelerator_blocked_reason=None,
_accelerator_flags=[],
_accelerator_env={},
_use_accelerator=True,
_executor=None,
_gapic_client=mock.Mock(
instance_path=lambda p, i: f"projects/{p}/instances/{i}"
),
_register_instance=mock.Mock(),
_remove_instance_registration=remove_reg,
)
mod = (
"google.cloud.bigtable.data._async.client"
if is_async
else "google.cloud.bigtable.data._sync_autogen.client"
)
task_fn = (
"CrossSync.create_task"
if is_async
else "CrossSync._Sync_Impl.create_task"
)
with (
mock.patch(f"{mod}.{task_fn}", return_value=fake_future),
mock.patch(f"{mod}.AcceleratorDaemon", return_value=fake_daemon),
mock.patch(f"{mod}.AcceleratorClientType", return_value=fake_client),
mock.patch(
f"{mod}.AcceleratorHealthMonitorType", return_value=fake_health
),
mock.patch.object(table_cls, "_verify_daemon_identity"),
):
table = table_cls(client, "inst-1", "tbl-1", app_profile_id="prof-1")
table_id = id(table)

if call_close:
if is_async:
asyncio.get_event_loop().run_until_complete(table.close())
else:
table.close()

del table
gc.collect()

if not call_close:
fake_health.request_stop.assert_called_once()
fake_daemon.close.assert_called_once()
fake_future.cancel.assert_called_once()
remove_reg.assert_called_once_with("inst-1", "prof-1", table_id)


# ---------------------------------------------------------------------------
# Event loop fixture so run_until_complete works
Expand Down
Loading