From 6dc54b0fef26391c81ca8a4160c6a259d491bb1e Mon Sep 17 00:00:00 2001
From: KyleWiese <kwiese@google.com>
Date: Wed, 10 Jan 2024 21:47:06 +0000
Subject: [PATCH] Add condition to bucket delete check to ensure an in-use
 bucket is not pre-maturely deleted

---
 quota/bucket.go      |  4 +++-
 quota/bucket_test.go | 28 +++++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/quota/bucket.go b/quota/bucket.go
index 934f34a..1ea3555 100644
--- a/quota/bucket.go
+++ b/quota/bucket.go
@@ -205,7 +205,9 @@ func (b *bucket) sync() error {
 func (b *bucket) needToDelete() bool {
 	b.lock.RLock()
 	defer b.lock.RUnlock()
-	return b.request.Weight == 0 && b.now().After(b.checked.Add(b.deleteAfter))
+	// Do not delete an in-use bucket to avoid allowing requests when a new bucket is created and is not yet synced to the remote service.
+	inUse := b.result != nil && b.result.Used != 0
+	return !inUse && b.request.Weight == 0 && b.now().After(b.checked.Add(b.deleteAfter))
 }
 
 func (b *bucket) needToSync() bool {
diff --git a/quota/bucket_test.go b/quota/bucket_test.go
index ab73b1c..8dc4802 100644
--- a/quota/bucket_test.go
+++ b/quota/bucket_test.go
@@ -146,6 +146,7 @@ func TestNeedToDelete(t *testing.T) {
 	cases := map[string]struct {
 		request *Request
 		checked time.Time
+		result  *Result
 		want    bool
 	}{
 		"empty": {
@@ -157,11 +158,35 @@ func TestNeedToDelete(t *testing.T) {
 			checked: now(),
 			want:    false,
 		},
-		"not recently checked": {
+		"not recently checked with empty result": {
 			request: &Request{},
 			checked: now().Add(-time.Hour),
 			want:    true,
 		},
+		"not recently checked with no quota allocated in result": {
+			request: &Request{},
+			checked: now().Add(-time.Hour),
+			result: &Result{
+				Allowed:    7,
+				Used:       0,
+				Exceeded:   0,
+				ExpiryTime: now().Unix(),
+				Timestamp:  now().Unix(),
+			},
+			want: true,
+		},
+		"not recently checked but quota has been allocated in result": {
+			request: &Request{},
+			checked: now().Add(-time.Hour),
+			result: &Result{
+				Allowed:    7,
+				Used:       7,
+				Exceeded:   1,
+				ExpiryTime: now().Unix(),
+				Timestamp:  now().Unix(),
+			},
+			want: false,
+		},
 		"has pending requests": {
 			request: &Request{Weight: 1},
 			checked: now().Add(-time.Hour),
@@ -176,6 +201,7 @@ func TestNeedToDelete(t *testing.T) {
 			deleteAfter: time.Minute,
 			request:     c.request,
 			checked:     c.checked,
+			result:      c.result,
 		}
 		if c.want != b.needToDelete() {
 			t.Errorf("want: %v got: %v", c.want, b.needToDelete())