feature: add diskquota support regular expression

Add diskquota support regular expression.

We can use pouch run --disk-quota="key=value" to set quota for
container. This pr can make it support to use regular expression key,
such as: --disk-quota=".=1g", it means all the quota of mount points
are'1g', the priority of key'.'is the lowest. If it has no'key', just
have'value', such as: --disk-quota=1g, it means the quota of rootfs is
1g. If the mountpoint is a volume, the quota of volume is base on
volume's size, the disk-quota is not take effect on a volume.

Signed-off-by: Rudy Zhang <rudyflyzhang@gmail.com>

apis/swagger.yml

@@ -1784,6 +1784,9 @@ definitions:
         type: "object"
         additionalProperties:
           type: "string"
+      QuotaID:
+        type: "string"
+        description: "set disk quota by specified quota id"
 
   ContainerCreateResp:
     description: "response returned by daemon when container create successfully"

daemon/mgr/container.go

@@ -7,6 +7,8 @@ import (
 	"os/exec"
 	"path"
 	"path/filepath"
+	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
@@ -20,6 +22,7 @@ import (
 	"github.com/alibaba/pouch/pkg/collect"
 	"github.com/alibaba/pouch/pkg/errtypes"
 	"github.com/alibaba/pouch/pkg/meta"
+	"github.com/alibaba/pouch/pkg/quota"
 	"github.com/alibaba/pouch/pkg/randomid"
 	"github.com/alibaba/pouch/pkg/reference"
 	"github.com/alibaba/pouch/pkg/utils"
@@ -437,6 +440,11 @@ func (mgr *ContainerManager) Create(ctx context.Context, name string, config *ty
 		return nil, errors.Wrap(err, "failed to parse volume argument")
 	}
 
+	// set mount point disk quota
+	if err := mgr.setMountPointDiskQuota(ctx, meta); err != nil {
+		return nil, errors.Wrap(err, "failed to set mount point disk quota")
+	}
+
 	// set container basefs
 	mgr.setBaseFS(ctx, meta, id)
 
@@ -1488,6 +1496,60 @@ func (mgr *ContainerManager) parseBinds(ctx context.Context, meta *ContainerMeta
 	return nil
 }
 
+func (mgr *ContainerManager) setMountPointDiskQuota(ctx context.Context, c *ContainerMeta) error {
+	qid := 0
+	if c.Config.QuotaID != "" {
+		var err error
+		qid, err = strconv.Atoi(c.Config.QuotaID)
+		if err != nil {
+			return errors.Wrapf(err, "invalid argument, QuotaID: %s", c.Config.QuotaID)
+		}
+	}
+
+	// parse diskquota regexe
+	quotas := c.Config.DiskQuota
+	res := make([]*quota.RegExp, 0)
+	for path, size := range quotas {
+		re := regexp.MustCompile(path)
+		res = append(res, &quota.RegExp{re, path, size})
+	}
+
+	for _, mp := range c.Mounts {
+		// skip volume mount or replace mode mount
+		if mp.Name != "" || mp.Replace != "" || mp.Source == "" || mp.Destination == "" {
+			continue
+		}
+
+		// skip non-directory path.
+		if fd, err := os.Stat(mp.Source); err != nil || !fd.IsDir() {
+			continue
+		}
+
+		matched := false
+		for _, re := range res {
+			findStr := re.Pattern.FindString(mp.Destination)
+			if findStr == mp.Destination {
+				quotas[mp.Destination] = re.Size
+				matched = true
+				if re.Path != ".*" {
+					break
+				}
+			}
+		}
+
+		if matched {
+			err := quota.SetDiskQuota(mp.Source, quotas[mp.Destination], uint32(qid))
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	c.Config.DiskQuota = quotas
+
+	return nil
+}
+
 func (mgr *ContainerManager) detachVolumes(ctx context.Context, c *ContainerMeta) error {
 	for name := range c.Config.Volumes {
 		v, err := mgr.VolumeMgr.Get(ctx, name)

daemon/mgr/spec_blkio.go

@@ -97,7 +97,11 @@ func setupDiskQuota(ctx context.Context, meta *ContainerMeta, spec *SpecWrapper)
 
 	rootFSQuota, ok := meta.Config.DiskQuota["/"]
 	if !ok || rootFSQuota == "" {
-		return nil
+		commonQuota, ok := meta.Config.DiskQuota[".*"]
+		if !ok || commonQuota == "" {
+			return nil
+		}
+		rootFSQuota = commonQuota
 	}
 
 	if s.Hooks == nil {
@@ -114,7 +118,7 @@ func setupDiskQuota(ctx context.Context, meta *ContainerMeta, spec *SpecWrapper)
 
 	quotaPrestart := specs.Hook{
 		Path: target,
-		Args: []string{"set-diskquota", meta.BaseFS, rootFSQuota},
+		Args: []string{"set-diskquota", meta.BaseFS, rootFSQuota, meta.Config.QuotaID},
 	}
 	s.Hooks.Prestart = append(s.Hooks.Prestart, quotaPrestart)
 

pkg/quota/grpquota.go

@@ -130,7 +130,7 @@ func (quota *GrpQuota) SetSubtree(dir string, qid uint32) (uint32, error) {
 }
 
 // SetDiskQuota is used to set quota for directory.
-func (quota *GrpQuota) SetDiskQuota(dir string, size string, quotaID int) error {
+func (quota *GrpQuota) SetDiskQuota(dir string, size string, quotaID uint32) error {
 	logrus.Debugf("set disk quota, dir: %s, size: %s, quotaID: %d", dir, size, quotaID)
 	if !UseQuota {
 		return nil
@@ -144,8 +144,8 @@ func (quota *GrpQuota) SetDiskQuota(dir string, size string, quotaID int) error
 		return fmt.Errorf("mountpoint not found: %s", dir)
 	}
 
-	id, err := quota.SetSubtree(dir, uint32(quotaID))
-	if id == 0 {
+	id, err := quota.SetSubtree(dir, quotaID)
+	if err != nil || id == 0 {
 		return fmt.Errorf("subtree not found: %s %v", dir, err)
 	}
 

pkg/quota/prjquota.go

@@ -100,7 +100,7 @@ func (quota *PrjQuota) SetSubtree(dir string, qid uint32) (uint32, error) {
 }
 
 // SetDiskQuota is used to set quota for directory.
-func (quota *PrjQuota) SetDiskQuota(dir string, size string, quotaID int) error {
+func (quota *PrjQuota) SetDiskQuota(dir string, size string, quotaID uint32) error {
 	logrus.Debugf("set disk quota, dir: %s, size: %s, quotaID: %d", dir, size, quotaID)
 	if !UseQuota {
 		return nil
@@ -114,8 +114,8 @@ func (quota *PrjQuota) SetDiskQuota(dir string, size string, quotaID int) error
 		return fmt.Errorf("mountpoint not found: %s", dir)
 	}
 
-	id, err := quota.SetSubtree(dir, uint32(quotaID))
-	if id == 0 {
+	id, err := quota.SetSubtree(dir, quotaID)
+	if err != nil || id == 0 {
 		return fmt.Errorf("subtree not found: %s %v", dir, err)
 	}
 

pkg/quota/quota.go

@@ -24,7 +24,7 @@ var (
 type BaseQuota interface {
 	StartQuotaDriver(dir string) (string, error)
 	SetSubtree(dir string, qid uint32) (uint32, error)
-	SetDiskQuota(dir string, size string, quotaID int) error
+	SetDiskQuota(dir string, size string, quotaID uint32) error
 	CheckMountpoint(devID uint64) (string, bool, string)
 	GetFileAttr(dir string) uint32
 	SetFileAttr(dir string, id uint32) error
@@ -92,7 +92,7 @@ func SetSubtree(dir string, qid uint32) (uint32, error) {
 }
 
 // SetDiskQuota is used to set quota for directory.
-func SetDiskQuota(dir string, size string, quotaID int) error {
+func SetDiskQuota(dir string, size string, quotaID uint32) error {
 	return Gquota.SetDiskQuota(dir, size, quotaID)
 }
 

pkg/quota/set_diskquota.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 
 	"github.com/docker/docker/pkg/reexec"
@@ -37,13 +38,15 @@ func processSetQuotaReexec() {
 		}
 	}()
 
-	if len(os.Args) != 3 {
-		err = fmt.Errorf("invalid arguments: %v, it should be: %s: <path> <size>", os.Args, os.Args[0])
+	if len(os.Args) != 4 {
+		err = fmt.Errorf("invalid arguments: %v, it should be: %s: <path> <size> <quota id>", os.Args, os.Args[0])
 		return
 	}
 
 	basefs := os.Args[1]
 	size := os.Args[2]
+	id, _ := strconv.Atoi(os.Args[3])
+	qid = uint32(id)
 
 	logrus.Infof("set diskquota: %v", os.Args)
 
@@ -60,18 +63,19 @@ func processSetQuotaReexec() {
 			return
 		}
 
-		qid, err = SetSubtree(dir, qid)
+		qid, err = SetSubtree(dir, uint32(qid))
 		if err != nil {
 			logrus.Errorf("failed to set subtree: %v", err)
 			return
 		}
 
-		err = SetDiskQuota(dir, size, int(qid))
+		err = SetDiskQuota(dir, size, qid)
 		if err != nil {
 			logrus.Errorf("failed to set disk quota: %v", err)
+			return
 		}
 
-		setQuotaForDir(dir, qid)
+		setQuotaForDir(dir, uint32(qid))
 	}
 
 	return

pkg/quota/type.go

@@ -0,0 +1,10 @@
+package quota
+
+import "regexp"
+
+// RegExp defines the regular expression of disk quota.
+type RegExp struct {
+	Pattern *regexp.Regexp
+	Path    string
+	Size    string
+}

test/cli_network_test.go

@@ -337,7 +337,7 @@ func (suite *PouchNetworkSuite) TestNetworkPortMapping(c *check.C) {
 		"-p", "9999:80",
 		image).Assert(c, icmd.Success)
 
-	time.Sleep(1 * time.Second)
+	time.Sleep(10 * time.Second)
 	err := icmd.RunCommand("curl", "localhost:9999").Compare(expct)
 	c.Assert(err, check.IsNil)
 

test/cli_run_test.go

@@ -808,6 +808,7 @@ func (suite *PouchRunSuite) TestRunWithDiskQuota(c *check.C) {
 	for _, line := range strings.Split(out, "\n") {
 		if strings.Contains(line, "/") && strings.Contains(line, "2048000") {
 			found = true
+			break
 		}
 	}
 
@@ -851,3 +852,65 @@ func (suite *PouchRunSuite) TestRunWithExitCode(c *check.C) {
 	}
 	c.Assert(result[0].State.ExitCode, check.Equals, int64(101))
 }
+
+// TestRunWithDiskQuotaRegular tests running container with --disk-quota.
+func (suite *PouchRunSuite) TestRunWithDiskQuotaRegular(c *check.C) {
+	if !environment.IsDiskQuota() {
+		c.Skip("Host does not support disk quota")
+	}
+
+	volumeName := "diskquota-volume"
+	containerName := "diskquota-regular"
+
+	ret := command.PouchRun("volume", "create", "-n", volumeName, "-o", "size=256m", "-o", "mount=/data/volume")
+	defer func() {
+		command.PouchRun("volume", "rm", volumeName).Assert(c, icmd.Success)
+	}()
+	ret.Assert(c, icmd.Success)
+
+	ret = command.PouchRun("run",
+		"--disk-quota=1024m",
+		`--disk-quota=".*=512m"`,
+		`--disk-quota="/mnt/mount1=768m"`,
+		"-v", "/data/mount1:/mnt/mount1",
+		"-v", "/data/mount2:/mnt/mount2",
+		"-v", "diskquota-volume:/mnt/mount3",
+		"--name", containerName, busyboxImage, "df")
+	defer func() {
+		command.PouchRun("rm", "-f", containerName).Assert(c, icmd.Success)
+	}()
+	ret.Assert(c, icmd.Success)
+
+	out := ret.Stdout()
+
+	rootFound := false
+	mount1Found := false
+	mount2Found := false
+	mount3Found := false
+	for _, line := range strings.Split(out, "\n") {
+		if strings.Contains(line, "/") && strings.Contains(line, "1048576") {
+			rootFound = true
+			continue
+		}
+
+		if strings.Contains(line, "/mnt/mount1") && strings.Contains(line, "786432") {
+			mount1Found = true
+			continue
+		}
+
+		if strings.Contains(line, "/mnt/mount2") && strings.Contains(line, "524288") {
+			mount2Found = true
+			continue
+		}
+
+		if strings.Contains(line, "/mnt/mount3") && strings.Contains(line, "262144") {
+			mount3Found = true
+			continue
+		}
+	}
+
+	c.Assert(rootFound, check.Equals, true)
+	c.Assert(mount1Found, check.Equals, true)
+	c.Assert(mount2Found, check.Equals, true)
+	c.Assert(mount3Found, check.Equals, true)
+}