REVERSE MERGE - v0.138.15-hotfix#1625
Merged
Merged
Conversation
…tkeeper and resolve sigsegv issues (#1616) * UN-2910 [FIX] Prevent task duplication during Kubernetes warm shutdown with HeartbeatKeeper Implemented HeartbeatKeeper bootstep to maintain RabbitMQ connection during Celery warm shutdown, preventing task requeuing and duplication during Kubernetes rolling updates and graceful pod terminations. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * structuretool-for-testing * handling envs * sonar issue addressing * forksafe gcs and gdrive connectors * updated grpc imports * fork safty in shared internal client session * reverted forksafty fropm baseclient, addressed coderabbitai comments * minor changes * reverted structure tool changes that was changed for tetsing * plugin registry changes for .so file comaptibility * added default value for heartbeat --------- Co-authored-by: Claude <noreply@anthropic.com>
Contributor
Summary by CodeRabbit
WalkthroughThis pull request introduces fork-safety improvements for Google Cloud integrations and Celery worker infrastructure. Changes include lazy initialization patterns for Google API clients to defer gRPC state creation after process fork, thread-safe locking mechanisms to prevent deadlocks in prefork scenarios, a new heartbeat keeper bootstep for warm shutdown, Celery pool type configuration support, and disabling API client singleton by default in prefork environments. Changes
Sequence Diagram(s)sequenceDiagram
actor Worker as Celery Worker
participant Init as __init__
participant Fork as fork()
participant Method as get_fsspec_fs()
participant Lock as Lock
participant Client as GCS Client
Worker->>Init: Create connector instance
Note over Init: Store credentials only,<br/>no client creation
Init-->>Worker: Return
Worker->>Fork: Fork process
Note over Fork: Child process spawned
Worker->>Method: Call get_fsspec_fs()
Method->>Lock: Acquire lock
Lock-->>Method: Lock acquired
alt Client already initialized
Method-->>Lock: Skip creation
else First call after fork
Method->>Client: Initialize credentials
Method->>Client: Create GCS client
Client-->>Method: Client ready
end
Method->>Lock: Release lock
Lock-->>Method: Lock released
Method-->>Worker: Return filesystem
sequenceDiagram
actor Start as Worker Startup
participant Bootstep as HeartbeatKeeper Bootstep
participant Signal as on_worker_process_init Signal
participant Task as gRPC State
participant Thread as Heartbeat Thread
Start->>Bootstep: Worker initialized
Bootstep->>Signal: Register process_init signal handler
Signal-->>Bootstep: Handler registered
Bootstep->>Thread: Create heartbeat thread
Note over Thread: Monitors warm shutdown<br/>standby during normal ops
Thread-->>Bootstep: Thread started
Note over Bootstep: When fork occurs
Signal->>Signal: on_worker_process_init triggered
Signal->>Task: Reinitialize gRPC state
Task->>Task: Clear gRPC channels
Task-->>Signal: State reinitialized
Signal->>Thread: Resume with clean state
Thread-->>Worker: Ready
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25-30 minutes
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
vishnuszipstack
approved these changes
Oct 31, 2025
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
workers/shared/api/internal_client.py (1)
149-172: Pass the shared WorkerConfig into WorkflowAPIClientIn both
_initialize_core_clients_optimizedand_initialize_core_clients_traditionalthe newWorkflowAPIClient()instances are created withoutself.config. If a caller provides a customizedWorkerConfig(different base URL, API key, timeouts, etc.), every other client respects it, but this new client will silently fall back to defaults, breaking those overrides and potentially sending requests to the wrong host. Please wire the shared config through so the workflow client stays consistent.- self.workflow_client = WorkflowAPIClient() + self.workflow_client = WorkflowAPIClient(self.config) self.workflow_client.session.close() self.workflow_client.session = InternalAPIClient._shared_session @@ - self.workflow_client = WorkflowAPIClient() + self.workflow_client = WorkflowAPIClient(self.config)Also applies to: 177-186
🧹 Nitpick comments (1)
workers/worker.py (1)
138-226: Usetyping.Anyin type hintsThe new helper methods return/accept
any, which references the builtin function rather than a type, so these annotations won’t pass type checking. Please importAnyfromtypingand swap the hints to keep static analysis green.-import logging -import os -import sys -import threading -import time +import logging +import os +import sys +import threading +import time +from typing import Any @@ - def _is_connection_ready(self) -> tuple[bool, any]: + def _is_connection_ready(self) -> tuple[bool, Any]: @@ - self, consumer: any, consecutive_errors: int, last_heartbeat: float + self, consumer: Any, consecutive_errors: int, last_heartbeat: float
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (11)
docker/dockerfiles/worker-unified.Dockerfile(1 hunks)unstract/connectors/src/unstract/connectors/databases/bigquery/bigquery.py(2 hunks)unstract/connectors/src/unstract/connectors/filesystems/google_cloud_storage/google_cloud_storage.py(4 hunks)unstract/connectors/src/unstract/connectors/filesystems/google_drive/google_drive.py(3 hunks)unstract/connectors/src/unstract/connectors/gcs_helper.py(4 hunks)workers/client_plugin_registry.py(2 hunks)workers/run-worker.sh(9 hunks)workers/sample.env(1 hunks)workers/shared/api/internal_client.py(1 hunks)workers/shared/infrastructure/config/worker_config.py(1 hunks)workers/worker.py(3 hunks)
🧰 Additional context used
🪛 Ruff (0.14.2)
workers/worker.py
20-20: Unused noqa directive (non-enabled: E402)
Remove unused noqa directive
(RUF100)
26-26: Unused noqa directive (non-enabled: E402)
Remove unused noqa directive
(RUF100)
27-27: Unused noqa directive (non-enabled: E402)
Remove unused noqa directive
(RUF100)
28-28: Unused noqa directive (non-enabled: E402)
Remove unused noqa directive
(RUF100)
29-29: Unused noqa directive (non-enabled: E402)
Remove unused noqa directive
(RUF100)
71-71: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
73-73: Unused method argument: kwargs
(ARG002)
80-80: Unused method argument: worker
(ARG002)
121-121: Unused method argument: worker
(ARG002)
130-130: Unused method argument: worker
(ARG002)
174-174: Consider moving this statement to an else block
(TRY300)
176-176: Do not catch blind exception: Exception
(BLE001)
222-222: Do not catch blind exception: Exception
(BLE001)
342-342: Unused function argument: kwargs
(ARG001)
unstract/connectors/src/unstract/connectors/filesystems/google_cloud_storage/google_cloud_storage.py
137-137: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
137-137: Use explicit conversion flag
Replace with conversion flag
(RUF010)
141-141: Use explicit conversion flag
Replace with conversion flag
(RUF010)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (11)
docker/dockerfiles/worker-unified.Dockerfile (1)
24-30: Addingbuild-essentialimproves reproducibility
Shipping the whole toolchain up front prevents missing headers during native builds. Looks good to me.unstract/connectors/src/unstract/connectors/databases/bigquery/bigquery.py (1)
9-235: Deferred import pattern looks solid
Thanks for moving the exceptions import intoexecute_query; this keeps the handler logic untouched while avoiding pre-fork gRPC init.unstract/connectors/src/unstract/connectors/filesystems/google_cloud_storage/google_cloud_storage.py (1)
16-145: Fork-safe lazy initialization implemented well
Double-checked locking plus JSON validation keeps the GCS filesystem creation both safe and deterministic. Nice work.unstract/connectors/src/unstract/connectors/filesystems/google_drive/google_drive.py (1)
14-210: Lazy secret loading + Drive init look correct
Appreciate the thorough locking and error mapping to keep Google Drive fork-safe without regressing UX.workers/shared/infrastructure/config/worker_config.py (1)
245-252: Defaulting the singleton off is a good callThanks for making the fork-safety stance explicit in both the comment and default value. This keeps prefork workers stable by default while still allowing explicit opt-in when needed. Nice alignment with the broader hotfix goals.
workers/client_plugin_registry.py (1)
137-193: Nice enhancement for compiled pluginsAppreciate the dual lookup for
.pyand Nuitka-produced.soartifacts with graceful fallback. This widens deployment options without disturbing the existing import flow.workers/sample.env (1)
92-101: Clear guidance in sample envThe expanded comment spells out exactly why the singleton is off by default and when it’s safe to re-enable it. That will save operators a lot of guesswork—thanks for keeping the docs in lockstep with the config.
workers/run-worker.sh (4)
330-333: Pool type option added to Celery command correctly.The conditional check for
pool_typeproperly gates the--pool=$pool_typeargument addition to the Celery command. The argument is inserted at the right place in the command structure—after logging and queue configuration, before concurrency settings—maintaining logical grouping of options.
89-89: Help text and argument parsing comprehensive and well‑documented.The usage documentation clearly describes the new
-P/--pooloption with supported pool types and practical examples. The argument parsing correctly handles the option withshift 2, and the examples (lines 115-117) demonstrate realistic usage patterns for operators.Also applies to: 115-117, 460-463
429-429: Pool type initialization aligns with other configuration variables.The
POOL_TYPE=""initialization follows the same pattern as other optional configuration variables (CUSTOM_QUEUES,HEALTH_PORT,CUSTOM_HOSTNAME), maintaining consistency in the script's variable setup.
269-270: Pool type parameter implementation is complete and consistent.Verification confirms both internal invocations of
run_worker()at lines 403 and 530 correctly pass all 8 parameters in the proper order. No external callers exist outside run-worker.sh (run-worker-docker.sh is an independent script with its own run_worker function). The signature change from 7 to 8 parameters is self-contained and poses no backward compatibility risks. The implementation is correct and complete.
Comment on lines
+48
to
+109
| # Lazy initialization - credentials created only when needed (after fork) | ||
| self._google_credentials: Any | None = None | ||
| self._credentials_lock = threading.Lock() | ||
|
|
||
| def _get_credentials(self) -> Any: | ||
| """Lazily create credentials object (fork-safe). | ||
|
|
||
| This method is called after fork, ensuring gRPC state is created | ||
| in the child process, not inherited from parent. | ||
| """ | ||
| if self._google_credentials is None: | ||
| with self._credentials_lock: | ||
| if self._google_credentials is None: | ||
| # Import inside method to avoid module-level initialization | ||
| from google.oauth2 import service_account | ||
|
|
||
| logger.debug("Creating Google credentials (lazy init after fork)") | ||
| self._google_credentials = ( | ||
| service_account.Credentials.from_service_account_info( | ||
| json.loads(self.google_service_json) | ||
| ) | ||
| ) | ||
| return self._google_credentials | ||
|
|
||
| def get_google_credentials(self) -> Any: | ||
| """Get Google credentials, creating them lazily if needed.""" | ||
| return self._get_credentials() | ||
|
|
||
| def get_secret(self, secret_name: str) -> str: | ||
| from google.cloud import secretmanager | ||
| """Get secret from Google Secret Manager (lazy init after fork). | ||
|
|
||
| Thread-safe: Uses global lock to prevent gRPC deadlock when multiple | ||
| threads create Secret Manager clients simultaneously. | ||
| """ | ||
| # Use global lock to serialize Secret Manager client creation ONLY | ||
| # This prevents deadlock when multiple threads hit this simultaneously | ||
| with _GOOGLE_API_INIT_LOCK: | ||
| from google.cloud import secretmanager | ||
|
|
||
| credentials = self._get_credentials() | ||
| google_secrets_client = secretmanager.SecretManagerServiceClient( | ||
| credentials=credentials | ||
| ) | ||
|
|
||
| google_secrets_client = secretmanager.SecretManagerServiceClient( | ||
| credentials=self.google_credentials | ||
| ) | ||
| # Network I/O happens OUTSIDE the lock for better concurrency | ||
| s = google_secrets_client.access_secret_version( | ||
| request={ | ||
| "name": f"projects/{self.google_project_id}/secrets/{secret_name}/versions/latest" # noqa: E501 | ||
| } | ||
| "name": f"projects/{self.google_project_id}/secrets/{secret_name}/versions/latest" | ||
| }, | ||
| timeout=30.0, # Add reasonable timeout to prevent indefinite hangs | ||
| ) | ||
| return s.payload.data.decode("UTF-8") | ||
|
|
||
| def _get_storage_client(self) -> Any: | ||
| """Lazily create GCS client (fork-safe and thread-safe).""" | ||
| # Use global lock to serialize GCS client creation | ||
| with _GOOGLE_API_INIT_LOCK: | ||
| from google.cloud.storage import Client | ||
|
|
||
| credentials = self._get_credentials() | ||
| return Client(credentials=credentials) | ||
|
|
Contributor
There was a problem hiding this comment.
Cache the storage client to avoid gRPC leaks
As written, every call to _get_storage_client() instantiates a brand-new storage.Client. Those gRPC clients spin up background threads and open channels that are never closed, so repeated uploads/downloads will steadily leak threads/file descriptors. We still need lazy, post-fork creation, but once created the client should be cached on the helper. Please store it on the instance and reuse it, e.g.:
@@
- self._google_credentials: Any | None = None
- self._credentials_lock = threading.Lock()
+ self._google_credentials: Any | None = None
+ self._credentials_lock = threading.Lock()
+ self._storage_client: Any | None = None
@@
- def _get_storage_client(self) -> Any:
- """Lazily create GCS client (fork-safe and thread-safe)."""
- # Use global lock to serialize GCS client creation
- with _GOOGLE_API_INIT_LOCK:
- from google.cloud.storage import Client
-
- credentials = self._get_credentials()
- return Client(credentials=credentials)
+ def _get_storage_client(self) -> Any:
+ """Lazily create and cache the GCS client (fork-safe and thread-safe)."""
+ if self._storage_client is None:
+ with _GOOGLE_API_INIT_LOCK:
+ if self._storage_client is None:
+ from google.cloud.storage import Client
+
+ credentials = self._get_credentials()
+ self._storage_client = Client(credentials=credentials)
+ return self._storage_clientThis keeps the initialization fork-safe while preventing the runaway resource growth.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # Lazy initialization - credentials created only when needed (after fork) | |
| self._google_credentials: Any | None = None | |
| self._credentials_lock = threading.Lock() | |
| def _get_credentials(self) -> Any: | |
| """Lazily create credentials object (fork-safe). | |
| This method is called after fork, ensuring gRPC state is created | |
| in the child process, not inherited from parent. | |
| """ | |
| if self._google_credentials is None: | |
| with self._credentials_lock: | |
| if self._google_credentials is None: | |
| # Import inside method to avoid module-level initialization | |
| from google.oauth2 import service_account | |
| logger.debug("Creating Google credentials (lazy init after fork)") | |
| self._google_credentials = ( | |
| service_account.Credentials.from_service_account_info( | |
| json.loads(self.google_service_json) | |
| ) | |
| ) | |
| return self._google_credentials | |
| def get_google_credentials(self) -> Any: | |
| """Get Google credentials, creating them lazily if needed.""" | |
| return self._get_credentials() | |
| def get_secret(self, secret_name: str) -> str: | |
| from google.cloud import secretmanager | |
| """Get secret from Google Secret Manager (lazy init after fork). | |
| Thread-safe: Uses global lock to prevent gRPC deadlock when multiple | |
| threads create Secret Manager clients simultaneously. | |
| """ | |
| # Use global lock to serialize Secret Manager client creation ONLY | |
| # This prevents deadlock when multiple threads hit this simultaneously | |
| with _GOOGLE_API_INIT_LOCK: | |
| from google.cloud import secretmanager | |
| credentials = self._get_credentials() | |
| google_secrets_client = secretmanager.SecretManagerServiceClient( | |
| credentials=credentials | |
| ) | |
| google_secrets_client = secretmanager.SecretManagerServiceClient( | |
| credentials=self.google_credentials | |
| ) | |
| # Network I/O happens OUTSIDE the lock for better concurrency | |
| s = google_secrets_client.access_secret_version( | |
| request={ | |
| "name": f"projects/{self.google_project_id}/secrets/{secret_name}/versions/latest" # noqa: E501 | |
| } | |
| "name": f"projects/{self.google_project_id}/secrets/{secret_name}/versions/latest" | |
| }, | |
| timeout=30.0, # Add reasonable timeout to prevent indefinite hangs | |
| ) | |
| return s.payload.data.decode("UTF-8") | |
| def _get_storage_client(self) -> Any: | |
| """Lazily create GCS client (fork-safe and thread-safe).""" | |
| # Use global lock to serialize GCS client creation | |
| with _GOOGLE_API_INIT_LOCK: | |
| from google.cloud.storage import Client | |
| credentials = self._get_credentials() | |
| return Client(credentials=credentials) | |
| # Lazy initialization - credentials created only when needed (after fork) | |
| self._google_credentials: Any | None = None | |
| self._credentials_lock = threading.Lock() | |
| self._storage_client: Any | None = None | |
| def _get_credentials(self) -> Any: | |
| """Lazily create credentials object (fork-safe). | |
| This method is called after fork, ensuring gRPC state is created | |
| in the child process, not inherited from parent. | |
| """ | |
| if self._google_credentials is None: | |
| with self._credentials_lock: | |
| if self._google_credentials is None: | |
| # Import inside method to avoid module-level initialization | |
| from google.oauth2 import service_account | |
| logger.debug("Creating Google credentials (lazy init after fork)") | |
| self._google_credentials = ( | |
| service_account.Credentials.from_service_account_info( | |
| json.loads(self.google_service_json) | |
| ) | |
| ) | |
| return self._google_credentials | |
| def get_google_credentials(self) -> Any: | |
| """Get Google credentials, creating them lazily if needed.""" | |
| return self._get_credentials() | |
| def get_secret(self, secret_name: str) -> str: | |
| """Get secret from Google Secret Manager (lazy init after fork). | |
| Thread-safe: Uses global lock to prevent gRPC deadlock when multiple | |
| threads create Secret Manager clients simultaneously. | |
| """ | |
| # Use global lock to serialize Secret Manager client creation ONLY | |
| # This prevents deadlock when multiple threads hit this simultaneously | |
| with _GOOGLE_API_INIT_LOCK: | |
| from google.cloud import secretmanager | |
| credentials = self._get_credentials() | |
| google_secrets_client = secretmanager.SecretManagerServiceClient( | |
| credentials=credentials | |
| ) | |
| # Network I/O happens OUTSIDE the lock for better concurrency | |
| s = google_secrets_client.access_secret_version( | |
| request={ | |
| "name": f"projects/{self.google_project_id}/secrets/{secret_name}/versions/latest" | |
| }, | |
| timeout=30.0, # Add reasonable timeout to prevent indefinite hangs | |
| ) | |
| return s.payload.data.decode("UTF-8") | |
| def _get_storage_client(self) -> Any: | |
| """Lazily create and cache the GCS client (fork-safe and thread-safe).""" | |
| if self._storage_client is None: | |
| with _GOOGLE_API_INIT_LOCK: | |
| if self._storage_client is None: | |
| from google.cloud.storage import Client | |
| credentials = self._get_credentials() | |
| self._storage_client = Client(credentials=credentials) | |
| return self._storage_client |
🤖 Prompt for AI Agents
In unstract/connectors/src/unstract/connectors/gcs_helper.py around lines 48 to
109, _get_storage_client() currently creates a new google.cloud.storage.Client
on every call which leaks gRPC threads; change it to lazily create and cache a
single client on the instance (e.g. self._google_storage_client = None at init),
use the existing _GOOGLE_API_INIT_LOCK to guard creation, return the cached
client on subsequent calls, and ensure you import and construct the Client only
once after the lock so the client is fork-safe and reused across calls.
ritwik-g
approved these changes
Oct 31, 2025
ritwik-g
changed the title
|
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Why
How
Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)
No, this PR will not break existing features:
Potential edge cases (not breaking, but worth noting):
Database Migrations
Env Config
New Environment Variables (all optional, have defaults):
HEARTBEAT_KEEPER_ENABLED(default:true)true(enabled), any other value (disabled)HEARTBEAT_KEEPER_SHUTDOWN_INTERVAL(default:CELERY_BROKER_HEARTBEAT / 5)HEARTBEAT_KEEPER_CHECK_INTERVAL(default:CELERY_BROKER_HEARTBEAT)HEARTBEAT_KEEPER_MAX_ERRORS(default:10)Existing Config Used:
CELERY_BROKER_HEARTBEAT(default: 10) - Used to calculate smart defaults and determines RabbitMQ connection timeoutRelevant Docs
Related Issues or PRs
Dependencies Versions
Notes on Testing
Manual Testing Performed:
Normal Operation Test:
Warm Shutdown Test (Task Completes):
Warm Shutdown Test (Connection Degraded):
Disable Feature Test:
Configuration Test:
RabbitMQ Logs Verified:
Expected Behavior:
Screenshots
Not applicable (backend infrastructure change)
Checklist
I have read and understood the Contribution Guidelines.
🤖 Generated with Claude Code