From 45d1ab1aa30fca41cf4013e5d3948ea4752767e5 Mon Sep 17 00:00:00 2001
From: Mehran Kholdi <semekh@hamravesh.com>
Date: Fri, 2 Jul 2021 20:31:01 +0430
Subject: [PATCH] Refuse to create/resize volumes in case of insufficient disk
 space

---
 README.md           |  1 -
 consts.py           |  1 +
 orchestrator/k8s.py |  6 +++++-
 rawfile_servicer.py | 34 ++++++++++++++++++++++++++--------
 remote.py           | 12 +++++++++++-
 5 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 352cea6..a9193dd 100644
--- a/README.md
+++ b/README.md
@@ -35,7 +35,6 @@ Features
 - [x] Direct I/O: Near-zero disk performance overhead
 - [x] Dynamic provisioning
 - [x] Enforced volume size limit
-- [x] Thin provisioned
 - [x] Access Modes
     - [x] ReadWriteOnce
     - ~~ReadOnlyMany~~
diff --git a/consts.py b/consts.py
index 9fd2164..7318145 100644
--- a/consts.py
+++ b/consts.py
@@ -4,3 +4,4 @@ PROVISIONER_NAME = os.getenv("PROVISIONER_NAME", "rawfile.csi.openebs.io")
 PROVISIONER_VERSION = "0.5.0"
 DATA_DIR = "/data"
 CONFIG = {}
+RESOURCE_EXHAUSTED_EXIT_CODE = 101
diff --git a/orchestrator/k8s.py b/orchestrator/k8s.py
index 13c3e88..f45057e 100644
--- a/orchestrator/k8s.py
+++ b/orchestrator/k8s.py
@@ -1,6 +1,7 @@
 import json
 import uuid
 from pathlib import Path
+from subprocess import CalledProcessError
 from time import sleep
 
 import pykube
@@ -66,5 +67,8 @@ def run_on_node(fn, node):
 
     wait_for(is_finished, "task to finish")
     if task_pod.obj["status"]["phase"] != "Succeeded":
-        raise Exception(f"Task {name} failed")
+        exit_code = task_pod.obj["status"]["containerStatuses"][0]["state"][
+            "terminated"
+        ]["exitCode"]
+        raise CalledProcessError(returncode=exit_code, cmd=f"Task: {name}")
     task_pod.delete()
diff --git a/rawfile_servicer.py b/rawfile_servicer.py
index 4fe44a4..614fc70 100644
--- a/rawfile_servicer.py
+++ b/rawfile_servicer.py
@@ -1,10 +1,11 @@
 from pathlib import Path
+from subprocess import CalledProcessError
 
 import grpc
 from google.protobuf.wrappers_pb2 import BoolValue
 
 import rawfile_util
-from consts import PROVISIONER_VERSION, PROVISIONER_NAME
+from consts import PROVISIONER_VERSION, PROVISIONER_NAME, RESOURCE_EXHAUSTED_EXIT_CODE
 from csi import csi_pb2, csi_pb2_grpc
 from declarative import be_symlink, be_absent
 from fs_util import device_stats, mountpoint_to_dev
@@ -184,10 +185,18 @@ class RawFileControllerServicer(csi_pb2_grpc.ControllerServicer):
                 grpc.StatusCode.INVALID_ARGUMENT, "Topology key not found... why?"
             )
 
-        run_on_node(
-            init_rawfile.as_cmd(volume_id=request.name, size=size),
-            node=node_name,
-        )
+        try:
+            run_on_node(
+                init_rawfile.as_cmd(volume_id=request.name, size=size),
+                node=node_name,
+            )
+        except CalledProcessError as exc:
+            if exc.returncode == RESOURCE_EXHAUSTED_EXIT_CODE:
+                context.abort(
+                    grpc.StatusCode.RESOURCE_EXHAUSTED, "Not enough disk space"
+                )
+            else:
+                raise exc
 
         return csi_pb2.CreateVolumeResponse(
             volume=csi_pb2.Volume(
@@ -210,9 +219,18 @@ class RawFileControllerServicer(csi_pb2_grpc.ControllerServicer):
         volume_id = request.volume_id
         node_name = volume_to_node(volume_id)
         size = request.capacity_range.required_bytes
-        run_on_node(
-            expand_rawfile.as_cmd(volume_id=volume_id, size=size), node=node_name
-        )
+
+        try:
+            run_on_node(
+                expand_rawfile.as_cmd(volume_id=volume_id, size=size), node=node_name
+            )
+        except CalledProcessError as exc:
+            if exc.returncode == RESOURCE_EXHAUSTED_EXIT_CODE:
+                context.abort(
+                    grpc.StatusCode.RESOURCE_EXHAUSTED, "Not enough disk space"
+                )
+            else:
+                raise exc
 
         return csi_pb2.ControllerExpandVolumeResponse(
             capacity_bytes=size,
diff --git a/remote.py b/remote.py
index ab00e01..72ab899 100644
--- a/remote.py
+++ b/remote.py
@@ -25,6 +25,10 @@ def init_rawfile(volume_id, size):
     from pathlib import Path
 
     from util import run
+    from consts import RESOURCE_EXHAUSTED_EXIT_CODE
+
+    if rawfile_util.get_capacity() < size:
+        exit(RESOURCE_EXHAUSTED_EXIT_CODE)
 
     img_dir = rawfile_util.img_dir(volume_id)
     img_dir.mkdir(exist_ok=True)
@@ -47,11 +51,17 @@ def init_rawfile(volume_id, size):
 @remote_fn
 def expand_rawfile(volume_id, size):
     import rawfile_util
+
     from util import run
+    from consts import RESOURCE_EXHAUSTED_EXIT_CODE
 
     img_file = rawfile_util.img_file(volume_id)
-    if rawfile_util.metadata(volume_id)["size"] >= size:
+    size_inc = size - rawfile_util.metadata(volume_id)["size"]
+    if size_inc <= 0:
         return
+    if rawfile_util.get_capacity() < size_inc:
+        exit(RESOURCE_EXHAUSTED_EXIT_CODE)
+
     rawfile_util.patch_metadata(
         volume_id,
         {"size": size},