View Javadoc

1   /*
2    *
3    * Licensed to the Apache Software Foundation (ASF) under one
4    * or more contributor license agreements.  See the NOTICE file
5    * distributed with this work for additional information
6    * regarding copyright ownership.  The ASF licenses this file
7    * to you under the Apache License, Version 2.0 (the
8    * "License"); you may not use this file except in compliance
9    * with the License.  You may obtain a copy of the License at
10   *
11   *     http://www.apache.org/licenses/LICENSE-2.0
12   *
13   * Unless required by applicable law or agreed to in writing, software
14   * distributed under the License is distributed on an "AS IS" BASIS,
15   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16   * See the License for the specific language governing permissions and
17   * limitations under the License.
18   */
19  
20  package org.apache.hadoop.hbase.io;
21  
22  import java.io.IOException;
23  import java.io.OutputStream;
24  import java.nio.ByteBuffer;
25  import java.nio.channels.Channels;
26  import java.nio.channels.WritableByteChannel;
27  
28  import org.apache.hadoop.classification.InterfaceAudience;
29  import org.apache.hadoop.classification.InterfaceStability;
30  import org.apache.hadoop.hbase.util.Bytes;
31  
32  /**
33   * Not thread safe!
34   */
35  @InterfaceAudience.Public
36  @InterfaceStability.Evolving
37  public class ByteBufferOutputStream extends OutputStream {
38  
39    protected ByteBuffer buf;
40  
41    public ByteBufferOutputStream(int capacity) {
42      this(capacity, false);
43    }
44  
45    public ByteBufferOutputStream(int capacity, boolean useDirectByteBuffer) {
46      if (useDirectByteBuffer) {
47        buf = ByteBuffer.allocateDirect(capacity);
48      } else {
49        buf = ByteBuffer.allocate(capacity);
50      }
51    }
52  
53    public int size() {
54      return buf.position();
55    }
56  
57    /**
58     * This flips the underlying BB so be sure to use it _last_!
59     * @return ByteBuffer
60     */
61    public ByteBuffer getByteBuffer() {
62      buf.flip();
63      return buf;
64    }
65  
66    private void checkSizeAndGrow(int extra) {
67      if ( (buf.position() + extra) > buf.limit()) {
68        // size calculation is complex, because we could overflow negative,
69        // and/or not allocate enough space. this fixes that.
70        int newSize = (int)Math.min((((long)buf.capacity()) * 2),
71            (long)(Integer.MAX_VALUE));
72        newSize = Math.max(newSize, buf.position() + extra);
73  
74        ByteBuffer newBuf = ByteBuffer.allocate(newSize);
75        buf.flip();
76        newBuf.put(buf);
77        buf = newBuf;
78      }
79    }
80  
81    // OutputStream
82    @Override
83    public void write(int b) throws IOException {
84      checkSizeAndGrow(Bytes.SIZEOF_BYTE);
85  
86      buf.put((byte)b);
87    }
88  
89   /**
90    * Writes the complete contents of this byte buffer output stream to
91    * the specified output stream argument.
92    *
93    * @param      out   the output stream to which to write the data.
94    * @exception  IOException  if an I/O error occurs.
95    */
96    public synchronized void writeTo(OutputStream out) throws IOException {
97      WritableByteChannel channel = Channels.newChannel(out);
98      ByteBuffer bb = buf.duplicate();
99      bb.flip();
100     channel.write(bb);
101   }
102 
103   @Override
104   public void write(byte[] b) throws IOException {
105     checkSizeAndGrow(b.length);
106 
107     buf.put(b);
108   }
109 
110   @Override
111   public void write(byte[] b, int off, int len) throws IOException {
112     checkSizeAndGrow(len);
113 
114     buf.put(b, off, len);
115   }
116 
117   @Override
118   public void flush() throws IOException {
119     // noop
120   }
121 
122   @Override
123   public void close() throws IOException {
124     // noop again. heh
125   }
126 
127   public byte[] toByteArray(int offset, int length) {
128     ByteBuffer bb = buf.duplicate();
129     bb.flip();
130 
131     byte[] chunk = new byte[length];
132 
133     bb.position(offset);
134     bb.get(chunk, 0, length);
135     return chunk;
136   }
137 }