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.hbase.classification.InterfaceAudience;
29  import org.apache.hadoop.hbase.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        ByteBuffer newBuf = null;
74        if (buf.isDirect()) {
75          newBuf = ByteBuffer.allocateDirect(newSize);
76        } else {
77          newBuf = ByteBuffer.allocate(newSize);
78        }
79        buf.flip();
80        newBuf.put(buf);
81        buf = newBuf;
82      }
83    }
84  
85    // OutputStream
86    @Override
87    public void write(int b) throws IOException {
88      checkSizeAndGrow(Bytes.SIZEOF_BYTE);
89  
90      buf.put((byte)b);
91    }
92  
93   /**
94    * Writes the complete contents of this byte buffer output stream to
95    * the specified output stream argument.
96    *
97    * @param      out   the output stream to which to write the data.
98    * @exception  IOException  if an I/O error occurs.
99    */
100   public synchronized void writeTo(OutputStream out) throws IOException {
101     WritableByteChannel channel = Channels.newChannel(out);
102     ByteBuffer bb = buf.duplicate();
103     bb.flip();
104     channel.write(bb);
105   }
106 
107   @Override
108   public void write(byte[] b) throws IOException {
109     checkSizeAndGrow(b.length);
110 
111     buf.put(b);
112   }
113 
114   @Override
115   public void write(byte[] b, int off, int len) throws IOException {
116     checkSizeAndGrow(len);
117 
118     buf.put(b, off, len);
119   }
120 
121   @Override
122   public void flush() throws IOException {
123     // noop
124   }
125 
126   @Override
127   public void close() throws IOException {
128     // noop again. heh
129   }
130 
131   public byte[] toByteArray(int offset, int length) {
132     ByteBuffer bb = buf.duplicate();
133     bb.flip();
134 
135     byte[] chunk = new byte[length];
136 
137     bb.position(offset);
138     bb.get(chunk, 0, length);
139     return chunk;
140   }
141 }