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

tanujnay112 · 2025-10-22T08:29:30Z

Description of changes

Summarize the changes made by this PR.

Improvements & Bug fixes
- Create heaptender client for rust to use in the compaction code within FinishTask.
New functionality
- ...

Test plan

How are these changes tested?

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

Migration plan

Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?

Observability plan

What is the plan to instrument and monitor this change?

Documentation Changes

Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the _docs section?_

github-actions · 2025-10-22T08:29:50Z

tanujnay112 · 2025-10-22T08:30:01Z

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

propel-code-bot · 2025-10-22T08:31:15Z

Centralize compaction & FinishTask flows through new Heaptender gRPC client

This PR introduces a reusable Rust gRPC client (HeaptenderClient) for the Heaptender S3-heap management service and rewires every compaction-related code path—including the FinishTask operator—to route through this client. By eliminating all ad-hoc direct S3 heap manipulation in the worker, the change establishes a single, typed, observable integration point for heap compaction and clean-up. New configuration knobs and service-discovery wiring are added so both local development and production clusters can locate the Heaptender service. The result is simplified code, improved observability, and a clear migration path for future heap-level optimizations and policies.

Because workers now hard-depend on a reachable Heaptender endpoint, the service must be deployed and configured before these binaries are rolled out. No data-plane or schema migration is required.

Key Changes

• Created s3heap-service client module exposing HeaptenderClient with TLS, retry, and timeout support (tonic/tower based)
• Replaced direct S3 heap calls in compactor/compaction_manager.rs and execution/operators/finish_task.rs with HeaptenderClient::finish_and_compact
• Added heaptender.{host,port,timeout} configuration sections to worker, frontend, and client crates
• Extended memberlist client_manager to advertise the Heaptender service for discovery, and updated Tiltfile to spin it up in local dev
• Integrated gRPC traffic into existing grpc_log tracing so Heaptender calls are automatically logged/observed

Affected Areas

• rust/s3heap-service (new client crate)
• rust/worker/compactor and execution operators
• rust/frontend executor config
• memberlist service discovery layer
• Dev-ops manifests (Tiltfile) and logging middleware

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

propel-code-bot · 2025-10-22T08:34:52Z

rust/worker/src/execution/operators/finish_task.rs

            );

-            // TODO: Schedule a new task for next nonce by pushing to the heap
+            // Schedule a new task for next nonce by pushing to the heap


[CriticalError]

Idempotency issue: The push operation to schedule the next task run lacks idempotency protection. If this function is called multiple times with the same task (e.g., due to retries or duplicate processing), it will schedule the same task multiple times in the heap, leading to duplicate task execution:

// Add idempotency check before scheduling if self.is_task_already_scheduled(&input.updated_task.id, &input.updated_task.next_nonce).await? { tracing::info!("Task already scheduled, skipping heap push"); return Ok(FinishTaskOutput::new()); }

Context for Agents

[**CriticalError**] Idempotency issue: The `push` operation to schedule the next task run lacks idempotency protection. If this function is called multiple times with the same task (e.g., due to retries or duplicate processing), it will schedule the same task multiple times in the heap, leading to duplicate task execution: ```rust // Add idempotency check before scheduling if self.is_task_already_scheduled(&input.updated_task.id, &input.updated_task.next_nonce).await? { tracing::info!("Task already scheduled, skipping heap push"); return Ok(FinishTaskOutput::new()); } ``` File: rust/worker/src/execution/operators/finish_task.rs Line: 135

rust/worker/src/execution/operators/finish_task.rs

rust/frontend/src/executor/distributed.rs

rust/log/src/grpc_log.rs

rust/worker/src/execution/orchestration/compact.rs

rust/worker/src/compactor/compaction_manager.rs

propel-code-bot · 2025-10-22T09:25:31Z

rust/worker/src/execution/operators/finish_task.rs

+        s3heap::client::GrpcHeapService::try_from_config(
+            &(config, system),
+            &registry,
+        )
+        .await
+        .expect("Failed to create test heap service client - ensure heap service is running on localhost:50052")


[BestPractice]

Test setup will panic if heap service is unavailable: The expect() will crash tests when the heap service isn't running on localhost:50052:

.expect("Failed to create test heap service client - ensure heap service is running on localhost:50052")

For tests, use a mock or make the test conditional:

let heap_service = s3heap::client::GrpcHeapService::try_from_config( &(config, system), &registry, ).await; match heap_service { Ok(service) => service, Err(_) => { eprintln!("Skipping test - heap service not available"); return; } }

Context for Agents

[**BestPractice**] Test setup will panic if heap service is unavailable: The `expect()` will crash tests when the heap service isn't running on localhost:50052: ```rust .expect("Failed to create test heap service client - ensure heap service is running on localhost:50052") ``` For tests, use a mock or make the test conditional: ```rust let heap_service = s3heap::client::GrpcHeapService::try_from_config( &(config, system), &registry, ).await; match heap_service { Ok(service) => service, Err(_) => { eprintln!("Skipping test - heap service not available"); return; } } ``` File: rust/worker/src/execution/operators/finish_task.rs Line: 222

rust/worker/src/execution/operators/finish_task.rs

propel-code-bot · 2025-10-22T10:00:55Z

rust/worker/src/execution/operators/finish_task.rs

+            };
+
+            heap_service
+                .push(
+                    vec![schedule],
+                    &input.updated_task.input_collection_id.to_string(),
+                )
+                .await?;


[BestPractice]

Missing error handling for heap service operations:

heap_service .push( vec![schedule], &input.updated_task.input_collection_id.to_string(), ) .await?;

If the heap service is unavailable or returns an error, this will cause the entire task finish operation to fail, potentially leaving the task in an inconsistent state. Consider handling heap service errors gracefully:

if let Err(e) = heap_service.push(vec![schedule], key).await { tracing::warn!("Failed to schedule next task run: {}", e); // Continue with task completion even if scheduling fails }

Context for Agents

[**BestPractice**] Missing error handling for heap service operations: ```rust heap_service .push( vec![schedule], &input.updated_task.input_collection_id.to_string(), ) .await?; ``` If the heap service is unavailable or returns an error, this will cause the entire task finish operation to fail, potentially leaving the task in an inconsistent state. Consider handling heap service errors gracefully: ```rust if let Err(e) = heap_service.push(vec![schedule], key).await { tracing::warn!("Failed to schedule next task run: {}", e); // Continue with task completion even if scheduling fails } ``` File: rust/worker/src/execution/operators/finish_task.rs Line: 150

rust/s3heap/Cargo.toml

rust/frontend/src/executor/distributed.rs

rust/s3heap/Cargo.toml

rust/worker/src/compactor/compaction_manager.rs

rust/worker/src/execution/operators/finish_task.rs

rust/s3heap-service/src/client/grpc.rs

propel-code-bot · 2025-10-24T21:27:06Z

rust/log/src/grpc_log.rs

        let endpoint_res = match Endpoint::from_shared(format!(
-            "grpc://rust-log-service-{ordinal}.rust-log-service:50051"
+            "grpc://rust-log-service-{ordinal}.rust-log-service:{port}"
        )) {


[BestPractice]

For consistency with Tonic best practices and other parts of the codebase that create gRPC endpoints, use the http:// scheme for plaintext connections. While grpc:// may work with some gRPC implementations, http:// is the standard scheme for plaintext gRPC connections in Tonic, as documented in official examples and used throughout the Tonic ecosystem. For TLS connections, use https://.

Suggested change

let endpoint_res = match Endpoint::from_shared(format!(

"grpc://rust-log-service-{ordinal}.rust-log-service:50051"

"grpc://rust-log-service-{ordinal}.rust-log-service:{port}"

)) {

let endpoint_res = match Endpoint::from_shared(format!(

"http://rust-log-service-{ordinal}.rust-log-service:{port}"

)) {

⚡ Committable suggestion

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

Context for Agents

[**BestPractice**] For consistency with Tonic best practices and other parts of the codebase that create gRPC endpoints, use the `http://` scheme for plaintext connections. While `grpc://` may work with some gRPC implementations, `http://` is the standard scheme for plaintext gRPC connections in Tonic, as documented in official examples and used throughout the Tonic ecosystem. For TLS connections, use `https://`. ```suggestion let endpoint_res = match Endpoint::from_shared(format!( "http://rust-log-service-{ordinal}.rust-log-service:{port}" )) { ``` ⚡ **Committable suggestion** Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. File: rust/log/src/grpc_log.rs Line: 647

Thank you, propel.

rust/worker/src/compactor/compaction_manager.rs

rescrv · 2025-10-26T17:41:08Z

rust/log/src/grpc_log.rs

        let endpoint_res = match Endpoint::from_shared(format!(
-            "grpc://rust-log-service-{ordinal}.rust-log-service:50051"
+            "grpc://rust-log-service-{ordinal}.rust-log-service:{port}"
        )) {


Thank you, propel.

rescrv · 2025-10-26T17:41:48Z

rust/memberlist/src/client_manager.rs

            1,
            1000,
            1000,
+            50051,


I'm fine with it.

rescrv · 2025-10-26T17:44:14Z

rust/s3heap-service/src/client/grpc.rs

+    }
+
+    /// Push schedules to the heap
+    #[tracing::instrument(skip(self, schedules))]


Our pattern in most places is to have the client do conversion to local types. These request and respond with chroma_proto types. Which is fine if we're not going to parse or validate them.

rust/worker/src/execution/orchestration/compact.rs

rescrv · 2025-10-26T17:48:47Z

rust/worker/src/execution/orchestration/compact.rs

+        let heap_service = GrpcHeapService::try_from_config(
+            &(GrpcHeapServiceConfig::default(), system.clone()),
+            &registry,
+        )
+        .await
+        .expect("Should connect to grpc heap service");


test_k8s_integration_ is a symbol that it's to be run with these services running.

propel-code-bot · 2025-10-27T18:31:53Z

rust/s3heap-service/src/client/grpc.rs

+        return Ok(GrpcHeapService::new(
+            my_config.clone(),
+            client_assigner,
+            client_manager_handle,
+            memberlist_provider_handle,
+        ));


[BestPractice]

The return keyword is redundant here as this is the last expression in the function. It can be removed for idiomatic Rust.

Context for Agents

[**BestPractice**] The `return` keyword is redundant here as this is the last expression in the function. It can be removed for idiomatic Rust. File: rust/s3heap-service/src/client/grpc.rs Line: 205

This was referenced Oct 22, 2025

[ENH]: Added SysDB functionality for TaskRunner support #5685

Merged

[ENH]: Add Rust task client and execution operators #5686

Merged

tanujnay112 mentioned this pull request Oct 22, 2025

[ENH]: Integrate task operators into compaction #5687

Merged

1 task

tanujnay112 marked this pull request as ready for review October 22, 2025 08:30