Add disk timeout configuration option (#6826)

* Add disk timeout setting
* Add tests for timeout
  - test_timeout_value which checks that every timeout call is done
     with the configuration value
  - test_timeout_warning which checks that a warning is raised when
     the timeout raises an exception

Author: mx-psi

disk/assets/configuration/spec.yaml

@@ -159,3 +159,9 @@ files:
             value:
               example: 0
               type: number
+          - name: timeout
+            description: Timeout of the disk query in seconds
+            value:
+              example: 5
+              default: 5
+              type: integer

disk/datadog_checks/disk/data/conf.yaml.example

@@ -142,3 +142,8 @@ instances:
     ## Exclude devices with a total disk size less than a minimum value (in MiB)
     #
     # min_disk_size: 0
+
+    ## @param timeout - integer - optional - default: 5
+    ## Timeout of the disk query in seconds
+    #
+    # timeout: 5

disk/datadog_checks/disk/disk.py

@@ -58,7 +58,7 @@ def __init__(self, name, init_config, instances):
         self._service_check_rw = is_affirmative(instance.get('service_check_rw', False))
         self._min_disk_size = instance.get('min_disk_size', 0) * 1024 * 1024
         self._blkid_cache_file = instance.get('blkid_cache_file')
-
+        self._timeout = instance.get('timeout', 5)
         self._compile_pattern_filters(instance)
         self._compile_tag_re()
         self._blkid_label_re = re.compile('LABEL=\"(.*?)\"', re.I)
@@ -78,10 +78,13 @@ def check(self, instance):
 
             # Get disk metrics here to be able to exclude on total usage
             try:
-                disk_usage = timeout(5)(psutil.disk_usage)(part.mountpoint)
+                disk_usage = timeout(self._timeout)(psutil.disk_usage)(part.mountpoint)
             except TimeoutException:
                 self.log.warning(
-                    u'Timeout while retrieving the disk usage of `%s` mountpoint. Skipping...', part.mountpoint
+                    u'Timeout after %d seconds while retrieving the disk usage of `%s` mountpoint. '
+                    u'You might want to change the timeout length in the settings.',
+                    self._timeout,
+                    part.mountpoint,
                 )
                 continue
             except Exception as e:
@@ -232,9 +235,14 @@ def _collect_inodes_metrics(self, mountpoint):
         metrics = {}
         # we need to timeout this, too.
         try:
-            inodes = timeout(5)(os.statvfs)(mountpoint)
+            inodes = timeout(self._timeout)(os.statvfs)(mountpoint)
         except TimeoutException:
-            self.log.warning(u'Timeout while retrieving the disk usage of `%s` mountpoint. Skipping...', mountpoint)
+            self.log.warning(
+                u'Timeout after %d seconds while retrieving the disk usage of `%s` mountpoint. '
+                u'You might want to change the timeout length in the settings.',
+                self._timeout,
+                mountpoint,
+            )
             return metrics
         except Exception as e:
             self.log.warning('Unable to get disk metrics for %s: %s', mountpoint, e)

disk/tests/test_unit.py

@@ -9,10 +9,11 @@
 from six import iteritems
 
 from datadog_checks.base.utils.platform import Platform
+from datadog_checks.base.utils.timeout import TimeoutException
 from datadog_checks.disk import Disk
 
 from .common import DEFAULT_DEVICE_BASE_NAME, DEFAULT_DEVICE_NAME, DEFAULT_FILE_SYSTEM, DEFAULT_MOUNT_POINT
-from .mocks import MockDiskMetrics, mock_blkid_output
+from .mocks import MockDiskMetrics, MockPart, mock_blkid_output
 
 
 def test_default_options():
@@ -30,6 +31,7 @@ def test_default_options():
     assert check._device_tag_re == []
     assert check._service_check_rw is False
     assert check._min_disk_size == 0
+    assert check._timeout == 5
 
 
 def test_bad_config():
@@ -211,3 +213,61 @@ def test_blkid_cache_file_contains_no_labels(
     c.check(instance_blkid_cache_file_no_label)
     for metric in chain(gauge_metrics, rate_metrics):
         aggregator.assert_metric(metric, tags=['device:/dev/sda1', 'device_name:sda1'])
+
+
+@pytest.mark.usefixtures('psutil_mocks')
+def test_timeout_config(aggregator, gauge_metrics, rate_metrics):
+    """Test timeout configuration value is used on every timeout on the check."""
+
+    # Arbitrary value
+    TIMEOUT_VALUE = 42
+    instance = {'timeout': TIMEOUT_VALUE}
+    c = Disk('disk', {}, [instance])
+
+    # Mock timeout version
+    def no_timeout(fun):
+        return lambda *args: fun(args)
+
+    with mock.patch('psutil.disk_partitions', return_value=[MockPart()]), mock.patch(
+        'datadog_checks.disk.disk.timeout', return_value=no_timeout
+    ) as mock_timeout:
+        c.check(instance)
+
+    mock_timeout.assert_called_with(TIMEOUT_VALUE)
+
+
+@pytest.mark.usefixtures('psutil_mocks')
+def test_timeout_warning(aggregator, gauge_metrics, rate_metrics):
+    """Test a warning is raised when there is a Timeout exception."""
+
+    # Raise exception for "/faulty" mountpoint
+    def faulty_timeout(fun):
+        def f(mountpoint):
+            if mountpoint == "/faulty":
+                raise TimeoutException
+            else:
+                return fun(mountpoint)
+
+        return f
+
+    c = Disk('disk', {}, [{}])
+    c.log = mock.MagicMock()
+    m = MockDiskMetrics()
+    m.total = 0
+
+    with mock.patch('psutil.disk_partitions', return_value=[MockPart(), MockPart(mountpoint="/faulty")]), mock.patch(
+        'psutil.disk_usage', return_value=m, __name__='disk_usage'
+    ), mock.patch('datadog_checks.disk.disk.timeout', return_value=faulty_timeout):
+        c.check({})
+
+    # Check that the warning is called once for the faulty disk
+    c.log.warning.assert_called_once()
+
+    for name in gauge_metrics:
+        aggregator.assert_metric(name, count=0)
+
+    for name in rate_metrics:
+        aggregator.assert_metric_has_tag(name, 'device:{}'.format(DEFAULT_DEVICE_NAME))
+        aggregator.assert_metric_has_tag(name, 'device_name:{}'.format(DEFAULT_DEVICE_BASE_NAME))
+
+    aggregator.assert_all_metrics_covered()