Release GIL for listing

Author: Kimahriman

python/src/lib.rs

@@ -71,7 +71,7 @@ impl PyFileStatus {
 
 #[pyclass(name = "FileStatusIter")]
 struct PyFileStatusIter {
-    inner: ListStatusIterator,
+    inner: Arc<ListStatusIterator>,
     rt: Arc<Runtime>,
 }
 
@@ -81,10 +81,11 @@ impl PyFileStatusIter {
         slf
     }
 
-    fn __next__(mut slf: PyRefMut<'_, Self>) -> PyHdfsResult<Option<PyFileStatus>> {
-        // This is dumb, figure out how to get around the double borrow here
+    fn __next__(slf: PyRefMut<'_, Self>) -> PyHdfsResult<Option<PyFileStatus>> {
+        // Kinda dumb, but lets us release the GIL while getting the next value
+        let inner = Arc::clone(&slf.inner);
         let rt = Arc::clone(&slf.rt);
-        if let Some(result) = rt.block_on(slf.inner.next()) {
+        if let Some(result) = slf.py().allow_threads(|| rt.block_on(inner.next())) {
             Ok(Some(PyFileStatus::from(result?)))
         } else {
             Ok(None)
@@ -305,7 +306,7 @@ impl RawClient {
     pub fn list_status(&self, path: &str, recursive: bool) -> PyFileStatusIter {
         let inner = self.inner.list_status_iter(path, recursive);
         PyFileStatusIter {
-            inner,
+            inner: Arc::new(inner),
             rt: Arc::clone(&self.rt),
         }
     }

rust/src/client.rs

@@ -588,7 +588,7 @@ impl DirListingIterator {
 pub struct ListStatusIterator {
     mount_table: Arc<MountTable>,
     recursive: bool,
-    iters: Vec<DirListingIterator>,
+    iters: Arc<tokio::sync::Mutex<Vec<DirListingIterator>>>,
 }
 
 impl ListStatusIterator {
@@ -598,20 +598,21 @@ impl ListStatusIterator {
         ListStatusIterator {
             mount_table,
             recursive,
-            iters: vec![initial],
+            iters: Arc::new(tokio::sync::Mutex::new(vec![initial])),
         }
     }
 
-    pub async fn next(&mut self) -> Option<Result<FileStatus>> {
+    pub async fn next(&self) -> Option<Result<FileStatus>> {
         let mut next_file: Option<Result<FileStatus>> = None;
+        let mut iters = self.iters.lock().await;
         while next_file.is_none() {
-            if let Some(iter) = self.iters.last_mut() {
+            if let Some(iter) = iters.last_mut() {
                 if let Some(file_result) = iter.next().await {
                     if let Ok(file) = file_result {
                         // Return the directory as the next result, but start traversing into that directory
                         // next if we're doing a recursive listing
                         if file.isdir && self.recursive {
-                            self.iters.push(DirListingIterator::new(
+                            iters.push(DirListingIterator::new(
                                 file.path.clone(),
                                 &self.mount_table,
                                 false,
@@ -624,7 +625,7 @@ impl ListStatusIterator {
                     }
                 } else {
                     // We've exhausted this directory
-                    self.iters.pop();
+                    iters.pop();
                 }
             } else {
                 // There's nothing left, just return None
@@ -636,7 +637,7 @@ impl ListStatusIterator {
     }
 
     pub fn into_stream(self) -> BoxStream<'static, Result<FileStatus>> {
-        let listing = stream::unfold(self, |mut state| async move {
+        let listing = stream::unfold(self, |state| async move {
             let next = state.next().await;
             next.map(|n| (n, state))
         });