[ENH] clean up component channels (#2429)

Author: codetheweb

rust/worker/Dockerfile

@@ -1,4 +1,4 @@
-FROM rust:1.77.2 as builder
+FROM rust:1.79.0 as builder
 ARG CHROMA_KUBERNETES_INTEGRATION=0
 ENV CHROMA_KUBERNETES_INTEGRATION $CHROMA_KUBERNETES_INTEGRATION
 

rust/worker/src/compactor/compaction_manager.rs

@@ -7,6 +7,7 @@ use crate::config::CompactionServiceConfig;
 use crate::config::Configurable;
 use crate::errors::ChromaError;
 use crate::errors::ErrorCodes;
+use crate::execution::dispatcher::Dispatcher;
 use crate::execution::operator::TaskMessage;
 use crate::execution::orchestration::CompactOrchestrator;
 use crate::execution::orchestration::CompactionResponse;
@@ -18,8 +19,8 @@ use crate::sysdb;
 use crate::sysdb::sysdb::SysDb;
 use crate::system::Component;
 use crate::system::ComponentContext;
+use crate::system::ComponentHandle;
 use crate::system::Handler;
-use crate::system::Receiver;
 use crate::system::System;
 use async_trait::async_trait;
 use futures::stream::FuturesUnordered;
@@ -42,7 +43,7 @@ pub(crate) struct CompactionManager {
     blockfile_provider: BlockfileProvider,
     hnsw_index_provider: HnswIndexProvider,
     // Dispatcher
-    dispatcher: Option<Box<dyn Receiver<TaskMessage>>>,
+    dispatcher: Option<ComponentHandle<Dispatcher>>,
     // Config
     compaction_manager_queue_size: usize,
     compaction_interval: Duration,
@@ -95,7 +96,7 @@ impl CompactionManager {
         compaction_job: &CompactionJob,
     ) -> Result<CompactionResponse, Box<dyn ChromaError>> {
         let dispatcher = match self.dispatcher {
-            Some(ref dispatcher) => dispatcher,
+            Some(ref dispatcher) => dispatcher.clone(),
             None => {
                 println!("No dispatcher found");
                 return Err(Box::new(CompactionError::FailedToCompact));
@@ -112,7 +113,7 @@ impl CompactionManager {
                     self.sysdb.clone(),
                     self.blockfile_provider.clone(),
                     self.hnsw_index_provider.clone(),
-                    dispatcher.clone(),
+                    dispatcher,
                     None,
                     None,
                     Arc::new(AtomicU32::new(0)),
@@ -161,7 +162,7 @@ impl CompactionManager {
         (num_completed_jobs, num_failed_jobs)
     }
 
-    pub(crate) fn set_dispatcher(&mut self, dispatcher: Box<dyn Receiver<TaskMessage>>) {
+    pub(crate) fn set_dispatcher(&mut self, dispatcher: ComponentHandle<Dispatcher>) {
         self.dispatcher = Some(dispatcher);
     }
 
@@ -253,12 +254,8 @@ impl Component for CompactionManager {
 
     async fn on_start(&mut self, ctx: &crate::system::ComponentContext<Self>) -> () {
         println!("Starting CompactionManager");
-        ctx.scheduler.schedule(
-            ctx.sender.clone(),
-            ScheduleMessage {},
-            self.compaction_interval,
-            ctx,
-        );
+        ctx.scheduler
+            .schedule(ScheduleMessage {}, self.compaction_interval, ctx);
     }
 }
 
@@ -279,12 +276,8 @@ impl Handler<ScheduleMessage> for CompactionManager {
         println!("CompactionManager: Performing compaction");
         self.compact_batch().await;
         // Compaction is done, schedule the next compaction
-        ctx.scheduler.schedule(
-            ctx.sender.clone(),
-            ScheduleMessage {},
-            self.compaction_interval,
-            ctx,
-        );
+        ctx.scheduler
+            .schedule(ScheduleMessage {}, self.compaction_interval, ctx);
     }
 }
 
@@ -508,7 +501,7 @@ mod tests {
 
         let dispatcher = Dispatcher::new(10, 10, 10);
         let dispatcher_handle = system.start_component(dispatcher);
-        manager.set_dispatcher(dispatcher_handle.receiver());
+        manager.set_dispatcher(dispatcher_handle);
         manager.set_system(system);
         let (num_completed, number_failed) = manager.compact_batch().await;
         assert_eq!(num_completed, 2);

rust/worker/src/execution/dispatcher.rs

@@ -3,7 +3,7 @@ use crate::execution::config::DispatcherConfig;
 use crate::{
     config::Configurable,
     errors::ChromaError,
-    system::{Component, ComponentContext, Handler, Receiver, System},
+    system::{Component, ComponentContext, Handler, ReceiverForMessage, System},
 };
 use async_trait::async_trait;
 use std::fmt::Debug;
@@ -83,7 +83,7 @@ impl Dispatcher {
     fn spawn_workers(
         &self,
         system: &mut System,
-        self_receiver: Box<dyn Receiver<TaskRequestMessage>>,
+        self_receiver: Box<dyn ReceiverForMessage<TaskRequestMessage>>,
     ) {
         for _ in 0..self.n_worker_threads {
             let worker = WorkerThread::new(self_receiver.clone(), self.worker_queue_size);
@@ -154,15 +154,15 @@ impl Configurable<DispatcherConfig> for Dispatcher {
 /// - reply_to: The receiver to send the task to, this is the worker thread
 #[derive(Debug)]
 pub(super) struct TaskRequestMessage {
-    reply_to: Box<dyn Receiver<TaskMessage>>,
+    reply_to: Box<dyn ReceiverForMessage<TaskMessage>>,
 }
 
 impl TaskRequestMessage {
     /// Create a new TaskRequestMessage
     /// # Parameters
     /// - reply_to: The receiver to send the task to, this is the worker thread
     /// that is requesting the task
-    pub(super) fn new(reply_to: Box<dyn Receiver<TaskMessage>>) -> Self {
+    pub(super) fn new(reply_to: Box<dyn ReceiverForMessage<TaskMessage>>) -> Self {
         TaskRequestMessage { reply_to }
     }
 }
@@ -180,7 +180,7 @@ impl Component for Dispatcher {
     }
 
     async fn on_start(&mut self, ctx: &ComponentContext<Self>) {
-        self.spawn_workers(&mut ctx.system.clone(), ctx.sender.as_receiver());
+        self.spawn_workers(&mut ctx.system.clone(), ctx.receiver());
     }
 }
 
@@ -207,7 +207,7 @@ mod tests {
     use super::*;
     use crate::{
         execution::operator::{wrap, Operator, TaskResult},
-        system::System,
+        system::{ComponentHandle, System},
     };
     use std::{
         collections::HashSet,
@@ -243,7 +243,7 @@ mod tests {
 
     #[derive(Debug)]
     struct MockDispatchUser {
-        pub dispatcher: Box<dyn Receiver<TaskMessage>>,
+        pub dispatcher: ComponentHandle<Dispatcher>,
         counter: Arc<AtomicUsize>, // We expect to recieve DISPATCH_COUNT messages
         sent_tasks: Arc<Mutex<HashSet<Uuid>>>,
         received_tasks: Arc<Mutex<HashSet<Uuid>>>,
@@ -261,13 +261,8 @@ mod tests {
         async fn on_start(&mut self, ctx: &ComponentContext<Self>) {
             // dispatch a new task every DISPATCH_FREQUENCY_MS for DISPATCH_COUNT times
             let duration = std::time::Duration::from_millis(DISPATCH_FREQUENCY_MS);
-            ctx.scheduler.schedule_interval(
-                ctx.sender.clone(),
-                (),
-                duration,
-                Some(DISPATCH_COUNT),
-                ctx,
-            );
+            ctx.scheduler
+                .schedule_interval((), duration, Some(DISPATCH_COUNT), ctx);
         }
     }
     #[async_trait]
@@ -290,7 +285,7 @@ mod tests {
     #[async_trait]
     impl Handler<()> for MockDispatchUser {
         async fn handle(&mut self, _message: (), ctx: &ComponentContext<MockDispatchUser>) {
-            let task = wrap(Box::new(MockOperator {}), 42.0, ctx.sender.as_receiver());
+            let task = wrap(Box::new(MockOperator {}), 42.0, ctx.receiver());
             let task_id = task.id();
             self.sent_tasks.lock().insert(task_id);
             let res = self.dispatcher.send(task, None).await;
@@ -306,7 +301,7 @@ mod tests {
         let sent_tasks = Arc::new(Mutex::new(HashSet::new()));
         let received_tasks = Arc::new(Mutex::new(HashSet::new()));
         let dispatch_user = MockDispatchUser {
-            dispatcher: dispatcher_handle.receiver(),
+            dispatcher: dispatcher_handle,
             counter: counter.clone(),
             sent_tasks: sent_tasks.clone(),
             received_tasks: received_tasks.clone(),

rust/worker/src/execution/operator.rs

@@ -1,4 +1,4 @@
-use crate::system::Receiver;
+use crate::system::ReceiverForMessage;
 use async_trait::async_trait;
 use std::fmt::Debug;
 use uuid::Uuid;
@@ -45,7 +45,7 @@ where
 {
     operator: Box<dyn Operator<Input, Output, Error = Error>>,
     input: Input,
-    reply_channel: Box<dyn Receiver<TaskResult<Output, Error>>>,
+    reply_channel: Box<dyn ReceiverForMessage<TaskResult<Output, Error>>>,
     task_id: Uuid,
 }
 
@@ -89,7 +89,7 @@ where
 pub(super) fn wrap<Input, Output, Error>(
     operator: Box<dyn Operator<Input, Output, Error = Error>>,
     input: Input,
-    reply_channel: Box<dyn Receiver<TaskResult<Output, Error>>>,
+    reply_channel: Box<dyn ReceiverForMessage<TaskResult<Output, Error>>>,
 ) -> TaskMessage
 where
     Error: Debug + 'static,

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

@@ -3,6 +3,7 @@ use crate::blockstore::provider::BlockfileProvider;
 use crate::compactor::CompactionJob;
 use crate::errors::ChromaError;
 use crate::execution::data::data_chunk::Chunk;
+use crate::execution::dispatcher::Dispatcher;
 use crate::execution::operator::TaskResult;
 use crate::execution::operators::flush_s3::FlushS3Input;
 use crate::execution::operators::flush_s3::FlushS3Operator;
@@ -33,8 +34,9 @@ use crate::sysdb::sysdb::GetCollectionsError;
 use crate::sysdb::sysdb::GetSegmentsError;
 use crate::sysdb::sysdb::SysDb;
 use crate::system::Component;
+use crate::system::ComponentHandle;
 use crate::system::Handler;
-use crate::system::Receiver;
+use crate::system::ReceiverForMessage;
 use crate::system::System;
 use crate::types::LogRecord;
 use crate::types::Segment;
@@ -92,7 +94,7 @@ pub struct CompactOrchestrator {
     pulled_log_offset: Option<i64>,
     record_segment: Option<Segment>,
     // Dispatcher
-    dispatcher: Box<dyn Receiver<TaskMessage>>,
+    dispatcher: ComponentHandle<Dispatcher>,
     // number of write segments tasks
     num_write_tasks: i32,
     // Result Channel
@@ -149,7 +151,7 @@ impl CompactOrchestrator {
         sysdb: Box<SysDb>,
         blockfile_provider: BlockfileProvider,
         hnsw_index_provider: HnswIndexProvider,
-        dispatcher: Box<dyn Receiver<TaskMessage>>,
+        dispatcher: ComponentHandle<Dispatcher>,
         result_channel: Option<
             tokio::sync::oneshot::Sender<Result<CompactionResponse, Box<dyn ChromaError>>>,
         >,
@@ -180,7 +182,7 @@ impl CompactOrchestrator {
     // of the segment, and not fully respect the offset_id from the compaction job
     async fn pull_logs(
         &mut self,
-        self_address: Box<dyn Receiver<TaskResult<PullLogsOutput, PullLogsError>>>,
+        self_address: Box<dyn ReceiverForMessage<TaskResult<PullLogsOutput, PullLogsError>>>,
     ) {
         self.state = ExecutionState::PullLogs;
         let operator = PullLogsOperator::new(self.log.clone());
@@ -215,7 +217,7 @@ impl CompactOrchestrator {
     async fn partition(
         &mut self,
         records: Chunk<LogRecord>,
-        self_address: Box<dyn Receiver<TaskResult<PartitionOutput, PartitionError>>>,
+        self_address: Box<dyn ReceiverForMessage<TaskResult<PartitionOutput, PartitionError>>>,
     ) {
         self.state = ExecutionState::Partition;
         // TODO: make this configurable
@@ -236,7 +238,7 @@ impl CompactOrchestrator {
         &mut self,
         partitions: Vec<Chunk<LogRecord>>,
         self_address: Box<
-            dyn Receiver<TaskResult<WriteSegmentsOutput, WriteSegmentsOperatorError>>,
+            dyn ReceiverForMessage<TaskResult<WriteSegmentsOutput, WriteSegmentsOperatorError>>,
         >,
     ) {
         self.state = ExecutionState::Write;
@@ -282,7 +284,7 @@ impl CompactOrchestrator {
         record_segment_writer: RecordSegmentWriter,
         hnsw_segment_writer: Box<DistributedHNSWSegmentWriter>,
         metadata_segment_writer: MetadataSegmentWriter<'static>,
-        self_address: Box<dyn Receiver<TaskResult<FlushS3Output, Box<dyn ChromaError>>>>,
+        self_address: Box<dyn ReceiverForMessage<TaskResult<FlushS3Output, Box<dyn ChromaError>>>>,
     ) {
         self.state = ExecutionState::Flush;
 
@@ -306,7 +308,7 @@ impl CompactOrchestrator {
         &mut self,
         log_position: i64,
         segment_flush_info: Arc<[SegmentFlushInfo]>,
-        self_address: Box<dyn Receiver<TaskResult<RegisterOutput, RegisterError>>>,
+        self_address: Box<dyn ReceiverForMessage<TaskResult<RegisterOutput, RegisterError>>>,
     ) {
         self.state = ExecutionState::Register;
         let operator = RegisterOperator::new();
@@ -494,7 +496,7 @@ impl Component for CompactOrchestrator {
     }
 
     async fn on_start(&mut self, ctx: &crate::system::ComponentContext<Self>) -> () {
-        self.pull_logs(ctx.sender.as_receiver()).await;
+        self.pull_logs(ctx.receiver()).await;
     }
 }
 
@@ -528,7 +530,7 @@ impl Handler<TaskResult<PullLogsOutput, PullLogsError>> for CompactOrchestrator
             Some(record) => {
                 self.pulled_log_offset = Some(record.log_offset);
                 println!("Pulled Logs Up To Offset: {:?}", self.pulled_log_offset);
-                self.partition(records, ctx.sender.as_receiver()).await;
+                self.partition(records, ctx.receiver()).await;
             }
             None => {
                 // Log an error and return
@@ -561,7 +563,7 @@ impl Handler<TaskResult<PartitionOutput, PartitionError>> for CompactOrchestrato
                 return;
             }
         };
-        self.write(records, _ctx.sender.as_receiver()).await;
+        self.write(records, _ctx.receiver()).await;
     }
 }
 
@@ -602,7 +604,7 @@ impl Handler<TaskResult<WriteSegmentsOutput, WriteSegmentsOperatorError>> for Co
                 output.record_segment_writer,
                 output.hnsw_segment_writer,
                 output.metadata_segment_writer,
-                _ctx.sender.as_receiver(),
+                _ctx.receiver(),
             )
             .await;
         }
@@ -623,7 +625,7 @@ impl Handler<TaskResult<FlushS3Output, Box<dyn ChromaError>>> for CompactOrchest
                 self.register(
                     self.pulled_log_offset.unwrap(),
                     msg.segment_flush_info,
-                    _ctx.sender.as_receiver(),
+                    _ctx.receiver(),
                 )
                 .await;
             }