[tor-commits] [stem/master] Adding proc.get_file_descriptors_used

[atagar at torproject.org]

commit 20d7f9d13cbc183d33a129324137b06d2aaeb313
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Sep 1 15:43:39 2014 -0700

    Adding proc.get_file_descriptors_used
    
    Just a little helper for getting the number of file descriptors a process is
    using.
---
 docs/change_log.rst    |    1 +
 stem/util/proc.py      |   38 +++++++++++++++++++++++++++++++++-----
 test/unit/util/proc.py |   31 +++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index 2b6964c..24e39f0 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -54,6 +54,7 @@ The following are only available within Stem's `git repository
   * Added support for directories to :func:`stem.util.conf.Config.load`.
   * Changed :func:`stem.util.conf.uses_settings` to only provide a 'config' keyword arument if the decorated function would accept it.
   * Added :func:`stem.util.str_tools.crop`
+  * Added :func:`stem.util.proc.get_file_descriptors_used`
 
  * **Interpreter**
 
diff --git a/stem/util/proc.py b/stem/util/proc.py
index 62b7a5e..f9b2a59 100644
--- a/stem/util/proc.py
+++ b/stem/util/proc.py
@@ -26,6 +26,7 @@ future, use them at your own risk.**
   get_uid - provides the user id a process is running under
   get_memory_usage - provides the memory usage of a process
   get_stats - queries statistics about a process
+  get_file_descriptors_used - number of file descriptors used by a process
   get_connections - provides the connections made by a process
 
 .. data:: Stat (enum)
@@ -292,6 +293,31 @@ def get_stats(pid, *stat_types):
   return tuple(results)
 
 
+def get_file_descriptors_used(pid):
+  """
+  Provides the number of file descriptors currently being used by a process.
+
+  :param int pid: process id of the process to be queried
+
+  :returns: **int** of the number of file descriptors used
+
+  :raises: **IOError** if it can't be determined
+  """
+
+  try:
+    pid = int(pid)
+
+    if pid < 0:
+      raise IOError("Process pids can't be negative: %s" % pid)
+  except (ValueError, TypeError):
+    raise IOError('Process pid was non-numeric: %s' % pid)
+
+  try:
+    return len(os.listdir('/proc/%i/fd' % pid))
+  except Exception as exc:
+    raise IOError("Unable to check number of file descriptors used: %s" % exc)
+
+
 def get_connections(pid):
   """
   Queries connection related information from the proc contents. This provides
@@ -307,11 +333,13 @@ def get_connections(pid):
   :raises: **IOError** if it can't be determined
   """
 
-  if isinstance(pid, str):
-    try:
-      pid = int(pid)
-    except ValueError:
-      raise IOError('Process pid was non-numeric: %s' % pid)
+  try:
+    pid = int(pid)
+
+    if pid < 0:
+      raise IOError("Process pids can't be negative: %s" % pid)
+  except (ValueError, TypeError):
+    raise IOError('Process pid was non-numeric: %s' % pid)
 
   if pid == 0:
     return []
diff --git a/test/unit/util/proc.py b/test/unit/util/proc.py
index c682853..1786a99 100644
--- a/test/unit/util/proc.py
+++ b/test/unit/util/proc.py
@@ -144,6 +144,37 @@ class TestProc(unittest.TestCase):
       self.assertEquals(response, proc.get_stats(0, *args))
 
   @patch('os.listdir')
+  def test_get_file_descriptors_used(self, listdir_mock):
+    """
+    Tests the get_file_descriptors_used function.
+    """
+
+    # check that we reject bad pids
+
+    for arg in (None, -100, 'hello',):
+      self.assertRaises(IOError, proc.get_file_descriptors_used, arg)
+
+    # when proc directory doesn't exist
+
+    error_msg = "OSError: [Errno 2] No such file or directory: '/proc/2118/fd'"
+    listdir_mock.side_effect = OSError(error_msg)
+
+    try:
+      proc.get_file_descriptors_used(2118)
+      self.fail("We should raise when listdir() fails")
+    except IOError as exc:
+      expected = "Unable to check number of file descriptors used: %s" % error_msg
+      self.assertEqual(expected, str(exc))
+
+    # successful calls
+
+    listdir_mock.return_value = ['0', '1', '2', '3', '4', '5']
+    listdir_mock.side_effect = None
+
+    self.assertEqual(6, proc.get_file_descriptors_used(2118))
+    self.assertEqual(6, proc.get_file_descriptors_used('2118'))
+
+  @patch('os.listdir')
   @patch('os.readlink')
   @patch('stem.util.proc.open', create = True)
   def test_get_connections(self, open_mock, readlink_mock, listdir_mock):