Release GIL when blocking on inner calls

Author: Kimahriman

python/src/lib.rs

@@ -150,21 +150,21 @@ impl RawFileReader {
         self.inner.tell()
     }
 
-    pub fn read(&mut self, len: i64) -> PyHdfsResult<Cow<[u8]>> {
+    pub fn read(&mut self, len: i64, py: Python) -> PyHdfsResult<Cow<[u8]>> {
         let read_len = if len < 0 {
             self.inner.remaining()
         } else {
             len as usize
         };
         Ok(Cow::from(
-            self.rt.block_on(self.inner.read(read_len))?.to_vec(),
+            py.allow_threads(|| self.rt.block_on(self.inner.read(read_len)))?
+                .to_vec(),
         ))
     }
 
-    pub fn read_range(&self, offset: usize, len: usize) -> PyHdfsResult<Cow<[u8]>> {
+    pub fn read_range(&self, offset: usize, len: usize, py: Python) -> PyHdfsResult<Cow<[u8]>> {
         Ok(Cow::from(
-            self.rt
-                .block_on(self.inner.read_range(offset, len))?
+            py.allow_threads(|| self.rt.block_on(self.inner.read_range(offset, len)))?
                 .to_vec(),
         ))
     }
@@ -254,12 +254,12 @@ struct RawFileWriter {
 
 #[pymethods]
 impl RawFileWriter {
-    pub fn write(&mut self, buf: Vec<u8>) -> PyHdfsResult<usize> {
-        Ok(self.rt.block_on(self.inner.write(Bytes::from(buf)))?)
+    pub fn write(&mut self, buf: Vec<u8>, py: Python) -> PyHdfsResult<usize> {
+        Ok(py.allow_threads(|| self.rt.block_on(self.inner.write(Bytes::from(buf))))?)
     }
 
-    pub fn close(&mut self) -> PyHdfsResult<()> {
-        Ok(self.rt.block_on(self.inner.close())?)
+    pub fn close(&mut self, py: Python) -> PyHdfsResult<()> {
+        Ok(py.allow_threads(|| self.rt.block_on(self.inner.close()))?)
     }
 }
 
@@ -294,11 +294,12 @@ impl RawClient {
         })
     }
 
-    pub fn get_file_info(&self, path: &str) -> PyHdfsResult<PyFileStatus> {
-        Ok(self
-            .rt
-            .block_on(self.inner.get_file_info(path))
-            .map(PyFileStatus::from)?)
+    pub fn get_file_info(&self, path: &str, py: Python) -> PyHdfsResult<PyFileStatus> {
+        Ok(py.allow_threads(|| {
+            self.rt
+                .block_on(self.inner.get_file_info(path))
+                .map(PyFileStatus::from)
+        })?)
     }
 
     pub fn list_status(&self, path: &str, recursive: bool) -> PyFileStatusIter {
@@ -309,78 +310,93 @@ impl RawClient {
         }
     }
 
-    pub fn read(&self, path: &str) -> PyHdfsResult<RawFileReader> {
-        let file_reader = self.rt.block_on(self.inner.read(path))?;
+    pub fn read(&self, path: &str, py: Python) -> PyHdfsResult<RawFileReader> {
+        let file_reader = py.allow_threads(|| self.rt.block_on(self.inner.read(path)))?;
 
         Ok(RawFileReader {
             inner: file_reader,
             rt: Arc::clone(&self.rt),
         })
     }
 
-    pub fn create(&self, src: &str, write_options: PyWriteOptions) -> PyHdfsResult<RawFileWriter> {
-        let file_writer = self
-            .rt
-            .block_on(self.inner.create(src, WriteOptions::from(write_options)))?;
+    pub fn create(
+        &self,
+        src: &str,
+        write_options: PyWriteOptions,
+        py: Python,
+    ) -> PyHdfsResult<RawFileWriter> {
+        let file_writer = py.allow_threads(|| {
+            self.rt
+                .block_on(self.inner.create(src, WriteOptions::from(write_options)))
+        })?;
 
         Ok(RawFileWriter {
             inner: file_writer,
             rt: Arc::clone(&self.rt),
         })
     }
 
-    pub fn append(&self, src: &str) -> PyHdfsResult<RawFileWriter> {
-        let file_writer = self.rt.block_on(self.inner.append(src))?;
+    pub fn append(&self, src: &str, py: Python) -> PyHdfsResult<RawFileWriter> {
+        let file_writer = py.allow_threads(|| self.rt.block_on(self.inner.append(src)))?;
 
         Ok(RawFileWriter {
             inner: file_writer,
             rt: Arc::clone(&self.rt),
         })
     }
 
-    pub fn mkdirs(&self, path: &str, permission: u32, create_parent: bool) -> PyHdfsResult<()> {
-        Ok(self
-            .rt
-            .block_on(self.inner.mkdirs(path, permission, create_parent))?)
+    pub fn mkdirs(
+        &self,
+        path: &str,
+        permission: u32,
+        create_parent: bool,
+        py: Python,
+    ) -> PyHdfsResult<()> {
+        Ok(py.allow_threads(|| {
+            self.rt
+                .block_on(self.inner.mkdirs(path, permission, create_parent))
+        })?)
     }
 
-    pub fn rename(&self, src: &str, dst: &str, overwrite: bool) -> PyHdfsResult<()> {
-        Ok(self.rt.block_on(self.inner.rename(src, dst, overwrite))?)
+    pub fn rename(&self, src: &str, dst: &str, overwrite: bool, py: Python) -> PyHdfsResult<()> {
+        Ok(py.allow_threads(|| self.rt.block_on(self.inner.rename(src, dst, overwrite)))?)
     }
 
-    pub fn delete(&self, path: &str, recursive: bool) -> PyHdfsResult<bool> {
-        Ok(self.rt.block_on(self.inner.delete(path, recursive))?)
+    pub fn delete(&self, path: &str, recursive: bool, py: Python) -> PyHdfsResult<bool> {
+        Ok(py.allow_threads(|| self.rt.block_on(self.inner.delete(path, recursive)))?)
     }
 
-    pub fn set_times(&self, path: &str, mtime: u64, atime: u64) -> PyHdfsResult<()> {
-        Ok(self.rt.block_on(self.inner.set_times(path, mtime, atime))?)
+    pub fn set_times(&self, path: &str, mtime: u64, atime: u64, py: Python) -> PyHdfsResult<()> {
+        Ok(py.allow_threads(|| self.rt.block_on(self.inner.set_times(path, mtime, atime)))?)
     }
 
     pub fn set_owner(
         &self,
         path: &str,
         owner: Option<&str>,
         group: Option<&str>,
+        py: Python,
     ) -> PyHdfsResult<()> {
-        Ok(self.rt.block_on(self.inner.set_owner(path, owner, group))?)
+        Ok(py.allow_threads(|| self.rt.block_on(self.inner.set_owner(path, owner, group)))?)
     }
 
-    pub fn set_permission(&self, path: &str, permission: u32) -> PyHdfsResult<()> {
-        Ok(self
-            .rt
-            .block_on(self.inner.set_permission(path, permission))?)
+    pub fn set_permission(&self, path: &str, permission: u32, py: Python) -> PyHdfsResult<()> {
+        Ok(py.allow_threads(|| {
+            self.rt
+                .block_on(self.inner.set_permission(path, permission))
+        })?)
     }
 
-    pub fn set_replication(&self, path: &str, replication: u32) -> PyHdfsResult<bool> {
-        Ok(self
-            .rt
-            .block_on(self.inner.set_replication(path, replication))?)
+    pub fn set_replication(&self, path: &str, replication: u32, py: Python) -> PyHdfsResult<bool> {
+        Ok(py.allow_threads(|| {
+            self.rt
+                .block_on(self.inner.set_replication(path, replication))
+        })?)
     }
 
-    pub fn get_content_summary(&self, path: &str) -> PyHdfsResult<PyContentSummary> {
-        Ok(self
-            .rt
-            .block_on(self.inner.get_content_summary(path))?
+    pub fn get_content_summary(&self, path: &str, py: Python) -> PyHdfsResult<PyContentSummary> {
+        Ok(py
+            .allow_threads(|| self.rt.block_on(self.inner.get_content_summary(path)))?
             .into())
     }
 }

python/tests/test_benchmark.py

@@ -0,0 +1,21 @@
+from concurrent.futures import ThreadPoolExecutor
+
+from pytest_benchmark.fixture import BenchmarkFixture
+
+from hdfs_native import Client
+
+
+def do_work(client: Client):
+    def delete(path: str):
+        client.delete(path)
+
+    with ThreadPoolExecutor(100) as executor:
+        for i in range(100):
+            executor.submit(delete, f"/bench{i}")
+
+
+def test_threading(client: Client, benchmark: BenchmarkFixture):
+    for i in range(100):
+        client.create(f"/bench{i}").close()
+
+    benchmark(do_work, client)