storaged-project · tbzatek · Nov 30, 2023
diff --git a/doc/udisks2-sections.txt.daemon.sections.in b/doc/udisks2-sections.txt.daemon.sections.in
@@ -143,6 +143,7 @@ udisks_linux_drive_object_should_include_device
 UDisksLinuxDrive
 udisks_linux_drive_new
 udisks_linux_drive_update
+udisks_linux_drive_recalculate_nvme_size
 <SUBSECTION Standard>
 UDISKS_LINUX_DRIVE
 UDISKS_IS_LINUX_DRIVE
@@ -283,6 +284,7 @@ udisks_linux_device_read_sysfs_attr_as_int
 udisks_linux_device_read_sysfs_attr_as_uint64
 udisks_linux_device_subsystem_is_nvme
 udisks_linux_device_nvme_is_fabrics
+udisks_linux_device_nvme_tnvmcap_supported
 <SUBSECTION Standard>
 UDISKS_TYPE_LINUX_DEVICE
 UDISKS_LINUX_DEVICE

diff --git a/src/tests/dbus-tests/test_nvme.py b/src/tests/dbus-tests/test_nvme.py
@@ -267,7 +267,7 @@ def test_controller_info(self):
         id = self.get_property_raw(drive_obj, '.Drive', 'Id')
         self.assertTrue(id.startswith('Linux-'))
         size = self.get_property_raw(drive_obj, '.Drive', 'Size')
-        self.assertEqual(size, 0)
+        self.assertEqual(size, self.NS_SIZE * self.NUM_NS)
 
         ctrl_id = self.get_property_raw(drive_obj, '.NVMe.Controller', 'ControllerID')
         self.assertGreater(ctrl_id, 0)

diff --git a/src/udiskslinuxblock.c b/src/udiskslinuxblock.c
@@ -49,6 +49,7 @@
 #include "udiskslinuxblock.h"
 #include "udiskslinuxblockobject.h"
 #include "udiskslinuxdriveobject.h"
+#include "udiskslinuxdrive.h"
 #include "udisksdaemon.h"
 #include "udisksstate.h"
 #include "udisksprivate.h"
@@ -220,20 +221,55 @@ find_block_device_by_sysfs_path (GDBusObjectManagerServer *object_manager,
 
 /* ---------------------------------------------------------------------------------------------------- */
 
+typedef struct
+{
+  UDisksDaemon *daemon;
+  gchar *obj_path;
+} PingDriveData;
+
+static void
+free_ping_drive_data (gpointer user_data)
+{
+  PingDriveData *data = user_data;
+
+  g_free (data->obj_path);
+  g_free (data);
+}
+
+static gboolean
+ping_drive_idle_cb (gpointer user_data)
+{
+  PingDriveData *data = user_data;
+  UDisksObject *object;
+  UDisksDrive *drive;
+
+  object = udisks_daemon_find_object (data->daemon, data->obj_path);
+  drive = object ? udisks_object_get_drive (object) : NULL;
+  if (object && drive)
+    {
+      udisks_linux_drive_recalculate_nvme_size (UDISKS_LINUX_DRIVE (drive),
+                                                UDISKS_LINUX_DRIVE_OBJECT (object));
+    }
+  g_clear_object (&object);
+  g_clear_object (&drive);
+
+  return G_SOURCE_REMOVE;
+}
+
 static gchar *
-find_drive (GDBusObjectManagerServer  *object_manager,
-            GUdevDevice               *block_device,
-            UDisksDrive              **out_drive)
+find_drive (UDisksDaemon  *daemon,
+            GUdevDevice   *block_device,
+            gboolean       update_size,
+            UDisksDrive  **out_drive)
 {
+  GDBusObjectManagerServer *object_manager;
   GUdevDevice *whole_disk_block_device;
   const gchar *whole_disk_block_device_sysfs_path;
   gchar **nvme_ctrls = NULL;
-  gchar *ret;
+  gchar *ret = NULL;
   GList *objects = NULL;
   GList *l;
 
-  ret = NULL;
-
   if (g_strcmp0 (g_udev_device_get_devtype (block_device), "disk") == 0)
     whole_disk_block_device = g_object_ref (block_device);
   else
@@ -267,6 +303,7 @@ find_drive (GDBusObjectManagerServer  *object_manager,
       g_clear_object (&parent_device);
     }
 
+  object_manager = udisks_daemon_get_object_manager (daemon);
   objects = g_dbus_object_manager_get_objects (G_DBUS_OBJECT_MANAGER (object_manager));
   for (l = objects; l != NULL; l = l->next)
     {
@@ -283,19 +320,48 @@ find_drive (GDBusObjectManagerServer  *object_manager,
           UDisksLinuxDevice *drive_device = UDISKS_LINUX_DEVICE (j->data);
           const gchar *drive_sysfs_path;
 
+          /* See if the drive object encloses our block device.
+           * For NVMe, see if the drive object representing a NVMe controller
+           * provides our namespace.
+           */
           drive_sysfs_path = g_udev_device_get_sysfs_path (drive_device->udev_device);
           if (g_strcmp0 (whole_disk_block_device_sysfs_path, drive_sysfs_path) == 0 ||
               (nvme_ctrls && g_strv_contains ((const gchar * const *) nvme_ctrls, drive_sysfs_path)))
             {
-              if (out_drive != NULL)
-                *out_drive = udisks_object_get_drive (UDISKS_OBJECT (object));
-              ret = g_strdup (g_dbus_object_get_object_path (G_DBUS_OBJECT (object)));
-              g_list_free_full (drive_devices, g_object_unref);
-              /* FIXME: NVMe namespace may be provided by multiple controllers within
-               *  a NVMe subsystem, however the org.freedesktop.UDisks2.Block.Drive
-               *  property may only contain single object path.
+              const gchar *obj_path;
+
+              /* FIXME: An NVMe namespace may be provided by multiple controllers within
+               *        an NVMe subsystem, however the org.freedesktop.UDisks2.Block.Drive
+               *        property may only contain a single object path.
                */
-              goto out;
+              if (out_drive != NULL && *out_drive == NULL)
+                *out_drive = udisks_object_get_drive (UDISKS_OBJECT (object));
+              obj_path = g_dbus_object_get_object_path (G_DBUS_OBJECT (object));
+              if (! ret)
+                ret = g_strdup (obj_path);
+              if (!nvme_ctrls || !update_size)
+                {
+                  g_list_free_full (drive_devices, g_object_unref);
+                  goto out;
+                }
+              else
+                {
+                  if (!udisks_linux_device_nvme_tnvmcap_supported (drive_device))
+                    {
+                      PingDriveData *data;
+
+                      /* ping the drive object to recalculate controller size
+                       * from all attached namespaces
+                       */
+                      data = g_new0 (PingDriveData, 1);
+                      data->daemon = daemon;
+                      data->obj_path = g_strdup (obj_path);
+                      g_idle_add_full (G_PRIORITY_DEFAULT_IDLE,
+                                       ping_drive_idle_cb,
+                                       data,
+                                       free_ping_drive_data);
+                    }
+                }
             }
         }
       g_list_free_full (drive_devices, g_object_unref);
@@ -1135,7 +1201,7 @@ udisks_linux_block_update (UDisksLinuxBlock       *block,
    * TODO: if this is slow we could have a cache or ensure that we
    * only do this once or something else
    */
-  drive_object_path = find_drive (object_manager, device->udev_device, &drive);
+  drive_object_path = find_drive (daemon, device->udev_device, TRUE, &drive);
   if (drive_object_path != NULL)
     {
       udisks_block_set_drive (iface, drive_object_path);
@@ -1978,7 +2044,7 @@ update_block_fstab (UDisksDaemon           *daemon,
 
   /* hints take fstab records in the calculation */
   device = udisks_linux_block_object_get_device (object);
-  drive_object_path = find_drive (udisks_daemon_get_object_manager (daemon), device->udev_device, &drive);
+  drive_object_path = find_drive (daemon, device->udev_device, FALSE, &drive);
   update_hints (daemon, block, device, drive);
   g_free (drive_object_path);
   g_clear_object (&device);

diff --git a/src/udiskslinuxdevice.c b/src/udiskslinuxdevice.c
@@ -463,3 +463,24 @@ udisks_linux_device_nvme_is_fabrics (UDisksLinuxDevice *device)
 
   return FALSE;
 }
+
+/**
+ * udisks_linux_device_nvme_tnvmcap_supported:
+ * @device: A #UDisksLinuxDevice.
+ *
+ * Determines whether @device supports Capacity information
+ * in the Identify Controller data structure.
+ *
+ * Returns: %TRUE if capacity reporting is supported, %FALSE otherwise.
+ */
+gboolean
+udisks_linux_device_nvme_tnvmcap_supported (UDisksLinuxDevice *device)
+{
+  if (device->nvme_ctrl_info == NULL)
+    return FALSE;
+
+  /* FIXME: find a more reliable way to detect controller
+   *        capacity reporting capability.
+   */
+  return device->nvme_ctrl_info->size_total > 0;
+}
diff --git a/src/udiskslinuxdevice.h b/src/udiskslinuxdevice.h
@@ -74,6 +74,7 @@ guint64            udisks_linux_device_read_sysfs_attr_as_uint64 (UDisksLinuxDev
 
 gboolean           udisks_linux_device_subsystem_is_nvme         (UDisksLinuxDevice  *device);
 gboolean           udisks_linux_device_nvme_is_fabrics           (UDisksLinuxDevice  *device);
+gboolean           udisks_linux_device_nvme_tnvmcap_supported    (UDisksLinuxDevice  *device);
 
 G_END_DECLS
 

diff --git a/src/udiskslinuxdrive.c b/src/udiskslinuxdrive.c
@@ -1016,6 +1016,67 @@ udisks_linux_drive_update (UDisksLinuxDrive       *drive,
 
 /* ---------------------------------------------------------------------------------------------------- */
 
+/**
+ * udisks_linux_drive_recalculate_nvme_size:
+ * @drive: A #UDisksLinuxDrive.
+ * @object: The enclosing #UDisksLinuxDriveObject instance.
+ *
+ * Find all block objects pointing to this drive, calculate
+ * NVMe namespace capacity numbers and update this interface.
+ */
+void
+udisks_linux_drive_recalculate_nvme_size (UDisksLinuxDrive       *drive,
+                                          UDisksLinuxDriveObject *object)
+{
+  UDisksDaemon *daemon;
+  GDBusObjectManagerServer *object_manager;
+  GList *objects = NULL;
+  GList *l;
+  const gchar *obj_path;
+  guint64 size_total = 0;
+
+  daemon = udisks_linux_drive_object_get_daemon (object);
+  object_manager = udisks_daemon_get_object_manager (daemon);
+  obj_path = g_dbus_object_get_object_path (G_DBUS_OBJECT (object));
+
+  objects = g_dbus_object_manager_get_objects (G_DBUS_OBJECT_MANAGER (object_manager));
+  for (l = objects; l != NULL; l = l->next)
+    {
+      UDisksObject *o = l->data;
+      UDisksBlock *block;
+      UDisksLinuxDevice *device;
+
+      if (!UDISKS_IS_LINUX_BLOCK_OBJECT (o))
+        continue;
+
+      block = udisks_object_get_block (o);
+      if (!block)
+          continue;
+
+      if (g_strcmp0 (udisks_block_get_drive (block), obj_path) != 0)
+        {
+          g_object_unref (block);
+          continue;
+        }
+
+      device = udisks_linux_block_object_get_device (UDISKS_LINUX_BLOCK_OBJECT (o));
+      if (device && device->nvme_ns_info &&
+          device->nvme_ns_info->current_lba_format.data_size > 0)
+        {
+          /* Namespace Size >= Namespace Capacity >= Namespace Utilization */
+          size_total += (guint64) device->nvme_ns_info->nsize *
+                        (guint64) device->nvme_ns_info->current_lba_format.data_size;
+        }
+      g_clear_object (&device);
+      g_object_unref (block);
+    }
+  g_list_free_full (objects, g_object_unref);
+
+  udisks_drive_set_size (UDISKS_DRIVE (drive), size_total);
+}
+
+/* ---------------------------------------------------------------------------------------------------- */
+
 static gboolean
 handle_eject (UDisksDrive           *_drive,
               GDBusMethodInvocation *invocation,

diff --git a/src/udiskslinuxdrive.h b/src/udiskslinuxdrive.h
@@ -34,6 +34,9 @@ UDisksDrive *udisks_linux_drive_new      (void);
 gboolean     udisks_linux_drive_update   (UDisksLinuxDrive       *drive,
                                           UDisksLinuxDriveObject *object);
 
+void         udisks_linux_drive_recalculate_nvme_size (UDisksLinuxDrive       *drive,
+                                                       UDisksLinuxDriveObject *object);
+
 G_END_DECLS
 
 #endif /* __UDISKS_LINUX_DRIVE_H__ */