fix(rust): forget to call report blockIds method await

rust/experimental/server/src/grpc.rs

@@ -446,7 +446,16 @@ impl ShuffleServer for DefaultShuffleServer {
  },
  blocks: partition_to_block_id.block_ids,
  };
- let _ = app.report_block_ids(ctx);
+
+ match app.report_block_ids(ctx).await {
+ Err(e) => {
+ return Ok(Response::new(ReportShuffleResultResponse {
+ status: StatusCode::INTERNAL_ERROR.into(),
+ ret_msg: e.to_string(),
+ }))
+ }
+ _ => (),
+ }
  }
 
  Ok(Response::new(ReportShuffleResultResponse {

rust/experimental/server/src/lib.rs

@@ -40,13 +40,15 @@ use crate::proto::uniffle::shuffle_server_client::ShuffleServerClient;
 use crate::proto::uniffle::shuffle_server_server::ShuffleServerServer;
 use crate::proto::uniffle::{
  GetLocalShuffleDataRequest, GetLocalShuffleIndexRequest, GetMemoryShuffleDataRequest,
- RequireBufferRequest, SendShuffleDataRequest, ShuffleBlock, ShuffleData,
- ShuffleRegisterRequest,
+ GetShuffleResultRequest, PartitionToBlockIds, ReportShuffleResultRequest, RequireBufferRequest,
+ SendShuffleDataRequest, ShuffleBlock, ShuffleData, ShuffleRegisterRequest,
 };
 use crate::runtime::manager::RuntimeManager;
 use crate::util::gen_worker_uid;
 use anyhow::Result;
 use bytes::{Buf, Bytes, BytesMut};
+use croaring::treemap::JvmSerializer;
+use croaring::Treemap;
 use log::info;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::time::Duration;
@@ -164,6 +166,20 @@ pub async fn write_read_for_one_time(mut client: ShuffleServerClient<Channel>) -
 
  let response = response.into_inner();
  assert_eq!(0, response.status);
+
+ // report the finished block ids
+ client
+ .report_shuffle_result(ReportShuffleResultRequest {
+ app_id: app_id.clone(),
+ shuffle_id: 0,
+ task_attempt_id: 0,
+ bitmap_num: 0,
+ partition_to_block_ids: vec![PartitionToBlockIds {
+ partition_id: idx,
+ block_ids: vec![idx as i64],
+ }],
+ })
+ .await?;
  }
 
  tokio::time::sleep(Duration::from_secs(1)).await;
@@ -173,6 +189,21 @@ pub async fn write_read_for_one_time(mut client: ShuffleServerClient<Channel>) -
 
  // firstly. read from the memory
  for idx in 0..batch_size {
+ let block_id_result = client
+ .get_shuffle_result(GetShuffleResultRequest {
+ app_id: app_id.clone(),
+ shuffle_id: 0,
+ partition_id: idx,
+ })
+ .await?
+ .into_inner();
+
+ assert_eq!(0, block_id_result.status);
+
+ let block_id_bitmap = Treemap::deserialize(&*block_id_result.serialized_bitmap)?;
+ assert_eq!(1, block_id_bitmap.iter().count());
+ assert!(block_id_bitmap.contains(idx as u64));
+
  let response_data = client
  .get_memory_shuffle_data(GetMemoryShuffleDataRequest {
  app_id: app_id.clone(),