io.bitrise.gradle.cache.connection.ClientBalancer.kt Maven / Gradle / Ivy

Go to download
/**
 * Copyright (C)2022 Bitrise
 * All rights reserved.
 */
package io.bitrise.gradle.cache.connection

import io.grpc.ManagedChannel
import kotlinx.coroutines.Job
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Semaphore
import kotlinx.coroutines.sync.withPermit
import java.time.LocalDateTime
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicInteger

/**
 * This interface and the associated wrapper provide a mechanism by which requests to a resource can
 * be limited (in this case a grpc client) and disposal can be guaranteed by use of `use()` fn thanks to
 * implementation of Closeable
 */
interface Client : Semaphore {
    suspend fun use(f: suspend (T) -> Unit): Job
}

class ClientWrapper(val c: T, maxConcurrency: Int, val debug: Boolean) : Client {
    private val semaphore = Semaphore(maxConcurrency, 0)

    override suspend fun acquire() {
        semaphore.acquire()
    }

    override suspend fun use(f: suspend (T) -> Unit): Job = coroutineScope {
        launch {
            if (semaphore.availablePermits == 0 && debug) {
                println("[${LocalDateTime.now()}] debug: Max concurrency reached, waiting for available permits")
            }
            semaphore.withPermit {
                f(c)
            }
        }
    }

    override fun release() {
        semaphore.release()
    }

    override fun tryAcquire(): Boolean {
        return semaphore.tryAcquire()
    }

    override val availablePermits: Int
        get() = semaphore.availablePermits
}

/**
 * load-balances requests across a collection of channels/factories in a uniform manner
 */
class ClientBalancer(
    numChannels: Int,
    maxConcurrencyPerChannel: Int,
    channelFactory: () -> ManagedChannel,
    clientFactory: (ManagedChannel) -> T,
    debug: Boolean,
) {
    private val channels: MutableList = mutableListOf()
    private val clients: MutableList> = mutableListOf()
    private val nextIndex: AtomicInteger = AtomicInteger()

    init {
        repeat(numChannels) {
            channels.add(channelFactory())
            clients.add(ClientWrapper(clientFactory(channels[it]), maxConcurrencyPerChannel, debug))
        }
    }

    val next: Client
        get() {
            return clients[nextIndex.getAndIncrement() % clients.size]
        }

    fun close() {
        channels.forEach {
            it.shutdownNow()
            it.awaitTermination(2, TimeUnit.SECONDS)
        }
    }
}