From 8fe578e7ab67947aee5f5737864d5aabeeb5f67f Mon Sep 17 00:00:00 2001
From: "rongfu.leng" <1275177125@qq.com>
Date: Mon, 5 Jun 2023 17:19:55 +0800
Subject: [PATCH] feat: Optimize quota checking when pushing images (#17392)

Signed-off-by: lengrongfu <1275177125@qq.com>
---
 src/controller/quota/controller_test.go       | 13 +++++++++++++
 src/pkg/quota/util.go                         |  2 +-
 .../quota/post_initiate_blob_upload.go        |  4 ++--
 .../quota/post_initiate_blob_upload_test.go   | 19 +++++++++++++++++++
 4 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/controller/quota/controller_test.go b/src/controller/quota/controller_test.go
index a98db9b3c..037d39bb6 100644
--- a/src/controller/quota/controller_test.go
+++ b/src/controller/quota/controller_test.go
@@ -143,6 +143,19 @@ func (suite *ControllerTestSuite) TestRequestFunctionFailed() {
 	suite.Error(suite.ctl.Request(ctx, suite.reference, referenceID, resources, func() error { return fmt.Errorf("error") }))
 }
 
+func (suite *ControllerTestSuite) TestRequestResourceIsZero() {
+	suite.PrepareForUpdate(suite.quota, nil)
+
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+	referenceID := uuid.New().String()
+	f := func() error {
+		return nil
+	}
+	res := types.ResourceList{types.ResourceStorage: 0}
+	err := suite.ctl.Request(ctx, suite.reference, referenceID, res, f)
+	suite.Nil(err)
+}
+
 func TestControllerTestSuite(t *testing.T) {
 	suite.Run(t, &ControllerTestSuite{})
 }
diff --git a/src/pkg/quota/util.go b/src/pkg/quota/util.go
index 8e3be434f..00d796d38 100644
--- a/src/pkg/quota/util.go
+++ b/src/pkg/quota/util.go
@@ -32,7 +32,7 @@ func IsSafe(hardLimits types.ResourceList, currentUsed types.ResourceList, newUs
 			continue
 		}
 
-		if hardLimit == types.UNLIMITED || value == currentUsed[resource] {
+		if hardLimit == types.UNLIMITED {
 			continue
 		}
 
diff --git a/src/server/middleware/quota/post_initiate_blob_upload.go b/src/server/middleware/quota/post_initiate_blob_upload.go
index f4653a9c3..efeb8dd33 100644
--- a/src/server/middleware/quota/post_initiate_blob_upload.go
+++ b/src/server/middleware/quota/post_initiate_blob_upload.go
@@ -38,8 +38,8 @@ func postInitiateBlobUploadResources(r *http.Request, reference, referenceID str
 	query := r.URL.Query()
 	mount := query.Get("mount")
 	if mount == "" {
-		// it is not mount blob http request, skip to request the resources
-		return nil, nil
+		// it is not mount blob http request, create length is zero resource to check quota is full
+		return types.ResourceList{types.ResourceStorage: 0}, nil
 	}
 
 	ctx := r.Context()
diff --git a/src/server/middleware/quota/post_initiate_blob_upload_test.go b/src/server/middleware/quota/post_initiate_blob_upload_test.go
index c6f0d1da5..d30e2cc93 100644
--- a/src/server/middleware/quota/post_initiate_blob_upload_test.go
+++ b/src/server/middleware/quota/post_initiate_blob_upload_test.go
@@ -70,6 +70,25 @@ func (suite *PostInitiateBlobUploadMiddlewareTestSuite) TestMiddleware() {
 		PostInitiateBlobUploadMiddleware()(next).ServeHTTP(rr, req)
 		suite.Equal(http.StatusOK, rr.Code)
 	}
+
+	{
+		url = "/v2/library/photon/blobs/uploads"
+		mock.OnAnything(suite.quotaController, "Request").Return(nil).Once().Run(func(args mock.Arguments) {
+			resources := args.Get(3).(types.ResourceList)
+			suite.Len(resources, 1)
+			suite.Equal(resources[types.ResourceStorage], int64(0))
+
+			f := args.Get(4).(func() error)
+			f()
+		})
+		mock.OnAnything(suite.quotaController, "GetByRef").Return(&quota.Quota{}, nil).Once()
+
+		req := httptest.NewRequest(http.MethodPost, url, nil)
+		rr := httptest.NewRecorder()
+
+		PostInitiateBlobUploadMiddleware()(next).ServeHTTP(rr, req)
+		suite.Equal(http.StatusOK, rr.Code)
+	}
 }
 
 func TestPostInitiateBlobUploadMiddlewareTestSuite(t *testing.T) {