co.cask.cdap.data2.transaction.stream.StreamConsumerStateStore.java Source code

Java tutorial

Introduction

Here is the source code for co.cask.cdap.data2.transaction.stream.StreamConsumerStateStore.java

Source

/*
 * Copyright  2014 Cask Data, Inc.
 *
 * Licensed 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 co.cask.cdap.data2.transaction.stream;

import co.cask.cdap.api.common.Bytes;
import co.cask.cdap.data.stream.StreamFileOffset;
import co.cask.cdap.data.stream.StreamUtils;
import co.cask.cdap.proto.Id;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSortedMap;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.io.ByteArrayDataOutput;
import com.google.common.io.ByteStreams;
import com.google.common.primitives.Ints;
import com.google.common.primitives.Longs;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutput;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.SortedMap;

/**
 * Represents storage for {@link ConsumerState} for stream consumers.
 */
public abstract class StreamConsumerStateStore
        implements ConsumerStateStore<StreamConsumerState, Iterable<StreamFileOffset>> {

    protected final StreamConfig streamConfig;
    protected final Id.Stream streamId;

    protected StreamConsumerStateStore(StreamConfig streamConfig) {
        this.streamConfig = streamConfig;
        this.streamId = streamConfig.getStreamId();
    }

    @Override
    public final void getAll(Collection<? super StreamConsumerState> result) throws IOException {
        SortedMap<byte[], byte[]> states = Maps.newTreeMap(Bytes.BYTES_COMPARATOR);
        fetchAll(streamId.toBytes(), states);

        for (Map.Entry<byte[], byte[]> entry : states.entrySet()) {
            byte[] column = entry.getKey();
            byte[] value = entry.getValue();
            if (value != null) {
                result.add(
                        new StreamConsumerState(getGroupId(column), getInstanceId(column), decodeOffsets(value)));
            }
        }
    }

    @Override
    public final void getByGroup(long groupId, Collection<? super StreamConsumerState> result) throws IOException {
        SortedMap<byte[], byte[]> states = Maps.newTreeMap(Bytes.BYTES_COMPARATOR);
        fetchAll(streamId.toBytes(), Bytes.toBytes(groupId), states);

        for (Map.Entry<byte[], byte[]> entry : states.entrySet()) {
            byte[] column = entry.getKey();
            if (getGroupId(column) != groupId) {
                continue;
            }
            byte[] value = entry.getValue();
            if (value != null) {
                result.add(new StreamConsumerState(groupId, getInstanceId(column), decodeOffsets(value)));
            }
        }
    }

    @Override
    public final StreamConsumerState get(long groupId, int instanceId) throws IOException {
        byte[] value = fetch(streamId.toBytes(), getColumn(groupId, instanceId));
        return value == null ? null : new StreamConsumerState(groupId, instanceId, decodeOffsets(value));
    }

    @Override
    public final void save(StreamConsumerState state) throws IOException {
        store(streamId.toBytes(), getColumn(state.getGroupId(), state.getInstanceId()),
                encodeOffsets(state.getState()));
    }

    @Override
    public final void save(Iterable<? extends StreamConsumerState> states) throws IOException {
        ImmutableSortedMap.Builder<byte[], byte[]> values = ImmutableSortedMap.orderedBy(Bytes.BYTES_COMPARATOR);
        ByteArrayOutputStream os = new ByteArrayOutputStream();
        DataOutput output = new DataOutputStream(os);
        for (StreamConsumerState state : states) {
            os.reset();
            encodeOffsets(state.getState(), output);
            values.put(getColumn(state.getGroupId(), state.getInstanceId()), os.toByteArray());
        }

        store(streamId.toBytes(), values.build());
    }

    @Override
    public final void remove(Iterable<? extends StreamConsumerState> states) throws IOException {
        Set<byte[]> columns = Sets.newTreeSet(Bytes.BYTES_COMPARATOR);
        for (StreamConsumerState state : states) {
            columns.add(getColumn(state.getGroupId(), state.getInstanceId()));
        }
        delete(streamId.toBytes(), columns);
    }

    /**
     * Fetches the cell value for the given row and column.
     * If no such value exists, {@code null} should be returned.
     */
    protected abstract byte[] fetch(byte[] row, byte[] column) throws IOException;

    /**
     * Fetches all values for the given row.
     *
     * @param row the row to fetch from.
     * @param result a map from column to value.
     */
    protected abstract void fetchAll(byte[] row, Map<byte[], byte[]> result) throws IOException;

    /**
     * Fetches all values for the given row. Children can optionally use the columnPrefix to do more efficient retrieval.
     *
     * @param row the row to fetch from.
     * @param columnPrefix the column prefix.
     * @param result a map from column to value. It's valid to contains columns that don't start with columnPrefix.
     */
    protected abstract void fetchAll(byte[] row, byte[] columnPrefix, Map<byte[], byte[]> result)
            throws IOException;

    /**
     * Stores the given value to cell identified by the given row and column.
     */
    protected abstract void store(byte[] row, byte[] column, byte[] value) throws IOException;

    /**
     * Stores all the values to the given row.
     * @param row the row to store to.
     * @param values Map from column name to value.
     */
    protected abstract void store(byte[] row, Map<byte[], byte[]> values) throws IOException;

    /**
     * Deletes the set of columns from the given row.
     * @param row the row to act on.
     * @param columns columns to get deleted.
     */
    protected abstract void delete(byte[] row, Set<byte[]> columns) throws IOException;

    /**
     * Encodes list of {@link StreamFileOffset} into bytes.
     */
    private byte[] encodeOffsets(Iterable<StreamFileOffset> offsets) throws IOException {
        // Assumption: Each offset encoded into ~40 bytes and there are 8 offsets (number of live files)
        ByteArrayDataOutput output = ByteStreams.newDataOutput(320);
        encodeOffsets(offsets, output);
        return output.toByteArray();
    }

    private void encodeOffsets(Iterable<StreamFileOffset> offsets, DataOutput output) throws IOException {
        for (StreamFileOffset offset : offsets) {
            StreamUtils.encodeOffset(output, offset);
        }
    }

    /**
     * Decodes encoded bytes back to list of {@link StreamFileOffset}.
     */
    private Iterable<StreamFileOffset> decodeOffsets(byte[] encoded) throws IOException {
        ImmutableList.Builder<StreamFileOffset> offsets = ImmutableList.builder();
        if (encoded != null && encoded.length > 0) {
            DataInputStream input = new DataInputStream(new ByteArrayInputStream(encoded));
            while (input.available() > 0) {
                offsets.add(StreamUtils.decodeOffset(streamConfig, input));
            }
        }
        return offsets.build();
    }

    private byte[] getColumn(long groupId, int instanceId) {
        byte[] column = new byte[Longs.BYTES + Ints.BYTES];
        Bytes.putLong(column, 0, groupId);
        Bytes.putInt(column, Longs.BYTES, instanceId);
        return column;
    }

    /**
     * Decodes the group id from the column name.
     */
    private long getGroupId(byte[] columnName) {
        return Bytes.toLong(columnName);
    }

    /**
     * Decodes the instance id from the column name.
     */
    private int getInstanceId(byte[] columnName) {
        return Bytes.toInt(columnName, Longs.BYTES);
    }
}