Skip to content

Commit

Permalink
Make multi subscriber optional (#185)
Browse files Browse the repository at this point in the history
* make optional

* fmt
  • Loading branch information
yoshidan authored Aug 14, 2023
1 parent e2d8530 commit e996765
Showing 1 changed file with 45 additions and 6 deletions.
51 changes: 45 additions & 6 deletions pubsub/src/subscription.rs
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,23 @@ pub struct SubscriptionConfigToUpdate {
pub retry_policy: Option<RetryPolicy>,
}

#[derive(Debug, Clone, Default)]
pub struct SubscribeConfig {
enable_multiple_subscriber: bool,
subscriber_config: SubscriberConfig,
}

impl SubscribeConfig {
pub fn with_enable_multiple_subscriber(mut self, v: bool) -> Self {
self.enable_multiple_subscriber = v;
self
}
pub fn with_subscriber_config(mut self, v: SubscriberConfig) -> Self {
self.subscriber_config = v;
self
}
}

#[derive(Debug, Clone)]
pub struct ReceiveConfig {
pub worker_count: usize,
Expand Down Expand Up @@ -353,13 +370,25 @@ impl Subscription {
/// Ok(())
/// }
/// ```
pub async fn subscribe(&self, opt: Option<SubscriberConfig>) -> Result<MessageStream, Status> {
pub async fn subscribe(&self, opt: Option<SubscribeConfig>) -> Result<MessageStream, Status> {
let (tx, rx) = async_channel::unbounded::<ReceivedMessage>();
let cancel = CancellationToken::new();

// spawn a separate subscriber task for each connection in the pool
for _ in 0..self.pool_size() {
Subscriber::start(cancel.clone(), self.fqsn.clone(), self.subc.clone(), tx.clone(), opt.clone());
let opt = opt.unwrap_or_default();
let subscribers = if opt.enable_multiple_subscriber {
self.pool_size()
} else {
1
};
for _ in 0..subscribers {
Subscriber::start(
cancel.clone(),
self.fqsn.clone(),
self.subc.clone(),
tx.clone(),
Some(opt.subscriber_config.clone()),
);
}

Ok(MessageStream { queue: rx, cancel })
Expand Down Expand Up @@ -584,7 +613,7 @@ mod tests {
use crate::apiv1::publisher_client::PublisherClient;
use crate::apiv1::subscriber_client::SubscriberClient;
use crate::subscriber::ReceivedMessage;
use crate::subscription::{SeekTo, Subscription, SubscriptionConfig, SubscriptionConfigToUpdate};
use crate::subscription::{SeekTo, SubscribeConfig, Subscription, SubscriptionConfig, SubscriptionConfigToUpdate};

const PROJECT_NAME: &str = "local-project";
const EMULATOR: &str = "localhost:8681";
Expand Down Expand Up @@ -966,12 +995,22 @@ mod tests {

#[tokio::test]
#[serial]
async fn test_subscribe() {
async fn test_subscribe_single_subscriber() {
test_subscribe(None).await
}

#[tokio::test]
#[serial]
async fn test_subscribe_multiple_subscriber() {
test_subscribe(Some(SubscribeConfig::default().with_enable_multiple_subscriber(true))).await
}

async fn test_subscribe(opt: Option<SubscribeConfig>) {
let subscription = create_subscription(false).await;
let received = Arc::new(Mutex::new(false));
let checking = received.clone();
let _handler = tokio::spawn(async move {
let mut iter = subscription.subscribe(None).await.unwrap();
let mut iter = subscription.subscribe(opt).await.unwrap();
while let Some(message) = iter.next().await {
*received.lock().unwrap() = true;
let _ = message.ack().await;
Expand Down

0 comments on commit e996765

Please sign in to comment.