[ENH]: Add heaptender client to compaction and make finishtask use this

Author: tanujnay112

Cargo.lock

@@ -74,6 +74,7 @@ impl Configurable<(config::DistributedExecutorConfig, System)> for DistributedEx
             config.connections_per_node,
             config.connect_timeout_ms,
             config.request_timeout_ms,
+            50051,
             ClientOptions::new(Some(config.max_query_service_response_size_bytes)),
         );
         let client_manager_handle = system.start_component(client_manager);

rust/frontend/src/executor/distributed.rs

@@ -263,6 +263,7 @@ impl Configurable<(GrpcLogConfig, System)> for GrpcLog {
             1,
             my_config.connect_timeout_ms,
             my_config.request_timeout_ms,
+            50051,
             ClientOptions::new(Some(my_config.max_decoding_message_size)),
         );
         let client_manager_handle = system.start_component(client_manager);

rust/log/src/grpc_log.rs

@@ -7,7 +7,8 @@ use chroma_error::ChromaError;
 use chroma_system::{Component, ComponentContext, Handler};
 use chroma_tracing::GrpcClientTraceService;
 use chroma_types::chroma_proto::{
-    log_service_client::LogServiceClient, query_executor_client::QueryExecutorClient,
+    heap_tender_service_client::HeapTenderServiceClient, log_service_client::LogServiceClient,
+    query_executor_client::QueryExecutorClient,
 };
 use parking_lot::RwLock;
 use std::{
@@ -171,6 +172,7 @@ pub struct ClientManager<T> {
     connections_per_node: usize,
     connect_timeout_ms: u64,
     request_timeout_ms: u64,
+    port: u16,
     old_memberlist: Memberlist,
     options: ClientOptions,
 }
@@ -184,6 +186,7 @@ where
         connections_per_node: usize,
         connect_timeout_ms: u64,
         request_timeout_ms: u64,
+        port: u16,
         options: ClientOptions,
     ) -> Self {
         Self {
@@ -192,6 +195,7 @@ where
             connections_per_node,
             connect_timeout_ms,
             request_timeout_ms,
+            port,
             old_memberlist: Memberlist::new(),
             options,
         }
@@ -233,8 +237,7 @@ where
     }
 
     async fn add_ip_for_node(&mut self, ip: String, node: &str) {
-        // TODO: Configure the port
-        let ip_with_port = format!("http://{}:{}", ip, 50051);
+        let ip_with_port = format!("http://{}:{}", ip, self.port);
         let endpoint = match Endpoint::from_shared(ip_with_port) {
             Ok(endpoint) => endpoint
                 .connect_timeout(std::time::Duration::from_millis(self.connect_timeout_ms))
@@ -417,6 +420,17 @@ impl ClientFactory
     }
 }
 
+impl ClientFactory
+    for HeapTenderServiceClient<chroma_tracing::GrpcClientTraceService<tonic::transport::Channel>>
+{
+    fn new_from_channel(channel: GrpcClientTraceService<Channel>) -> Self {
+        HeapTenderServiceClient::new(channel)
+    }
+    fn max_decoding_message_size(self, max_size: usize) -> Self {
+        self.max_decoding_message_size(max_size)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::super::memberlist_provider::Member;
@@ -436,6 +450,7 @@ mod test {
             1,
             1000,
             1000,
+            50051,
             ClientOptions::default(),
         );
         (client_manager, client_assigner)

rust/memberlist/src/client_manager.rs

@@ -14,10 +14,18 @@ parquet = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 thiserror = { workspace = true }
+tonic = { workspace = true }
+tower = { workspace = true }
+tracing = { workspace = true }
 uuid = { workspace = true }
 
+chroma-config = { workspace = true }
 chroma-error = { workspace = true }
+chroma-memberlist = { workspace = true }
 chroma-storage = { workspace = true }
+chroma-system = { workspace = true }
+chroma-tracing = { workspace = true }
+chroma-types = { workspace = true }
 wal3 = { workspace = true }
 
 [dev-dependencies]

rust/s3heap/Cargo.toml

@@ -0,0 +1,109 @@
+use chroma_memberlist::config::CustomResourceMemberlistProviderConfig;
+use serde::{Deserialize, Serialize};
+
+/// Configuration for the gRPC heap service client.
+///
+/// This configures how to connect to the heap tender service instances via memberlist.
+#[derive(Deserialize, Clone, Serialize, Debug)]
+pub struct GrpcHeapServiceConfig {
+    /// Whether the heap service client is enabled. Defaults to false.
+    #[serde(default = "GrpcHeapServiceConfig::default_enabled")]
+    pub enabled: bool,
+    /// Connection timeout in milliseconds. Defaults to 5000ms.
+    #[serde(default = "GrpcHeapServiceConfig::default_connect_timeout_ms")]
+    pub connect_timeout_ms: u64,
+    /// Request timeout in milliseconds. Defaults to 5000ms.
+    #[serde(default = "GrpcHeapServiceConfig::default_request_timeout_ms")]
+    pub request_timeout_ms: u64,
+    /// Maximum message size for encoding. Defaults to 32MB.
+    #[serde(default = "GrpcHeapServiceConfig::default_max_encoding_message_size")]
+    pub max_encoding_message_size: usize,
+    /// Maximum message size for decoding. Defaults to 32MB.
+    #[serde(default = "GrpcHeapServiceConfig::default_max_decoding_message_size")]
+    pub max_decoding_message_size: usize,
+    /// Memberlist provider configuration. Defaults to rust-log-service-memberlist (colocated).
+    #[serde(default = "GrpcHeapServiceConfig::default_memberlist_provider")]
+    pub memberlist_provider: chroma_memberlist::config::MemberlistProviderConfig,
+    /// Assignment policy. Must match log service (RendezvousHashing + Murmur3) for data locality.
+    #[serde(default = "GrpcHeapServiceConfig::default_assignment")]
+    pub assignment: chroma_config::assignment::config::AssignmentPolicyConfig,
+    /// Port the heap service listens on. Defaults to 50052.
+    #[serde(default = "GrpcHeapServiceConfig::default_port")]
+    pub port: u16,
+}
+
+impl GrpcHeapServiceConfig {
+    fn default_enabled() -> bool {
+        false
+    }
+
+    fn default_connect_timeout_ms() -> u64 {
+        5000
+    }
+
+    fn default_request_timeout_ms() -> u64 {
+        5000
+    }
+
+    fn default_max_encoding_message_size() -> usize {
+        32_000_000
+    }
+
+    fn default_max_decoding_message_size() -> usize {
+        32_000_000
+    }
+
+    fn default_memberlist_provider() -> chroma_memberlist::config::MemberlistProviderConfig {
+        chroma_memberlist::config::MemberlistProviderConfig::CustomResource(
+            CustomResourceMemberlistProviderConfig {
+                kube_namespace: "chroma".to_string(),
+                memberlist_name: "rust-log-service-memberlist".to_string(), // Colocated with log service
+                queue_size: 100,
+            },
+        )
+    }
+
+    fn default_assignment() -> chroma_config::assignment::config::AssignmentPolicyConfig {
+        // IMPORTANT: Must match log service assignment policy (RendezvousHashing + Murmur3)
+        // since heap and log services are colocated. This ensures that a collection's
+        // heap operations go to the same node as its log operations.
+        chroma_config::assignment::config::AssignmentPolicyConfig::RendezvousHashing(
+            chroma_config::assignment::config::RendezvousHashingAssignmentPolicyConfig {
+                hasher: chroma_config::assignment::config::HasherType::Murmur3,
+            },
+        )
+    }
+
+    fn default_port() -> u16 {
+        50052
+    }
+}
+
+impl Default for GrpcHeapServiceConfig {
+    fn default() -> Self {
+        GrpcHeapServiceConfig {
+            enabled: GrpcHeapServiceConfig::default_enabled(),
+            connect_timeout_ms: GrpcHeapServiceConfig::default_connect_timeout_ms(),
+            request_timeout_ms: GrpcHeapServiceConfig::default_request_timeout_ms(),
+            max_encoding_message_size: GrpcHeapServiceConfig::default_max_encoding_message_size(),
+            max_decoding_message_size: GrpcHeapServiceConfig::default_max_decoding_message_size(),
+            memberlist_provider: GrpcHeapServiceConfig::default_memberlist_provider(),
+            assignment: GrpcHeapServiceConfig::default_assignment(),
+            port: GrpcHeapServiceConfig::default_port(),
+        }
+    }
+}
+
+/// Configuration for heap service client.
+#[derive(Deserialize, Clone, Serialize, Debug)]
+pub enum HeapServiceConfig {
+    /// gRPC-based heap service configuration.
+    #[serde(alias = "grpc")]
+    Grpc(GrpcHeapServiceConfig),
+}
+
+impl Default for HeapServiceConfig {
+    fn default() -> Self {
+        HeapServiceConfig::Grpc(GrpcHeapServiceConfig::default())
+    }
+}