HDDS-3728. Bucket space: check quotaUsageInBytes when write key and allocate block.

hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/client/rpc/TestOzoneRpcClientAbstract.java

@@ -813,6 +813,7 @@ public void testPutKey() throws IOException {
   }
 
   @Test
+  @SuppressWarnings("methodlength")
   public void testCheckUsedBytesQuota() throws IOException {
     String volumeName = UUID.randomUUID().toString();
     String bucketName = UUID.randomUUID().toString();
@@ -826,13 +827,15 @@ public void testCheckUsedBytesQuota() throws IOException {
 
     store.createVolume(volumeName);
     volume = store.getVolume(volumeName);
+
+    // Test volume quota.
     // Set quota In Bytes for a smaller value
     store.getVolume(volumeName).setQuota(
         OzoneQuota.parseQuota("1 Bytes", 100));
     volume.createBucket(bucketName);
     OzoneBucket bucket = volume.getBucket(bucketName);
 
-    // Test write key.
+    // Test volume quota: write key.
     // The remaining quota does not satisfy a block size, so the write fails.
     try {
       writeKey(bucket, UUID.randomUUID().toString(), ONE, value, valueLength);
@@ -843,7 +846,7 @@ public void testCheckUsedBytesQuota() throws IOException {
     // Write failed, volume usedBytes should be 0
     Assert.assertEquals(0L, store.getVolume(volumeName).getUsedBytes());
 
-    // Test write file.
+    // Test volume quota: write file.
     // The remaining quota does not satisfy a block size, so the write fails.
     try {
       writeFile(bucket, UUID.randomUUID().toString(), ONE, value, 0);
@@ -854,7 +857,7 @@ public void testCheckUsedBytesQuota() throws IOException {
     // Write failed, volume usedBytes should be 0
     Assert.assertEquals(0L, store.getVolume(volumeName).getUsedBytes());
 
-    // Write a key(with two blocks), test allocateBlock fails.
+    // Test volume quota: write key(with two blocks), test allocateBlock fails.
     store.getVolume(volumeName).setQuota(
         OzoneQuota.parseQuota(blockSize + "Bytes", 100));
     try {
@@ -871,8 +874,8 @@ public void testCheckUsedBytesQuota() throws IOException {
     // AllocateBlock failed, volume usedBytes should be 1 * blockSize.
     Assert.assertEquals(blockSize, store.getVolume(volumeName).getUsedBytes());
 
-    // Write large key(with five blocks), the first four blocks will succeed，
-    // while the later block will fail.
+    // Test volume quota: write large key(with five blocks), the first four
+    // blocks will succeed，while the later block will fail.
     store.getVolume(volumeName).setQuota(
         OzoneQuota.parseQuota(5 * blockSize + "Bytes", 100));
     try {
@@ -890,7 +893,59 @@ public void testCheckUsedBytesQuota() throws IOException {
     Assert.assertEquals(5 * blockSize,
         store.getVolume(volumeName).getUsedBytes());
 
-    Assert.assertEquals(4, countException);
+    // Test bucket quota.
+    // Set quota In Bytes for a smaller value
+    store.getVolume(volumeName).setQuota(
+        OzoneQuota.parseQuota(Long.MAX_VALUE + " Bytes", 100));
+    bucketName = UUID.randomUUID().toString();
+    volume.createBucket(bucketName);
+    bucket = volume.getBucket(bucketName);
+    bucket.setQuota(OzoneQuota.parseQuota("1 Bytes", 100));
+
+    // Test bucket quota: write key.
+    // The remaining quota does not satisfy a block size, so the write fails.
+    try {
+      writeKey(bucket, UUID.randomUUID().toString(), ONE, value, valueLength);
+    } catch (IOException ex) {
+      countException++;
+      GenericTestUtils.assertExceptionContains("QUOTA_EXCEEDED", ex);
+    }
+    // Write failed, bucket usedBytes should be 0
+    Assert.assertEquals(0L,
+        store.getVolume(volumeName).getBucket(bucketName).getUsedBytes());
+
+    // Test bucket quota: write file.
+    // The remaining quota does not satisfy a block size, so the write fails.
+    try {
+      writeFile(bucket, UUID.randomUUID().toString(), ONE, value, 0);
+    } catch (IOException ex) {
+      countException++;
+      GenericTestUtils.assertExceptionContains("QUOTA_EXCEEDED", ex);
+    }
+    // Write failed, bucket usedBytes should be 0
+    Assert.assertEquals(0L,
+        store.getVolume(volumeName).getBucket(bucketName).getUsedBytes());
+
+    // Test bucket quota: write large key(with five blocks), the first four
+    // blocks will succeed，while the later block will fail.
+    bucket.setQuota(OzoneQuota.parseQuota(
+        4 * blockSize + " Bytes", 100));
+    try {
+      OzoneOutputStream out = bucket.createKey(UUID.randomUUID().toString(),
+          valueLength, STAND_ALONE, ONE, new HashMap<>());
+      for (int i = 0; i <= (4 * blockSize) / value.length(); i++) {
+        out.write(value.getBytes());
+      }
+      out.close();
+    } catch (IOException ex) {
+      countException++;
+      GenericTestUtils.assertExceptionContains("QUOTA_EXCEEDED", ex);
+    }
+    // AllocateBlock failed, bucket usedBytes should be 4 * blockSize
+    Assert.assertEquals(4 * blockSize,
+        store.getVolume(volumeName).getBucket(bucketName).getUsedBytes());
+
+    Assert.assertEquals(7, countException);
   }
 
   @Test

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/file/OMFileCreateRequest.java

@@ -279,10 +279,11 @@ public OMClientResponse validateAndUpdateCache(OzoneManager ozoneManager,
 
       omVolumeArgs = getVolumeInfo(omMetadataManager, volumeName);
       omBucketInfo = getBucketInfo(omMetadataManager, volumeName, bucketName);
-      // check volume quota
+      // check bucket and volume quota
       long preAllocatedSpace = newLocationList.size()
           * ozoneManager.getScmBlockSize()
           * omKeyInfo.getFactor().getNumber();
+      checkBucketQuotaInBytes(omBucketInfo, preAllocatedSpace);
       checkVolumeQuotaInBytes(omVolumeArgs, preAllocatedSpace);
 
       // Add to cache entry can be done outside of lock for this openKey.

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/key/OMAllocateBlockRequest.java

@@ -196,10 +196,11 @@ public OMClientResponse validateAndUpdateCache(OzoneManager ozoneManager,
           OmKeyLocationInfo.getFromProtobuf(blockLocation));
       omVolumeArgs = getVolumeInfo(omMetadataManager, volumeName);
       omBucketInfo = getBucketInfo(omMetadataManager, volumeName, bucketName);
-      // check volume quota
+      // check bucket and volume quota
       long preAllocatedSpace = newLocationList.size()
           * ozoneManager.getScmBlockSize()
           * openKeyInfo.getFactor().getNumber();
+      checkBucketQuotaInBytes(omBucketInfo, preAllocatedSpace);
       checkVolumeQuotaInBytes(omVolumeArgs, preAllocatedSpace);
       // Append new block
       openKeyInfo.appendNewBlocks(newLocationList, false);

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/key/OMKeyCreateRequest.java

@@ -298,7 +298,8 @@ public OMClientResponse validateAndUpdateCache(OzoneManager ozoneManager,
       long preAllocatedSpace = newLocationList.size()
           * ozoneManager.getScmBlockSize()
           * omKeyInfo.getFactor().getNumber();
-      // check volume quota
+      // check bucket and volume quota
+      checkBucketQuotaInBytes(omBucketInfo, preAllocatedSpace);
       checkVolumeQuotaInBytes(omVolumeArgs, preAllocatedSpace);
 
       // Add to cache entry can be done outside of lock for this openKey.

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/key/OMKeyRequest.java

@@ -555,6 +555,27 @@ protected void checkVolumeQuotaInBytes(OmVolumeArgs omVolumeArgs,
     }
   }
 
+  /**
+   * Check bucket quota in bytes.
+   * @param omBucketInfo
+   * @param allocateSize
+   * @throws IOException
+   */
+  protected void checkBucketQuotaInBytes(OmBucketInfo omBucketInfo,
+      long allocateSize) throws IOException {
+    if (omBucketInfo.getQuotaInBytes() > OzoneConsts.QUOTA_RESET) {
+      long usedBytes = omBucketInfo.getUsedBytes().sum();
+      long quotaInBytes = omBucketInfo.getQuotaInBytes();
+      if (quotaInBytes - usedBytes < allocateSize) {
+        throw new OMException("The DiskSpace quota of bucket:"
+            + omBucketInfo.getBucketName() + "exceeded: quotaInBytes: "
+            + quotaInBytes + " Bytes but diskspace consumed: " + (usedBytes
+            + allocateSize) + " Bytes.",
+            OMException.ResultCodes.QUOTA_EXCEEDED);
+      }
+    }
+  }
+
   /**
    * Check directory exists. If exists return true, else false.
    * @param volumeName