money.rave.common.backend.redis.lettuce.StringRedis.kt Maven / Gradle / Ivy

Go to download
package money.rave.common.backend.redis.lettuce

import io.lettuce.core.RedisClient
import io.lettuce.core.RedisURI
import io.lettuce.core.pubsub.RedisPubSubListener
import io.lettuce.core.pubsub.StatefulRedisPubSubConnection
import io.lettuce.core.pubsub.api.sync.RedisPubSubCommands
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ReceiveChannel
import kotlinx.coroutines.withContext
import java.util.concurrent.ConcurrentHashMap

class StringRedis(
    private val host: String,
    private val port: Int,
    private val username: String,
    private val password: String,
    private val database: Int,
) {

    private lateinit var redisClient: RedisClient
    private lateinit var connection: StatefulRedisPubSubConnection
    private lateinit var command: RedisPubSubCommands
    private val subscriptions = ConcurrentHashMap>>()
    private val psubscriptions = ConcurrentHashMap>>()

    suspend fun connect() {
        withContext(Dispatchers.IO) {
            redisClient = RedisClient.create(
                RedisURI.Builder
                    .redis(host, port)
                    .withAuthentication(username, password)
                    .withDatabase(database)
                    .build()
            )

            connection = redisClient.connectPubSub()
            command = connection.sync()
        }
    }

    suspend fun  execute(block: RedisPubSubCommands.() -> T): T = withContext(Dispatchers.IO) { command.block() }

    suspend fun subscribe(key: String, capacity: Int = Channel.RENDEZVOUS): ReceiveChannel {
        val channel = Channel(capacity)
        withContext(Dispatchers.IO) { command.subscribe(key) }
        if (!subscriptions.containsKey(key)) {
            subscriptions[key] = mutableListOf()
        }
        subscriptions.getValue(key).add(channel)
        return channel
    }

    suspend fun psubscribe(key: String, capacity: Int = Channel.RENDEZVOUS): ReceiveChannel {
        val channel = Channel(capacity)
        withContext(Dispatchers.IO) { command.psubscribe(key) }
        if (!psubscriptions.containsKey(key)) {
            psubscriptions[key] = mutableListOf()
        }
        psubscriptions.getValue(key).add(channel)
        return channel
    }

    suspend fun publish(key: String, value: String) { withContext(Dispatchers.IO) { command.publish(key, value) } }

    suspend fun unsubscribe(key: String, channel: ReceiveChannel) {
        subscriptions[key]?.remove(channel)
        withContext(Dispatchers.IO) { command.unsubscribe(key) }
    }

    suspend fun punsubscribe(key: String, channel: ReceiveChannel) {
        psubscriptions[key]?.remove(channel)
        withContext(Dispatchers.IO) { command.punsubscribe(key) }
    }

    suspend fun addListener(listener: RedisPubSubListener) = withContext(Dispatchers.IO) { connection.addListener(listener) }

    suspend fun removeListener(listener: RedisPubSubListener) = withContext(Dispatchers.IO) { connection.removeListener(listener) }

    fun close() {
        connection.close()
    }
}