fix(batch-exports): Handle cancellation of activities

[rossgray]

Problem

Currently whenever batch export Temporal activities timeout or are cancelled the SPMC consumer task continues to run in the background. This consumes additional resources in ClickHouse and whatever external destination it is exporting data to.

Changes

Ensure the record batch consumer is cancelled if the activity itself is cancelled by Temporal.

Does this work well for both Cloud and self-hosted?

Yes

How did you test this code?

Extensive testing locally

[greptile-apps]

PR Summary

This PR improves cancellation handling in the batch exports system by implementing proper propagation of cancellation signals to all pending tasks.

• Added asyncio.TaskGroup to manage consumer tasks, ensuring cancellation is properly propagated to all pending tasks
• Modified Consumer.create_consumer_task() to accept a TaskGroup parameter instead of directly using asyncio.create_task()
• Added cancellation detection in the consumer_done_callback function to log when a consumer task is cancelled
• Restructured run_consumer() function to use TaskGroup for better cancellation handling and cleanup

Greptile AI

_{1 file(s) reviewed, 1 comment(s)
Edit PR Review Bot Settings | Greptile}