/** * Adds a new partition to the working set. */ public void addPartition(PartitionKey partitionKey) { partitions.add(new DefaultConsumablePartition(partitionKey)); }
@Override public void take() { assertState(ProcessState.AVAILABLE); processState = ProcessState.IN_PROGRESS; }
@Override public void retry() { untake(); 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(); }
public byte[] toBytes() { // first byte for serialization format version int numBytes = 1; numBytes += Bytes.SIZEOF_INT; for (ConsumablePartition partition : partitions) { byte[] partitionBytes = ((DefaultConsumablePartition) partition).toBytes(); numBytes += Bytes.SIZEOF_INT; numBytes += partitionBytes.length; } byte[] markerBytes = partitionConsumerState.toBytes(); numBytes += Bytes.SIZEOF_INT; numBytes += markerBytes.length; ByteBuffer bb = ByteBuffer.allocate(numBytes); bb.put((byte) VERSION); bb.putInt(partitions.size()); for (ConsumablePartition partition : partitions) { byte[] partitionBytes = ((DefaultConsumablePartition) partition).toBytes(); bb.putInt(partitionBytes.length); bb.put(partitionBytes); } bb.putInt(markerBytes.length); bb.put(markerBytes); return bb.array(); }
public static ConsumerWorkingSet fromBytes(byte[] bytes) { ByteBuffer bb = ByteBuffer.wrap(bytes); byte serializationFormatVersion = bb.get(); if (serializationFormatVersion != VERSION) { throw new IllegalArgumentException("Unsupported serialization format: " + serializationFormatVersion); } int numPartitions = bb.getInt(); List<ConsumablePartition> partitions = new ArrayList<>(numPartitions); for (int i = 0; i < numPartitions; i++) { int consumablePartitionBytesLength = bb.getInt(); byte[] consumablePartitionBytes = new byte[consumablePartitionBytesLength]; bb.get(consumablePartitionBytes, 0, consumablePartitionBytesLength); partitions.add(DefaultConsumablePartition.fromBytes(consumablePartitionBytes)); } int sizeOfMarker = bb.getInt(); byte[] markerBytes = new byte[sizeOfMarker]; bb.get(markerBytes); return new ConsumerWorkingSet(PartitionConsumerState.fromBytes(markerBytes), partitions); }
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); }
@Override public void complete() { assertState(ProcessState.IN_PROGRESS); processState = ProcessState.COMPLETED; }
@Test(expected = IllegalStateException.class) public void testRetryWithoutTakenTransition() { // cannot take retry a partition without it being taken first ConsumablePartition partition = new DefaultConsumablePartition(generateUniqueKey()); partition.retry(); }
@Override public void untake() { assertState(ProcessState.IN_PROGRESS); processState = ProcessState.AVAILABLE; timestamp = 0; }
@Test(expected = IllegalStateException.class) public void testAlreadyTakenTransition() { // cannot take a partition that's already taken ConsumablePartition partition = new DefaultConsumablePartition(generateUniqueKey()); partition.take(); partition.take(); }
@Override public void discard() { assertState(ProcessState.IN_PROGRESS); processState = ProcessState.DISCARDED; }
@Test public void testByteSerialization() { ConsumerWorkingSet workingSet = new ConsumerWorkingSet(); // test with empty partitions lists testSerDe(workingSet); // test with two elements in AVAILABLE and none in IN_PROGRESS workingSet.getPartitions().add(new DefaultConsumablePartition(generateUniqueKey())); workingSet.getPartitions().add(new DefaultConsumablePartition(generateUniqueKey())); testSerDe(workingSet); // test with three elements in partitions and none in inProgressPartitions workingSet.getPartitions().add(new DefaultConsumablePartition(generateUniqueKey())); testSerDe(workingSet); // mark the first element as IN_PROGRESS workingSet.getPartitions().get(0).take(); workingSet.getPartitions().get(0).setTimestamp(System.currentTimeMillis()); testSerDe(workingSet); }
@Test(expected = IllegalStateException.class) public void testAlreadyCompletedTransition() { // cannot complete a partition that has already been completed ConsumablePartition partition = new DefaultConsumablePartition(generateUniqueKey()); partition.take(); partition.complete(); partition.complete(); }
@Test public void testSimpleTransitions() { // tests simple success case ConsumablePartition partition = new DefaultConsumablePartition(generateUniqueKey()); Assert.assertEquals(0, partition.getNumFailures()); partition.take(); partition.retry(); Assert.assertEquals(1, partition.getNumFailures()); partition.take(); // test that untake doesn't increment failure count partition.untake(); Assert.assertEquals(1, partition.getNumFailures()); partition.take(); partition.complete(); }