From 8a40fbd4180635e91fc8e79d4392b9b4cb7b4488 Mon Sep 17 00:00:00 2001
From: cina_pm <cinapm375@gmail.com>
Date: Thu, 13 Jun 2024 13:55:35 +0330
Subject: [PATCH] feat(zfs-localpv): add option for choosing between refquota
 and quota

Signed-off-by: cina_pm <cinapm375@gmail.com>
---
 cmd/main.go          | 15 ++++++++++++---
 pkg/config/config.go |  7 +++++++
 pkg/zfs/zfs_util.go  |  9 +++++----
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index f5bb5652d..2b7b40eb2 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -22,7 +22,8 @@ import (
 	"log"
 	"os"
 
-	config "github.com/openebs/zfs-localpv/pkg/config"
+	configs "github.com/openebs/zfs-localpv/pkg/config"
+
 	"github.com/openebs/zfs-localpv/pkg/driver"
 	"github.com/openebs/zfs-localpv/pkg/version"
 	zfs "github.com/openebs/zfs-localpv/pkg/zfs"
@@ -40,7 +41,7 @@ import (
  */
 func main() {
 	_ = flag.CommandLine.Parse([]string{})
-	var config = config.Default()
+	var config = configs.Default()
 
 	cmd := &cobra.Command{
 		Use:   "zfs-driver",
@@ -74,6 +75,10 @@ func main() {
 		&config.PluginType, "plugin", "csi-plugin", "Type of this driver i.e. controller or node",
 	)
 
+	cmd.PersistentFlags().StringVar(
+		&configs.QuotaType, "quota-type", "quota", "quota type: refquota or quota",
+	)
+
 	err := cmd.Execute()
 	if err != nil {
 		_, _ = fmt.Fprintf(os.Stderr, "%s", err.Error())
@@ -81,11 +86,15 @@ func main() {
 	}
 }
 
-func run(config *config.Config) {
+func run(config *configs.Config) {
 	if config.Version == "" {
 		config.Version = version.Current()
 	}
 
+	if configs.QuotaType != configs.Quota && configs.QuotaType != configs.RefQuota {
+		log.Fatalln(fmt.Errorf("quota-type should be quota or refquota"))
+	}
+
 	klog.Infof("ZFS Driver Version :- %s - commit :- %s", version.Current(), version.GetGitCommit())
 	klog.Infof(
 		"DriverName: %s Plugin: %s EndPoint: %s Node Name: %s",
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 97788cd93..dc58f7cc5 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -43,6 +43,13 @@ type Config struct {
 	Nodename string
 }
 
+const (
+	Quota    = "quota"    // This option sets a limit on the amount of disk space a dataset can use
+	RefQuota = "refquota" // This option sets a limit on the amount of disk space a dataset and all its snapshots can use
+)
+
+var QuotaType string
+
 // Default returns a new instance of config
 // required to initialize a driver instance
 func Default() *Config {
diff --git a/pkg/zfs/zfs_util.go b/pkg/zfs/zfs_util.go
index 2c9b331e2..de6404b97 100644
--- a/pkg/zfs/zfs_util.go
+++ b/pkg/zfs/zfs_util.go
@@ -18,6 +18,7 @@ package zfs
 
 import (
 	"bufio"
+	configs "github.com/openebs/zfs-localpv/pkg/config"
 	"os/exec"
 	"path/filepath"
 	"strconv"
@@ -143,7 +144,7 @@ func buildCloneCreateArgs(vol *apis.ZFSVolume) []string {
 
 	if vol.Spec.VolumeType == VolTypeDataset {
 		if len(vol.Spec.Capacity) != 0 {
-			quotaProperty := "quota=" + vol.Spec.Capacity
+			quotaProperty := configs.QuotaType + "=" + vol.Spec.Capacity
 			ZFSVolArg = append(ZFSVolArg, "-o", quotaProperty)
 		}
 		if len(vol.Spec.RecordSize) != 0 {
@@ -216,7 +217,7 @@ func buildDatasetCreateArgs(vol *apis.ZFSVolume) []string {
 	ZFSVolArg = append(ZFSVolArg, ZFSCreateArg)
 
 	if len(vol.Spec.Capacity) != 0 {
-		quotaProperty := "quota=" + vol.Spec.Capacity
+		quotaProperty := configs.QuotaType + "=" + vol.Spec.Capacity
 		ZFSVolArg = append(ZFSVolArg, "-o", quotaProperty)
 	}
 	if len(vol.Spec.RecordSize) != 0 {
@@ -292,7 +293,7 @@ func buildVolumeResizeArgs(vol *apis.ZFSVolume) []string {
 	ZFSVolArg = append(ZFSVolArg, ZFSSetArg)
 
 	if vol.Spec.VolumeType == VolTypeDataset {
-		quotaProperty := "quota=" + vol.Spec.Capacity
+		quotaProperty := configs.QuotaType + "=" + vol.Spec.Capacity
 		ZFSVolArg = append(ZFSVolArg, quotaProperty)
 	} else {
 		volsizeProperty := "volsize=" + vol.Spec.Capacity
@@ -350,7 +351,7 @@ func buildVolumeRestoreArgs(rstr *apis.ZFSRestore) ([]string, error) {
 
 	if rstr.VolSpec.VolumeType == VolTypeDataset {
 		if len(rstr.VolSpec.Capacity) != 0 {
-			ZFSRecvParam += " -o quota=" + rstr.VolSpec.Capacity
+			ZFSRecvParam += " -o " + configs.QuotaType + "=" + rstr.VolSpec.Capacity
 		}
 		if len(rstr.VolSpec.RecordSize) != 0 {
 			ZFSRecvParam += " -o recordsize=" + rstr.VolSpec.RecordSize