org.apache.cassandra.io.util.TeeDataInputPlus Maven / Gradle / Ivy

Go to download
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.cassandra.io.util;

import java.io.EOFException;
import java.io.IOException;

import org.apache.cassandra.db.TypeSizes;
import org.apache.cassandra.utils.Throwables;

/**
 * DataInput that also stores the raw inputs into an output buffer
 * This is useful for storing serialized buffers as they are deserialized.
 *
 * Note: If a non-zero limit is included it is important to for callers to check {@link #isLimitReached()}
 * before using the tee buffer as it could be cropped.
 */
public class TeeDataInputPlus implements DataInputPlus
{
    private final DataInputPlus source;
    private final DataOutputPlus teeBuffer;

    private final long limit;
    private boolean limitReached;

    public TeeDataInputPlus(DataInputPlus source, DataOutputPlus teeBuffer)
    {
        this(source, teeBuffer, 0);
    }

    public TeeDataInputPlus(DataInputPlus source, DataOutputPlus teeBuffer, long limit)
    {
        assert source != null && teeBuffer != null;
        this.source = source;
        this.teeBuffer = teeBuffer;
        this.limit = limit;
        this.limitReached = false;
    }

    private void maybeWrite(int length, Throwables.DiscreteAction writeAction) throws IOException
    {
        if (limit <= 0 || (!limitReached && (teeBuffer.position() + length) < limit))
            writeAction.perform();
        else
            limitReached = true;
    }

    @Override
    public void readFully(byte[] bytes) throws IOException
    {
        source.readFully(bytes);
        maybeWrite(bytes.length, () -> teeBuffer.write(bytes));
    }

    @Override
    public void readFully(byte[] bytes, int offset, int length) throws IOException
    {
        source.readFully(bytes, offset, length);
        maybeWrite(length, () -> teeBuffer.write(bytes, offset, length));
    }

    @Override
    public int skipBytes(int n) throws IOException
    {
        for (int i = 0; i < n; i++)
        {
            try
            {
                byte v = source.readByte();
                maybeWrite(TypeSizes.BYTE_SIZE, () -> teeBuffer.writeByte(v));
            }
            catch (EOFException eof)
            {
                return i;
            }
        }
        return n;
    }

    @Override
    public boolean readBoolean() throws IOException
    {
        boolean v = source.readBoolean();
        maybeWrite(TypeSizes.BOOL_SIZE, () -> teeBuffer.writeBoolean(v));
        return v;
    }

    @Override
    public byte readByte() throws IOException
    {
        byte v = source.readByte();
        maybeWrite(TypeSizes.BYTE_SIZE, () -> teeBuffer.writeByte(v));
        return v;
    }

    @Override
    public int readUnsignedByte() throws IOException
    {
        int v = source.readUnsignedByte();
        maybeWrite(TypeSizes.BYTE_SIZE, () -> teeBuffer.writeByte(v));
        return v;
    }

    @Override
    public short readShort() throws IOException
    {
        short v = source.readShort();
        maybeWrite(TypeSizes.SHORT_SIZE, () -> teeBuffer.writeShort(v));
        return v;
    }

    @Override
    public int readUnsignedShort() throws IOException
    {
        int v = source.readUnsignedShort();
        maybeWrite(TypeSizes.SHORT_SIZE, () -> teeBuffer.writeShort(v));
        return v;
    }

    @Override
    public char readChar() throws IOException
    {
        char v = source.readChar();
        maybeWrite(TypeSizes.BYTE_SIZE, () -> teeBuffer.writeChar(v));
        return v;
    }

    @Override
    public int readInt() throws IOException
    {
        int v = source.readInt();
        maybeWrite(TypeSizes.INT_SIZE, () -> teeBuffer.writeInt(v));
        return v;
    }

    @Override
    public long readLong() throws IOException
    {
        long v = source.readLong();
        maybeWrite(TypeSizes.LONG_SIZE, () -> teeBuffer.writeLong(v));
        return v;
    }

    @Override
    public float readFloat() throws IOException
    {
        float v = source.readFloat();
        maybeWrite(TypeSizes.FLOAT_SIZE, () -> teeBuffer.writeFloat(v));
        return v;
    }

    @Override
    public double readDouble() throws IOException
    {
        double v = source.readDouble();
        maybeWrite(TypeSizes.DOUBLE_SIZE, () -> teeBuffer.writeDouble(v));
        return v;
    }

    @Override
    public String readLine() throws IOException
    {
        //This one isn't safe since we know the actual line termination type
        throw new UnsupportedOperationException();
    }

    @Override
    public String readUTF() throws IOException
    {
        String v = source.readUTF();
        maybeWrite(TypeSizes.sizeof(v), () -> teeBuffer.writeUTF(v));
        return v;
    }

    @Override
    public long readVInt() throws IOException
    {
        long v = source.readVInt();
        maybeWrite(TypeSizes.sizeofVInt(v), () -> teeBuffer.writeVInt(v));
        return v;
    }

    @Override
    public long readUnsignedVInt() throws IOException
    {
        long v = source.readUnsignedVInt();
        maybeWrite(TypeSizes.sizeofUnsignedVInt(v), () -> teeBuffer.writeUnsignedVInt(v));
        return v;
    }

    @Override
    public void skipBytesFully(int n) throws IOException
    {
        source.skipBytesFully(n);
        maybeWrite(n, () -> {
            for (int i = 0; i < n; i++)
                teeBuffer.writeByte(0);
        });
    }

    /**
     * Used to detect if the teeBuffer hit the supplied limit.
     * If true this means the teeBuffer does not contain the full input.
     */
    public boolean isLimitReached()
    {
        return limitReached;
    }
}