View Javadoc

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