refact: refact disk quota api

for pouch cli command, support two api: ```bash --disk-quota []string --quota-id string ``` --disk-quota can be set: 1. --disk-quota=60G 2. --disk-quota=/abc=60G 3. --disk-quota=/&/abc=60G 4. --disk-quota=.*=60G --quota-id can be set: 1. --quota-id=0, or haven't set quota id(--quota-id="") 2. --quota-id=-1 3. --quota-id=16777216, or more than 16777216 detail rules you can see disk quota documents. validate disk quota config when create container, add test case for disk quota. Signed-off-by: Rudy Zhang <rudyflyzhang@gmail.com>
AliyunContainerService · Feb 15, 2019 · 37c682a · 37c682a
1 parent 7403638
commit 37c682a
Show file tree

Hide file tree

Showing 9 changed files with 197 additions and 169 deletions.
diff --git a/apis/opts/diskquota.go b/apis/opts/diskquota.go
@@ -17,7 +17,7 @@ func ParseDiskQuota(quotas []string) (map[string]string, error) {
 		parts := strings.Split(quota, "=")
 		switch len(parts) {
 		case 1:
-			quotaMaps["/"] = parts[0]
+			quotaMaps[".*"] = parts[0]
 		case 2:
 			quotaMaps[parts[0]] = parts[1]
 		default:

diff --git a/apis/opts/diskquota_test.go b/apis/opts/diskquota_test.go
@@ -18,7 +18,7 @@ func TestParseDiskQuota(t *testing.T) {
 		// TODO: Add test cases.
 		{name: "test1", args: args{diskquota: []string{""}}, want: nil, wantErr: true},
 		{name: "test2", args: args{diskquota: []string{"foo=foo=foo"}}, want: nil, wantErr: true},
-		{name: "test3", args: args{diskquota: []string{"foo"}}, want: map[string]string{"/": "foo"}, wantErr: false},
+		{name: "test3", args: args{diskquota: []string{"foo"}}, want: map[string]string{".*": "foo"}, wantErr: false},
 		{name: "test4", args: args{diskquota: []string{"foo=foo"}}, want: map[string]string{"foo": "foo"}, wantErr: false},
 		{name: "test5", args: args{diskquota: []string{"foo=foo", "bar=bar"}}, want: map[string]string{"foo": "foo", "bar": "bar"}, wantErr: false},
 	}

diff --git a/daemon/mgr/container.go b/daemon/mgr/container.go
@@ -9,7 +9,6 @@ import (
 	"os/exec"
 	"path"
 	"path/filepath"
-	"strconv"
 	"strings"
 	"time"
 
@@ -29,7 +28,6 @@ import (
 	mountutils "github.com/alibaba/pouch/pkg/mount"
 	"github.com/alibaba/pouch/pkg/streams"
 	"github.com/alibaba/pouch/pkg/utils"
-	"github.com/alibaba/pouch/storage/quota"
 	volumetypes "github.com/alibaba/pouch/storage/volume/types"
 
 	"github.com/containerd/cgroups"
@@ -358,6 +356,11 @@ func (mgr *ContainerManager) Create(ctx context.Context, name string, config *ty
 		return nil, errors.Wrapf(errtypes.ErrInvalidParam, "NetworkingConfig cannot be empty")
 	}
 
+	// validate disk quota
+	if err := mgr.validateDiskQuota(config); err != nil {
+		return nil, errors.Wrapf(err, "invalid disk quota config")
+	}
+
 	id, err := mgr.generateContainerID(config.SpecificID)
 	if err != nil {
 		return nil, err
@@ -1262,75 +1265,35 @@ func (mgr *ContainerManager) Remove(ctx context.Context, name string, options *t
 	return nil
 }
 
-func (mgr *ContainerManager) updateContainerDiskQuota(ctx context.Context, c *Container, diskQuota map[string]string) error {
+func (mgr *ContainerManager) updateContainerDiskQuota(ctx context.Context, c *Container, diskQuota map[string]string) (err error) {
 	if diskQuota == nil {
 		return nil
 	}
 
+	// backup diskquota
+	origDiskQuota := c.Config.DiskQuota
+	defer func() {
+		if err != nil {
+			c.Lock()
+			c.Config.DiskQuota = origDiskQuota
+			c.Unlock()
+		}
+	}()
+
 	c.Lock()
+	if c.Config.DiskQuota == nil {
+		c.Config.DiskQuota = make(map[string]string)
+	}
 	for dir, quota := range diskQuota {
 		c.Config.DiskQuota[dir] = quota
 	}
 	c.Unlock()
 
 	// set mount point disk quota
-	if err := mgr.setMountPointDiskQuota(ctx, c); err != nil {
+	if err = mgr.setDiskQuota(ctx, c, false); err != nil {
 		return errors.Wrapf(err, "failed to set mount point disk quota")
 	}
 
-	c.Lock()
-	var qid uint32
-	if c.Config.QuotaID != "" {
-		id, err := strconv.Atoi(c.Config.QuotaID)
-		if err != nil {
-			return errors.Wrapf(err, "failed to convert QuotaID %s", c.Config.QuotaID)
-		}
-
-		qid = uint32(id)
-		if id < 0 {
-			// QuotaID is < 0, it means pouchd alloc a unique quota id.
-			qid, err = quota.GetNextQuotaID()
-			if err != nil {
-				return errors.Wrap(err, "failed to get next quota id")
-			}
-
-			// update QuotaID
-			c.Config.QuotaID = strconv.Itoa(int(qid))
-		}
-	}
-	c.Unlock()
-
-	// get rootfs quota
-	defaultQuota := quota.GetDefaultQuota(c.Config.DiskQuota)
-	if qid > 0 && defaultQuota == "" {
-		return fmt.Errorf("set quota id but have no set default quota size")
-	}
-	// update container rootfs disk quota
-	// TODO: add lock for container?
-	rootfs := ""
-	if c.IsRunningOrPaused() && c.Snapshotter != nil {
-		basefs, ok := c.Snapshotter.Data["MergedDir"]
-		if !ok || basefs == "" {
-			return fmt.Errorf("Container is running, but MergedDir is missing")
-		}
-		rootfs = basefs
-	} else {
-		if err := mgr.Mount(ctx, c); err != nil {
-			return errors.Wrapf(err, "failed to mount rootfs: (%s)", c.MountFS)
-		}
-		rootfs = c.MountFS
-
-		defer func() {
-			if err := mgr.Unmount(ctx, c); err != nil {
-				logrus.Errorf("failed to umount rootfs: (%s), err: (%v)", c.MountFS, err)
-			}
-		}()
-	}
-	_, err := quota.SetRootfsDiskQuota(rootfs, defaultQuota, qid)
-	if err != nil {
-		return errors.Wrapf(err, "failed to set container rootfs diskquota")
-	}
-
 	return nil
 }
 

diff --git a/daemon/mgr/container_storage.go b/daemon/mgr/container_storage.go
@@ -471,96 +471,60 @@ func (mgr *ContainerManager) setMountTab(ctx context.Context, c *Container) erro
 	return nil
 }
 
-func (mgr *ContainerManager) setRootfsQuota(ctx context.Context, c *Container) error {
-	logrus.Debugf("start to set rootfs quota, dir(%s)", c.MountFS)
+func (mgr *ContainerManager) setDiskQuota(ctx context.Context, c *Container, mounted bool) error {
+	var globalQuotaID uint32
 
-	if c.MountFS == "" {
-		return nil
-	}
-
-	rootfsQuota := quota.GetDefaultQuota(c.Config.DiskQuota)
-	if rootfsQuota == "" {
-		return nil
-	}
-
-	qid := "0"
-	if c.Config.QuotaID != "" {
-		qid = c.Config.QuotaID
-	}
-
-	id, err := strconv.Atoi(qid)
-	if err != nil {
-		return errors.Wrapf(err, "failed to change quota id(%s) from string to int", qid)
-	}
-
-	// set rootfs quota
-	_, err = quota.SetRootfsDiskQuota(c.MountFS, rootfsQuota, uint32(id))
-	if err != nil {
-		return errors.Wrapf(err, "failed to set rootfs quota, mountfs(%s), quota(%s), quota id(%d)",
-			c.MountFS, rootfsQuota, id)
-	}
-
-	return nil
-}
-
-func (mgr *ContainerManager) setMountPointDiskQuota(ctx context.Context, c *Container) error {
-	if c.Config.DiskQuota == nil {
-		if c.Config.QuotaID != "" && c.Config.QuotaID != "0" {
-			return fmt.Errorf("invalid argument, set quota-id without disk-quota")
-		}
-		return nil
-	}
-
-	var (
-		qid        uint32
-		setQuotaID bool
-	)
-
-	if c.Config.QuotaID != "" {
+	if c.Config.QuotaID != "" && c.Config.QuotaID != "0" {
 		id, err := strconv.Atoi(c.Config.QuotaID)
 		if err != nil {
 			return errors.Wrapf(err, "invalid argument, QuotaID(%s)", c.Config.QuotaID)
 		}
 
 		// if QuotaID is < 0, it means pouchd alloc a unique quota id.
 		if id < 0 {
-			qid, err = quota.GetNextQuotaID()
+			globalQuotaID, err = quota.GetNextQuotaID()
 			if err != nil {
 				return errors.Wrap(err, "failed to get next quota id")
 			}
 
 			// update QuotaID
-			c.Config.QuotaID = strconv.Itoa(int(qid))
+			c.Config.QuotaID = strconv.Itoa(int(globalQuotaID))
 		} else {
-			qid = uint32(id)
+			globalQuotaID = uint32(id)
 		}
 	}
 
-	if qid > 0 {
-		setQuotaID = true
-	}
-
-	// get rootfs quota
+	// get default quota
 	quotas := c.Config.DiskQuota
 	defaultQuota := quota.GetDefaultQuota(quotas)
-	if setQuotaID && defaultQuota == "" {
-		return fmt.Errorf("set quota id but have no set default quota size")
-	}
 
-	// parse diskquota regexe
+	// parse disk quota regexp
 	var res []*quota.RegExp
-	for path, size := range quotas {
-		re := regexp.MustCompile(path)
-		res = append(res, &quota.RegExp{Pattern: re, Path: path, Size: size})
+	for key, size := range quotas {
+		var err error
+		id := globalQuotaID
+		if id == 0 {
+			id, err = quota.GetNextQuotaID()
+			if err != nil {
+				return errors.Wrap(err, "failed to get next quota id")
+			}
+		}
+		paths := strings.Split(key, "&")
+		for _, path := range paths {
+			re := regexp.MustCompile(path)
+			res = append(res, &quota.RegExp{Pattern: re, Path: path, Size: size, QuotaID: id})
+		}
 	}
 
+	var mounts []*types.MountPoint
 	for _, mp := range c.Mounts {
 		// skip volume mount or replace mode mount
 		if mp.Replace != "" || mp.Source == "" || mp.Destination == "" {
 			logrus.Debugf("skip volume mount or replace mode mount")
 			continue
 		}
 
+		// skip volume that has set size
 		if mp.Name != "" {
 			v, err := mgr.VolumeMgr.Get(ctx, mp.Name)
 			if err != nil {
@@ -580,33 +544,60 @@ func (mgr *ContainerManager) setMountPointDiskQuota(ctx context.Context, c *Cont
 			continue
 		}
 
-		matched := false
+		mounts = append(mounts, mp)
+	}
+
+	// add rootfs mountpoint
+	rootfs, err := mgr.getRootfs(ctx, c, mounted)
+	if err != nil {
+		return errors.Wrapf(err, "failed to get rootfs")
+	}
+	mounts = append(mounts, &types.MountPoint{
+		Source:      rootfs,
+		Destination: "/",
+	})
+
+	for _, mp := range mounts {
+		var (
+			size string
+			id   uint32
+		)
+
 		for _, re := range res {
 			findStr := re.Pattern.FindString(mp.Destination)
 			if findStr == mp.Destination {
-				quotas[mp.Destination] = re.Size
-				matched = true
+				size = re.Size
+				id = re.QuotaID
 				if re.Path != ".*" {
 					break
 				}
 			}
 		}
 
-		size := ""
-		if matched && !setQuotaID {
-			size = quotas[mp.Destination]
-		} else {
-			size = defaultQuota
+		if size == "" {
+			if defaultQuota != "" {
+				size = defaultQuota
+			} else {
+				continue
+			}
 		}
-		err := quota.SetDiskQuota(mp.Source, size, qid)
-		if err != nil {
-			// just ignore set disk quota fail
-			logrus.Warnf("failed to set disk quota, directory(%s), size(%s), quotaID(%d), err(%v)",
-				mp.Source, size, qid, err)
+
+		if mp.Destination == "/" {
+			// set rootfs quota
+			_, err = quota.SetRootfsDiskQuota(mp.Source, size, id)
+			if err != nil {
+				return errors.Wrapf(err, "failed to set rootfs quota, mountfs(%s), size(%s), quota id(%d)",
+					mp.Source, size, id)
+			}
+		} else {
+			err := quota.SetDiskQuota(mp.Source, size, id)
+			if err != nil {
+				return errors.Wrapf(err, "failed to set disk quota, directory(%s), size(%s), quota id(%d)",
+					mp.Source, size, id)
+			}
 		}
-	}
 
-	c.Config.DiskQuota = quotas
+	}
 
 	return nil
 }
@@ -699,13 +690,9 @@ func (mgr *ContainerManager) initContainerStorage(ctx context.Context, c *Contai
 	}
 
 	// set mount point disk quota
-	if err = mgr.setMountPointDiskQuota(ctx, c); err != nil {
-		return errors.Wrap(err, "failed to set mount point disk quota")
-	}
-
-	// set rootfs disk quota
-	if err = mgr.setRootfsQuota(ctx, c); err != nil {
-		logrus.Warnf("failed to set rootfs disk quota, err(%v)", err)
+	if err = mgr.setDiskQuota(ctx, c, true); err != nil {
+		// just ignore failed to set disk quota
+		logrus.Warnf("failed to set disk quota, err(%v)", err)
 	}
 
 	// set volumes into /etc/mtab in container
@@ -738,6 +725,35 @@ func (mgr *ContainerManager) SetupWorkingDirectory(ctx context.Context, c *Conta
 	return nil
 }
 
+func (mgr *ContainerManager) getRootfs(ctx context.Context, c *Container, mounted bool) (string, error) {
+	var (
+		rootfs string
+		err    error
+	)
+	if c.IsRunningOrPaused() && c.Snapshotter != nil {
+		basefs, ok := c.Snapshotter.Data["MergedDir"]
+		if !ok || basefs == "" {
+			return "", fmt.Errorf("Container is running, but MergedDir is missing")
+		}
+		rootfs = basefs
+	} else if !mounted {
+		if err = mgr.Mount(ctx, c); err != nil {
+			return "", errors.Wrapf(err, "failed to mount rootfs: (%s)", c.MountFS)
+		}
+		rootfs = c.MountFS
+
+		defer func() {
+			if err = mgr.Unmount(ctx, c); err != nil {
+				logrus.Errorf("failed to umount rootfs: (%s), err: (%v)", c.MountFS, err)
+			}
+		}()
+	} else {
+		rootfs = c.MountFS
+	}
+
+	return rootfs, nil
+}
+
 func sortMountPoint(mounts []*types.MountPoint) []*types.MountPoint {
 	sort.Slice(mounts, func(i, j int) bool {
 		if len(mounts[i].Destination) < len(mounts[j].Destination) {