CountingRequestBody.java

package net.callumtaylor.asynchttp.obj;

import java.io.IOException;

import okhttp3.MediaType;
import okhttp3.RequestBody;
import okio.Buffer;
import okio.BufferedSink;
import okio.ForwardingSink;
import okio.Okio;
import okio.Sink;

/**
 * Decorates an OkHttp request body to count the number of bytes written when writing it. Can
 * decorate any request body, but is most useful for tracking the upload progress of large
 * multipart requests.
 *
 * @author Leo Nikkilä
 */
public class CountingRequestBody extends RequestBody
{
	protected RequestBody delegate;
	protected Listener listener;

	protected CountingSink countingSink;

	public CountingRequestBody(RequestBody delegate, Listener listener)
	{
		this.delegate = delegate;
		this.listener = listener;
	}

	@Override public MediaType contentType()
	{
		return delegate.contentType();
	}

	@Override public long contentLength()
	{
		try
		{
			return delegate.contentLength();
		}
		catch (IOException e)
		{
			e.printStackTrace();
		}

		return 0;
	}

	@Override public void writeTo(BufferedSink sink) throws IOException
	{
		BufferedSink bufferedSink;

		countingSink = new CountingSink(sink);
		bufferedSink = Okio.buffer(countingSink);

		delegate.writeTo(bufferedSink);

		bufferedSink.flush();
	}

	protected final class CountingSink extends ForwardingSink
	{
		private long bytesWritten = 0;

		public CountingSink(Sink delegate)
		{
			super(delegate);
		}

		@Override public void write(Buffer source, long byteCount) throws IOException
		{
			Buffer copy = new Buffer();
			source.copyTo(copy, 0, byteCount);

			super.write(source, byteCount);

			bytesWritten += byteCount;
			listener.onRequestProgress(copy.readByteArray(), byteCount, bytesWritten, contentLength());
		}
	}

	public static interface Listener
	{
		public void onRequestProgress(byte[] buffer, long bufferCount, long bytesWritten, long contentLength);
	}
}