@Override public S getSchema() throws IOException { return this.embeddedExtractor.getSchema(); }
@Override public D readRecord(D reuse) throws DataRecordException, IOException { return actualExtractor.readRecord(reuse); }
@Override public long getExpectedRecordCount() { return this.embeddedExtractor.getExpectedRecordCount(); }
Assert.assertEquals(extractor.getSchema().toString(), jsonSchema); Assert.assertEquals(extractor.readRecord(null).toString(), json);
@Test public void testSourceExtractor() throws DataRecordException, IOException { SourceState state = new SourceState(); state.setProp(HelloWorldSource.NUM_HELLOS_FULL_KEY, 10); HelloWorldSource source = new HelloWorldSource(); List<WorkUnit> wus = source.getWorkunits(state); Assert.assertEquals(wus.size(), 10); for (int i = 0; i < wus.size(); ++i) { WorkUnit wu = wus.get(i); Assert.assertEquals(wu.getPropAsInt(HelloWorldSource.HELLO_ID_FULL_KEY), i + 1); WorkUnitState wuState = new WorkUnitState(wu, state); Extractor<String, String> extr = source.getExtractor(wuState); Assert.assertEquals(extr.getExpectedRecordCount(), 1); Assert.assertEquals(extr.readRecord(null), "Hello world "+ (i+1) + " !"); } }
/** * @param shutdownRequest an {@link AtomicBoolean} that becomes true when a shutdown has been requested. * @return a {@link Flowable} with the records from this source. Note the flowable should honor downstream backpressure. */ default RecordStreamWithMetadata<D, S> recordStream(AtomicBoolean shutdownRequest) throws IOException { S schema = getSchema(); Flowable<StreamEntity<D>> recordStream = Flowable.generate(() -> shutdownRequest, (BiConsumer<AtomicBoolean, Emitter<StreamEntity<D>>>) (state, emitter) -> { if (state.get()) { emitter.onComplete(); } try { StreamEntity<D> record = readStreamEntity(); if (record != null) { emitter.onNext(record); } else { emitter.onComplete(); } } catch (DataRecordException | IOException exc) { emitter.onError(exc); } }); recordStream = recordStream.doFinally(this::close); return new RecordStreamWithMetadata<>(recordStream, GlobalMetadata.<S>builder().schema(schema).build()); }
@Override public long getHighWatermark() { return actualExtractor.getHighWatermark(); } }
/** * Read an {@link StreamEntity}. By default, just return result of {@link #readRecordEnvelope()}. */ default StreamEntity<D> readStreamEntity() throws DataRecordException, IOException { return readRecordEnvelope(); }
@Override public RecordStreamWithMetadata<D, S> recordStream(AtomicBoolean shutdownRequest) throws IOException { if (this.isEmbeddedInstrumented) { return this.embeddedExtractor.recordStream(shutdownRequest); } RecordStreamWithMetadata<D, S> stream = this.embeddedExtractor.recordStream(shutdownRequest); stream = stream.mapRecords(r -> { if (this.lastRecordTime == 0) { this.lastRecordTime = System.nanoTime(); } afterRead(r.getRecord(), this.lastRecordTime); this.lastRecordTime = System.nanoTime(); return r; }); return stream; }
@Override public void close() throws IOException { try { this.extractor.close(); } finally { this.limiter.stop(); } } }
@Test public void testSourceExtractor() throws DataRecordException, IOException { final int MEM_ALLOC_BYTES = 100; final int NUM_WORK_UNITS = 10; final int COMPUTE_TIME_MICRO = 10; final int NUM_RECORDS = 10000; SourceState state = new SourceState(); state.setProp(StressTestingSource.NUM_WORK_UNITS_KEY, NUM_WORK_UNITS); state.setProp(StressTestingSource.MEM_ALLOC_BYTES_KEY, MEM_ALLOC_BYTES); state.setProp(StressTestingSource.COMPUTE_TIME_MICRO_KEY, COMPUTE_TIME_MICRO); state.setProp(StressTestingSource.NUM_RECORDS_KEY, NUM_RECORDS); StressTestingSource source = new StressTestingSource(); List<WorkUnit> wus = source.getWorkunits(state); Assert.assertEquals(wus.size(), NUM_WORK_UNITS); for (int i = 0; i < wus.size(); ++i) { WorkUnit wu = wus.get(i); WorkUnitState wuState = new WorkUnitState(wu, state); Extractor<String, byte[]> extractor = source.getExtractor(wuState); Assert.assertEquals(extractor.getExpectedRecordCount(), NUM_RECORDS); Assert.assertEquals(extractor.readRecord(null).length, 100); } }
/** * @param shutdownRequest an {@link AtomicBoolean} that becomes true when a shutdown has been requested. * @return a {@link Flowable} with the records from this source. Note the flowable should honor downstream backpressure. */ default RecordStreamWithMetadata<D, S> recordStream(AtomicBoolean shutdownRequest) throws IOException { S schema = getSchema(); Flowable<StreamEntity<D>> recordStream = Flowable.generate(() -> shutdownRequest, (BiConsumer<AtomicBoolean, Emitter<StreamEntity<D>>>) (state, emitter) -> { if (state.get()) { emitter.onComplete(); } try { StreamEntity<D> record = readStreamEntity(); if (record != null) { emitter.onNext(record); } else { emitter.onComplete(); } } catch (DataRecordException | IOException exc) { emitter.onError(exc); } }); recordStream = recordStream.doFinally(this::close); return new RecordStreamWithMetadata<>(recordStream, GlobalMetadata.<S>builder().schema(schema).build()); }
@Override public long getHighWatermark() { return this.embeddedExtractor.getHighWatermark(); }
@Override protected RecordEnvelope<D> readRecordEnvelopeImpl() throws DataRecordException, IOException { return this.embeddedExtractor.readRecordEnvelope(); }
RecordStreamWithMetadata<?, ?> stream = this.extractor.recordStream(this.shutdownRequested); ConnectableFlowable connectableStream = stream.getRecordStream().publish(); stream = stream.withRecordStream(connectableStream);
@Override public void close() throws IOException { try { this.extractor.close(); } finally { this.limiter.stop(); } } }
/** * Read an {@link RecordEnvelope}. By default, just wrap {@link #readRecord(Object)} in a {@link RecordEnvelope}. */ @SuppressWarnings(value = "RCN_REDUNDANT_NULLCHECK_OF_NONNULL_VALUE", justification = "Findbugs believes readRecord(null) is non-null. This is not true.") default RecordEnvelope<D> readRecordEnvelope() throws DataRecordException, IOException { D record = readRecord(null); return record == null ? null : new RecordEnvelope<>(record); }
@Override public S getSchema() throws IOException { return actualExtractor.getSchema(); }
@Override public long getExpectedRecordCount() { return actualExtractor.getExpectedRecordCount(); }
@Deprecated @Override public long getHighWatermark() { return this.extractor.getHighWatermark(); }