From 6e57d8d0de6863fd8488eceee621f52d149374b2 Mon Sep 17 00:00:00 2001
From: chrislu <chris.lu@gmail.com>
Date: Fri, 21 Jan 2022 02:15:27 -0800
Subject: [PATCH] s3: check bucket usage and adjust read only according to
 quota

---
 weed/shell/command_s3_bucket_quota_check.go | 141 ++++++++++++++++++++
 1 file changed, 141 insertions(+)
 create mode 100644 weed/shell/command_s3_bucket_quota_check.go

diff --git a/weed/shell/command_s3_bucket_quota_check.go b/weed/shell/command_s3_bucket_quota_check.go
new file mode 100644
index 000000000..fc8291148
--- /dev/null
+++ b/weed/shell/command_s3_bucket_quota_check.go
@@ -0,0 +1,141 @@
+package shell
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"github.com/chrislusf/seaweedfs/weed/filer"
+	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
+	"io"
+	"math"
+)
+
+func init() {
+	Commands = append(Commands, &commandS3BucketQuotaCheck{})
+}
+
+type commandS3BucketQuotaCheck struct {
+}
+
+func (c *commandS3BucketQuotaCheck) Name() string {
+	return "s3.bucket.quota.check"
+}
+
+func (c *commandS3BucketQuotaCheck) Help() string {
+	return `check quota for all buckets, make the bucket read only if over the limit
+
+	Example:
+		s3.bucket.quota.check -force
+`
+}
+
+func (c *commandS3BucketQuotaCheck) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {
+
+	bucketCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
+	enforceQuotaLimit := bucketCommand.Bool("force", false, "actually change the buckets readonly attribute")
+	if err = bucketCommand.Parse(args); err != nil {
+		return nil
+	}
+
+	// collect collection information
+	topologyInfo, _, err := collectTopologyInfo(commandEnv)
+	if err != nil {
+		return err
+	}
+	collectionInfos := make(map[string]*CollectionInfo)
+	collectCollectionInfo(topologyInfo, collectionInfos)
+
+	// read buckets path
+	var filerBucketsPath string
+	filerBucketsPath, err = readFilerBucketsPath(commandEnv)
+	if err != nil {
+		return fmt.Errorf("read buckets: %v", err)
+	}
+
+	// read existing filer configuration
+	fc, err := filer.ReadFilerConf(commandEnv.option.FilerAddress, commandEnv.option.GrpcDialOption, commandEnv.MasterClient)
+	if err != nil {
+		return err
+	}
+
+	// process each bucket
+	hasConfChanges := false
+	err = filer_pb.List(commandEnv, filerBucketsPath, "", func(entry *filer_pb.Entry, isLast bool) error {
+		if !entry.IsDirectory {
+			return nil
+		}
+		collection := entry.Name
+		var collectionSize uint64
+		if collectionInfo, found := collectionInfos[collection]; found {
+			collectionSize = collectionInfo.Size
+		}
+		if c.processEachBucket(fc, filerBucketsPath, entry, writer, collectionSize) {
+			hasConfChanges = true
+		}
+		return nil
+	}, "", false, math.MaxUint32)
+	if err != nil {
+		return fmt.Errorf("list buckets under %v: %v", filerBucketsPath, err)
+	}
+
+	// apply the configuration changes
+	if hasConfChanges && *enforceQuotaLimit {
+
+		var buf2 bytes.Buffer
+		fc.ToText(&buf2)
+
+		if err = commandEnv.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error {
+			return filer.SaveInsideFiler(client, filer.DirectoryEtcSeaweedFS, filer.FilerConfName, buf2.Bytes())
+		}); err != nil && err != filer_pb.ErrNotFound {
+			return err
+		}
+	}
+
+	return err
+
+}
+
+func (c *commandS3BucketQuotaCheck) processEachBucket(fc *filer.FilerConf, filerBucketsPath string, entry *filer_pb.Entry, writer io.Writer, collectionSize uint64) (hasConfChanges bool) {
+
+	locPrefix := filerBucketsPath + "/" + entry.Name + "/"
+	locConf := fc.MatchStorageRule(locPrefix)
+	if locConf == nil {
+		locConf = &filer_pb.FilerConf_PathConf{
+			LocationPrefix: locPrefix,
+			Collection:     entry.Name,
+		}
+	}
+
+	if entry.Quota > 0 {
+		if locConf.ReadOnly {
+			if collectionSize < uint64(entry.Quota) {
+				locConf.ReadOnly = false
+				hasConfChanges = true
+			}
+		} else {
+			if collectionSize > uint64(entry.Quota) {
+				locConf.ReadOnly = true
+				hasConfChanges = true
+			}
+		}
+	} else {
+		if locConf.ReadOnly {
+			locConf.ReadOnly = false
+			hasConfChanges = true
+		}
+	}
+
+	if hasConfChanges {
+		fmt.Fprintf(writer, "  %s\tsize:%d", entry.Name, collectionSize)
+		fmt.Fprintf(writer, "\tquota:%d\tusage:%.2f%%", entry.Quota, float64(collectionSize)*100/float64(entry.Quota))
+		fmt.Fprintln(writer)
+		if locConf.ReadOnly {
+			fmt.Fprintf(writer, "    changing bucket %s to read only!\n", entry.Name)
+		} else {
+			fmt.Fprintf(writer, "    changing bucket %s to writable.\n", entry.Name)
+		}
+		fc.AddLocationConf(locConf)
+	}
+
+	return
+}