Examples with Input org.talend.sdk.component.runtime.input.Input used on opensource projects

Example 11 with Input

use of org.talend.sdk.component.runtime.input.Input in project component-runtime by Talend.

the class BaseComponentsHandler method collect.

/**
 * Collects data emitted from this mapper. If the split creates more than one
 * mapper, it will create as much threads as mappers otherwise it will use the
 * caller thread.
 *
 * IMPORTANT: don't forget to consume all the stream to ensure the underlying
 * { @see org.talend.sdk.component.runtime.input.Input} is closed.
 *
 * @param recordType the record type to use to type the returned type.
 * @param mapper the mapper to go through.
 * @param maxRecords maximum number of records, allows to stop the source when
 * infinite.
 * @param concurrency requested (1 can be used instead if <= 0) concurrency for the reader execution.
 * @param <T> the returned type of the records of the mapper.
 * @return all the records emitted by the mapper.
 */
@Override
public <T> Stream<T> collect(final Class<T> recordType, final Mapper mapper, final int maxRecords, final int concurrency) {
    mapper.start();
    final State state = STATE.get();
    final long assess = mapper.assess();
    final int proc = Math.max(1, concurrency);
    final List<Mapper> mappers = mapper.split(Math.max(assess / proc, 1));
    switch(mappers.size()) {
        case 0:
            return Stream.empty();
        case 1:
            return StreamDecorator.decorate(asStream(asIterator(mappers.iterator().next().create(), new AtomicInteger(maxRecords))), collect -> {
                try {
                    collect.run();
                } finally {
                    mapper.stop();
                }
            });
        default:
            // N producers-1 consumer pattern
            final AtomicInteger threadCounter = new AtomicInteger(0);
            final ExecutorService es = Executors.newFixedThreadPool(mappers.size(), r -> new Thread(r) {

                {
                    setName(BaseComponentsHandler.this.getClass().getSimpleName() + "-pool-" + abs(mapper.hashCode()) + "-" + threadCounter.incrementAndGet());
                }
            });
            final AtomicInteger recordCounter = new AtomicInteger(maxRecords);
            final Semaphore permissions = new Semaphore(0);
            final Queue<T> records = new ConcurrentLinkedQueue<>();
            final CountDownLatch latch = new CountDownLatch(mappers.size());
            final List<? extends Future<?>> tasks = mappers.stream().map(Mapper::create).map(input -> (Iterator<T>) asIterator(input, recordCounter)).map(it -> es.submit(() -> {
                try {
                    while (it.hasNext()) {
                        final T next = it.next();
                        records.add(next);
                        permissions.release();
                    }
                } finally {
                    latch.countDown();
                }
            })).collect(toList());
            es.shutdown();
            final int timeout = Integer.getInteger("talend.component.junit.timeout", 5);
            new Thread() {

                {
                    setName(BaseComponentsHandler.class.getSimpleName() + "-monitor_" + abs(mapper.hashCode()));
                }

                @Override
                public void run() {
                    try {
                        latch.await(timeout, MINUTES);
                    } catch (final InterruptedException e) {
                        Thread.interrupted();
                    } finally {
                        permissions.release();
                    }
                }
            }.start();
            return StreamDecorator.decorate(asStream(new Iterator<T>() {

                @Override
                public boolean hasNext() {
                    try {
                        permissions.acquire();
                    } catch (final InterruptedException e) {
                        Thread.interrupted();
                        fail(e.getMessage());
                    }
                    return !records.isEmpty();
                }

                @Override
                public T next() {
                    T poll = records.poll();
                    if (poll != null) {
                        return mapRecord(state, recordType, poll);
                    }
                    return null;
                }
            }), task -> {
                try {
                    task.run();
                } finally {
                    tasks.forEach(f -> {
                        try {
                            f.get(5, SECONDS);
                        } catch (final InterruptedException e) {
                            Thread.interrupted();
                        } catch (final ExecutionException | TimeoutException e) {
                        // no-op
                        } finally {
                            if (!f.isDone() && !f.isCancelled()) {
                                f.cancel(true);
                            }
                        }
                    });
                }
            });
    }
}