update

Author: luoluoyuyu

src/runtime/streaming/api/context.rs

@@ -10,19 +10,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::runtime::streaming::memory::MemoryPool;
-use crate::runtime::streaming::network::endpoint::PhysicalSender;
-use crate::runtime::streaming::protocol::event::StreamEvent;
-use crate::runtime::streaming::protocol::event::TrackedEvent;
+use std::sync::Arc;
+use std::time::{Duration, SystemTime};
 
+use anyhow::{Context, Result};
 use arrow_array::RecordBatch;
-use std::sync::Arc;
-use std::time::Duration;
 
-/// 与单个子任务绑定的运行时参数（可由 `TaskContext::new` 默认填充，后续可扩展为从 Job 配置注入）。
+use crate::runtime::streaming::memory::MemoryPool;
+use crate::runtime::streaming::network::endpoint::PhysicalSender;
+use crate::runtime::streaming::protocol::event::{StreamEvent, TrackedEvent};
+
 #[derive(Debug, Clone)]
 pub struct TaskContextConfig {
-    /// Source 在无数据（`SourceEvent::Idle`）时的退避休眠时长。
     pub source_idle_timeout: Duration,
 }
 
@@ -34,110 +33,168 @@ impl Default for TaskContextConfig {
     }
 }
 
+/// Task execution context.
+///
+/// Acts as the sole bridge between operators and engine infrastructure (network, memory,
+/// configuration) for a single subtask.
 pub struct TaskContext {
+    /// Job identifier.
     pub job_id: String,
-    pub vertex_id: u32,
-    pub subtask_idx: u32,
+    /// Logical pipeline (vertex) index within the job graph.
+    pub pipeline_id: u32,
+    /// This subtask's index within the pipeline's parallelism.
+    pub subtask_index: u32,
+    /// Number of parallel subtasks for this pipeline.
     pub parallelism: u32,
 
-    pub outboxes: Vec<PhysicalSender>,
+    /// Precomputed display string for high-frequency logging without per-call allocation.
+    task_name: String,
 
+    /// Downstream physical senders (outbound edges).
+    downstream_senders: Vec<PhysicalSender>,
+
+    /// Global memory pool for back-pressure and accounting.
     memory_pool: Arc<MemoryPool>,
 
-    current_watermark: Option<std::time::SystemTime>,
+    /// Latest aligned event-time watermark for this subtask.
+    current_watermark: Option<SystemTime>,
 
+    /// Subtask-level tunables.
     config: TaskContextConfig,
 }
 
 impl TaskContext {
     pub fn new(
         job_id: String,
-        vertex_id: u32,
-        subtask_idx: u32,
+        pipeline_id: u32,
+        subtask_index: u32,
         parallelism: u32,
-        outboxes: Vec<PhysicalSender>,
+        downstream_senders: Vec<PhysicalSender>,
         memory_pool: Arc<MemoryPool>,
     ) -> Self {
+        let task_name = format!(
+            "Task-[{}]-Pipe[{}]-Sub[{}/{}]",
+            job_id, pipeline_id, subtask_index, parallelism
+        );
+
         Self {
             job_id,
-            vertex_id,
-            subtask_idx,
+            pipeline_id,
+            subtask_index,
             parallelism,
-            outboxes,
+            task_name,
+            downstream_senders,
             memory_pool,
             current_watermark: None,
             config: TaskContextConfig::default(),
         }
     }
 
+    #[inline]
     pub fn config(&self) -> &TaskContextConfig {
         &self.config
     }
 
-    // ========================================================================
-    // ========================================================================
+    #[inline]
+    pub fn task_name(&self) -> &str {
+        &self.task_name
+    }
+
+    // -------------------------------------------------------------------------
+    // Watermark
+    // -------------------------------------------------------------------------
 
-    pub fn last_present_watermark(&self) -> Option<std::time::SystemTime> {
+    #[inline]
+    pub fn current_watermark(&self) -> Option<SystemTime> {
         self.current_watermark
     }
 
-    pub fn advance_watermark(&mut self, watermark: std::time::SystemTime) {
+    pub fn advance_watermark(&mut self, watermark: SystemTime) {
         if let Some(current) = self.current_watermark {
-            if watermark > current {
-                self.current_watermark = Some(watermark);
-            }
+            self.current_watermark = Some(current.max(watermark));
         } else {
             self.current_watermark = Some(watermark);
         }
     }
 
-    // ========================================================================
-    // ========================================================================
-
-    pub fn task_identity(&self) -> String {
-        format!(
-            "Job[{}], Vertex[{}], Subtask[{}/{}]",
-            self.job_id, self.vertex_id, self.subtask_idx, self.parallelism
-        )
-    }
-
-    // ========================================================================
-    // ========================================================================
+    // -------------------------------------------------------------------------
+    // Data emission
+    // -------------------------------------------------------------------------
 
-    pub async fn collect(&self, batch: RecordBatch) -> anyhow::Result<()> {
-        if self.outboxes.is_empty() {
+    /// Fan-out a data batch to all downstreams (forward / broadcast).
+    pub async fn collect(&self, batch: RecordBatch) -> Result<()> {
+        if self.downstream_senders.is_empty() {
             return Ok(());
         }
 
         let bytes_required = batch.get_array_memory_size();
         let ticket = self.memory_pool.request_memory(bytes_required).await;
         let tracked_event = TrackedEvent::new(StreamEvent::Data(batch), Some(ticket));
 
-        for outbox in &self.outboxes {
-            outbox.send(tracked_event.clone()).await?;
-        }
-        Ok(())
+        self.broadcast_event(tracked_event).await
     }
 
-    pub async fn collect_keyed(&self, key_hash: u64, batch: RecordBatch) -> anyhow::Result<()> {
-        if self.outboxes.is_empty() {
+    /// Route a batch to one downstream by hash partitioning (shuffle).
+    pub async fn collect_keyed(&self, key_hash: u64, batch: RecordBatch) -> Result<()> {
+        let num_downstreams = self.downstream_senders.len();
+        if num_downstreams == 0 {
             return Ok(());
         }
 
         let bytes_required = batch.get_array_memory_size();
         let ticket = self.memory_pool.request_memory(bytes_required).await;
-        let tracked_event = TrackedEvent::new(StreamEvent::Data(batch), Some(ticket));
+        let event = TrackedEvent::new(StreamEvent::Data(batch), Some(ticket));
+
+        let target_idx = (key_hash as usize) % num_downstreams;
+
+        self.downstream_senders[target_idx]
+            .send(event)
+            .await
+            .with_context(|| {
+                format!(
+                    "{} failed to route keyed data to downstream index {}",
+                    self.task_name, target_idx
+                )
+            })?;
 
-        let target_idx = (key_hash as usize) % self.outboxes.len();
-        self.outboxes[target_idx].send(tracked_event).await?;
         Ok(())
     }
 
-    pub async fn broadcast(&self, event: StreamEvent) -> anyhow::Result<()> {
+    /// Broadcast a control event (watermark, barrier, end-of-stream).
+    pub async fn broadcast(&self, event: StreamEvent) -> Result<()> {
+        if self.downstream_senders.is_empty() {
+            return Ok(());
+        }
         let tracked_event = TrackedEvent::control(event);
-        for outbox in &self.outboxes {
-            outbox.send(tracked_event.clone()).await?;
+        self.broadcast_event(tracked_event).await
+    }
+
+    // -------------------------------------------------------------------------
+    // Internal dispatch
+    // -------------------------------------------------------------------------
+
+    async fn broadcast_event(&self, event: TrackedEvent) -> Result<()> {
+        let mut iter = self.downstream_senders.iter().enumerate().peekable();
+
+        while let Some((idx, sender)) = iter.next() {
+            if iter.peek().is_some() {
+                sender.send(event.clone()).await.with_context(|| {
+                    format!(
+                        "{} failed to broadcast event to downstream index {}",
+                        self.task_name, idx
+                    )
+                })?;
+            } else {
+                sender.send(event).await.with_context(|| {
+                    format!(
+                        "{} failed to send final event to downstream index {}",
+                        self.task_name, idx
+                    )
+                })?;
+                break;
+            }
         }
+
         Ok(())
     }
 }

src/runtime/streaming/execution/pipeline.rs

@@ -68,7 +68,7 @@ impl Pipeline {
         let span = info_span!(
             "pipeline_run",
             job_id = %self.ctx.job_id,
-            vertex = self.ctx.vertex_id
+            pipeline_id = self.ctx.pipeline_id
         );
 
         async move {

src/runtime/streaming/execution/source_driver.rs

@@ -50,7 +50,7 @@ impl SourceDriver {
         let span = info_span!(
             "source_run",
             job_id = %self.ctx.job_id,
-            vertex = self.ctx.vertex_id,
+            pipeline_id = self.ctx.pipeline_id,
             op = self.operator.name()
         );
 

src/runtime/streaming/operators/joins/join_instance.rs

@@ -152,7 +152,7 @@ impl InstantJoinOperator {
         let min_timestamp = min(time_column).ok_or_else(|| anyhow!("empty timestamp column"))?;
         let max_timestamp = max(time_column).ok_or_else(|| anyhow!("empty timestamp column"))?;
 
-        if let Some(watermark) = ctx.last_present_watermark()
+        if let Some(watermark) = ctx.current_watermark()
             && watermark > from_nanos(min_timestamp as u128)
         {
             warn!("Dropped late batch from {:?} before watermark", side);

src/runtime/streaming/operators/joins/join_with_expiration.rs

@@ -132,7 +132,7 @@ impl JoinWithExpirationOperator {
         batch: RecordBatch,
         ctx: &mut TaskContext,
     ) -> Result<Vec<StreamOutput>> {
-        let current_time = ctx.last_present_watermark().unwrap_or_else(SystemTime::now);
+        let current_time = ctx.current_watermark().unwrap_or_else(SystemTime::now);
 
         self.left_state.expire(current_time);
         self.right_state.expire(current_time);

src/runtime/streaming/operators/sink/kafka/mod.rs

@@ -117,7 +117,7 @@ impl KafkaSinkOperator {
             config.set("enable.idempotence", "true");
             let transactional_id = format!(
                 "fs-tx-{}-{}-{}-{}",
-                ctx.job_id, self.topic, ctx.subtask_idx, idx
+                ctx.job_id, self.topic, ctx.subtask_index, idx
             );
             config.set("transactional.id", &transactional_id);
 

src/runtime/streaming/operators/source/kafka/mod.rs

@@ -189,9 +189,9 @@ impl KafkaSourceOperator {
         let group_id = match (&self.group_id, &self.group_id_prefix) {
             (Some(gid), _) => gid.clone(),
             (None, Some(prefix)) => {
-                format!("{}-fs-{}-{}", prefix, ctx.job_id, ctx.subtask_idx)
+                format!("{}-fs-{}-{}", prefix, ctx.job_id, ctx.subtask_index)
             }
-            (None, None) => format!("fs-{}-{}-consumer", ctx.job_id, ctx.subtask_idx),
+            (None, None) => format!("fs-{}-{}-consumer", ctx.job_id, ctx.subtask_index),
         };
 
         for (key, value) in &self.client_configs {
@@ -223,7 +223,7 @@ impl KafkaSourceOperator {
         let pmax = ctx.parallelism.max(1) as i32;
 
         for p in partitions {
-            if p.id().rem_euclid(pmax) == ctx.subtask_idx as i32 {
+            if p.id().rem_euclid(pmax) == ctx.subtask_index as i32 {
                 let offset = state_map
                     .get(&p.id())
                     .map(|s| Offset::Offset(s.offset))
@@ -241,7 +241,7 @@ impl KafkaSourceOperator {
         if our_partitions.is_empty() {
             warn!(
                 "[Task {}] Subscribed to no partitions. Entering idle mode.",
-                ctx.subtask_idx
+                ctx.subtask_index
             );
             self.is_empty_assignment = true;
         } else {
@@ -366,7 +366,7 @@ impl SourceOperator for KafkaSourceOperator {
         _barrier: CheckpointBarrier,
         ctx: &mut TaskContext,
     ) -> Result<()> {
-        debug!("Source [{}] executing checkpoint", ctx.subtask_idx);
+        debug!("Source [{}] executing checkpoint", ctx.subtask_index);
 
         let mut topic_partitions = TopicPartitionList::new();
         for (&partition, &offset) in &self.current_offsets {

src/runtime/streaming/operators/watermark/watermark_generator.rs

@@ -141,7 +141,7 @@ impl Operator for WatermarkGeneratorOperator {
         if self.is_idle || time_since_last_emit > self.interval {
             debug!(
                 "[{}] emitting expression watermark {}",
-                ctx.subtask_idx,
+                ctx.subtask_index,
                 to_millis(self.state.max_watermark)
             );
 
@@ -174,7 +174,7 @@ impl Operator for WatermarkGeneratorOperator {
             if !self.is_idle && elapsed > idle_timeout {
                 info!(
                     "task [{}] entering Idle after {:?}",
-                    ctx.subtask_idx, idle_timeout
+                    ctx.subtask_index, idle_timeout
                 );
                 self.is_idle = true;
                 return Ok(vec![StreamOutput::Watermark(Watermark::Idle)]);

src/runtime/streaming/operators/windows/session_aggregating_window.rs

@@ -676,7 +676,7 @@ impl Operator for SessionWindowOperator {
         batch: RecordBatch,
         ctx: &mut TaskContext,
     ) -> Result<Vec<StreamOutput>> {
-        let watermark_time = ctx.last_present_watermark();
+        let watermark_time = ctx.current_watermark();
 
         let filtered_batch = self.filter_batch_by_time(batch, watermark_time)?;
         if filtered_batch.num_rows() == 0 {

src/runtime/streaming/operators/windows/sliding_aggregating_window.rs

@@ -339,7 +339,7 @@ impl Operator for SlidingWindowOperator {
             .ok_or_else(|| anyhow!("binning function must produce TimestampNanosecond"))?;
         let partition_ranges = partition(std::slice::from_ref(&sorted_bins))?.ranges();
 
-        let watermark = ctx.last_present_watermark();
+        let watermark = ctx.current_watermark();
 
         for range in partition_ranges {
             let bin_start = from_nanos(typed_bin.value(range.start) as u128);