use of org.apache.samza.system.inmemory.InMemorySystemFactory in project samza by apache.
the class TestRunner method initializeInMemoryInputStream.
/**
* Creates an in memory stream with {@link InMemorySystemFactory} and feeds its partition with stream of messages
* @param partitionData key of the map represents partitionId and value represents messages in the partition
* @param descriptor describes a stream to initialize with the in memory system
*/
private <StreamMessageType> void initializeInMemoryInputStream(InMemoryInputDescriptor<?> descriptor, Map<Integer, Iterable<StreamMessageType>> partitionData) {
String systemName = descriptor.getSystemName();
String streamName = (String) descriptor.getPhysicalName().orElse(descriptor.getStreamId());
if (this.app instanceof LegacyTaskApplication) {
// for legacy applications that only specify task.class.
if (configs.containsKey(TaskConfig.INPUT_STREAMS)) {
configs.put(TaskConfig.INPUT_STREAMS, configs.get(TaskConfig.INPUT_STREAMS).concat("," + systemName + "." + streamName));
} else {
configs.put(TaskConfig.INPUT_STREAMS, systemName + "." + streamName);
}
}
InMemorySystemDescriptor imsd = (InMemorySystemDescriptor) descriptor.getSystemDescriptor();
imsd.withInMemoryScope(this.inMemoryScope);
addConfig(descriptor.toConfig());
addConfig(descriptor.getSystemDescriptor().toConfig());
addSerdeConfigs(descriptor);
StreamSpec spec = new StreamSpec(descriptor.getStreamId(), streamName, systemName, partitionData.size());
SystemFactory factory = new InMemorySystemFactory();
Config config = new MapConfig(descriptor.toConfig(), descriptor.getSystemDescriptor().toConfig());
factory.getAdmin(systemName, config).createStream(spec);
InMemorySystemProducer producer = (InMemorySystemProducer) factory.getProducer(systemName, config, null);
SystemStream sysStream = new SystemStream(systemName, streamName);
partitionData.forEach((partitionId, partition) -> {
partition.forEach(e -> {
Object key = e instanceof KV ? ((KV) e).getKey() : null;
Object value = e instanceof KV ? ((KV) e).getValue() : e;
if (value instanceof IncomingMessageEnvelope) {
producer.send((IncomingMessageEnvelope) value);
} else {
producer.send(systemName, new OutgoingMessageEnvelope(sysStream, Integer.valueOf(partitionId), key, value));
}
});
producer.send(systemName, new OutgoingMessageEnvelope(sysStream, Integer.valueOf(partitionId), null, new EndOfStreamMessage(null)));
});
}
use of org.apache.samza.system.inmemory.InMemorySystemFactory in project samza by apache.
the class TestRunner method consumeStream.
/**
* Gets the contents of the output stream represented by {@code outputDescriptor} after {@link TestRunner#run(Duration)}
* has completed
*
* @param outputDescriptor describes the stream to be consumed
* @param timeout timeout for consumption of stream in Ms
* @param <StreamMessageType> type of message
*
* @return a map whose key is {@code partitionId} and value is messages in partition
* @throws SamzaException Thrown when a poll is incomplete
*/
public static <StreamMessageType> Map<Integer, List<StreamMessageType>> consumeStream(InMemoryOutputDescriptor outputDescriptor, Duration timeout) throws SamzaException {
Preconditions.checkNotNull(outputDescriptor);
String streamId = outputDescriptor.getStreamId();
String systemName = outputDescriptor.getSystemName();
Set<SystemStreamPartition> ssps = new HashSet<>();
Set<String> streamIds = new HashSet<>();
streamIds.add(streamId);
SystemFactory factory = new InMemorySystemFactory();
Config config = new MapConfig(outputDescriptor.toConfig(), outputDescriptor.getSystemDescriptor().toConfig());
Map<String, SystemStreamMetadata> metadata = factory.getAdmin(systemName, config).getSystemStreamMetadata(streamIds);
SystemConsumer consumer = factory.getConsumer(systemName, config, null);
String name = (String) outputDescriptor.getPhysicalName().orElse(streamId);
metadata.get(name).getSystemStreamPartitionMetadata().keySet().forEach(partition -> {
SystemStreamPartition temp = new SystemStreamPartition(systemName, streamId, partition);
ssps.add(temp);
consumer.register(temp, "0");
});
long t = System.currentTimeMillis();
Map<SystemStreamPartition, List<IncomingMessageEnvelope>> output = new HashMap<>();
HashSet<SystemStreamPartition> didNotReachEndOfStream = new HashSet<>(ssps);
while (System.currentTimeMillis() < t + timeout.toMillis()) {
Map<SystemStreamPartition, List<IncomingMessageEnvelope>> currentState = null;
try {
currentState = consumer.poll(ssps, 10);
} catch (InterruptedException e) {
throw new SamzaException("Timed out while consuming stream \n" + e.getMessage());
}
for (Map.Entry<SystemStreamPartition, List<IncomingMessageEnvelope>> entry : currentState.entrySet()) {
SystemStreamPartition ssp = entry.getKey();
output.computeIfAbsent(ssp, k -> new LinkedList<IncomingMessageEnvelope>());
List<IncomingMessageEnvelope> currentBuffer = entry.getValue();
int totalMessagesToFetch = Integer.valueOf(metadata.get(outputDescriptor.getStreamId()).getSystemStreamPartitionMetadata().get(ssp.getPartition()).getUpcomingOffset());
if (output.get(ssp).size() + currentBuffer.size() == totalMessagesToFetch) {
didNotReachEndOfStream.remove(entry.getKey());
ssps.remove(entry.getKey());
}
output.get(ssp).addAll(currentBuffer);
}
if (didNotReachEndOfStream.isEmpty()) {
break;
}
}
if (!didNotReachEndOfStream.isEmpty()) {
throw new IllegalStateException("Could not poll for all system stream partitions");
}
return output.entrySet().stream().collect(Collectors.toMap(entry -> entry.getKey().getPartition().getPartitionId(), entry -> entry.getValue().stream().map(e -> (StreamMessageType) e.getMessage()).collect(Collectors.toList())));
}
use of org.apache.samza.system.inmemory.InMemorySystemFactory in project samza by apache.
the class TestRunner method addOutputStream.
/**
* Adds the provided output stream to the test application. Default configs and user added configs have a higher
* precedence over system and stream descriptor generated configs.
* @param streamDescriptor describes the stream that is supposed to be output for the Samza application
* @param partitionCount partition count of output stream
* @return this {@link TestRunner}
*/
public TestRunner addOutputStream(InMemoryOutputDescriptor<?> streamDescriptor, int partitionCount) {
Preconditions.checkNotNull(streamDescriptor);
Preconditions.checkState(partitionCount >= 1);
InMemorySystemDescriptor imsd = (InMemorySystemDescriptor) streamDescriptor.getSystemDescriptor();
imsd.withInMemoryScope(this.inMemoryScope);
Config config = new MapConfig(streamDescriptor.toConfig(), streamDescriptor.getSystemDescriptor().toConfig());
InMemorySystemFactory factory = new InMemorySystemFactory();
String physicalName = (String) streamDescriptor.getPhysicalName().orElse(streamDescriptor.getStreamId());
StreamSpec spec = new StreamSpec(streamDescriptor.getStreamId(), physicalName, streamDescriptor.getSystemName(), partitionCount);
factory.getAdmin(streamDescriptor.getSystemName(), config).createStream(spec);
addConfig(streamDescriptor.toConfig());
addConfig(streamDescriptor.getSystemDescriptor().toConfig());
addSerdeConfigs(streamDescriptor);
return this;
}
use of org.apache.samza.system.inmemory.InMemorySystemFactory in project beam by apache.
the class TranslationContext method createDummyStreamDescriptor.
/**
* The dummy stream created will only be used in Beam tests.
*/
private static InputDescriptor<OpMessage<String>, ?> createDummyStreamDescriptor(String id) {
final GenericSystemDescriptor dummySystem = new GenericSystemDescriptor(id, InMemorySystemFactory.class.getName());
final GenericInputDescriptor<OpMessage<String>> dummyInput = dummySystem.getInputDescriptor(id, new NoOpSerde<>());
dummyInput.withOffsetDefault(SystemStreamMetadata.OffsetType.OLDEST);
final Config config = new MapConfig(dummyInput.toConfig(), dummySystem.toConfig());
final SystemFactory factory = new InMemorySystemFactory();
final StreamSpec dummyStreamSpec = new StreamSpec(id, id, id, 1);
factory.getAdmin(id, config).createStream(dummyStreamSpec);
final SystemProducer producer = factory.getProducer(id, config, null);
final SystemStream sysStream = new SystemStream(id, id);
final Consumer<Object> sendFn = (msg) -> {
producer.send(id, new OutgoingMessageEnvelope(sysStream, 0, null, msg));
};
final WindowedValue<String> windowedValue = WindowedValue.timestampedValueInGlobalWindow("dummy", new Instant());
sendFn.accept(OpMessage.ofElement(windowedValue));
sendFn.accept(new WatermarkMessage(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
sendFn.accept(new EndOfStreamMessage(null));
return dummyInput;
}
Aggregations