1// Copyright 2014 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5package org.chromium.net;
6
7import java.io.IOException;
8import java.nio.ByteBuffer;
9import java.nio.channels.ClosedChannelException;
10import java.nio.channels.WritableByteChannel;
11import java.util.ArrayList;
12
13/**
14 * A writable byte channel that is optimized for chunked writing. Each call to
15 * {@link #write} results in a ByteBuffer being created and remembered. Then all
16 * of those byte buffers are combined on demand. This approach allows to avoid
17 * the cost of reallocating a byte buffer.
18 */
19public class ChunkedWritableByteChannel implements WritableByteChannel {
20
21    private final ArrayList<ByteBuffer> mBuffers = new ArrayList<ByteBuffer>();
22
23    private ByteBuffer mInitialBuffer;
24
25    private ByteBuffer mBuffer;
26
27    private int mSize;
28
29    private boolean mClosed;
30
31    public void setCapacity(int capacity) {
32        if (!mBuffers.isEmpty() || mInitialBuffer != null) {
33            throw new IllegalStateException();
34        }
35
36        mInitialBuffer = ByteBuffer.allocateDirect(capacity);
37    }
38
39    @Override
40    public int write(ByteBuffer buffer) throws IOException {
41        if (mClosed) {
42            throw new ClosedChannelException();
43        }
44
45        int size = buffer.remaining();
46        mSize += size;
47
48        if (mInitialBuffer != null) {
49            if (size <= mInitialBuffer.remaining()) {
50                mInitialBuffer.put(buffer);
51                return size;
52            }
53
54            // The supplied initial size was incorrect. Keep the accumulated
55            // data and switch to the usual "sequence of buffers" mode.
56            mInitialBuffer.flip();
57            mBuffers.add(mInitialBuffer);
58            mInitialBuffer = null;
59        }
60
61        // We can't hold a reference to this buffer, because it may wrap native
62        // memory and is not guaranteed to be immutable.
63        ByteBuffer tmpBuf = ByteBuffer.allocateDirect(size);
64        tmpBuf.put(buffer).rewind();
65        mBuffers.add(tmpBuf);
66        return size;
67    }
68
69    /**
70     * Returns the entire content accumulated by the channel as a ByteBuffer.
71     */
72    public ByteBuffer getByteBuffer() {
73        if (mInitialBuffer != null) {
74            mInitialBuffer.flip();
75            mBuffer = mInitialBuffer;
76            mInitialBuffer = null;
77        } else if (mBuffer != null && mSize == mBuffer.capacity()) {
78            // Cache hit
79        } else if (mBuffer == null && mBuffers.size() == 1) {
80            mBuffer = mBuffers.get(0);
81        } else {
82            mBuffer = ByteBuffer.allocateDirect(mSize);
83            int count = mBuffers.size();
84            for (int i = 0; i < count; i++) {
85                mBuffer.put(mBuffers.get(i));
86            }
87            mBuffer.rewind();
88        }
89        return mBuffer;
90    }
91
92    /**
93     * Returns the entire content accumulated by the channel as a byte array.
94     */
95    public byte[] getBytes() {
96        byte[] bytes = new byte[mSize];
97        if (mInitialBuffer != null) {
98            mInitialBuffer.flip();
99            mInitialBuffer.get(bytes);
100        } else {
101            int bufferCount = mBuffers.size();
102            int offset = 0;
103            for (int i = 0; i < bufferCount; i++) {
104                ByteBuffer buffer = mBuffers.get(i);
105                int bufferSize = buffer.remaining();
106                buffer.get(bytes, offset, bufferSize);
107                buffer.rewind();
108                offset += bufferSize;
109            }
110        }
111        return bytes;
112    }
113
114    @Override
115    public void close() {
116        mClosed = true;
117    }
118
119    @Override
120    public boolean isOpen() {
121        return !mClosed;
122    }
123}
124