dalbodeule/chibot-chzzk-bot
1
0
Fork 0
mirror of https://github.com/dalbodeule/chibot-chzzk-bot.git synced 2025-06-08 23:08:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor WebSocket handlers and add ACK-based message flow

Browse Source 
Consolidated coroutine scopes into `songListScope` and `timerScope` for better management across WebSocket routes. Introduced ACK (acknowledgment) handling for reliable message delivery with retries and timeouts. Updated session handling for multiple WebSocket routes to improve code maintainability and consistency.
This commit is contained in:
dalbodeule 2025-04-24 16:56:49 +09:00
parent d07cdb6ae8
commit 8230762053
No known key found for this signature in database
GPG Key ID: EFA860D069C9FA65
5 changed files with 190 additions and 161 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
common/src/main/kotlin/space/mori/chzzk_bot/common/events/TimerEvents.kt
View File

@ -5,7 +5,8 @@ enum class TimerType(var value: Int) {
TIMER(1), TIMER(1),
REMOVE(2), REMOVE(2),
STREAM_OFF(50) STREAM_OFF(50),
ACK(51)
} }
class TimerEvent( class TimerEvent(

2
webserver/src/main/kotlin/space/mori/chzzk_bot/webserver/Main.kt
View File

@ -22,11 +22,9 @@ import io.ktor.server.websocket.*
import kotlinx.coroutines.delay import kotlinx.coroutines.delay
import kotlinx.serialization.Serializable import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json import kotlinx.serialization.json.Json
import org.slf4j.LoggerFactory
import space.mori.chzzk_bot.common.services.UserService import space.mori.chzzk_bot.common.services.UserService
import space.mori.chzzk_bot.webserver.routes.* import space.mori.chzzk_bot.webserver.routes.*
import space.mori.chzzk_bot.webserver.utils.DiscordRatelimits import space.mori.chzzk_bot.webserver.utils.DiscordRatelimits
import wsSongListRoutes
import java.math.BigInteger import java.math.BigInteger
import java.security.SecureRandom import java.security.SecureRandom
import java.time.Duration import java.time.Duration

118
webserver/src/main/kotlin/space/mori/chzzk_bot/webserver/routes/WSSongListRoutes.kt
View File

@ -1,3 +1,5 @@
package space.mori.chzzk_bot.webserver.routes
import io.ktor.client.plugins.websocket.WebSocketException import io.ktor.client.plugins.websocket.WebSocketException
import io.ktor.server.application.ApplicationStopped import io.ktor.server.application.ApplicationStopped
import io.ktor.server.routing.* import io.ktor.server.routing.*
@ -6,7 +8,7 @@ import io.ktor.server.websocket.*
import io.ktor.util.logging.Logger import io.ktor.util.logging.Logger
import io.ktor.utils.io.CancellationException import io.ktor.utils.io.CancellationException
import io.ktor.websocket.* import io.ktor.websocket.*
import io.ktor.websocket.Frame.* import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.SupervisorJob
@ -14,8 +16,10 @@ import kotlinx.coroutines.cancel
import kotlinx.coroutines.channels.ClosedReceiveChannelException import kotlinx.coroutines.channels.ClosedReceiveChannelException
import kotlinx.coroutines.delay import kotlinx.coroutines.delay
import kotlinx.coroutines.launch import kotlinx.coroutines.launch
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withTimeoutOrNull
import kotlinx.io.IOException import kotlinx.io.IOException
import kotlinx.serialization.Serializable import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json import kotlinx.serialization.json.Json
@ -30,35 +34,29 @@ import space.mori.chzzk_bot.common.services.UserService
import space.mori.chzzk_bot.common.utils.YoutubeVideo import space.mori.chzzk_bot.common.utils.YoutubeVideo
import space.mori.chzzk_bot.common.utils.getYoutubeVideo import space.mori.chzzk_bot.common.utils.getYoutubeVideo
import space.mori.chzzk_bot.webserver.UserSession import space.mori.chzzk_bot.webserver.UserSession
import space.mori.chzzk_bot.webserver.routes.SongResponse
import space.mori.chzzk_bot.webserver.routes.toSerializable
import space.mori.chzzk_bot.webserver.utils.CurrentSong import space.mori.chzzk_bot.webserver.utils.CurrentSong
import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.ConcurrentHashMap
val routeScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) val songListScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
fun Routing.wsSongListRoutes() { fun Routing.wsSongListRoutes() {
val sessions = ConcurrentHashMap<String, WebSocketServerSession>() val sessions = ConcurrentHashMap<String, WebSocketServerSession>()
val status = ConcurrentHashMap<String, SongType>() val status = ConcurrentHashMap<String, SongType>()
val logger = LoggerFactory.getLogger("WSSongListRoutes") val logger = LoggerFactory.getLogger("WSSongListRoutes")
val dispatcher: CoroutinesEventBus by inject(CoroutinesEventBus::class.java) val dispatcher: CoroutinesEventBus by inject(CoroutinesEventBus::class.java)
// 세션 관련 작업을 위한 Mutex 추가
val sessionMutex = Mutex() val sessionMutex = Mutex()
environment.monitor.subscribe(ApplicationStopped) { environment.monitor.subscribe(ApplicationStopped) {
routeScope.cancel() songListScope.cancel()
} }
suspend fun addSession(uid: String, session: WebSocketServerSession) { suspend fun addSession(uid: String, session: WebSocketServerSession) {
val oldSession = sessionMutex.withLock { val oldSession = sessionMutex.withLock {
val old = sessions[uid] val old = sessions[uid]
sessions[uid] = session sessions[uid] = session
old old
} }
if(oldSession != null) { if(oldSession != null) {
routeScope.launch { songListScope.launch {
try { try {
oldSession.close(CloseReason( oldSession.close(CloseReason(
CloseReason.Codes.VIOLATED_POLICY, "Another session is already active.")) CloseReason.Codes.VIOLATED_POLICY, "Another session is already active."))
@ -69,87 +67,54 @@ fun Routing.wsSongListRoutes() {
} }
} }
} }
suspend fun removeSession(uid: String) { suspend fun removeSession(uid: String) {
sessionMutex.withLock { sessionMutex.withLock {
sessions.remove(uid) sessions.remove(uid)
} }
} }
suspend fun waitForAck(ws: WebSocketServerSession, expectedType: Int): Boolean { val ackMap = ConcurrentHashMap<String, CompletableDeferred<Boolean>>()
try {
for (frame in ws.incoming) {
if (frame is Text) {
val message = frame.readText()
if (message == "ping") {
continue // Keep the loop running if a ping is received
}
val data = Json.decodeFromString<SongRequest>(message)
if (data.type == SongType.ACK.value) {
return true // ACK received
}
}
}
} catch (e: Exception) {
logger.warn("Error waiting for ACK: ${e.message}")
}
return false // Return false if no ACK received
}
suspend fun sendWithRetry(uid: String, res: SongResponse, maxRetries: Int = 5, delayMillis: Long = 3000L) { suspend fun sendWithRetry(uid: String, res: SongResponse, maxRetries: Int = 5, delayMillis: Long = 3000L) {
var attempt = 0 var attempt = 0
var sentSuccessfully = false while (attempt < maxRetries) {
val ws: WebSocketServerSession? = sessionMutex.withLock { sessions[uid] }
while (attempt < maxRetries && !sentSuccessfully) { if (ws == null) {
val ws: WebSocketServerSession? = sessionMutex.withLock { sessions[uid] } ?: run {
logger.debug("No active session for $uid. Retrying in $delayMillis ms.") logger.debug("No active session for $uid. Retrying in $delayMillis ms.")
delay(delayMillis) delay(delayMillis)
attempt++ attempt++
continue
null
} }
if (ws == null) continue
try { try {
// 메시지 전송 시도
ws.sendSerialized(res) ws.sendSerialized(res)
logger.debug("Message sent successfully to $uid on attempt $attempt") logger.debug("Message sent successfully to $uid on attempt $attempt")
val ackDeferred = CompletableDeferred<Boolean>()
// ACK 대기 ackMap[res.uid] = ackDeferred
val ackReceived = waitForAck(ws, res.type) val ackReceived = withTimeoutOrNull(delayMillis) { ackDeferred.await() } ?: false
if (ackReceived) { if (ackReceived) {
logger.debug("ACK received for message to $uid on attempt $attempt.") logger.debug("ACK received for message to $uid on attempt $attempt.")
sentSuccessfully = true return
} else { } else {
logger.warn("ACK not received for message to $uid on attempt $attempt.") logger.warn("ACK not received for message to $uid on attempt $attempt.")
attempt++ attempt++
} }
} catch (e: CancellationException) { } catch (e: CancellationException) {
// 코루틴 취소는 다시 throw
throw e throw e
} catch (e: Exception) { } catch (e: Exception) {
attempt++ attempt++
logger.warn("Failed to send message to $uid on attempt $attempt: ${e.message}") logger.warn("Failed to send message to $uid on attempt $attempt: ${e.message}")
if (e is WebSocketException || e is IOException) { if (e is WebSocketException || e is IOException) {
logger.warn("Connection issue detected, session may be invalid") logger.warn("Connection issue detected, session may be invalid")
// 연결 문제로 보이면 세션을 제거할 수도 있음
removeSession(uid) removeSession(uid)
} }
} }
if (attempt < maxRetries) {
if (!sentSuccessfully && attempt < maxRetries) {
delay(delayMillis) delay(delayMillis)
} }
} }
logger.error("Failed to send message to $uid after $maxRetries attempts.")
if (!sentSuccessfully) {
logger.error("Failed to send message to $uid after $maxRetries attempts.")
}
} }
webSocket("/songlist") { webSocket("/songlist") {
val session = call.sessions.get<UserSession>() val session = call.sessions.get<UserSession>()
val user = session?.id?.let { UserService.getUser(it) } val user = session?.id?.let { UserService.getUser(it) }
@ -157,13 +122,10 @@ fun Routing.wsSongListRoutes() {
close(CloseReason(CloseReason.Codes.CANNOT_ACCEPT, "Invalid SID")) close(CloseReason(CloseReason.Codes.CANNOT_ACCEPT, "Invalid SID"))
return@webSocket return@webSocket
} }
val uid = user.token val uid = user.token
addSession(uid, this) addSession(uid, this)
if (status[uid] == SongType.STREAM_OFF) { if (status[uid] == SongType.STREAM_OFF) {
routeScope.launch { songListScope.launch {
sendSerialized(SongResponse( sendSerialized(SongResponse(
SongType.STREAM_OFF.value, SongType.STREAM_OFF.value,
uid, uid,
@ -175,23 +137,32 @@ fun Routing.wsSongListRoutes() {
removeSession(uid) removeSession(uid)
} }
try { songListScope.launch {
for (frame in incoming) { for (frame in incoming) {
when (frame) { when (frame) {
is Text -> { is Frame.Text -> {
val text = frame.readText() val text = frame.readText()
if (text.trim() == "ping") { if (text.trim() == "ping") {
send("pong") send("pong")
} else { } else {
val data = Json.decodeFromString<SongRequest>(text) val data = Json.decodeFromString<SongRequest>(text)
// Handle song requests if (data.type == SongType.ACK.value) {
handleSongRequest(data, user, dispatcher, logger) ackMap[data.uid]?.complete(true)
ackMap.remove(data.uid)
} else {
handleSongRequest(data, user, dispatcher, logger)
}
} }
} }
is Ping -> send(Pong(frame.data)) is Frame.Ping -> send(Frame.Pong(frame.data))
else -> "" else -> {}
} }
} }
}
try {
// Keep the connection alive
suspendCancellableCoroutine<Unit> {}
} catch (e: ClosedReceiveChannelException) { } catch (e: ClosedReceiveChannelException) {
logger.error("WebSocket connection closed: ${e.message}") logger.error("WebSocket connection closed: ${e.message}")
} catch(e: Exception) { } catch(e: Exception) {
@ -203,7 +174,7 @@ fun Routing.wsSongListRoutes() {
dispatcher.subscribe(SongEvent::class) { dispatcher.subscribe(SongEvent::class) {
logger.debug("SongEvent: {} / {} {}", it.uid, it.type, it.current?.name) logger.debug("SongEvent: {} / {} {}", it.uid, it.type, it.current?.name)
routeScope.launch { songListScope.launch {
try { try {
val user = UserService.getUser(it.uid) val user = UserService.getUser(it.uid)
if (user != null) { if (user != null) {
@ -223,10 +194,9 @@ fun Routing.wsSongListRoutes() {
} }
} }
} }
dispatcher.subscribe(TimerEvent::class) { dispatcher.subscribe(TimerEvent::class) {
if (it.type == TimerType.STREAM_OFF) { if (it.type == TimerType.STREAM_OFF) {
routeScope.launch { songListScope.launch {
try { try {
val user = UserService.getUser(it.uid) val user = UserService.getUser(it.uid)
if (user != null) { if (user != null) {
@ -247,10 +217,8 @@ fun Routing.wsSongListRoutes() {
} }
} }
} }
// 노래 처리를 위한 Mutex 추가 // 노래 처리를 위한 Mutex 추가
private val songMutex = Mutex() private val songMutex = Mutex()
suspend fun handleSongRequest( suspend fun handleSongRequest(
data: SongRequest, data: SongRequest,
user: User, user: User,
@ -261,14 +229,13 @@ suspend fun handleSongRequest(
if (data.maxUserLimit != null && data.maxUserLimit > 0) SongConfigService.updatePersonalLimit(user, data.maxUserLimit) if (data.maxUserLimit != null && data.maxUserLimit > 0) SongConfigService.updatePersonalLimit(user, data.maxUserLimit)
if (data.isStreamerOnly != null) SongConfigService.updateStreamerOnly(user, data.isStreamerOnly) if (data.isStreamerOnly != null) SongConfigService.updateStreamerOnly(user, data.isStreamerOnly)
if (data.isDisabled != null) SongConfigService.updateDisabled(user, data.isDisabled) if (data.isDisabled != null) SongConfigService.updateDisabled(user, data.isDisabled)
when (data.type) { when (data.type) {
SongType.ADD.value -> { SongType.ADD.value -> {
data.url?.let { url -> data.url?.let { url ->
try { try {
val youtubeVideo = getYoutubeVideo(url) val youtubeVideo = getYoutubeVideo(url)
if (youtubeVideo != null) { if (youtubeVideo != null) {
routeScope.launch { songListScope.launch {
songMutex.withLock { songMutex.withLock {
SongListService.saveSong( SongListService.saveSong(
user, user,
@ -298,7 +265,7 @@ suspend fun handleSongRequest(
} }
SongType.REMOVE.value -> { SongType.REMOVE.value -> {
data.url?.let { url -> data.url?.let { url ->
routeScope.launch { songListScope.launch {
songMutex.withLock { songMutex.withLock {
val songs = SongListService.getSong(user) val songs = SongListService.getSong(user)
val exactSong = songs.firstOrNull { it.url == url } val exactSong = songs.firstOrNull { it.url == url }
@ -320,17 +287,15 @@ suspend fun handleSongRequest(
} }
} }
SongType.NEXT.value -> { SongType.NEXT.value -> {
routeScope.launch { songListScope.launch {
songMutex.withLock { songMutex.withLock {
val songList = SongListService.getSong(user) val songList = SongListService.getSong(user)
var song: SongList? = null var song: SongList? = null
var youtubeVideo: YoutubeVideo? = null var youtubeVideo: YoutubeVideo? = null
if (songList.isNotEmpty()) { if (songList.isNotEmpty()) {
song = songList[0] song = songList[0]
SongListService.deleteSong(user, song.uid, song.name) SongListService.deleteSong(user, song.uid, song.name)
} }
song?.let { song?.let {
youtubeVideo = YoutubeVideo( youtubeVideo = YoutubeVideo(
song.url, song.url,
@ -347,7 +312,6 @@ suspend fun handleSongRequest(
youtubeVideo youtubeVideo
) )
) )
CurrentSong.setSong(user, youtubeVideo) CurrentSong.setSong(user, youtubeVideo)
} }
} }

72
webserver/src/main/kotlin/space/mori/chzzk_bot/webserver/routes/WSSongRoutes.kt
View File

@ -4,6 +4,7 @@ import io.ktor.server.application.ApplicationStopped
import io.ktor.server.routing.* import io.ktor.server.routing.*
import io.ktor.server.websocket.* import io.ktor.server.websocket.*
import io.ktor.websocket.* import io.ktor.websocket.*
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.SupervisorJob
@ -11,7 +12,9 @@ import kotlinx.coroutines.cancel
import kotlinx.coroutines.channels.ClosedReceiveChannelException import kotlinx.coroutines.channels.ClosedReceiveChannelException
import kotlinx.coroutines.delay import kotlinx.coroutines.delay
import kotlinx.coroutines.launch import kotlinx.coroutines.launch
import kotlinx.coroutines.withTimeoutOrNull
import kotlinx.serialization.Serializable import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import org.koin.java.KoinJavaComponent.inject import org.koin.java.KoinJavaComponent.inject
import org.slf4j.LoggerFactory import org.slf4j.LoggerFactory
import space.mori.chzzk_bot.common.events.* import space.mori.chzzk_bot.common.events.*
@ -20,21 +23,20 @@ import space.mori.chzzk_bot.common.utils.YoutubeVideo
import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.ConcurrentLinkedQueue import java.util.concurrent.ConcurrentLinkedQueue
val routeScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) val songScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
fun Routing.wsSongRoutes() { fun Routing.wsSongRoutes() {
environment.monitor.subscribe(ApplicationStopped) { environment.monitor.subscribe(ApplicationStopped) {
routeScope.cancel() songListScope.cancel()
} }
val sessions = ConcurrentHashMap<String, ConcurrentLinkedQueue<WebSocketServerSession>>() val sessions = ConcurrentHashMap<String, ConcurrentLinkedQueue<WebSocketServerSession>>()
val status = ConcurrentHashMap<String, SongType>() val status = ConcurrentHashMap<String, SongType>()
val logger = LoggerFactory.getLogger("WSSongRoutes") val logger = LoggerFactory.getLogger("WSSongRoutes")
val dispatcher: CoroutinesEventBus by inject(CoroutinesEventBus::class.java)
val ackMap = ConcurrentHashMap<String, ConcurrentHashMap<WebSocketServerSession, CompletableDeferred<Boolean>>>()
fun addSession(uid: String, session: WebSocketServerSession) { fun addSession(uid: String, session: WebSocketServerSession) {
sessions.computeIfAbsent(uid) { ConcurrentLinkedQueue() }.add(session) sessions.computeIfAbsent(uid) { ConcurrentLinkedQueue() }.add(session)
} }
fun removeSession(uid: String, session: WebSocketServerSession) { fun removeSession(uid: String, session: WebSocketServerSession) {
sessions[uid]?.remove(session) sessions[uid]?.remove(session)
if(sessions[uid]?.isEmpty() == true) { if(sessions[uid]?.isEmpty() == true) {
@ -51,27 +53,35 @@ fun Routing.wsSongRoutes() {
var attempt = 0 var attempt = 0
while (attempt < maxRetries) { while (attempt < maxRetries) {
try { try {
session.sendSerialized(message) // 메시지 전송 시도 session.sendSerialized(message)
return true // 성공하면 true 반환 val ackDeferred = CompletableDeferred<Boolean>()
ackMap.computeIfAbsent(message.uid) { ConcurrentHashMap() }[session] = ackDeferred
val ackReceived = withTimeoutOrNull(delayMillis) { ackDeferred.await() } ?: false
if (ackReceived) {
ackMap[message.uid]?.remove(session)
return true
} else {
attempt++
logger.warn("ACK not received for message to ${message.uid} on attempt $attempt.")
}
} catch (e: Exception) { } catch (e: Exception) {
attempt++ attempt++
logger.info("Failed to send message on attempt $attempt. Retrying in $delayMillis ms.") logger.info("Failed to send message on attempt $attempt. Retrying in $delayMillis ms.")
e.printStackTrace() e.printStackTrace()
delay(delayMillis) // 재시도 전 대기 delay(delayMillis)
} }
} }
return false // 재시도 실패 시 false 반환 return false
} }
fun broadcastMessage(userId: String, message: SongResponse) { fun broadcastMessage(userId: String, message: SongResponse) {
val userSessions = sessions[userId] val userSessions = sessions[userId]
userSessions?.forEach { session -> userSessions?.forEach { session ->
routeScope.launch { songListScope.launch {
val success = sendWithRetry(session, message) val success = sendWithRetry(session, message)
if (!success) { if (!success) {
println("Removing session for user $userId due to repeated failures.") logger.info("Removing session for user $userId due to repeated failures.")
userSessions.remove(session) // 실패 시 세션 제거 removeSession(userId, session)
} }
} }
} }
@ -80,19 +90,13 @@ fun Routing.wsSongRoutes() {
webSocket("/song/{uid}") { webSocket("/song/{uid}") {
val uid = call.parameters["uid"] val uid = call.parameters["uid"]
val user = uid?.let { UserService.getUser(it) } val user = uid?.let { UserService.getUser(it) }
if (uid == null) { if (uid == null || user == null) {
close(CloseReason(CloseReason.Codes.CANNOT_ACCEPT, "Invalid UID")) close(CloseReason(CloseReason.Codes.CANNOT_ACCEPT, "Invalid UID"))
return@webSocket return@webSocket
} }
if (user == null) {
close(CloseReason(CloseReason.Codes.CANNOT_ACCEPT, "Invalid UID"))
return@webSocket
}
addSession(uid, this) addSession(uid, this)
if(status[uid] == SongType.STREAM_OFF) { if(status[uid] == SongType.STREAM_OFF) {
routeScope.launch { songListScope.launch {
sendSerialized(SongResponse( sendSerialized(SongResponse(
SongType.STREAM_OFF.value, SongType.STREAM_OFF.value,
uid, uid,
@ -102,33 +106,36 @@ fun Routing.wsSongRoutes() {
)) ))
} }
} }
try { try {
for (frame in incoming) { for (frame in incoming) {
when(frame) { when(frame) {
is Frame.Text -> { is Frame.Text -> {
if(frame.readText().trim() == "ping") { val text = frame.readText().trim()
if(text == "ping") {
send("pong") send("pong")
} else {
val data = Json.decodeFromString<SongRequest>(text)
if (data.type == SongType.ACK.value) {
ackMap[data.uid]?.get(this)?.complete(true)
ackMap[data.uid]?.remove(this)
}
} }
} }
is Frame.Ping -> send(Frame.Pong(frame.data)) is Frame.Ping -> send(Frame.Pong(frame.data))
else -> { else -> {}
}
} }
} }
} catch(e: ClosedReceiveChannelException) { } catch(e: ClosedReceiveChannelException) {
logger.error("Error in WebSocket: ${e.message}") logger.error("Error in WebSocket: ${e.message}")
} finally { } finally {
removeSession(uid, this) removeSession(uid, this)
ackMap[uid]?.remove(this)
} }
} }
val dispatcher: CoroutinesEventBus by inject(CoroutinesEventBus::class.java)
dispatcher.subscribe(SongEvent::class) { dispatcher.subscribe(SongEvent::class) {
logger.debug("SongEvent: {} / {} {}", it.uid, it.type, it.current?.name) logger.debug("SongEvent: {} / {} {}", it.uid, it.type, it.current?.name)
routeScope.launch { songListScope.launch {
broadcastMessage(it.uid, SongResponse( broadcastMessage(it.uid, SongResponse(
it.type.value, it.type.value,
it.uid, it.uid,
@ -141,7 +148,7 @@ fun Routing.wsSongRoutes() {
} }
dispatcher.subscribe(TimerEvent::class) { dispatcher.subscribe(TimerEvent::class) {
if(it.type == TimerType.STREAM_OFF) { if(it.type == TimerType.STREAM_OFF) {
routeScope.launch { songListScope.launch {
broadcastMessage(it.uid, SongResponse( broadcastMessage(it.uid, SongResponse(
it.type.value, it.type.value,
it.uid, it.uid,
@ -153,7 +160,6 @@ fun Routing.wsSongRoutes() {
} }
} }
} }
@Serializable @Serializable
data class SerializableYoutubeVideo( data class SerializableYoutubeVideo(
val url: String, val url: String,
@ -161,9 +167,7 @@ data class SerializableYoutubeVideo(
val author: String, val author: String,
val length: Int val length: Int
) )
fun YoutubeVideo.toSerializable() = SerializableYoutubeVideo(url, name, author, length) fun YoutubeVideo.toSerializable() = SerializableYoutubeVideo(url, name, author, length)
@Serializable @Serializable
data class SongResponse( data class SongResponse(
val type: Int, val type: Int,
@ -172,4 +176,4 @@ data class SongResponse(
val current: SerializableYoutubeVideo? = null, val current: SerializableYoutubeVideo? = null,
val next: SerializableYoutubeVideo? = null, val next: SerializableYoutubeVideo? = null,
val delUrl: String? = null val delUrl: String? = null
) )

156
webserver/src/main/kotlin/space/mori/chzzk_bot/webserver/routes/WSTimerRoutes.kt
View File

@ -1,13 +1,20 @@
package space.mori.chzzk_bot.webserver.routes package space.mori.chzzk_bot.webserver.routes
import io.ktor.server.application.ApplicationStopped
import io.ktor.server.routing.* import io.ktor.server.routing.*
import io.ktor.server.websocket.* import io.ktor.server.websocket.*
import io.ktor.websocket.* import io.ktor.websocket.*
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.channels.ClosedReceiveChannelException import kotlinx.coroutines.channels.ClosedReceiveChannelException
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch import kotlinx.coroutines.launch
import kotlinx.coroutines.withTimeoutOrNull
import kotlinx.serialization.Serializable import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import org.koin.java.KoinJavaComponent.inject import org.koin.java.KoinJavaComponent.inject
import org.slf4j.LoggerFactory import org.slf4j.LoggerFactory
import space.mori.chzzk_bot.common.events.* import space.mori.chzzk_bot.common.events.*
@ -17,14 +24,19 @@ import space.mori.chzzk_bot.webserver.utils.CurrentTimer
import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.ConcurrentLinkedQueue import java.util.concurrent.ConcurrentLinkedQueue
val timerScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
fun Routing.wsTimerRoutes() { fun Routing.wsTimerRoutes() {
environment.monitor.subscribe(ApplicationStopped) {
songListScope.cancel()
}
val sessions = ConcurrentHashMap<String, ConcurrentLinkedQueue<WebSocketServerSession>>() val sessions = ConcurrentHashMap<String, ConcurrentLinkedQueue<WebSocketServerSession>>()
val logger = LoggerFactory.getLogger("WSTimerRoutes") val logger = LoggerFactory.getLogger("WSTimerRoutes")
val dispatcher: CoroutinesEventBus by inject(CoroutinesEventBus::class.java)
val ackMap = ConcurrentHashMap<String, ConcurrentHashMap<WebSocketServerSession, CompletableDeferred<Boolean>>>()
fun addSession(uid: String, session: WebSocketServerSession) { fun addSession(uid: String, session: WebSocketServerSession) {
sessions.computeIfAbsent(uid) { ConcurrentLinkedQueue() }.add(session) sessions.computeIfAbsent(uid) { ConcurrentLinkedQueue() }.add(session)
} }
fun removeSession(uid: String, session: WebSocketServerSession) { fun removeSession(uid: String, session: WebSocketServerSession) {
sessions[uid]?.remove(session) sessions[uid]?.remove(session)
if(sessions[uid]?.isEmpty() == true) { if(sessions[uid]?.isEmpty() == true) {
@ -32,82 +44,132 @@ fun Routing.wsTimerRoutes() {
} }
} }
webSocket("/timer/{uid}") { suspend fun sendWithRetry(
val uid = call.parameters["uid"] session: WebSocketServerSession,
val user = uid?.let { UserService.getUser(it) } message: TimerResponse,
if (uid == null) { maxRetries: Int = 3,
close(CloseReason(CloseReason.Codes.CANNOT_ACCEPT, "Invalid UID")) delayMillis: Long = 2000L
return@webSocket ): Boolean {
} var attempt = 0
if (user == null) { while (attempt < maxRetries) {
close(CloseReason(CloseReason.Codes.CANNOT_ACCEPT, "Invalid UID")) try {
return@webSocket session.sendSerialized(message)
} val ackDeferred = CompletableDeferred<Boolean>()
ackMap.computeIfAbsent(message.uid) { ConcurrentHashMap() }[session] = ackDeferred
addSession(uid, this) val ackReceived = withTimeoutOrNull(delayMillis) { ackDeferred.await() } ?: false
val timer = CurrentTimer.getTimer(user) if (ackReceived) {
ackMap[message.uid]?.remove(session)
if(timer?.type == TimerType.STREAM_OFF) { return true
CoroutineScope(Dispatchers.Default).launch {
sendSerialized(TimerResponse(TimerType.STREAM_OFF.value, null))
}
} else {
CoroutineScope(Dispatchers.Default).launch {
if (timer == null) {
sendSerialized(
TimerResponse(
TimerConfigService.getConfig(user)?.option ?: TimerType.REMOVE.value,
null
)
)
} else { } else {
sendSerialized( attempt++
TimerResponse( logger.warn("ACK not received for message to ${message.uid} on attempt $attempt.")
timer.type.value, }
timer.time } catch (e: Exception) {
) attempt++
) logger.info("Failed to send message on attempt $attempt. Retrying in $delayMillis ms.")
e.printStackTrace()
delay(delayMillis)
}
}
return false
}
fun broadcastMessage(uid: String, message: TimerResponse) {
val userSessions = sessions[uid]
userSessions?.forEach { session ->
songListScope.launch {
val success = sendWithRetry(session, message.copy(uid = uid))
if (!success) {
logger.info("Removing session for user $uid due to repeated failures.")
removeSession(uid, session)
} }
} }
} }
}
webSocket("/timer/{uid}") {
val uid = call.parameters["uid"]
val user = uid?.let { UserService.getUser(it) }
if (uid == null || user == null) {
close(CloseReason(CloseReason.Codes.CANNOT_ACCEPT, "Invalid UID"))
return@webSocket
}
addSession(uid, this)
val timer = CurrentTimer.getTimer(user)
if (timer?.type == TimerType.STREAM_OFF) {
songListScope.launch {
sendSerialized(TimerResponse(TimerType.STREAM_OFF.value, null, uid))
}
} else {
songListScope.launch {
if(timer?.type == TimerType.STREAM_OFF) {
sendSerialized(TimerResponse(TimerType.STREAM_OFF.value, null, uid))
} else {
if (timer == null) {
sendSerialized(
TimerResponse(
TimerConfigService.getConfig(user)?.option ?: TimerType.REMOVE.value,
null,
uid
)
)
} else {
sendSerialized(
TimerResponse(
timer.type.value,
timer.time,
uid
)
)
}
}
}
}
try { try {
for (frame in incoming) { for (frame in incoming) {
when(frame) { when(frame) {
is Frame.Text -> { is Frame.Text -> {
if(frame.readText().trim() == "ping") { val text = frame.readText().trim()
if(text == "ping") {
send("pong") send("pong")
} else {
val data = Json.decodeFromString<TimerRequest>(text)
if (data.type == TimerType.ACK.value) {
ackMap[data.uid]?.get(this)?.complete(true)
ackMap[data.uid]?.remove(this)
}
} }
} }
is Frame.Ping -> send(Frame.Pong(frame.data)) is Frame.Ping -> send(Frame.Pong(frame.data))
else -> { else -> {}
}
} }
} }
} catch(e: ClosedReceiveChannelException) { } catch(e: ClosedReceiveChannelException) {
logger.error("Error in WebSocket: ${e.message}") logger.error("Error in WebSocket: ${e.message}")
} finally { } finally {
removeSession(uid, this) removeSession(uid, this)
ackMap[uid]?.remove(this)
} }
} }
val dispatcher: CoroutinesEventBus by inject(CoroutinesEventBus::class.java)
dispatcher.subscribe(TimerEvent::class) { dispatcher.subscribe(TimerEvent::class) {
logger.debug("TimerEvent: {} / {}", it.uid, it.type) logger.debug("TimerEvent: {} / {}", it.uid, it.type)
val user = UserService.getUser(it.uid) val user = UserService.getUser(it.uid)
CurrentTimer.setTimer(user!!, it) CurrentTimer.setTimer(user!!, it)
CoroutineScope(Dispatchers.Default).launch { songListScope.launch {
sessions[it.uid]?.forEach { ws -> broadcastMessage(it.uid, TimerResponse(it.type.value, it.time ?: "", it.uid))
ws.sendSerialized(TimerResponse(it.type.value, it.time ?: ""))
}
} }
} }
} }
@Serializable @Serializable
data class TimerResponse( data class TimerResponse(
val type: Int, val type: Int,
val time: String? val time: String?,
val uid: String
)
@Serializable
data class TimerRequest(
val type: Int,
val uid: String
) )