Okio Throttler integration with OkHttp

Question

My team is suffering from this issue with slack integration to upload files, so following the comments in that issue I would like to throttle the requests in our Kotlin implementation.

I am trying to integrate Okio Throttler within an OkHttp interceptor, so I have the setup:

val client = OkHttpClient.Builder()
            .retryOnConnectionFailure(false)
            .addInterceptor { chain ->
                val request = chain.request()
                val originalRequestBody = request.body
                val newRequest = if (originalRequestBody != null) {
                    val wrappedRequestBody = ThrottledRequestBody(originalRequestBody)
                    request.newBuilder()
                        .method(request.method, wrappedRequestBody)
                        .build()
                } else {
                    request
                }
                chain.proceed(newRequest)
            }
            .build()

class ThrottledRequestBody(private val delegate: RequestBody) : RequestBody() {
    private val throttler = Throttler().apply {
        bytesPerSecond(1024, 1024 * 4, 1024 * 8)
    }


    override fun contentType(): MediaType? {
        return delegate.contentType()
    }

    override fun writeTo(sink: BufferedSink) {
        delegate.writeTo(throttler.sink(sink).buffer())
    }
}

It seems throttler.sink returns a Sink, but a BufferedSink is required to the method delegate.writeTo, so I called buffer() to get that BufferedSink. Am I doing it wrong ? Is the call for .buffer() breaking the integration?

Answer 1

It's almost perfect. You just need to flush the buffer when you're done otherwise it'll finish with a few bytes inside.

override fun writeTo(sink: BufferedSink) {
  throttler.sink(sink).buffer().use {
    delegate.writeTo(it)
  }
}