Skip to content

Commit 05c7f37

Browse files
s3riuss3rius
committed
Implemented better subsccription interface.
1 parent 550762e commit 05c7f37

7 files changed

Lines changed: 245 additions & 16 deletions

File tree

‎python/natsrpy/__init__.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
from ._natsrpy_rs import Message, Nats, Subscription
1+
from ._natsrpy_rs import CallbackSubscription, IteratorSubscription, Message, Nats
22

33
__all__ = [
4+
"CallbackSubscription",
5+
"IteratorSubscription",
46
"Message",
57
"Nats",
6-
"Subscription",
78
]

‎python/natsrpy/_natsrpy_rs/__init__.pyi‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,20 @@
11
from collections.abc import Awaitable, Callable
22
from datetime import timedelta
3-
from typing import Any
3+
from typing import Any, overload
44

55
from natsrpy._natsrpy_rs.js import JetStream
66
from natsrpy._natsrpy_rs.message import Message
77

8-
class Subscription:
9-
def __aiter__(self) -> Subscription: ...
8+
class IteratorSubscription:
9+
def __aiter__(self) -> IteratorSubscription: ...
1010
async def __anext__(self) -> Message: ...
1111
async def unsubscribe(self, limit: int | None = None) -> None: ...
1212
async def drain(self) -> None: ...
1313

14+
class CallbackSubscription:
15+
async def unsubscribe(self, limit: int | None = None) -> None: ...
16+
async def drain(self) -> None: ...
17+
1418
class Nats:
1519
def __init__(
1620
self,
@@ -40,11 +44,18 @@ class Nats:
4044
async def request(self, subject: str, payload: bytes) -> None: ...
4145
async def drain(self) -> None: ...
4246
async def flush(self) -> None: ...
47+
@overload
48+
async def subscribe(
49+
self,
50+
subject: str,
51+
callback: Callable[[Message], Awaitable[None]],
52+
) -> CallbackSubscription: ...
53+
@overload
4354
async def subscribe(
4455
self,
4556
subject: str,
46-
callback: Callable[[Message], Awaitable[None]] | None = None,
47-
) -> Subscription: ...
57+
callback: None = None,
58+
) -> IteratorSubscription: ...
4859
async def jetstream(self) -> JetStream: ...
4960

50-
__all__ = ["Message", "Nats", "Subscription"]
61+
__all__ = ["CallbackSubscription", "IteratorSubscription", "Message", "Nats"]

‎src/lib.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ pub mod exceptions;
2424
pub mod js;
2525
pub mod message;
2626
pub mod nats_cls;
27-
pub mod subscription;
27+
pub mod subscriptions;
2828
pub mod utils;
2929

3030
#[pyo3::pymodule]
@@ -38,7 +38,7 @@ pub mod _natsrpy_rs {
3838
#[pymodule_export]
3939
use super::nats_cls::NatsCls;
4040
#[pymodule_export]
41-
use super::subscription::Subscription;
41+
use super::subscriptions::{callback::CallbackSubscription, iterator::IteratorSubscription};
4242

4343
#[pymodule_export]
4444
use super::js::pymod as js;

‎src/nats_cls.rs‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
use async_nats::{Subject, client::traits::Publisher, message::OutboundMessage};
22
use pyo3::{
3-
Bound, Py, PyAny, Python,
3+
Bound, IntoPyObjectExt, Py, PyAny, Python,
44
types::{PyBytes, PyBytesMethods, PyDict},
55
};
66
use std::{sync::Arc, time::Duration};
77
use tokio::sync::RwLock;
88

99
use crate::{
1010
exceptions::rust_err::{NatsrpyError, NatsrpyResult},
11-
subscription::Subscription,
11+
subscriptions::{callback::CallbackSubscription, iterator::IteratorSubscription},
1212
utils::{
1313
futures::natsrpy_future_with_timeout,
1414
headers::NatsrpyHeadermapExt,
@@ -208,10 +208,13 @@ impl NatsCls {
208208
let session = self.nats_session.clone();
209209
natsrpy_future(py, async move {
210210
if let Some(session) = session.read().await.as_ref() {
211-
Ok(Subscription::new(
212-
session.subscribe(subject).await?,
213-
callback,
214-
)?)
211+
if let Some(cb) = callback {
212+
let sub = CallbackSubscription::new(session.subscribe(subject).await?, cb)?;
213+
Ok(Python::attach(|gil| sub.into_py_any(gil))?)
214+
} else {
215+
let sub = IteratorSubscription::new(session.subscribe(subject).await?);
216+
Ok(Python::attach(|gil| sub.into_py_any(gil))?)
217+
}
215218
} else {
216219
Err(NatsrpyError::NotInitialized)
217220
}

‎src/subscriptions/callback.rs‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
use std::sync::Arc;
2+
3+
use futures_util::StreamExt;
4+
use pyo3::{Bound, Py, PyAny, Python};
5+
use tokio::sync::Mutex;
6+
7+
use crate::{exceptions::rust_err::NatsrpyResult, utils::natsrpy_future};
8+
9+
#[pyo3::pyclass]
10+
pub struct CallbackSubscription {
11+
inner: Option<Arc<Mutex<async_nats::Subscriber>>>,
12+
reading_task: tokio::task::AbortHandle,
13+
}
14+
15+
async fn process_message(message: async_nats::message::Message, py_callback: Py<PyAny>) {
16+
let task = async || -> NatsrpyResult<()> {
17+
let message = crate::message::Message::try_from(&message)?;
18+
let awaitable = Python::attach(|gil| -> NatsrpyResult<_> {
19+
let res = py_callback.call1(gil, (message,))?;
20+
let rust_task = pyo3_async_runtimes::tokio::into_future(res.into_bound(gil))?;
21+
Ok(rust_task)
22+
})?;
23+
awaitable.await?;
24+
Ok(())
25+
};
26+
if let Err(err) = task().await {
27+
log::error!("Cannot process message {message:?}. Error: {err}");
28+
}
29+
}
30+
31+
async fn start_py_sub(
32+
sub: Arc<Mutex<async_nats::Subscriber>>,
33+
py_callback: Py<PyAny>,
34+
locals: pyo3_async_runtimes::TaskLocals,
35+
) {
36+
while let Some(message) = sub.lock().await.next().await {
37+
let py_cb = Python::attach(|py| py_callback.clone_ref(py));
38+
tokio::spawn(pyo3_async_runtimes::tokio::scope(
39+
locals.clone(),
40+
process_message(message, py_cb),
41+
));
42+
}
43+
}
44+
45+
impl CallbackSubscription {
46+
pub fn new(sub: async_nats::Subscriber, callback: Py<PyAny>) -> NatsrpyResult<Self> {
47+
let sub = Arc::new(Mutex::new(sub));
48+
let cb_sub = sub.clone();
49+
let task_locals = Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?;
50+
let task_handle = tokio::task::spawn(pyo3_async_runtimes::tokio::scope(
51+
task_locals.clone(),
52+
start_py_sub(cb_sub, callback, task_locals),
53+
))
54+
.abort_handle();
55+
Ok(Self {
56+
inner: Some(sub),
57+
reading_task: task_handle,
58+
})
59+
}
60+
}
61+
62+
#[pyo3::pymethods]
63+
impl CallbackSubscription {
64+
#[pyo3(signature=(limit=None))]
65+
pub fn unsubscribe<'py>(
66+
&self,
67+
py: Python<'py>,
68+
limit: Option<u64>,
69+
) -> NatsrpyResult<Bound<'py, PyAny>> {
70+
let Some(inner) = self.inner.clone() else {
71+
unreachable!("Subscription used after del")
72+
};
73+
natsrpy_future(py, async move {
74+
if let Some(limit) = limit {
75+
inner.lock().await.unsubscribe_after(limit).await?;
76+
} else {
77+
inner.lock().await.unsubscribe().await?;
78+
}
79+
Ok(())
80+
})
81+
}
82+
83+
pub fn drain<'py>(&self, py: Python<'py>) -> NatsrpyResult<Bound<'py, PyAny>> {
84+
let Some(inner) = self.inner.clone() else {
85+
unreachable!("Subscription used after del")
86+
};
87+
natsrpy_future(py, async move {
88+
inner.lock().await.drain().await?;
89+
Ok(())
90+
})
91+
}
92+
}
93+
94+
/// This is required only because
95+
/// in nats library they run async operation on Drop.
96+
///
97+
/// Because of that we need to execute drop in async
98+
/// runtime's context.
99+
///
100+
/// And because we want to perform a drop,
101+
/// we need somehow drop the inner variable,
102+
/// but leave self intouch. That is exactly why we have
103+
/// Option<Arc<...>>. So we can just assign it to None
104+
/// and it will perform a drop.
105+
impl Drop for CallbackSubscription {
106+
fn drop(&mut self) {
107+
pyo3_async_runtimes::tokio::get_runtime().block_on(async move {
108+
self.inner = None;
109+
self.reading_task.abort();
110+
});
111+
}
112+
}

‎src/subscriptions/iterator.rs‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
use std::sync::Arc;
2+
3+
use futures_util::StreamExt;
4+
use pyo3::{Bound, PyAny, PyRef, Python};
5+
use tokio::sync::Mutex;
6+
7+
use crate::exceptions::rust_err::{NatsrpyError, NatsrpyResult};
8+
use crate::utils::futures::natsrpy_future_with_timeout;
9+
use crate::utils::natsrpy_future;
10+
use crate::utils::py_types::TimeValue;
11+
12+
#[pyo3::pyclass]
13+
pub struct IteratorSubscription {
14+
inner: Option<Arc<Mutex<async_nats::Subscriber>>>,
15+
}
16+
17+
impl IteratorSubscription {
18+
#[must_use]
19+
pub fn new(sub: async_nats::Subscriber) -> Self {
20+
Self {
21+
inner: Some(Arc::new(Mutex::new(sub))),
22+
}
23+
}
24+
}
25+
26+
#[pyo3::pymethods]
27+
impl IteratorSubscription {
28+
#[must_use]
29+
pub const fn __aiter__(slf: PyRef<Self>) -> PyRef<Self> {
30+
slf
31+
}
32+
33+
pub fn next<'py>(
34+
&self,
35+
py: Python<'py>,
36+
timeout: Option<TimeValue>,
37+
) -> NatsrpyResult<Bound<'py, PyAny>> {
38+
let Some(inner) = self.inner.clone() else {
39+
unreachable!("Subscription used after del")
40+
};
41+
natsrpy_future_with_timeout(py, timeout, async move {
42+
let Some(message) = inner.lock().await.next().await else {
43+
return Err(NatsrpyError::AsyncStopIteration);
44+
};
45+
crate::message::Message::try_from(message)
46+
})
47+
}
48+
49+
pub fn __anext__<'py>(&self, py: Python<'py>) -> NatsrpyResult<Bound<'py, PyAny>> {
50+
self.next(py, None)
51+
}
52+
53+
#[pyo3(signature=(limit=None))]
54+
pub fn unsubscribe<'py>(
55+
&self,
56+
py: Python<'py>,
57+
limit: Option<u64>,
58+
) -> NatsrpyResult<Bound<'py, PyAny>> {
59+
let Some(inner) = self.inner.clone() else {
60+
unreachable!("Subscription used after del")
61+
};
62+
natsrpy_future(py, async move {
63+
if let Some(limit) = limit {
64+
inner.lock().await.unsubscribe_after(limit).await?;
65+
} else {
66+
inner.lock().await.unsubscribe().await?;
67+
}
68+
Ok(())
69+
})
70+
}
71+
72+
pub fn drain<'py>(&self, py: Python<'py>) -> NatsrpyResult<Bound<'py, PyAny>> {
73+
let Some(inner) = self.inner.clone() else {
74+
unreachable!("Subscription used after del")
75+
};
76+
natsrpy_future(py, async move {
77+
inner.lock().await.drain().await?;
78+
Ok(())
79+
})
80+
}
81+
}
82+
83+
/// This is required only because
84+
/// in nats library they run async operation on Drop.
85+
///
86+
/// Because of that we need to execute drop in async
87+
/// runtime's context.
88+
///
89+
/// And because we want to perform a drop,
90+
/// we need somehow drop the inner variable,
91+
/// but leave self intouch. That is exactly why we have
92+
/// Option<Arc<...>>. So we can just assign it to None
93+
/// and it will perform a drop.
94+
impl Drop for IteratorSubscription {
95+
fn drop(&mut self) {
96+
pyo3_async_runtimes::tokio::get_runtime().block_on(async move {
97+
self.inner = None;
98+
});
99+
}
100+
}

‎src/subscriptions/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
pub mod callback;
2+
pub mod iterator;

0 commit comments

Comments
 (0)