Springboot 整合 WebSocket

├─kotlin
│  └─ski
│      └─mashiro
│          │  WsDemoApplication.kt
│          ├─annotation
│          │      Slf4j.kt
│          ├─config
│          │      WebMvcConfig.kt
│          │      WebSocketConfig.kt
│          ├─controller
│          │      WebSocketController.kt
│          └─ws
│                  WebSocketServer.kt
└─resources
        application.yml

dependencies {
    implementation("org.springframework.boot:spring-boot-starter-web")
    // https://mvnrepository.com/artifact/org.springframework.boot/spring-boot-starter-websocket
    implementation("org.springframework.boot:spring-boot-starter-websocket")
    implementation("com.fasterxml.jackson.module:jackson-module-kotlin")
    implementation("com.github.yitter:yitter-idgenerator:1.0.6")
    implementation("org.jetbrains.kotlin:kotlin-reflect")
    testImplementation("org.springframework.boot:spring-boot-starter-test")
}

@Configuration
class WebMvcConfig: WebMvcConfigurer {
    override fun addCorsMappings(registry: CorsRegistry) {
        registry.addMapping("/**")
            .allowedOriginPatterns("*")
            .allowedHeaders("*")
            .allowedMethods("*")
            .allowCredentials(true)
    }
}

@Configuration
class WebSocketConfig {
    @Bean
    fun serverEndpointExporter(): ServerEndpointExporter {
        return ServerEndpointExporter()
    }
}

@Service
@Slf4j
@ServerEndpoint("/api/ws/{userId}")
class WebSocketServer {
    private var session: Session? = null
    private var userId: Long? = null

    // 里面的相当于Java的static变量
    companion object {
        private val webSockets = CopyOnWriteArraySet< WebSocketServer>()
        private val sessionPool = ConcurrentHashMap<Long, Session>()
    }

    @OnOpen
    fun onOpen(session: Session, @PathParam("userId") userId: Long) {
        this.session = session
        this.userId = userId
        webSockets.add(this)
        sessionPool[userId] = session
        logger.info("【websocket消息】有新的连接, userId = $userId, 总数为: ${webSockets.size}")
    }

    @OnClose
    fun onClose() {
        webSockets.remove(this)
        sessionPool.remove(userId)
        logger.info("【websocket消息】有连接断开, userId = $userId, 总数为: ${webSockets.size}")
    }

    @OnMessage
    fun onMessage(message: String) {
        logger.info("【websocket消息】收到客户端消息: $message");
    }

    @OnError
    fun onError(session: Session, error: Throwable) {
        logger.error("【websocket消息】错误, userId: $userId, 原因: ${error.message}")
    }

    fun broadcast(msg: String) {
        logger.info("【websocket消息】广播消息: $msg")
        webSockets.forEach {
            try {
                if (it.session!!.isOpen) {
                    it.session!!.asyncRemote.sendText(msg)
                }
            } catch (e: Exception) {
                e.printStackTrace()
            }
        }
    }

    fun sendMsg2User(userId: Long, msg: String) {
        val session = sessionPool[userId] ?: run {
            logger.error("userId: $userId 不在sessionPool中")
            return
        }
        try {
            session.userPrincipal
            if (session.isOpen) {
                logger.info("【websocket消息】单点消息: $msg")
                session.asyncRemote.sendText(msg)
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    fun sendMsg2Users(userIds: LongArray, msg: String) {
        userIds.forEach {userId ->
            sendMsg2User(userId, msg)
        }
    }
}

@RestController
@RequestMapping("/api/rest/ws")
class WebSocketController(
    private val webSocketServer: WebSocketServer,
) {

    @PostMapping("/broadcast")
    fun broadcast(@RequestBody msg: String) {
        webSocketServer.broadcast(msg)
    }

    @PostMapping("/sendMsg/user/{userId}")
    fun sendMsg2User(@PathVariable userId: Long, @RequestBody msg: String) {
        webSocketServer.sendMsg2User(userId, msg)
    }

    @PostMapping("/sendMsg/users")
    fun sendMsg2Users(@RequestParam userIds: LongArray, @RequestBody msg: String) {
        webSocketServer.sendMsg2Users(userIds, msg)
    }

    @GetMapping("/userId")
    fun getUserId(): String {
        return YitIdHelper.nextId().toString()
    }
}

<template>
  <h3>UserId: {{ userId }}</h3>
  <div class="operation">
    <div>
      <button type="button" @click="getUserId">获取Id</button>
    </div>
    <div>
      <button type="button" @click="connect2Server">连接ws</button>
      <button type="button" @click="disconnect2Server">断开ws</button>
      <span>状态: {{ connectStatus }}</span>
    </div>
    <div>
      <button type="button" @click="broadcast">发送广播</button>
      <span>Content: </span><input type="text" v-model="broadcastContent">
    </div>
    <div>
      <button type="button" @click="privateChat">指定UserId发送</button>
      <span>Receiver's userId: </span><input type="text" v-model="privateChatUserId">
      <span>Content: </span><input type="text" v-model="privateChatContent">
    </div>
  </div>
  <hr/>
  <div>
    <div v-for="(msg, idx) in msgQueue" :key="idx">
      {{ msg }}
    </div>
  </div>
</template>

<script setup>
import axios from "axios";
import {reactive, ref} from "vue";

const uri = "127.0.0.1:8080"
const httpUri = `http://${uri}`
const wsEndpoint = `ws://${uri}/api/ws/`

const userId = ref(0);
const connectStatus = ref("断开")
const broadcastContent = ref("")
const privateChatUserId = ref("")
const privateChatContent = ref("")
const msgQueue = reactive([])
let webSocket;

const getUserId = () => axios.get(`${httpUri}/api/rest/ws/userId`)
    .then(resp => {
      userId.value = resp.data
    })
    .catch(e => console.log(e));

function connect2Server() {
  // 详细可看：https://cloud.tencent.com/developer/article/1887095
  webSocket = new WebSocket(wsEndpoint + userId.value);
  webSocket.onopen = onOpen
  webSocket.onclose = onClose
  webSocket.onerror = onErr
  webSocket.onmessage = onMessage
}

function disconnect2Server() {
  webSocket.close()
}

const send2Server = (msg) => {
  webSocket.send(msg);
}

const onOpen = () => {
  sysBroadcast(`【广播】userId: ${userId.value} 连接到Server`)
  connectStatus.value = "连接"
}

const onClose = (e) => {
  sysBroadcast(`【广播】userId: ${userId.value} 断开连接, ${e.code}`)
  connectStatus.value = "断开"
}

const onMessage = (e) => {
  // data的数据是像这样： {"msg":"【广播】userId: 472661222139909 连接到Server"}
  // 似乎是通过eval将json转为对象，我的前端半桶水，欢迎斧正
  const data = eval("(" + e.data + ")")
  msgQueue.push(data.msg)
}

const onErr = () => console.log("WebSocket连接发生错误")

const broadcast = () => {
  axios.post(`${httpUri}/api/rest/ws/broadcast`, {
    msg: "【广播】" + userId.value + ": " + broadcastContent.value
  })
  broadcastContent.value = ""
}

const sysBroadcast = (msg) => axios.post(`${httpUri}/api/rest/ws/broadcast`, {
  msg: msg
})

const privateChat = () => {
  axios.post(`${httpUri}/api/rest/ws/sendMsg/user/${privateChatUserId.value}`, {
    msg: "【私聊】" + userId.value + ": " + privateChatContent.value
  })
  msgQueue.push(`【私聊】To ${privateChatUserId.value}: ${privateChatContent.value}`)
  privateChatUserId.value = ""
  privateChatContent.value = ""
}
</script>