/*
* Copyright 2008-2013 LinkedIn, 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 voldemort.client.protocol.pb;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import voldemort.VoldemortException;
import voldemort.client.protocol.pb.VAdminProto.ROStoreVersionDirMap;
import voldemort.client.protocol.pb.VAdminProto.RebalanceTaskInfoMap;
import voldemort.client.protocol.pb.VAdminProto.StoreToPartitionsIds;
import voldemort.client.protocol.pb.VProto.KeyedVersions;
import voldemort.client.rebalance.RebalanceTaskInfo;
import voldemort.store.ErrorCodeMapper;
import voldemort.utils.ByteArray;
import voldemort.versioning.ClockEntry;
import voldemort.versioning.VectorClock;
import voldemort.versioning.Version;
import voldemort.versioning.Versioned;
import voldemort.xml.ClusterMapper;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.protobuf.ByteString;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.CodedOutputStream;
import com.google.protobuf.Message;
/**
* Helper functions for serializing or deserializing client requests in protocol
* buffers
*
*
*/
public class ProtoUtils {
/*
* Begin decode RebalanceTaskInfoMap methods
*/
/**
* Given a protobuf rebalance-partition info, converts it into our
* rebalance-partition info
*
* @param rebalanceTaskInfoMap Proto-buff version of
* RebalanceTaskInfoMap
* @return RebalanceTaskInfo object.
*/
public static RebalanceTaskInfo decodeRebalanceTaskInfoMap(VAdminProto.RebalanceTaskInfoMap rebalanceTaskInfoMap) {
RebalanceTaskInfo rebalanceTaskInfo = new RebalanceTaskInfo(
rebalanceTaskInfoMap.getStealerId(),
rebalanceTaskInfoMap.getDonorId(),
decodeStoreToPartitionIds(rebalanceTaskInfoMap.getPerStorePartitionIdsList()),
new ClusterMapper().readCluster(new StringReader(rebalanceTaskInfoMap.getInitialCluster())));
return rebalanceTaskInfo;
}
public static HashMap<String, List<Integer>> decodeStoreToPartitionIds(List<StoreToPartitionsIds> storeToPartitionIds) {
HashMap<String, List<Integer>> storeToPartitionIdsList = Maps.newHashMap();
for(StoreToPartitionsIds tuple: storeToPartitionIds) {
storeToPartitionIdsList.put(tuple.getStoreName(), tuple.getPartitionIdsList());
}
return storeToPartitionIdsList;
}
/*
* End decode RebalanceTaskInfoMap methods
*/
/*
* Begin encode RebalanceTaskInfo methods
*/
/**
* Given a rebalance-task info, convert it into the protobuf equivalent
*
* @param stealInfo Rebalance task info
* @return Protobuff equivalent of the same
*/
public static RebalanceTaskInfoMap encodeRebalanceTaskInfoMap(RebalanceTaskInfo stealInfo) {
return RebalanceTaskInfoMap.newBuilder()
.setStealerId(stealInfo.getStealerId())
.setDonorId(stealInfo.getDonorId())
.addAllPerStorePartitionIds(ProtoUtils.encodeStoreToPartitionsTuple(stealInfo.getStoreToPartitionIds()))
.setInitialCluster(new ClusterMapper().writeCluster(stealInfo.getInitialCluster()))
.build();
}
public static List<StoreToPartitionsIds> encodeStoreToPartitionsTuple(HashMap<String, List<Integer>> storeToPartitionIds) {
List<StoreToPartitionsIds> perStorePartitionTuples = Lists.newArrayList();
for(Entry<String, List<Integer>> entry: storeToPartitionIds.entrySet()) {
StoreToPartitionsIds.Builder tupleBuilder = StoreToPartitionsIds.newBuilder();
tupleBuilder.setStoreName(entry.getKey());
tupleBuilder.addAllPartitionIds(entry.getValue());
perStorePartitionTuples.add(tupleBuilder.build());
}
return perStorePartitionTuples;
}
/*
* End encode RebalanceTaskInfo methods
*/
public static Map<String, String> encodeROMap(List<ROStoreVersionDirMap> metadataMap) {
Map<String, String> storeToValue = Maps.newHashMap();
for(ROStoreVersionDirMap currentStore: metadataMap) {
storeToValue.put(currentStore.getStoreName(), currentStore.getStoreDir());
}
return storeToValue;
}
public static VProto.Error.Builder encodeError(ErrorCodeMapper mapper, VoldemortException e) {
return VProto.Error.newBuilder()
.setErrorCode(mapper.getCode(e))
.setErrorMessage(e.getMessage());
}
public static VProto.Versioned.Builder encodeVersioned(Versioned<byte[]> versioned) {
return VProto.Versioned.newBuilder()
.setValue(ByteString.copyFrom(versioned.getValue()))
.setVersion(ProtoUtils.encodeClock(versioned.getVersion()));
}
public static Versioned<byte[]> decodeVersioned(VProto.Versioned versioned) {
return new Versioned<byte[]>(versioned.getValue().toByteArray(),
decodeClock(versioned.getVersion()));
}
/**
* Given a list of value versions for the metadata keys we are just
* interested in the value at index 0 This is because even if we have to
* update the cluster.xml we marshall a single key into a versioned list
* Hence we just look at the value at index 0
*
*/
public static Versioned<byte[]> decodeVersionedMetadataKeyValue(KeyedVersions keyValue) {
return decodeVersioned(keyValue.getVersions(0));
}
public static List<Versioned<byte[]>> decodeVersions(List<VProto.Versioned> versioned) {
List<Versioned<byte[]>> values = new ArrayList<Versioned<byte[]>>(versioned.size());
for(VProto.Versioned v: versioned)
values.add(decodeVersioned(v));
return values;
}
public static VectorClock decodeClock(VProto.VectorClock encoded) {
List<ClockEntry> entries = new ArrayList<ClockEntry>(encoded.getEntriesCount());
for(VProto.ClockEntry entry: encoded.getEntriesList())
entries.add(new ClockEntry((short) entry.getNodeId(), entry.getVersion()));
return new VectorClock(entries, encoded.getTimestamp());
}
public static VProto.VectorClock.Builder encodeClock(Version version) {
VectorClock clock = (VectorClock) version;
VProto.VectorClock.Builder encoded = VProto.VectorClock.newBuilder();
encoded.setTimestamp(clock.getTimestamp());
for(ClockEntry entry: clock.getEntries())
encoded.addEntries(VProto.ClockEntry.newBuilder()
.setNodeId(entry.getNodeId())
.setVersion(entry.getVersion()));
return encoded;
}
public static ByteArray decodeBytes(ByteString string) {
return new ByteArray(string.toByteArray());
}
public static ByteString encodeBytes(ByteArray array) {
return ByteString.copyFrom(array.get());
}
public static ByteString encodeTransform(byte[] transform) {
return ByteString.copyFrom(transform);
}
public static void writeMessage(DataOutputStream output, Message message) throws IOException {
/*
* We don't use varints here because the c++ version of the protocol
* buffer classes seem to be buggy requesting more data than necessary
* from the underlying stream causing it to block forever
*/
output.writeInt(message.getSerializedSize());
CodedOutputStream codedOut = CodedOutputStream.newInstance(output);
message.writeTo(codedOut);
codedOut.flush();
}
public static void writeEndOfStream(DataOutputStream output) throws IOException {
output.writeInt(-1);
}
public static <T extends Message.Builder> T readToBuilder(DataInputStream input, T builder)
throws IOException {
int size = input.readInt();
CodedInputStream codedIn = CodedInputStream.newInstance(input);
codedIn.pushLimit(size);
builder.mergeFrom(codedIn);
return builder;
}
}