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

Go to download
package io.bitrise.gradle.cache.connection

import com.google.bytestream.ByteStreamProto.ReadRequest
import com.google.bytestream.ByteStreamProto.ReadResponse
import com.google.bytestream.ByteStreamProto.WriteRequest
import com.google.bytestream.ByteStreamProto.WriteResponse
import io.grpc.CallOptions
import io.grpc.ManagedChannel
import io.grpc.Metadata
import io.grpc.stub.MetadataUtils
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kv_storage.KVStorageGrpcKt
import java.io.Closeable
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicReference

class KVStorageClient(
    private val channel: ManagedChannel,
    private val callOptions: CallOptions,
    private val metadata: Map> = emptyMap(),
) : Closeable {

    private val headers by lazy {
        val header = Metadata()
        for ((k, v) in metadata) {
            val key = Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER)
            v.forEach { header.put(key, it) }
        }
        header
    }

    private val stub = KVStorageGrpcKt.KVStorageCoroutineStub(channel, callOptions)

    fun withMetadata(additionalMetadata: Map>): KVStorageClient {
        return KVStorageClient(channel, callOptions, metadata.plus(additionalMetadata))
    }

    suspend fun put(requests: Flow, onResponseMetadataHandler: (Metadata) -> Unit = {}): WriteResponse {
        val responseHeaders = AtomicReference(Metadata())
        val trailerHeaders = AtomicReference(Metadata())

        val resp = stub.withInterceptors(
            // LoggingInterceptor(), Add to log headers
            MetadataUtils.newCaptureMetadataInterceptor(
                responseHeaders,
                trailerHeaders,
            ),
        )
            .put(requests, headers)

        val joinedHeaders = responseHeaders.get()
        joinedHeaders.merge(trailerHeaders.get())
        onResponseMetadataHandler(joinedHeaders)

        return resp
    }

    fun get(request: ReadRequest, onResponseMetadataHandler: (Metadata) -> Unit = {}): Flow {
        val responseHeaders = AtomicReference(Metadata())
        val trailerHeaders = AtomicReference(Metadata())
        val resp = stub.withInterceptors(
            // LoggingInterceptor(), Add to log headers
            MetadataUtils.newCaptureMetadataInterceptor(
                responseHeaders,
                trailerHeaders,
            ),
        )
            .get(request, headers)

        return flow {
            resp.collect {
                emit(it)
            }
            val joinedHeaders = responseHeaders.get()
            joinedHeaders.merge(trailerHeaders.get())
            onResponseMetadataHandler(joinedHeaders)
        }
    }

    override fun close() {
        channel.shutdown().awaitTermination(2, TimeUnit.SECONDS)
    }
}