use crate::client_factory::ClientFactory;
use crate::event_reader::EventReader;
use crate::reader_group::reader_group_state::Offset;
use crate::reader_group_config::ReaderGroupConfig;
use pravega_client_shared::{Reader, Scope, ScopedSegment, ScopedStream};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
cfg_if::cfg_if! {
if #[cfg(test)] {
use crate::reader_group::reader_group_state::MockReaderGroupState as ReaderGroupState;
} else {
use crate::reader_group::reader_group_state::ReaderGroupState;
}
}
#[derive(new)]
pub struct ReaderGroup {
name: String,
config: ReaderGroupConfig,
pub state: Arc<Mutex<ReaderGroupState>>,
client_factory: ClientFactory,
}
impl ReaderGroup {
cfg_if::cfg_if! {
if #[cfg(test)] {
async fn create_rg_state(
_scope: Scope,
_name: String,
_rg_config: ReaderGroupConfig,
_client_factory: &ClientFactory,
_init_segments: HashMap<ScopedSegment, Offset>,
) -> ReaderGroupState {
ReaderGroupState::default()
}
} else {
async fn create_rg_state (
scope: Scope,
name: String,
rg_config: ReaderGroupConfig,
client_factory: &ClientFactory,
init_segments: HashMap<ScopedSegment, Offset>,
) -> ReaderGroupState {
ReaderGroupState::new(scope, name, client_factory, rg_config.config, init_segments).await
}
}
}
pub async fn create(
scope: Scope,
name: String,
rg_config: ReaderGroupConfig,
client_factory: ClientFactory,
) -> ReaderGroup {
let streams: Vec<ScopedStream> = rg_config.get_streams();
let mut init_segments: HashMap<ScopedSegment, Offset> = HashMap::new();
for stream in streams {
let segments = client_factory
.get_controller_client()
.get_head_segments(&stream)
.await
.expect("Error while fetching stream's starting segments to read from ");
init_segments.extend(segments.iter().map(|(seg, off)| {
(
ScopedSegment {
scope: stream.scope.clone(),
stream: stream.stream.clone(),
segment: seg.clone(),
},
Offset::new(*off),
)
}));
}
let rg_state = ReaderGroup::create_rg_state(
scope,
name.clone(),
rg_config.clone(),
&client_factory,
init_segments,
)
.await;
ReaderGroup {
name: name.clone(),
config: rg_config.clone(),
state: Arc::new(Mutex::new(rg_state)),
client_factory,
}
}
pub fn get_name(&self) -> String {
self.name.clone()
}
pub async fn create_reader(&self, reader_id: String) -> EventReader {
let r: Reader = Reader::from(reader_id.clone());
self.state
.lock()
.await
.add_reader(&r)
.await
.expect("Error while creating the reader");
EventReader::init_reader(reader_id, self.state.clone(), self.client_factory.clone()).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::SynchronizerError::SyncUpdateError;
use crate::reader_group::reader_group_state::ReaderGroupStateError;
use crate::reader_group_config::ReaderGroupConfigBuilder;
use pravega_client_config::ClientConfigBuilder;
use pravega_client_config::MOCK_CONTROLLER_URI;
#[test]
#[should_panic]
fn test_create_reader_error() {
let client_factory = ClientFactory::new(
ClientConfigBuilder::default()
.controller_uri(MOCK_CONTROLLER_URI)
.build()
.unwrap(),
);
let mut mock_rg_state = ReaderGroupState::default();
let err: Result<(), ReaderGroupStateError> = Result::Err(ReaderGroupStateError::SyncError {
error_msg: "Reader already exists".to_string(),
source: SyncUpdateError {
error_msg: "update failed".to_string(),
},
});
mock_rg_state.expect_add_reader().return_once(move |_| err);
let rg = ReaderGroup {
name: "rg".to_string(),
config: ReaderGroupConfigBuilder::default()
.add_stream(ScopedStream::from("scope/s1"))
.build(),
state: Arc::new(Mutex::new(mock_rg_state)),
client_factory: client_factory.clone(),
};
client_factory
.get_runtime()
.block_on(rg.create_reader("r1".to_string()));
}
}