🚧 Remove use of static DATA_DIR

Signed-off-by: Muhammed Hussain Karimi <[email protected]>

Author: mhkarimi1383

rawfile/consts.py

@@ -5,7 +5,6 @@
 PROVISIONER_VERSION = os.getenv("PROVISIONER_VERSION") or importlib.metadata.version(
     "rawfile"
 )
-DATA_DIR = "/data"
 D_PERMS = 0o700
 F_PERMS = 0o600
 OWNER_UMASK = 0o077
@@ -15,4 +14,3 @@
 CSI_K8S_PVC_NAME_KEY = "csi.storage.k8s.io/pvc/name"
 GA_ID = "Ry1UWkdQNDY2MThX"
 GA_KEY = "OTFKR2RUZzlRd0duN1ktdnZ1TTR6QQ=="
-COW_SUPPORTED: bool | None = None

rawfile/internal_svc.py

@@ -57,8 +57,8 @@ def ExpandRawFile(
                     is_attached=volume_manager.is_attached(request.volume_id),
                     status=internal_pb2.ExpandRawFileStatus.OK,
                 )
-
-            if get_capacity() < size_inc:
+            meta = metadata(request.volume_id)
+            if get_capacity(meta["storage_pool"]) < request.new_size:
                 return internal_pb2.ExpandRawFileResponse(
                     is_attached=volume_manager.is_attached(request.volume_id),
                     status=internal_pb2.ExpandRawFileStatus.RESOURCE_EXHAUSTED,

rawfile/metrics.py

@@ -3,6 +3,14 @@
 from prometheus_client.metrics_core import GaugeMetricFamily
 from utils.storage_pool import get_capacity
 from utils.volume_manager import manager as volume_manager
+from config import config
+
+
+def get_remaining_capacity():
+    val = 0
+    for pool in config.csi_driver.storage_pools.values():
+        val += get_capacity(pool.path.as_posix())
+    return val
 
 
 class VolumeStatsCollector(object):
@@ -28,7 +36,7 @@ def collect(self):
             labels=["node", "volume"],
             unit="bytes",
         )
-        remaining_capacity.add_metric([self.node], get_capacity())
+        remaining_capacity.add_metric([self.node], get_remaining_capacity())
         for volume_id, stats in volume_manager.get_all_volumes_stats().items():
             volume_used.add_metric([self.node, volume_id], stats["used"])
             volume_total.add_metric([self.node, volume_id], stats["total"])

rawfile/rawfile.py

@@ -12,7 +12,6 @@
 from internal import internal_pb2_grpc
 from metrics import expose_metrics
 from utils import task_manager
-from utils.rawfile import is_cow_supported
 from utils.logs import init as init_logging, logger
 from internal_svc import InternalServicer, SignatureInterceptor
 import consts
@@ -39,17 +38,16 @@ def node_driver_preflight_checks(task_manager: task_manager.TaskManager):
         raise RuntimeError("CSI Driver configuration is missing")
     if not config.csi_driver.metadata_dir:
         raise RuntimeError("Metadata directory is not set for node plugin")
-    data_dir = Path(consts.DATA_DIR)
-    dirs = (
+    dirs = [
         config.csi_driver.metadata_dir,
-        data_dir,
-    )
+    ]
+    dirs.extend([pool.path for pool in config.csi_driver.storage_pools.values()])
     for dir in dirs:
         __create_and_check_directory(dir)
     volume_manager.migrate_metadata_dir()
     volume_manager.migrate_all_volume_schemas()
     task_manager.migrate_tasks_file_path()
-    consts.COW_SUPPORTED = is_cow_supported(data_dir)
+    # consts.COW_SUPPORTED = is_cow_supported(data_dir)
 
 
 def csi_driver(driver_config: CSIDriverCmd):

rawfile/rawfile_servicer.py

@@ -192,6 +192,7 @@ def CreateVolume(self, request: csi_pb2.CreateVolumeRequest, context):
         if copy_on_write_param is not None:
             copy_on_write = str_to_bool(copy_on_write_param)
         freezefs = str_to_bool(params.get("freezefs", "no"))
+        storage_pool = params.get("storagepool", config.csi_driver.default_pool)
         source_type = None
         source_id = None
         node_name = None
@@ -240,6 +241,7 @@ def CreateVolume(self, request: csi_pb2.CreateVolumeRequest, context):
                 TaskName.CREATE_VOLUME,
                 request.name,
                 size,
+                storage_pool,
                 thin_provision,
                 freezefs,
                 copy_on_write,

rawfile/utils/rawfile.py

@@ -51,9 +51,14 @@ def img_file(volume_id):
 def snapshots_dir(volume_id: str, temporary: bool = False):
     from utils.volume_manager import manager as volume_manager
 
+    meta = metadata(volume_id)
     if temporary:
-        return Path(f"{volume_manager._get_volume_path(volume_id)}/snapshots/temp")
-    return Path(f"{volume_manager._get_volume_path(volume_id)}/snapshots")
+        return Path(
+            f"{volume_manager._get_volume_path(meta['storage_pool'], volume_id)}/snapshots/temp"
+        )
+    return Path(
+        f"{volume_manager._get_volume_path(meta['storage_pool'], volume_id)}/snapshots"
+    )
 
 
 def img_size(volume_id) -> int:
@@ -94,7 +99,8 @@ def update_permissions(volume_id: str) -> None:
     from utils.volume_manager import manager as volume_manager
     from itertools import chain
 
-    _img_dir = volume_manager._get_volume_path(volume_id)
+    meta = metadata(volume_id)
+    _img_dir = volume_manager._get_volume_path(meta["storage_pool"], volume_id)
     if not _img_dir.exists():
         return
     _img_dir.chmod(D_PERMS)

rawfile/utils/snapshot_manager.py

@@ -2,9 +2,8 @@
 from pathlib import Path
 import time
 from dataclasses import dataclass
-import consts
 from utils.commands import run
-
+from config import config
 from utils.errors import FsFreezeNotSupportedOnBlockVolumes, SnapshotCreateVolumeInUse
 from utils.lock import VolLock
 from utils.rawfile import (
@@ -139,10 +138,21 @@ def list_snapshots(
     ) -> SnapshotList:
         """List available snapshots"""
         subdir = "snapshots/temp" if temporary else "snapshots"
-        pattern = f"{consts.DATA_DIR}/{volume_id if volume_id else '**'}/{subdir}/{snapshot_name if snapshot_name else '*'}.img"
-
+        patterns = []
+        if volume_id:
+            patterns = [
+                f"{metadata(volume_id)['storage_pool']}/{volume_id if volume_id else '**'}/{subdir}/{snapshot_name if snapshot_name else '*'}.img"
+            ]
+        else:
+            for pool in config.csi_driver.storage_pools.values():
+                patterns.append(
+                    f"{pool.path}/**/{subdir}/{snapshot_name if snapshot_name else '*'}.img"
+                )
         snapshots = []
-        snapshot_files = sorted(glob(pattern, recursive=True))
+        snapshot_files = []
+        for pattern in patterns:
+            snapshot_files.extend(sorted(glob(pattern, recursive=True)))
+
         count = 0
         idx = 0
         for idx, snap_filename in enumerate(snapshot_files):

rawfile/utils/storage_pool.py

@@ -1,16 +1,18 @@
 from utils.devices import path_stats
 from utils.volume_manager import manager as volume_manager
 from config import config
-from consts import DATA_DIR
 
 
-def get_capacity():
-    disk_free_size = path_stats(DATA_DIR, config.capacity_override)["fs_avail"]
+def get_capacity(storage_pool: str | None = None):
+    if not storage_pool:
+        storage_pool = config.csi_driver.default_pool
+    pool = config.csi_driver.storage_pools[storage_pool]
+    disk_free_size = path_stats(pool.path, pool.capacity_override)["fs_avail"]
     capacity = disk_free_size
     for volume_stat in volume_manager.get_all_volumes_stats().values():
         capacity -= volume_stat["total"] - volume_stat["used"]
-    if isinstance(config.reserved_capacity, str):
-        capacity -= capacity * int(config.reserved_capacity[:-1]) / 100
+    if isinstance(pool.reserved_capacity, str):
+        capacity -= capacity * int(pool.reserved_capacity[:-1]) / 100
     else:
-        capacity -= config.reserved_capacity.to("B")
+        capacity -= pool.reserved_capacity.to("B")
     return max(capacity, 0)

rawfile/utils/task_manager.py

@@ -3,7 +3,7 @@
 import shutil
 from typing import Callable, TypedDict
 import time
-import consts
+from config import config
 from utils.snapshot_manager import manager as snapshot_manager
 from utils.volume_manager import manager as volume_manager
 from pathlib import Path
@@ -203,7 +203,9 @@ def shutdown(self, timeout: int):
         self._executor.shutdown(wait=False, cancel_futures=True)
 
     def migrate_tasks_file_path(self):
-        src = Path(f"{consts.DATA_DIR}/tasks.json")
+        src = Path(
+            f"{config.csi_driver.storage_pools[config.csi_driver.default_pool].path}/tasks.json"
+        )
         dst = self._tasks_store_path
         if src.exists():
             shutil.move(src, dst)

rawfile/utils/volume_manager.py

@@ -39,14 +39,15 @@ class VolumeStats(TypedDict):
 
 
 class VolumeManager:
-    def _get_volume_path(self, volume_id: str) -> Path:
-        return Path(f"{consts.DATA_DIR}/{volume_id}")
+    def _get_volume_path(self, storage_pool: str, volume_id: str) -> Path:
+        return Path(f"{config.csi_driver.storage_pools[storage_pool].path}/{volume_id}")
 
     def create_volume(
         self,
         volume_id: str,
         size: int,
         thin_provision: bool,
+        storage_pool: str,
         freezefs: bool | None = None,
         copy_on_write: bool | None = None,
         source_type: VolumeSource | None = None,
@@ -56,11 +57,16 @@ def create_volume(
             raise SourceTypeRequired(source_id)
         snapshot_name = None
         source_volume_id = None
-        img_data_dir = self._get_volume_path(volume_id)
+        img_data_dir = self._get_volume_path(storage_pool, volume_id)
         img_data_dir.mkdir(mode=consts.D_PERMS, exist_ok=True)
         meta_dir(volume_id).mkdir(exist_ok=True, parents=True)
         patch_metadata(
-            volume_id, {"ready": False, "schema_version": LATEST_SCHEMA_VERSION}
+            volume_id,
+            {
+                "schema_version": LATEST_SCHEMA_VERSION,
+                "storage_pool": storage_pool,
+                "ready": False,
+            },
         )
         try:
             if source_type == VolumeSource.snapshot:
@@ -185,7 +191,7 @@ def rmdir(path: Path):
             for file in meta_dir(volume_id).glob("*"):
                 file.unlink(missing_ok=True)
             rmdir(meta_dir(volume_id))
-            rmdir(self._get_volume_path(volume_id))
+            rmdir(self._get_volume_path(meta["storage_pool"], volume_id))
 
     def gc_if_needed(self, volume_id, dry_run=True):
         with VolLock(volume_id):
@@ -201,7 +207,8 @@ def gc_if_needed(self, volume_id, dry_run=True):
         return False
 
     def delete_volume(self, volume_id):
-        img_data_dir = self._get_volume_path(volume_id)
+        meta = metadata_or(volume_id)
+        img_data_dir = self._get_volume_path(meta["storage_pool"], volume_id)
         if not img_data_dir.exists():
             return 0
         vol_img_file = img_file(volume_id)
@@ -259,14 +266,18 @@ def migrate_metadata(self, volume_id, target_version):
         return update_metadata(volume_id, new_data)
 
     def migrate_metadata_dir(self):
-        for old_meta in glob(f"{consts.DATA_DIR}/**/disk.meta"):
+        for old_meta in glob(
+            f"{config.csi_driver.storage_pools[config.csi_driver.default_pool].path}/**/disk.meta"
+        ):
             volume_id = basename(dirname(old_meta))
             for f in (
                 "disk.meta",
                 "disk.meta.tmp",
                 "disk.lock",
             ):
-                src = Path(f"{consts.DATA_DIR}/{volume_id}/{f}")
+                src = Path(
+                    f"{config.csi_driver.storage_pools[config.csi_driver.default_pool].path}/{volume_id}/{f}"
+                )
                 if src.exists():
                     dst = meta_dir(volume_id) / f
                     dst.parent.mkdir(parents=True, exist_ok=True)
@@ -278,7 +289,8 @@ def migrate_all_volume_schemas(self):
             self.migrate_metadata(volume_id, target_version)
 
     def is_attached(self, volume_id):
-        vol_img_dir = self._get_volume_path(volume_id)
+        meta = metadata_or(volume_id)
+        vol_img_dir = self._get_volume_path(meta["storage_pool"], volume_id)
         if not vol_img_dir.exists():
             return False