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 */
018package org.apache.hadoop.hbase.procedure2.util;
019
020import java.io.IOException;
021import java.io.OutputStream;
022import java.util.Arrays;
023import org.apache.yetus.audience.InterfaceAudience;
024import org.apache.yetus.audience.InterfaceStability;
025
026/**
027 * Similar to the ByteArrayOutputStream, with the exception that we can prepend an header. e.g. you
028 * write some data and you want to prepend an header that contains the data len or cksum. <code>
029 * ByteSlot slot = new ByteSlot();
030 * // write data
031 * slot.write(...);
032 * slot.write(...);
033 * // write header with the size of the written data
034 * slot.markHead();
035 * slot.write(Bytes.toBytes(slot.size()));
036 * // flush to stream as [header, data]
037 * slot.writeTo(stream);
038 * </code>
039 */
040@InterfaceAudience.Private
041@InterfaceStability.Evolving
042public class ByteSlot extends OutputStream {
043  private static final int LARGE_GROW_SIZE_THRESHOLD = 8 << 20;
044  private static final int LARGE_GROW_SIZE = 1 << 20;
045  private static final int RESET_THRESHOLD = 64 << 20;
046  private static final int GROW_ALIGN = 128;
047
048  private byte[] buf;
049  private int head;
050  private int size;
051
052  public void reset() {
053    if (buf != null && buf.length > RESET_THRESHOLD) {
054      buf = null;
055    }
056    head = 0;
057    size = 0;
058  }
059
060  public void markHead() {
061    head = size;
062  }
063
064  public int getHead() {
065    return head;
066  }
067
068  public int size() {
069    return size;
070  }
071
072  public byte[] getBuffer() {
073    return buf;
074  }
075
076  public void writeAt(int offset, int b) {
077    head = Math.min(head, offset);
078    buf[offset] = (byte) b;
079  }
080
081  @Override
082  public void write(int b) {
083    ensureCapacity(size + 1);
084    buf[size++] = (byte) b;
085  }
086
087  @Override
088  public void write(byte[] b, int off, int len) {
089    ensureCapacity(size + len);
090    System.arraycopy(b, off, buf, size, len);
091    size += len;
092  }
093
094  public void writeTo(final OutputStream stream) throws IOException {
095    if (head != 0) {
096      stream.write(buf, head, size - head);
097      stream.write(buf, 0, head);
098    } else {
099      stream.write(buf, 0, size);
100    }
101  }
102
103  private void ensureCapacity(int minCapacity) {
104    minCapacity = (minCapacity + (GROW_ALIGN - 1)) & -GROW_ALIGN;
105    if (buf == null) {
106      buf = new byte[minCapacity];
107    } else if (minCapacity > buf.length) {
108      int newCapacity;
109      if (buf.length <= LARGE_GROW_SIZE_THRESHOLD) {
110        newCapacity = buf.length << 1;
111      } else {
112        newCapacity = buf.length + LARGE_GROW_SIZE;
113      }
114      if (minCapacity > newCapacity) {
115        newCapacity = minCapacity;
116      }
117      buf = Arrays.copyOf(buf, newCapacity);
118    }
119  }
120}