feat:增加websocket消息大小限制配置项 #2829

src/backend/websocket/biz-websocket/src/main/kotlin/com/tencent/bkrepo/websocket/config/WebSocketMetrics.kt

@@ -0,0 +1,51 @@
+/*
+ * Tencent is pleased to support the open source community by making BK-CI 蓝鲸持续集成平台 available.
+ *
+ * Copyright (C) 2019 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.websocket.config
+
+import io.micrometer.core.instrument.Gauge
+import io.micrometer.core.instrument.MeterRegistry
+import io.micrometer.core.instrument.binder.MeterBinder
+import org.springframework.stereotype.Component
+import java.util.concurrent.atomic.AtomicInteger
+
+@Component
+class WebSocketMetrics: MeterBinder {
+
+    var connectionCount: AtomicInteger = AtomicInteger(0)
+
+    override fun bindTo(registry: MeterRegistry) {
+        Gauge.builder(WEBSOCKET_CONNECTION_COUNT, connectionCount) { it.get().toDouble() }
+            .description(WEBSOCKET_CONNECTION_COUNT_DESC)
+            .register(registry)
+    }
+
+    companion object {
+        private const val WEBSOCKET_CONNECTION_COUNT = "websocket_connection_count"
+        private const val WEBSOCKET_CONNECTION_COUNT_DESC = "websocket connection count"
+    }
+}

src/backend/websocket/biz-websocket/src/main/kotlin/com/tencent/bkrepo/websocket/config/WebSocketProperties.kt

@@ -33,5 +33,8 @@ import org.springframework.boot.context.properties.ConfigurationProperties
 data class WebSocketProperties(
     var cacheLimit: Int = 3600,
     var minThread: Int = 8,
-    var transfer: Boolean = false
+    var transfer: Boolean = false,
+    var messageSizeLimit: Int = 8*1024*1024,
+    var sendTimeLimit: Int = 10*1000,
+    var sendBufferSizeLimit: Int = 1024*1024,
 )

src/backend/websocket/biz-websocket/src/main/kotlin/com/tencent/bkrepo/websocket/config/WebsocketConfiguration.kt

@@ -56,6 +56,7 @@ class WebsocketConfiguration(
     private val websocketService: WebsocketService,
     private val jwtAuthProperties: JwtAuthProperties,
     private val authenticationManager: AuthenticationManager,
+    private val webSocketMetrics: WebSocketMetrics
 ) : WebSocketMessageBrokerConfigurer {
 
     override fun configureMessageBroker(config: MessageBrokerRegistry) {
@@ -94,6 +95,9 @@ class WebsocketConfiguration(
 
     override fun configureWebSocketTransport(registration: WebSocketTransportRegistration) {
         registration.addDecoratorFactory(wsHandlerDecoratorFactory())
+        registration.setMessageSizeLimit(webSocketProperties.messageSizeLimit)
+        registration.setSendTimeLimit(webSocketProperties.sendTimeLimit)
+        registration.setSendBufferSizeLimit(webSocketProperties.sendBufferSizeLimit)
         super.configureWebSocketTransport(registration)
     }
 
@@ -102,7 +106,8 @@ class WebsocketConfiguration(
         return SessionWebSocketHandlerDecoratorFactory(
             websocketService = websocketService,
             authenticationManager = authenticationManager,
-            jwtAuthProperties = jwtAuthProperties
+            jwtAuthProperties = jwtAuthProperties,
+            webSocketMetrics = webSocketMetrics
         )
     }
 

src/backend/websocket/biz-websocket/src/main/kotlin/com/tencent/bkrepo/websocket/handler/SessionHandler.kt

@@ -38,6 +38,7 @@ import com.tencent.bkrepo.common.security.exception.AuthenticationException
 import com.tencent.bkrepo.common.security.http.jwt.JwtAuthProperties
 import com.tencent.bkrepo.common.security.manager.AuthenticationManager
 import com.tencent.bkrepo.common.security.util.JwtUtils
+import com.tencent.bkrepo.websocket.config.WebSocketMetrics
 import com.tencent.bkrepo.websocket.constant.APP_ENDPOINT
 import com.tencent.bkrepo.websocket.constant.DESKTOP_ENDPOINT
 import com.tencent.bkrepo.websocket.constant.SESSION_ID
@@ -59,6 +60,7 @@ class SessionHandler(
     delegate: WebSocketHandler,
     private val websocketService: WebsocketService,
     private val authenticationManager: AuthenticationManager,
+    private val webSocketMetrics: WebSocketMetrics,
     jwtProperties: JwtAuthProperties
 ) : WebSocketHandlerDecorator(delegate) {
 
@@ -73,10 +75,10 @@ class SessionHandler(
         val sessionId = HostUtils.getRealSession(session.uri?.query)
         if (sessionId.isNullOrEmpty()) {
             logger.warn("connection closed can not find sessionId, $uri| ${session.remoteAddress}")
-            super.afterConnectionClosed(session, closeStatus)
+        } else {
+            websocketService.removeCacheSession(sessionId)
         }
-        websocketService.removeCacheSession(sessionId!!)
-
+        webSocketMetrics.connectionCount.decrementAndGet()
         super.afterConnectionClosed(session, closeStatus)
     }
 
@@ -123,6 +125,7 @@ class SessionHandler(
         websocketService.addCacheSession(sessionId!!)
         session.attributes[SESSION_ID] = sessionId
         logger.info("connection success: |$sessionId| $uri | $remoteId | ${session.attributes[USER_KEY]} ")
+        webSocketMetrics.connectionCount.incrementAndGet()
         super.afterConnectionEstablished(session)
     }
 

src/backend/websocket/biz-websocket/src/main/kotlin/com/tencent/bkrepo/websocket/handler/SessionWebSocketHandlerDecoratorFactory.kt

@@ -29,6 +29,7 @@ package com.tencent.bkrepo.websocket.handler
 
 import com.tencent.bkrepo.common.security.http.jwt.JwtAuthProperties
 import com.tencent.bkrepo.common.security.manager.AuthenticationManager
+import com.tencent.bkrepo.websocket.config.WebSocketMetrics
 import com.tencent.bkrepo.websocket.service.WebsocketService
 import org.springframework.web.socket.WebSocketHandler
 import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory
@@ -37,9 +38,10 @@ class SessionWebSocketHandlerDecoratorFactory (
     private val websocketService: WebsocketService,
     private val authenticationManager: AuthenticationManager,
     private val jwtAuthProperties: JwtAuthProperties,
+    private val webSocketMetrics: WebSocketMetrics
 ) : WebSocketHandlerDecoratorFactory {
 
     override fun decorate(handler: WebSocketHandler): WebSocketHandler {
-        return SessionHandler(handler, websocketService, authenticationManager, jwtAuthProperties)
+        return SessionHandler(handler, websocketService, authenticationManager, webSocketMetrics, jwtAuthProperties)
     }
 }