crackthecodeabhi / kreds Goto Github PK

@Suppress("TooGenericExceptionThrown", "TooManyFunctions")
@Service
class NettyRedisClient(
    @Value("\${redis.host}") private val host: String,
    @Value("\${redis.port}") private val port: Int,
    @Value("\${redis.password}") private val password: String
) : ExclusiveObject, DisposableBean {
    private val logger = KotlinLogging.logger {}

    final override val mutex: Mutex = Mutex()
    override val key: ReentrantMutexContextKey = ReentrantMutexContextKey(mutex)

    private val group = NioEventLoopGroup()
    private val bootstrap = Bootstrap()
        .group(group)
        .remoteAddress(host, port)
        .option(ChannelOption.SO_KEEPALIVE, true)
        .channel(NioSocketChannel::class.java)

    val sslContext: SslContext = SslContextBuilder
        .forClient()
        .trustManager(InsecureTrustManagerFactory.INSTANCE)
        .build()

    private var writeChannel: SocketChannel? = null
    private var readChannel: KChannel<RedisMessage>? = null

    fun pipelined(): Pipelined = Pipelined(this)

    suspend fun connect() = withReentrantLock {
        if (!isConnected()) {
            val newReadChannel = KChannel<RedisMessage>(KChannel.UNLIMITED)
            val newWriteChannel = bootstrap
                .handler(LoggingHandler(LogLevel.INFO))
                .handler(channelInitializer(newReadChannel))
                .connect()
                .suspendableAwait() as SocketChannel

            readChannel = newReadChannel
            writeChannel = newWriteChannel

            auth(newWriteChannel)
        }
    }

    private fun channelInitializer(newReadChannel: KChannel<RedisMessage>): ChannelInitializer<SocketChannel> {
        return object : ChannelInitializer<SocketChannel>() {
            override fun initChannel(channel: SocketChannel) {
                val pipeline: ChannelPipeline = channel.pipeline()
                pipeline.addLast(sslContext.newHandler(channel.alloc(), host, port))
                pipeline.addLast(RedisDecoder())
                pipeline.addLast(RedisBulkStringAggregator())
                pipeline.addLast(RedisArrayAggregator())
                pipeline.addLast(RedisEncoder())
                pipeline.addLast(commandHandler(newReadChannel))
            }
        }
    }

    private fun commandHandler(newReadChannel: KChannel<RedisMessage>) = object : ChannelDuplexHandler() {
        override fun write(
            handlerContext: ChannelHandlerContext,
            message: Any,
            promise: ChannelPromise
        ) {
            val commands = (message as String)
                .trim()
                .split(Regex("\\s+"))
                .map { command ->
                    FullBulkStringRedisMessage(
                        ByteBufUtil.writeUtf8(
                            handlerContext.alloc(),
                            command
                        )
                    )
                }
            val request: RedisMessage = ArrayRedisMessage(commands)
            handlerContext.write(request, promise)
        }

        override fun channelRead(handlerContext: ChannelHandlerContext, message: Any) {
            message as RedisMessage
            newReadChannel.trySend(message)
        }

        override fun exceptionCaught(handlerContext: ChannelHandlerContext, cause: Throwable) {
            handlerContext.close()
            newReadChannel.close(cause)
        }
    }

    suspend fun executeCommands(commands: List<String>): List<RedisMessage> = withReentrantLock {
        connect()

        commands.forEach {
            write(it)
        }

        flush()

        commands.map {
            read()
        }
    }

    suspend fun executeCommand(command: String): RedisMessage = withReentrantLock {
        connect()
        writeAndFlush(command)
        read()
    }

    suspend fun info(): String {
        return decode(executeCommand("INFO"))
    }

    suspend fun flushAll(): String {
        return decode(executeCommand("FLUSHALL"))
    }

    suspend fun dbSize(): Long {
        return decode(executeCommand("DBSIZE")).toLong()
    }

    suspend fun xread(streamNames: List<String>, block: Long?): String {
        val streamNamesString = streamNames.joinToString(" ")
        val streamOffsetsString = streamNames.joinToString(" ") { "0-0" }
        val string = "$streamNamesString $streamOffsetsString"
        val command = if (block == null) {
            "XREAD STREAMS $string"
        } else {
            "XREAD BLOCK $block STREAMS $string"
        }

        return decode(executeCommand((command)))
    }

    private suspend fun auth(writeChannel: SocketChannel) = withReentrantLock {
        writeChannel.writeAndFlush("AUTH $password")
        val response = (read() as SimpleStringRedisMessage).content()
        if (response == "OK") {
            logger.info("AUTH successful")
        } else {
            throw RuntimeException("AUTH failed")
        }
    }

    private suspend fun isConnected(): Boolean = withReentrantLock {
        if (writeChannel == null || readChannel == null) {
            false
        } else {
            writeChannel!!.isActive
        }
    }

    private suspend fun write(message: String): Unit = withReentrantLock {
        if (!isConnected()) {
            throw RuntimeException("Not yet connected")
        } else {
            writeChannel!!.write(message)
        }
    }

    private suspend fun writeAndFlush(message: String): Unit = withReentrantLock {
        if (!isConnected()) {
            throw RuntimeException("Not yet connected")
        } else {
            writeChannel!!.writeAndFlush(message)
        }
    }

    private suspend fun flush(): Unit = withReentrantLock {
        if (!isConnected()) {
            throw RuntimeException("Not yet connected")
        } else {
            writeChannel!!.flush()
        }
    }

    private suspend fun read(): RedisMessage = withReentrantLock {
        if (!isConnected()) {
            throw RuntimeException("Not yet connected")
        } else {
            readChannel!!.receive()
        }
    }

    override fun destroy() {
        runBlocking {
            withReentrantLock {
                readChannel?.close()
                writeChannel?.close()
                group.shutdownGracefully()
            }
        }
    }
}

@Suppress("TooGenericExceptionThrown")
class Pipelined(private val client: NettyRedisClient) : ExclusiveObject {
    override val mutex: Mutex = Mutex()
    override val key: ReentrantMutexContextKey = ReentrantMutexContextKey(mutex)

    private var done = false
    private val responseFlow = MutableSharedFlow<List<String>>(1)
    private val sharedResponseFlow: Flow<List<String>> = responseFlow.asSharedFlow()
    private val commands = mutableListOf<String>()
    private val commandResponse = mutableListOf<String>()

    suspend fun xadd(streamName: String, keyValues: Map<String, String>): Response {
        val keyValuesString = keyValues.map { "${it.key} ${it.value}" }.joinToString(" ")
        val command = "XADD $streamName * $keyValuesString"

        return add(command)
    }

    suspend fun expire(streamName: String, seconds: Long): Response {
        return add("EXPIRE $streamName $seconds")
    }

    suspend fun execute(): Unit = withReentrantLock {
        if (!done) {
            commandResponse.addAll(executePipeline(commands))
            done = true
            responseFlow.tryEmit(commandResponse.toMutableList())
        }
    }

    private suspend fun add(command: String): Response = withReentrantLock {
        commands.add(command)
        Response(sharedResponseFlow, commands.lastIndex)
    }

    private suspend fun executePipeline(commands: List<String>): List<String> = withReentrantLock {
        val responseMessages = client.executeCommands(commands)

        responseMessages.map { message ->
            decode(message)
        }
    }
}

@Suppress("TooGenericExceptionThrown")
internal fun decode(message: RedisMessage): String {
    return when (message) {
        is ErrorRedisMessage -> message.content()
        is SimpleStringRedisMessage -> message.content()
        is IntegerRedisMessage -> message.value().toString()
        is FullBulkStringRedisMessage -> {
            if (message.isNull) {
                throw RuntimeException("Stream response is null")
            } else {
                message.content().toString(Charset.defaultCharset())
            }
        }

        is ArrayRedisMessage -> {
            message.children().joinToString(" ") { child ->
                decode(child)
            }
        }

        else -> throw NotImplementedError("Message type not implemented")
    }
}

@Suppress("TooGenericExceptionThrown")
class Response internal constructor(
    private val flow: Flow<List<String>>,
    private val index: Int
) {
    suspend operator fun invoke(): String {
        return flow.first().ifEmpty { throw RuntimeException("Operation was cancelled.") }[index]
    }

    suspend fun get(): String = invoke()
}

internal interface ExclusiveObject {
    val mutex: Mutex
    val key: ReentrantMutexContextKey
}

data class ReentrantMutexContextKey(val mutex: Mutex) : CoroutineContext.Key<ReentrantMutexContextElement>
internal class ReentrantMutexContextElement(override val key: ReentrantMutexContextKey) : CoroutineContext.Element

internal suspend inline fun <R> ExclusiveObject.withReentrantLock(crossinline block: suspend () -> R): R {
    if (coroutineContext[key] != null) return block()

    return withContext(ReentrantMutexContextElement(key)) {
        [email protected] {
            block()
        }
    }
}

internal suspend fun ChannelFuture.suspendableAwait(): Channel {
    return suspendCoroutine { continuation ->
        addListener(object : ChannelFutureListener {
            override fun operationComplete(future: ChannelFuture) {
                if (future.isDone && future.isSuccess) {
                    continuation.resume(future.channel())
                } else {
                    continuation.resumeWithException(future.cause())
                }
            }
        })
    }
}