[ENH] Clean up client manager into manager/assigner - make log client use it #4640

HammadB · 2025-05-28T03:12:58Z

Description of changes

Summarize the changes made by this PR.

Remove set_system pattern in lieu of passing system into try_from_config - aligning with the common. pattern since we touched this
Refactor client manager into client manager + assigner thats generic over its clients
Use this in the old call site and in the log, the log uses assignment for write path, and reads all for read path

Test plan

How are these changes tested?

Tests pass locally with pytest for python, yarn test for js, cargo test for rust

Documentation Changes

None

HammadB · 2025-05-28T03:13:14Z

This stack of pull requests is managed by Graphite. Learn more about stacking.

github-actions · 2025-05-28T03:13:19Z

github-actions · 2025-05-28T03:13:20Z

Please tag your PR title with one of: [ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas

github-actions · 2025-05-28T03:14:31Z

Please tag your PR title with one of: [ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas

HammadB · 2025-05-28T03:14:58Z

Cargo.toml

@@ -69,7 +69,7 @@ tokio = { version = "1.41", features = ["fs", "macros", "rt-multi-thread"] }
 tokio-util = "0.7.12"
 tonic = "0.12"
 tonic-health = "0.12.3"
-tower = "0.5"
+tower = { version = "0.4.13", features = ["discover"] }


Some users of this were on an older version and depend on it, keep that version for now.

HammadB · 2025-05-28T03:16:30Z

rust/frontend/src/executor/distributed.rs

        let client_manager = ClientManager::new(
-            node_name_to_client.clone(),
+            client_assigner.clone(),
            config.connections_per_node,


I will move these other config into the ClientOptions

will address in another PR

HammadB · 2025-05-28T03:17:08Z

rust/frontend/src/executor/distributed.rs

+                    .collection_id
+                    .to_string(),
+            )
+            .map_err(|e| ExecutorError::Internal(e.boxed()))?;


Forced to box this error due to a cyclic dep since ExecutorError is in types and the crate that defines this error depends on types

HammadB · 2025-05-28T03:18:09Z

rust/log/src/grpc_log.rs

+        let client_assigner = ClientAssigner::new(assignment_policy, 1);
+        let client_manager = ClientManager::new(
+            client_assigner.clone(),
+            1,


replication factor is always one for log

HammadB · 2025-05-28T03:18:26Z

rust/log/src/grpc_log.rs

-        } else {
-            None
+
+        // Alt client manager setup


This could be made a bit cleaner, will leave that for another pr since the other use already did it this way

HammadB · 2025-05-28T03:18:32Z

rust/log/src/grpc_log.rs

+        memberlist_provider.subscribe(client_manager_handle.receiver());
+        let _memberlist_provider_handle = system.start_component(memberlist_provider);
+
+        // let alt_client = if let Some(alt_host) = my_config.alt_host.as_ref() {


will remove

HammadB · 2025-05-28T03:18:51Z

rust/log/src/grpc_log.rs

+                    .unwrap_or_default()
+                    .drain(..)
+                    .next()
+                    // fallback to the standard client if no alt client is available,


this is a temporary choice, when we remove alt client we can remove this

HammadB · 2025-05-28T03:19:02Z

rust/log/src/grpc_log.rs

            }
        }
        tracing::info!("using standard client for {collection_id}");
-        &mut self.client
+        self.client.clone()


grpc client clones are cheap

HammadB · 2025-05-28T03:19:25Z

rust/log/src/grpc_log.rs

-                        },
-                    )
-                    .await
+        let mut combined_response = Vec::new();


Please ensure this is correct

This looks correct to me. I read it twice over.

propel-code-bot · 2025-05-28T03:22:25Z

Major Refactor: Client Manager Abstraction, Memberlist Integration, & System Hydration Cleanup

This PR significantly refactors and centralizes client (gRPC) management logic for both the distributed executor and log components by introducing a generic, reusable ClientManager and ClientAssigner abstraction within the chroma-memberlist crate. The PR eliminates legacy patterns such as set_system, streamlines system hydration across components, and migrates all systems that relied on ad hoc or redundant client management (including log and executor) to use the new abstractions. It also adapts interfaces and configs to support the new generic lifecycle and dependency requirements, and updates configuration to facilitate easier swapping of assignment policies and memberlist providers. Support and wiring for System everywhere is modernized, and integration test and config changes reflect all these improvements.

Key Changes:
• Introduced a new, generic ClientManager<T> and ClientAssigner<T> system in chroma-memberlist for type-safe management/assignment of gRPC clients.
• Refactored the distributed executor, log client, and test logic to rely on the new managed client system instead of direct references or custom node-to-client maps.
• Made client management generic over client type, allowing both log and query services to use the same infra.
• Eliminated legacy set_system/Option<System> patterns in favor of constructor injection everywhere.
• Adapted configs, especially GrpcLogConfig, to hold assignment and memberlist provider settings (with sensible defaults).
• Upgraded dependency handling/tower versioning and workspace usage.
• Removed legacy client_manager module from frontend (now in memberlist crate).
• Reworked and expanded tests to work with the new manager/assigner.
• Updated Cargo.toml/Cargo.lock as appropriate for new dependencies and organization.

Affected Areas:
• chroma-memberlist (new client_manager)
• chroma-frontend (executor/distributed)
• chroma-log (grpc_log, config, lib)
• chroma-worker (compactor/compaction_manager, server, lib, config)
• types/api_types.rs (ExecutorError simplification)
• Cargo.toml, Cargo.lock (dependency updates)
• Various component and config wiring across log, worker, frontend

Potential Impact:

Functionality: Should improve reliability and flexibility in rolling out new node membership, assignment policies, and dynamic membership scenarios for both log and query services. Removes old system hydration footguns.

Performance: No expected regression; centralization could improve efficiency, but extra indirection may have minor effect. gRPC client clone noted as cheap.

Security: No direct security changes, but more robust system wiring and centralized error paths could reduce some edge error handling flaws.

Scalability: Enables more robust scaling with clean assignment and memberlist policy swapping. Generalizes client replica management and flexibility.

Review Focus:
• Correctness of the client manager/assigner generic abstractions: race/locking, type safety.
• All sites now using new manager; no direct NodeNameToClient usages left.
• Construction (hydration) of all major components (worker, frontend, log, compactor) -- especially System injection consistency.
• Backward compatibility: configuration defaults, assignment/memberlist config override points.
• Cargo.toml and lockfile workspace sanity, no duplicate tower/tonic/futures versions creeping in.
• Log clients and distributed executor: assignment and fallback logic (no client found, alt client selection).

Testing Needed

• Ensure that all tests pass and that membership changes propagate correctly to clients in log and executor paths.
• Test assignment policy swaps, especially with non-default configs.
• Run distributed queries and log operations (write and read paths) across node churn, multi-node clusters.
• Regression test for system startup/shutdown and graceful teardown (signal handling etc.).
• Test client fallback paths for missing alt log clients (log client assignment logic).

Code Quality Assessment

rust/memberlist/src/client_manager.rs: High complexity, but clearly structured; heavy use of generics and trait bounds is correctly applied.

rust/log/src/grpc_log.rs: Now cleaner with manager/assigner usage; alt client handing more robust but could use future cleanup.

rust/frontend/src/executor/distributed.rs: Greatly simplified, clearer separation of assignment vs. concrete GRPC client handling.

rust/worker/src/compactor/compaction_manager.rs: Simplified system wiring, less boilerplate and set_system/set_dispatcher hacks.

Cargo.toml/Cargo.lock: Workspace improvements and dependency pinning all look correct for new crate structure.

Best Practices

Dependency Injection:
• Moves all hydration to constructor style, eliminates Options in stateful core types

Error Handling:
• Boxing errors where cyclic deps exist; error paths are clear and most new APIs are Result-returning

Testing:
• Extensive test updates, including multi-member handling, churn, and unit tests for new manager/assigner

Configuration:
• Config structs now use sensible defaults and allow flexible overrides

Abstraction:
• Extracts reusable generic patterns (ClientManager, ClientAssigner)
• Reduces code duplication and centralizes relevant ownership

Potential Issues

• If any user relies on direct low-level node_name-to-client manipulation, those code paths will now be broken.
• Some logging and assignment-fallback code is intentionally temporary (marked as such) and will need cleanup in a future pass.
• Any custom deployment scenario with non-default memberlist/assignment policies will need to validate the config path now passes through correctly.
• Backward compatibility with older versions of tower in downstream consumers could require followup (see review comment).
• Some APIs now require (Config, System) tuples vs. just Config -- upgrading is required across the board.

This summary was automatically generated by @propel-code-bot

rescrv · 2025-05-28T03:28:28Z

rust/memberlist/src/client_manager.rs

+    fn new_from_channel(channel: GrpcTraceService<Channel>) -> Self {
+        QueryExecutorClient::new(channel)
+    }
+    fn max_decoding_message_size(self, max_size: usize) -> Self {


I'm surprised this works with a trait, but I guess it works only on concrete types.

rust/worker/src/test.rs

rust/log/src/grpc_log.rs

rescrv · 2025-05-28T03:43:49Z

rust/log/src/grpc_log.rs

-                        },
-                    )
-                    .await
+        let mut combined_response = Vec::new();


This looks correct to me. I read it twice over.

HammadB added 9 commits May 27, 2025 20:12

pro

77d9f83

[ENH] Make client manager generic

0e3c854

cln

9861a11

cln

cb06be2

[CLN] Move client manager out of FE into ML

86c4e9e

[ENH] Use client manager in log client

0d7fbb2

system hydration cleanup

dd6944c

cln

ea40c39

[ENH] Use client manager / client assigner

fe3519d

This was referenced May 28, 2025

[CLN] Move client manager out of FE into ML #4641

Merged

[ENH] Use client manager in log client #4642

Closed

HammadB commented May 28, 2025

View reviewed changes

HammadB changed the title ~~pro~~ [ENH] Clean up client manager into manager/assigner - make log client use it May 28, 2025

HammadB commented May 28, 2025

View reviewed changes

HammadB marked this pull request as ready for review May 28, 2025 03:21

rescrv approved these changes May 28, 2025

View reviewed changes

cln

4c44012

HammadB enabled auto-merge (squash) May 29, 2025 00:09

HammadB merged commit 258f7d6 into main May 29, 2025
72 checks passed

[ENH] Clean up client manager into manager/assigner - make log client use it #4640

[ENH] Clean up client manager into manager/assigner - make log client use it #4640

Uh oh!

Conversation

HammadB commented May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description of changes

Test plan

Documentation Changes

Uh oh!

HammadB commented May 28, 2025

Uh oh!

github-actions bot commented May 28, 2025

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

Uh oh!

github-actions bot commented May 28, 2025

Uh oh!

github-actions bot commented May 28, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HammadB May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

propel-code-bot bot commented May 28, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

HammadB commented May 28, 2025 •

edited

Loading

HammadB May 28, 2025 •

edited

Loading