goodmetrics.downstream.GrpcTrailerLogger.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of goodmetrics-kotlin Show documentation

A metrics recording library that is good

The newest version!

package goodmetrics.downstream

import io.grpc.CallOptions
import io.grpc.Channel
import io.grpc.ClientCall
import io.grpc.ClientInterceptor
import io.grpc.ForwardingClientCall.SimpleForwardingClientCall
import io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener
import io.grpc.MethodDescriptor
import io.grpc.Status
import io.grpc.Metadata as GrpcMetadata

class GrpcTrailerLoggerInterceptor(
    private val onTrailers: (Status, GrpcMetadata) -> Unit
) : ClientInterceptor {
    override fun  interceptCall(
        method: MethodDescriptor,
        callOptions: CallOptions,
        next: Channel
    ): ClientCall {
        return TrailersLoggingClientCall(next.newCall(method, callOptions), onTrailers)
    }

    private class TrailersLoggingClientCall (
        call: ClientCall?,
        private val onTrailers: (Status, GrpcMetadata) -> Unit,
    ) : SimpleForwardingClientCall(call) {
        override fun start(responseListener: Listener, headers: GrpcMetadata) {
            super.start(MetadataCapturingClientCallListener(responseListener), headers)
        }

        private inner class MetadataCapturingClientCallListener(
            responseListener: Listener?
        ) :
            SimpleForwardingClientCallListener(responseListener) {
            override fun onHeaders(headers: GrpcMetadata) {
                super.onHeaders(headers)
            }

            override fun onClose(status: Status, trailers: GrpcMetadata?) {
                if (trailers != null && trailers.keys().isNotEmpty()) {
                    onTrailers(status, trailers)
                }
                super.onClose(status, trailers)
            }
        }
    }
}