channel try
diff --git a/rust/experimental/server/Cargo.toml b/rust/experimental/server/Cargo.toml
index 860ad43..0b55654 100644
--- a/rust/experimental/server/Cargo.toml
+++ b/rust/experimental/server/Cargo.toml
@@ -67,7 +67,6 @@
log = "0.4.17"
env_logger = "0.10.0"
crossbeam = "0.8.2"
-crossbeam-channel = "0.5"
tempdir = "0.3.7"
async-trait = "0.1.68"
futures = "0.3"
@@ -98,6 +97,7 @@
cap = "0.1.2"
spin = "0.9.8"
opendal = { version = "0.44.0", features = ["services-fs"]}
+crossbeam-channel = "0.5.8"
[dependencies.hdfs-native]
version = "0.7.0"
diff --git a/rust/experimental/server/src/channel.rs b/rust/experimental/server/src/channel.rs
new file mode 100644
index 0000000..13d9994
--- /dev/null
+++ b/rust/experimental/server/src/channel.rs
@@ -0,0 +1,35 @@
+use std::sync::Arc;
+use crossbeam_channel::{Receiver, Sender, unbounded};
+use crate::app::{AppManager, WritingViewContext};
+use crate::error::WorkerError;
+use crate::metric::TOTAL_RECEIVED_DATA;
+use crate::runtime::manager::RuntimeManager;
+
+#[derive(Clone)]
+pub struct Channel {
+ app_manager: Arc<AppManager>,
+ runtime_manager: RuntimeManager,
+
+ sender: Sender<WritingViewContext>,
+ receiver: Receiver<WritingViewContext>,
+}
+
+impl Channel {
+ pub fn new(app_manager: Arc<AppManager>, runtime_manager: RuntimeManager) -> Self {
+ let (s, r) = unbounded();
+ Self {
+ app_manager,
+ runtime_manager,
+ sender: s,
+ receiver: r,
+ }
+ }
+
+ pub fn send(&self, ctx: WritingViewContext) -> anyhow::Result<i32, WorkerError> {
+ let len: i32 = ctx.data_blocks.iter().map(|block| block.length).sum();
+ TOTAL_RECEIVED_DATA.inc_by(len as u64);
+
+ self.sender.send(ctx).map_err(|_| WorkerError::WRITING_TO_CHANNEL_FAIL)?;
+ Ok(len)
+ }
+}
\ No newline at end of file
diff --git a/rust/experimental/server/src/error.rs b/rust/experimental/server/src/error.rs
index 7a70408..cfbe608 100644
--- a/rust/experimental/server/src/error.rs
+++ b/rust/experimental/server/src/error.rs
@@ -16,7 +16,6 @@
// under the License.
use anyhow::Error;
-
use log::error;
use poem::error::ParseQueryError;
use thiserror::Error;
@@ -63,6 +62,9 @@
#[error("Spill event has been retried exceed the max limit for app: {0}")]
SPILL_EVENT_EXCEED_RETRY_MAX_LIMIT(String),
+
+ #[error("Failed to writing data to channel.")]
+ WRITING_TO_CHANNEL_FAIL,
}
impl From<AcquireError> for WorkerError {
diff --git a/rust/experimental/server/src/grpc.rs b/rust/experimental/server/src/grpc.rs
index 1ff4bd4..658b94f 100644
--- a/rust/experimental/server/src/grpc.rs
+++ b/rust/experimental/server/src/grpc.rs
@@ -48,6 +48,7 @@
};
use crate::util;
use tonic::{Request, Response, Status};
+use crate::channel::Channel;
/// Use the maximum value for HTTP/2 connection window size to avoid deadlock among multiplexed
/// streams on the same connection.
@@ -77,11 +78,12 @@
pub struct DefaultShuffleServer {
app_manager_ref: AppManagerRef,
+ writing_channel: Channel,
}
impl DefaultShuffleServer {
- pub fn from(app_manager_ref: AppManagerRef) -> DefaultShuffleServer {
- DefaultShuffleServer { app_manager_ref }
+ pub fn from(app_manager_ref: AppManagerRef, writing_channel: Channel) -> DefaultShuffleServer {
+ DefaultShuffleServer { app_manager_ref, writing_channel }
}
}
@@ -217,10 +219,12 @@
};
let ctx = WritingViewContext::from(uid.clone(), blocks);
- let inserted = app
- .insert(ctx)
- .instrument_await(format!("insert data for app. uid: {:?}", &uid))
- .await;
+ let inserted = self.writing_channel.send(ctx);
+
+ // let inserted = app
+ // .insert(ctx)
+ // .instrument_await(format!("insert data for app. uid: {:?}", &uid))
+ // .await;
if inserted.is_err() {
let err = format!(
"Errors on putting data. app_id: {}, err: {:?}",
diff --git a/rust/experimental/server/src/lib.rs b/rust/experimental/server/src/lib.rs
index 1219285..54d3f18 100644
--- a/rust/experimental/server/src/lib.rs
+++ b/rust/experimental/server/src/lib.rs
@@ -31,6 +31,7 @@
pub mod signal;
pub mod store;
pub mod util;
+pub mod channel;
use crate::app::AppManager;
use crate::grpc::DefaultShuffleServer;
@@ -69,14 +70,16 @@
HTTP_SERVICE.start(runtime_manager.clone(), http_port);
let (tx, rx) = oneshot::channel::<()>();
+ let cloned_runtime_manager = runtime_manager.clone();
+ let app_manager_ref = AppManager::get_ref(cloned_runtime_manager, config.clone());
+
+ let writing_channel = channel::Channel::new(app_manager_ref.clone(), runtime_manager.clone());
// implement server startup
- let cloned_runtime_manager = runtime_manager.clone();
runtime_manager.grpc_runtime.spawn(async move {
- let app_manager_ref = AppManager::get_ref(cloned_runtime_manager, config.clone());
let rpc_port = config.grpc_port.unwrap_or(19999);
info!("Starting GRpc server with port:[{}] ......", rpc_port);
- let shuffle_server = DefaultShuffleServer::from(app_manager_ref);
+ let shuffle_server = DefaultShuffleServer::from(app_manager_ref, writing_channel.clone());
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), rpc_port as u16);
let service = ShuffleServerServer::new(shuffle_server)
.max_decoding_message_size(usize::MAX)
diff --git a/rust/experimental/server/src/main.rs b/rust/experimental/server/src/main.rs
index f5e93c5..86c7ba9 100644
--- a/rust/experimental/server/src/main.rs
+++ b/rust/experimental/server/src/main.rs
@@ -62,6 +62,7 @@
pub mod signal;
pub mod store;
mod util;
+pub mod channel;
const MAX_MEMORY_ALLOCATION_SIZE_ENV_KEY: &str = "MAX_MEMORY_ALLOCATION_SIZE";
const DEFAULT_SHUFFLE_SERVER_TAG: &str = "ss_v4";
@@ -233,8 +234,10 @@
let available_cores = std::thread::available_parallelism()?;
debug!("GRpc service with parallelism: [{}]", &available_cores);
+ let writing_channel = channel::Channel::new(app_manager_ref.clone(), runtime_manager.clone());
+
for _ in 0..available_cores.into() {
- let shuffle_server = DefaultShuffleServer::from(app_manager_ref.clone());
+ let shuffle_server = DefaultShuffleServer::from(app_manager_ref.clone(), writing_channel.clone());
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), rpc_port as u16);
let service = ShuffleServerServer::new(shuffle_server)
.max_decoding_message_size(usize::MAX)
diff --git a/rust/experimental/server/src/store/memory.rs b/rust/experimental/server/src/store/memory.rs
index e93e4ae..fad27e7 100644
--- a/rust/experimental/server/src/store/memory.rs
+++ b/rust/experimental/server/src/store/memory.rs
@@ -256,12 +256,13 @@
async fn insert(&self, ctx: WritingViewContext) -> Result<(), WorkerError> {
let uid = ctx.uid;
- let buffer = self.get_or_create_underlying_staging_buffer(uid.clone());
- let mut buffer_guarded = buffer.lock();
-
- let blocks = ctx.data_blocks;
- let inserted_size = buffer_guarded.add(blocks)?;
- drop(buffer_guarded);
+ let inserted_size = {
+ let buffer = self.get_or_create_underlying_staging_buffer(uid);
+ let mut buffer_guarded = buffer.lock();
+ let blocks = ctx.data_blocks;
+ let inserted_size = buffer_guarded.add(blocks)?;
+ inserted_size
+ };
self.budget.allocated_to_used(inserted_size)?;
