001/**
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018
019package org.apache.hadoop.hbase.procedure2.util;
020
021import java.io.IOException;
022import java.io.OutputStream;
023import java.util.Arrays;
024
025import org.apache.yetus.audience.InterfaceAudience;
026import org.apache.yetus.audience.InterfaceStability;
027
028/**
029 * Similar to the ByteArrayOutputStream, with the exception that we can prepend an header.
030 * e.g. you write some data and you want to prepend an header that contains the data len or cksum.
031 * <code>
032 * ByteSlot slot = new ByteSlot();
033 * // write data
034 * slot.write(...);
035 * slot.write(...);
036 * // write header with the size of the written data
037 * slot.markHead();
038 * slot.write(Bytes.toBytes(slot.size()));
039 * // flush to stream as [header, data]
040 * slot.writeTo(stream);
041 * </code>
042 */
043@InterfaceAudience.Private
044@InterfaceStability.Evolving
045public class ByteSlot extends OutputStream {
046  private static final int LARGE_GROW_SIZE_THRESHOLD = 8 << 20;
047  private static final int LARGE_GROW_SIZE = 1 << 20;
048  private static final int RESET_THRESHOLD = 64 << 20;
049  private static final int GROW_ALIGN = 128;
050
051  private byte[] buf;
052  private int head;
053  private int size;
054
055  public void reset() {
056    if (buf != null && buf.length > RESET_THRESHOLD) {
057      buf = null;
058    }
059    head = 0;
060    size = 0;
061  }
062
063  public void markHead() {
064    head = size;
065  }
066
067  public int getHead() {
068    return head;
069  }
070
071  public int size() {
072    return size;
073  }
074
075  public byte[] getBuffer() {
076    return buf;
077  }
078
079  public void writeAt(int offset, int b) {
080    head = Math.min(head, offset);
081    buf[offset] = (byte)b;
082  }
083
084  @Override
085  public void write(int b) {
086    ensureCapacity(size + 1);
087    buf[size++] = (byte)b;
088  }
089
090  @Override
091  public void write(byte[] b, int off, int len) {
092    ensureCapacity(size + len);
093    System.arraycopy(b, off, buf, size, len);
094    size += len;
095  }
096
097  public void writeTo(final OutputStream stream) throws IOException {
098    if (head != 0) {
099      stream.write(buf, head, size - head);
100      stream.write(buf, 0, head);
101    } else {
102      stream.write(buf, 0, size);
103    }
104  }
105
106  private void ensureCapacity(int minCapacity) {
107    minCapacity = (minCapacity + (GROW_ALIGN - 1)) & -GROW_ALIGN;
108    if (buf == null) {
109      buf = new byte[minCapacity];
110    } else if (minCapacity > buf.length) {
111      int newCapacity;
112      if (buf.length <= LARGE_GROW_SIZE_THRESHOLD) {
113        newCapacity = buf.length << 1;
114      } else {
115        newCapacity = buf.length + LARGE_GROW_SIZE;
116      }
117      if (minCapacity > newCapacity) {
118        newCapacity = minCapacity;
119      }
120      buf = Arrays.copyOf(buf, newCapacity);
121    }
122  }
123}