/*
* Copyright © 2015-2016 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.api.dataset.lib.partitioned;
import co.cask.cdap.api.common.Bytes;
import co.cask.cdap.api.dataset.lib.PartitionKey;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import java.nio.ByteBuffer;
/**
* Default implementation of {@link ConsumablePartition}.
*/
public final class DefaultConsumablePartition implements ConsumablePartition {
private static final Gson GSON =
new GsonBuilder().registerTypeAdapter(PartitionKey.class, new PartitionKeyCodec()).create();
private final PartitionKey partitionKey;
private ProcessState processState;
private int numFailures;
// the timestamp when it was marked IN_PROGRESS. it will be 0 if the ProcessState is AVAILABLE
private long timestamp;
public DefaultConsumablePartition(PartitionKey partitionKey) {
this(partitionKey, ProcessState.AVAILABLE, 0, 0);
}
public DefaultConsumablePartition(PartitionKey partitionKey, ProcessState processState,
long timestamp, int numFailures) {
this.partitionKey = partitionKey;
this.processState = processState;
this.timestamp = timestamp;
this.numFailures = numFailures;
}
@Override
public PartitionKey getPartitionKey() {
return partitionKey;
}
@Override
public long getTimestamp() {
return timestamp;
}
@Override
public void setTimestamp(long timestamp) {
this.timestamp = timestamp;
}
@Override
public int getNumFailures() {
return numFailures;
}
@Override
public int incrementNumFailures() {
return ++numFailures;
}
@Override
public ProcessState getProcessState() {
return processState;
}
@Override
public void setProcessState(ProcessState processState) {
this.processState = processState;
}
@Override
public void take() {
if (processState != ProcessState.AVAILABLE) {
throw new IllegalStateException();
}
processState = ProcessState.IN_PROGRESS;
}
@Override
public void retry() {
if (processState != ProcessState.IN_PROGRESS) {
throw new IllegalStateException();
}
processState = ProcessState.AVAILABLE;
numFailures++;
timestamp = 0;
}
@Override
public void complete() {
if (processState != ProcessState.IN_PROGRESS) {
throw new IllegalStateException();
}
processState = ProcessState.COMPLETED;
}
@Override
public void discard() {
if (processState != ProcessState.IN_PROGRESS) {
throw new IllegalStateException();
}
processState = ProcessState.DISCARDED;
}
static DefaultConsumablePartition fromBytes(byte[] bytes) {
ByteBuffer bb = ByteBuffer.wrap(bytes);
ProcessState processState = ProcessState.fromByte(bb.get());
int keyLength = bb.getInt();
byte[] stringBytes = new byte[keyLength];
bb.get(stringBytes, 0, keyLength);
long timestamp = bb.getLong();
int numFailures = bb.getInt();
return new DefaultConsumablePartition(GSON.fromJson(Bytes.toString(stringBytes), PartitionKey.class),
processState, timestamp, numFailures);
}
public byte[] toBytes() {
byte[] partitionKeyBytes = Bytes.toBytes(GSON.toJson(partitionKey));
// 1 byte for the ProcessState
int numBytes = 1;
numBytes += Bytes.SIZEOF_INT;
numBytes += partitionKeyBytes.length;
numBytes += Bytes.SIZEOF_LONG;
numBytes += Bytes.SIZEOF_INT;
ByteBuffer bb = ByteBuffer.allocate(numBytes);
bb.put(processState.toByte());
bb.putInt(partitionKeyBytes.length);
bb.put(partitionKeyBytes);
bb.putLong(getTimestamp());
bb.putInt(getNumFailures());
return bb.array();
}
@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
DefaultConsumablePartition that = (DefaultConsumablePartition) o;
return numFailures == that.numFailures && timestamp == that.timestamp
&& partitionKey.equals(that.partitionKey) && processState == that.processState;
}
@Override
public int hashCode() {
int result = partitionKey.hashCode();
result = 31 * result + processState.hashCode();
result = 31 * result + numFailures;
result = 31 * result + (int) (timestamp ^ (timestamp >>> 32));
return result;
}
}