feat:增加仓库级别下载限速 #1092 (#1094)

* feat: 增加仓库级别限速 #1092 * feat: 剔除上传限制 #1092 * feat: 代码还原 #1092 * feat: 增加缓存 #1092 * feat: 当仓库配置下载限速小于等于最低限速时则直接将请求断开, 避免占用过多连接 #1092 * feat: 增加配置 #1092 * feat: 配置调整#1092
TencentBlueKing · Aug 29, 2023 · eee7445 · eee7445
1 parent ea60e25
commit eee7445
Show file tree

Hide file tree

Showing 11 changed files with 170 additions and 3 deletions.
diff --git a/...n-api/src/main/kotlin/com/tencent/bkrepo/common/api/exception/TooManyRequestsException.kt b/...n-api/src/main/kotlin/com/tencent/bkrepo/common/api/exception/TooManyRequestsException.kt
@@ -0,0 +1,38 @@
+/*
+ * Tencent is pleased to support the open source community by making BK-CI 蓝鲸持续集成平台 available.
+ *
+ * Copyright (C) 2022 THL A29 Limited, a Tencent company.  All rights reserved.
+ *
+ * BK-CI 蓝鲸持续集成平台 is licensed under the MIT license.
+ *
+ * A copy of the MIT License is included in this file.
+ *
+ *
+ * Terms of the MIT License:
+ * ---------------------------------------------------
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+ * documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
+ * permit persons to whom the Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all copies or substantial portions of
+ * the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
+ * LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
+ * NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+ * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+package com.tencent.bkrepo.common.api.exception
+
+import com.tencent.bkrepo.common.api.constant.HttpStatus
+import com.tencent.bkrepo.common.api.message.CommonMessageCode
+
+/**
+ * 请求过多，超出了频次限制
+ */
+class TooManyRequestsException(
+    parameter: String = "Too Many Requests"
+) : ErrorCodeException(HttpStatus.TOO_MANY_REQUESTS, CommonMessageCode.TOO_MANY_REQUESTS, arrayOf(parameter))
diff --git a/...mon/common-api/src/main/kotlin/com/tencent/bkrepo/common/api/message/CommonMessageCode.kt b/...mon/common-api/src/main/kotlin/com/tencent/bkrepo/common/api/message/CommonMessageCode.kt
@@ -55,6 +55,7 @@ enum class CommonMessageCode(private val key: String) : MessageCode {
     MODIFY_PASSWORD_FAILED("modify.password.failed"),
     OPERATION_CROSS_CLUSTER_NOT_ALLOWED("operation.cross-cluster.not-allowed"),
     MEDIA_TYPE_UNACCEPTABLE("system.media-type.unacceptable"),
+    TOO_MANY_REQUESTS("too.many.requests")
     ;
 
     override fun getBusinessCode() = ordinal + 1

diff --git a/...fact-api/src/main/kotlin/com/tencent/bkrepo/common/artifact/constant/ArtifactConstants.kt b/...fact-api/src/main/kotlin/com/tencent/bkrepo/common/artifact/constant/ArtifactConstants.kt
@@ -51,6 +51,11 @@ const val ARTIFACT_CONFIGURER = "artifact-configurer"
  */
 const val NODE_DETAIL_KEY = "node-detail"
 
+/**
+ * 查询仓库上传下载限速配置后后将其写入request attributes的key
+ */
+const val REPO_RATE_LIMIT_KEY = "repo-rate-limit"
+
 /**
  * 项目id字段
  */

diff --git a/...ain/kotlin/com/tencent/bkrepo/common/artifact/repository/context/ArtifactContextHolder.kt b/...ain/kotlin/com/tencent/bkrepo/common/artifact/repository/context/ArtifactContextHolder.kt
@@ -41,18 +41,21 @@ import com.tencent.bkrepo.common.artifact.constant.NODE_DETAIL_KEY
 import com.tencent.bkrepo.common.artifact.constant.PROJECT_ID
 import com.tencent.bkrepo.common.artifact.constant.REPO_KEY
 import com.tencent.bkrepo.common.artifact.constant.REPO_NAME
+import com.tencent.bkrepo.common.artifact.constant.REPO_RATE_LIMIT_KEY
 import com.tencent.bkrepo.common.artifact.exception.RepoNotFoundException
 import com.tencent.bkrepo.common.artifact.pojo.RepositoryCategory
 import com.tencent.bkrepo.common.artifact.pojo.RepositoryType
 import com.tencent.bkrepo.common.artifact.repository.composite.CompositeRepository
 import com.tencent.bkrepo.common.artifact.repository.core.ArtifactRepository
 import com.tencent.bkrepo.common.security.http.core.HttpAuthSecurity
 import com.tencent.bkrepo.common.service.util.HttpContextHolder
+import com.tencent.bkrepo.common.storage.core.config.RateLimitProperties
 import com.tencent.bkrepo.repository.api.NodeClient
 import com.tencent.bkrepo.repository.api.RepositoryClient
 import com.tencent.bkrepo.repository.pojo.node.NodeDetail
 import com.tencent.bkrepo.repository.pojo.repo.RepositoryDetail
 import org.springframework.beans.factory.ObjectProvider
+import org.springframework.util.unit.DataSize
 import org.springframework.web.servlet.HandlerMapping
 import java.util.concurrent.TimeUnit
 import javax.servlet.http.HttpServletRequest
@@ -85,6 +88,10 @@ class ArtifactContextHolder(
         private lateinit var nodeClient: NodeClient
         private lateinit var httpAuthSecurity: ObjectProvider<HttpAuthSecurity>
 
+
+        private const val RECEIVE_RATE_LIMIT_OF_REPO = "receiveRateLimit"
+        private const val RESPONSE_RATE_LIMIT_OF_REPO = "responseRateLimit"
+
         private val artifactConfigurerMap = mutableMapOf<RepositoryType, ArtifactConfigurer>()
         private val repositoryDetailCache = CacheBuilder.newBuilder()
             .maximumSize(1000)
@@ -277,6 +284,34 @@ class ArtifactContextHolder(
             nodeDetail?.let { request.setAttribute(attrKey, nodeDetail) }
             return nodeDetail
         }
+
+        /**
+         * 获取仓库级别的限速配置
+         */
+        fun getRateLimitOfRepo(): RateLimitProperties {
+            val request = HttpContextHolder.getRequestOrNull() ?: return RateLimitProperties()
+            val repoRateLimitAttribute = request.getAttribute(REPO_RATE_LIMIT_KEY)
+            if (repoRateLimitAttribute != null) {
+                require(repoRateLimitAttribute is RateLimitProperties)
+                return repoRateLimitAttribute
+            }
+            val repo = getRepoDetail() ?: return RateLimitProperties()
+            val receiveRateLimit = convertToDataSize(repo.configuration.getStringSetting(RECEIVE_RATE_LIMIT_OF_REPO))
+            val responseRateLimit = convertToDataSize(repo.configuration.getStringSetting(RESPONSE_RATE_LIMIT_OF_REPO))
+            val rateLimitProperties = RateLimitProperties(receiveRateLimit, responseRateLimit)
+            request.setAttribute(REPO_RATE_LIMIT_KEY, rateLimitProperties)
+            return rateLimitProperties
+
+        }
+
+        private fun convertToDataSize(dataSize: String?): DataSize {
+            if (dataSize.isNullOrEmpty()) return DataSize.ofBytes(-1)
+            return try {
+                DataSize.parse(dataSize)
+            } catch (e: Exception) {
+                DataSize.ofBytes(-1)
+            }
+        }
     }
 
     /**

diff --git a/...tlin/com/tencent/bkrepo/common/artifact/resolve/response/DefaultArtifactResourceWriter.kt b/...tlin/com/tencent/bkrepo/common/artifact/resolve/response/DefaultArtifactResourceWriter.kt
@@ -31,11 +31,13 @@ import com.tencent.bkrepo.common.api.constant.HttpHeaders
 import com.tencent.bkrepo.common.api.constant.HttpStatus
 import com.tencent.bkrepo.common.api.constant.MediaTypes
 import com.tencent.bkrepo.common.api.constant.StringPool
+import com.tencent.bkrepo.common.api.exception.TooManyRequestsException
 import com.tencent.bkrepo.common.artifact.constant.X_CHECKSUM_MD5
 import com.tencent.bkrepo.common.artifact.constant.X_CHECKSUM_SHA256
 import com.tencent.bkrepo.common.artifact.exception.ArtifactResponseException
 import com.tencent.bkrepo.common.artifact.metrics.RecordAbleInputStream
 import com.tencent.bkrepo.common.artifact.path.PathUtils
+import com.tencent.bkrepo.common.artifact.repository.context.ArtifactContextHolder
 import com.tencent.bkrepo.common.artifact.stream.ArtifactInputStream
 import com.tencent.bkrepo.common.artifact.stream.Range
 import com.tencent.bkrepo.common.artifact.stream.STREAM_BUFFER_SIZE
@@ -50,6 +52,7 @@ import com.tencent.bkrepo.common.storage.monitor.Throughput
 import com.tencent.bkrepo.common.storage.monitor.measureThroughput
 import com.tencent.bkrepo.repository.pojo.node.NodeDetail
 import org.springframework.http.HttpMethod
+import org.springframework.util.unit.DataSize
 import java.io.IOException
 import java.time.LocalDateTime
 import java.time.ZoneOffset
@@ -68,6 +71,7 @@ open class DefaultArtifactResourceWriter(
 
     @Throws(ArtifactResponseException::class)
     override fun write(resource: ArtifactResource): Throughput {
+        responseRateLimitCheck()
         return if (resource.containsMultiArtifact()) {
             writeMultiArtifact(resource)
         } else {
@@ -185,7 +189,7 @@ open class DefaultArtifactResourceWriter(
         val recordAbleInputStream = RecordAbleInputStream(inputStream)
         try {
             return measureThroughput {
-                recordAbleInputStream.rateLimit(storageProperties.response.rateLimit.toBytes()).use {
+                recordAbleInputStream.rateLimit(responseRateLimitWrapper(storageProperties.response.rateLimit)).use {
                     it.copyTo(
                         out = response.outputStream,
                         bufferSize = getBufferSize(inputStream.range.length.toInt())
@@ -222,7 +226,9 @@ open class DefaultArtifactResourceWriter(
                     resource.artifactMap.forEach { (name, inputStream) ->
                         val recordAbleInputStream = RecordAbleInputStream(inputStream)
                         zipOutput.putNextEntry(generateZipEntry(name, inputStream))
-                        recordAbleInputStream.rateLimit(storageProperties.response.rateLimit.toBytes()).use {
+                        recordAbleInputStream.rateLimit(
+                            responseRateLimitWrapper(storageProperties.response.rateLimit)
+                        ).use {
                             it.copyTo(
                                 out = zipOutput,
                                 bufferSize = getBufferSize(inputStream.range.length.toInt())
@@ -242,6 +248,31 @@ open class DefaultArtifactResourceWriter(
         }
     }
 
+    /**
+     * 将仓库级别的限速配置导入
+     * 当同时存在全局限速配置以及仓库级别限速配置时，以仓库级别配置优先
+     */
+    private fun responseRateLimitWrapper(rateLimit: DataSize): Long {
+        val rateLimitOfRepo = ArtifactContextHolder.getRateLimitOfRepo()
+        if (rateLimitOfRepo.responseRateLimit != DataSize.ofBytes(-1)) {
+            return rateLimitOfRepo.responseRateLimit.toBytes()
+        }
+        return rateLimit.toBytes()
+    }
+
+    /**
+     * 当仓库配置下载限速小于等于最低限速时则直接将请求断开, 避免占用过多连接
+     */
+    private fun responseRateLimitCheck() {
+        val rateLimitOfRepo = ArtifactContextHolder.getRateLimitOfRepo()
+        if (rateLimitOfRepo.responseRateLimit != DataSize.ofBytes(-1) &&
+            rateLimitOfRepo.responseRateLimit <= storageProperties.response.circuitBreakerThreshold) {
+            throw TooManyRequestsException(
+                "The circuit breaker is activated when too many download requests are made to the service!"
+            )
+        }
+    }
+
     /**
      * 判断charset,一些媒体类型设置了charset会影响其表现，如application/vnd.android.package-archive
      * */

diff --git a/...ice/src/main/kotlin/com/tencent/bkrepo/common/service/exception/GlobalExceptionHandler.kt b/...ice/src/main/kotlin/com/tencent/bkrepo/common/service/exception/GlobalExceptionHandler.kt
@@ -34,6 +34,7 @@ package com.tencent.bkrepo.common.service.exception
 import com.fasterxml.jackson.module.kotlin.MissingKotlinParameterException
 import com.tencent.bkrepo.common.api.constant.HttpStatus
 import com.tencent.bkrepo.common.api.exception.ErrorCodeException
+import com.tencent.bkrepo.common.api.exception.TooManyRequestsException
 import com.tencent.bkrepo.common.api.message.CommonMessageCode
 import com.tencent.bkrepo.common.api.pojo.Response
 import org.springframework.core.Ordered
@@ -137,6 +138,11 @@ class GlobalExceptionHandler : AbstractExceptionHandler() {
         return deferredResult
     }
 
+    @ExceptionHandler(TooManyRequestsException::class)
+    fun handleException(exception: TooManyRequestsException): Response<Void> {
+        return response(exception)
+    }
+
     @ExceptionHandler(Exception::class)
     fun handleException(exception: Exception): Response<Void> {
         return response(exception)

diff --git a/src/backend/common/common-service/src/main/resources/i18n/messages_en.properties b/src/backend/common/common-service/src/main/resources/i18n/messages_en.properties
@@ -51,3 +51,4 @@ permission.project.denied = user[{0}] does not have [{1}] permission in project[
 permission.repo.denied = user[{0}] does not have [{1}] permission in project[{2}] repo[{3}]
 operation.cross-cluster.not-allowed=Cross location operation is not allowed
 system.media-type.unacceptable=Unacceptable Media Type
+too.many.requests=Too Many Requests: {0}
diff --git a/src/backend/common/common-service/src/main/resources/i18n/messages_zh_CN.properties b/src/backend/common/common-service/src/main/resources/i18n/messages_zh_CN.properties
@@ -51,3 +51,4 @@ permission.project.denied=用户[{0}]没有[{1}]权限在[{2}]项目上
 permission.repo.denied = 用户[{0}]没有[{1}]权限在[{2}]项目的[{3}]仓库上
 operation.cross-cluster.not-allowed=不允许跨地点操作
 system.media-type.unacceptable=不接受的Media Type
+too.many.requests=请求过多: {0}
diff --git a/src/backend/common/common-service/src/main/resources/i18n/messages_zh_TW.properties b/src/backend/common/common-service/src/main/resources/i18n/messages_zh_TW.properties
@@ -51,3 +51,4 @@ permission.project.denied=用戶[{0}]沒有[{1}]權限在[{2}]項目上
 permission.repo.denied = 用戶[{0}]沒有[{1}]權限在[{2}]項目的[{3}]倉庫上
 operation.cross-cluster.not-allowed=不允許跨地點操作
 system.media-type.unacceptable=不接受的Media Type
+too.many.requests=請求過多: {0}
diff --git a/...vice/src/main/kotlin/com/tencent/bkrepo/common/storage/core/config/RateLimitProperties.kt b/...vice/src/main/kotlin/com/tencent/bkrepo/common/storage/core/config/RateLimitProperties.kt
@@ -0,0 +1,44 @@
+/*
+ * Tencent is pleased to support the open source community by making BK-CI 蓝鲸持续集成平台 available.
+ *
+ * Copyright (C) 2022 THL A29 Limited, a Tencent company.  All rights reserved.
+ *
+ * BK-CI 蓝鲸持续集成平台 is licensed under the MIT license.
+ *
+ * A copy of the MIT License is included in this file.
+ *
+ *
+ * Terms of the MIT License:
+ * ---------------------------------------------------
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+ * documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
+ * permit persons to whom the Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all copies or substantial portions of
+ * the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
+ * LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
+ * NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+ * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+package com.tencent.bkrepo.common.storage.core.config
+
+import org.springframework.util.unit.DataSize
+
+/**
+ * 上传、下载限速配置
+ */
+data class RateLimitProperties(
+    /**
+     * 每秒接收数据量
+     */
+    var receiveRateLimit: DataSize = DataSize.ofBytes(-1),
+    /**
+     * 每秒传输数据量
+     */
+    var responseRateLimit: DataSize = DataSize.ofBytes(-1)
+)
diff --git a/...rvice/src/main/kotlin/com/tencent/bkrepo/common/storage/core/config/ResponseProperties.kt b/...rvice/src/main/kotlin/com/tencent/bkrepo/common/storage/core/config/ResponseProperties.kt
@@ -54,5 +54,9 @@ data class ResponseProperties(
     /**
      * 二进制媒体类型，不指定编码。
      * */
-    var binaryMediaTypes: Set<String> = emptySet()
+    var binaryMediaTypes: Set<String> = emptySet(),
+    /**
+     * 限速熔断阈值，当仓库配置的rateLimit小于等于限速熔断阈值时则直接将请求断开
+     */
+    var circuitBreakerThreshold: DataSize = DataSize.ofKilobytes(1),
 )