Examples with QueueConsumer co.cask.cdap.data2.queue.QueueConsumer used on opensource projects

Example 6 with QueueConsumer

use of co.cask.cdap.data2.queue.QueueConsumer in project cdap by caskdata.

the class HBaseQueueTest method testReconfigure.

@Test(timeout = 30000L)
public void testReconfigure() throws Exception {
    final QueueName queueName = QueueName.fromFlowlet(NamespaceId.DEFAULT.getEntityName(), "app", "flow", "flowlet", "changeinstances");
    ConsumerGroupConfig groupConfig = new ConsumerGroupConfig(0L, 2, DequeueStrategy.HASH, "key");
    configureGroups(queueName, ImmutableList.of(groupConfig));
    // Enqueue 10 items
    createEnqueueRunnable(queueName, 10, 1, null).run();
    // Map from instance id to items dequeued
    final Multimap<Integer, Integer> dequeued = ArrayListMultimap.create();
    // Consume 2 items for each consumer instances
    for (int instanceId = 0; instanceId < groupConfig.getGroupSize(); instanceId++) {
        final ConsumerConfig consumerConfig = new ConsumerConfig(groupConfig, instanceId);
        try (QueueConsumer consumer = queueClientFactory.createConsumer(queueName, consumerConfig, 1)) {
            Transactions.createTransactionExecutor(executorFactory, (TransactionAware) consumer).execute(new TransactionExecutor.Subroutine() {

                @Override
                public void apply() throws Exception {
                    DequeueResult<byte[]> result = consumer.dequeue(2);
                    Assert.assertEquals(2, result.size());
                    for (byte[] data : result) {
                        dequeued.put(consumerConfig.getInstanceId(), Bytes.toInt(data));
                    }
                }
            });
        }
    }
    // Increase number of instances to 3
    changeInstances(queueName, 0L, 3);
    // Enqueue 10 more items
    createEnqueueRunnable(queueName, 10, 1, null).run();
    groupConfig = new ConsumerGroupConfig(0L, 3, DequeueStrategy.HASH, "key");
    // Dequeue everything
    while (dequeued.size() != 20) {
        for (int instanceId = 0; instanceId < groupConfig.getGroupSize(); instanceId++) {
            final ConsumerConfig consumerConfig = new ConsumerConfig(groupConfig, instanceId);
            try (QueueConsumer consumer = queueClientFactory.createConsumer(queueName, consumerConfig, 1)) {
                Transactions.createTransactionExecutor(executorFactory, (TransactionAware) consumer).execute(new TransactionExecutor.Subroutine() {

                    @Override
                    public void apply() throws Exception {
                        for (byte[] data : consumer.dequeue(20)) {
                            dequeued.put(consumerConfig.getInstanceId(), Bytes.toInt(data));
                        }
                    }
                });
            }
        }
    }
    // Instance 0 should see all evens before change instances
    Assert.assertEquals(ImmutableList.of(0, 2, 4, 6, 8, 0, 3, 6, 9), dequeued.get(0));
    // Instance 1 should see all odds before change instances
    Assert.assertEquals(ImmutableList.of(1, 3, 5, 7, 9, 1, 4, 7), dequeued.get(1));
    // Instance 2 should only see entries after change instances
    Assert.assertEquals(ImmutableList.of(2, 5, 8), dequeued.get(2));
    // All consumers should have empty dequeue now
    for (int instanceId = 0; instanceId < groupConfig.getGroupSize(); instanceId++) {
        final ConsumerConfig consumerConfig = new ConsumerConfig(groupConfig, instanceId);
        try (QueueConsumer consumer = queueClientFactory.createConsumer(queueName, consumerConfig, 1)) {
            Transactions.createTransactionExecutor(executorFactory, (TransactionAware) consumer).execute(new TransactionExecutor.Subroutine() {

                @Override
                public void apply() throws Exception {
                    DequeueResult<byte[]> result = consumer.dequeue(20);
                    Assert.assertTrue(result.isEmpty());
                }
            });
        }
    }
    // Enqueue 6 more items for the 3 instances
    createEnqueueRunnable(queueName, 6, 1, null).run();
    // Reduce to 1 consumer
    changeInstances(queueName, 0L, 1);
    // The consumer 0 should be able to consume all 10 new items
    dequeued.clear();
    final ConsumerConfig consumerConfig = new ConsumerConfig(0L, 0, 1, DequeueStrategy.HASH, "key");
    try (final QueueConsumer consumer = queueClientFactory.createConsumer(queueName, consumerConfig, 1)) {
        while (dequeued.size() != 6) {
            Transactions.createTransactionExecutor(executorFactory, (TransactionAware) consumer).execute(new TransactionExecutor.Subroutine() {

                @Override
                public void apply() throws Exception {
                    for (byte[] data : consumer.dequeue(1)) {
                        dequeued.put(consumerConfig.getInstanceId(), Bytes.toInt(data));
                    }
                }
            });
        }
    }
    Assert.assertEquals(ImmutableList.of(0, 1, 2, 3, 4, 5), dequeued.get(0));
}

Example 7 with QueueConsumer

use of co.cask.cdap.data2.queue.QueueConsumer in project cdap by caskdata.

the class FlowletProgramRunner method processSpecificationFactory.

private ProcessSpecificationFactory processSpecificationFactory(final BasicFlowletContext flowletContext, final DataFabricFacade dataFabricFacade, final QueueReaderFactory queueReaderFactory, final String flowletName, final Table<Node, String, Set<QueueSpecification>> queueSpecs, final ImmutableList.Builder<ConsumerSupplier<?>> queueConsumerSupplierBuilder, final SchemaCache schemaCache) {
    return new ProcessSpecificationFactory() {

        @Override
        public <T> ProcessSpecification create(Set<String> inputNames, Schema schema, TypeToken<T> dataType, ProcessMethod<T> method, ConsumerConfig consumerConfig, int batchSize, Tick tickAnnotation) throws Exception {
            List<QueueReader<T>> queueReaders = Lists.newLinkedList();
            for (Map.Entry<Node, Set<QueueSpecification>> entry : queueSpecs.column(flowletName).entrySet()) {
                for (QueueSpecification queueSpec : entry.getValue()) {
                    final QueueName queueName = queueSpec.getQueueName();
                    if (queueSpec.getInputSchema().equals(schema) && (inputNames.contains(queueName.getSimpleName()) || inputNames.contains(FlowletDefinition.ANY_INPUT))) {
                        Node sourceNode = entry.getKey();
                        if (sourceNode.getType() == FlowletConnection.Type.STREAM) {
                            ConsumerSupplier<StreamConsumer> consumerSupplier = ConsumerSupplier.create(flowletContext.getOwners(), runtimeUsageRegistry, dataFabricFacade, queueName, consumerConfig);
                            queueConsumerSupplierBuilder.add(consumerSupplier);
                            // No decoding is needed, as a process method can only have StreamEvent as type for consuming stream
                            Function<StreamEvent, T> decoder = wrapInputDecoder(flowletContext, null, queueName, new Function<StreamEvent, T>() {

                                @Override
                                @SuppressWarnings("unchecked")
                                public T apply(StreamEvent input) {
                                    return (T) input;
                                }
                            });
                            queueReaders.add(queueReaderFactory.createStreamReader(queueName.toStreamId(), consumerSupplier, batchSize, decoder));
                        } else {
                            int numGroups = getNumGroups(Iterables.concat(queueSpecs.row(entry.getKey()).values()), queueName);
                            Function<ByteBuffer, T> decoder = wrapInputDecoder(// the producer flowlet,
                            flowletContext, // the producer flowlet,
                            entry.getKey().getName(), queueName, createInputDatumDecoder(dataType, schema, schemaCache));
                            ConsumerSupplier<QueueConsumer> consumerSupplier = ConsumerSupplier.create(flowletContext.getOwners(), runtimeUsageRegistry, dataFabricFacade, queueName, consumerConfig, numGroups);
                            queueConsumerSupplierBuilder.add(consumerSupplier);
                            queueReaders.add(queueReaderFactory.createQueueReader(consumerSupplier, batchSize, decoder));
                        }
                    }
                }
            }
            // If inputs is needed but there is no available input queue, return null
            if (!inputNames.isEmpty() && queueReaders.isEmpty()) {
                return null;
            }
            return new ProcessSpecification<>(new RoundRobinQueueReader<>(queueReaders), method, tickAnnotation);
        }
    };
}

Example 8 with QueueConsumer

use of co.cask.cdap.data2.queue.QueueConsumer in project cdap by caskdata.

the class HBaseQueueClientFactory method createConsumer.

@Override
public QueueConsumer createConsumer(final QueueName queueName, final ConsumerConfig consumerConfig, int numGroups) throws IOException {
    final HBaseQueueAdmin admin = ensureTableExists(queueName);
    try {
        final long groupId = consumerConfig.getGroupId();
        // A callback for create a list of HBaseQueueConsumer
        // based on the current queue consumer state of the given group
        Callable<List<HBaseQueueConsumer>> consumerCreator = new Callable<List<HBaseQueueConsumer>>() {

            @Override
            public List<HBaseQueueConsumer> call() throws Exception {
                List<HBaseConsumerState> states;
                try (HBaseConsumerStateStore stateStore = admin.getConsumerStateStore(queueName)) {
                    TransactionExecutor txExecutor = Transactions.createTransactionExecutor(txExecutorFactory, stateStore);
                    // Find all consumer states for consumers that need to be created based on current state
                    states = txExecutor.execute(new Callable<List<HBaseConsumerState>>() {

                        @Override
                        public List<HBaseConsumerState> call() throws Exception {
                            List<HBaseConsumerState> consumerStates = Lists.newArrayList();
                            HBaseConsumerState state = stateStore.getState(groupId, consumerConfig.getInstanceId());
                            if (state.getPreviousBarrier() == null) {
                                // Old HBase consumer (Salted based, not sharded)
                                consumerStates.add(state);
                                return consumerStates;
                            }
                            // Find the smallest start barrier that has something to consume for this instance.
                            // It should always exists since we assume the queue is configured before this method is called
                            List<QueueBarrier> queueBarriers = stateStore.getAllBarriers(groupId);
                            if (queueBarriers.isEmpty()) {
                                throw new IllegalStateException(String.format("No consumer information available. Queue: %s, GroupId: %d, InstanceId: %d", queueName, groupId, consumerConfig.getInstanceId()));
                            }
                            QueueBarrier startBarrier = Iterables.find(Lists.reverse(queueBarriers), new Predicate<QueueBarrier>() {

                                @Override
                                public boolean apply(QueueBarrier barrier) {
                                    return barrier.getGroupConfig().getGroupSize() > consumerConfig.getInstanceId() && stateStore.isAllConsumed(consumerConfig, barrier.getStartRow());
                                }
                            }, queueBarriers.get(0));
                            int groupSize = startBarrier.getGroupConfig().getGroupSize();
                            for (int i = consumerConfig.getInstanceId(); i < groupSize; i += consumerConfig.getGroupSize()) {
                                consumerStates.add(stateStore.getState(groupId, i));
                            }
                            return consumerStates;
                        }
                    });
                }
                List<HBaseQueueConsumer> consumers = Lists.newArrayList();
                for (HBaseConsumerState state : states) {
                    QueueType queueType = (state.getPreviousBarrier() == null) ? QueueType.QUEUE : QueueType.SHARDED_QUEUE;
                    HTable hTable = createHTable(admin.getDataTableId(queueName, queueType));
                    int distributorBuckets = getDistributorBuckets(hTable.getTableDescriptor());
                    HBaseQueueStrategy strategy = (state.getPreviousBarrier() == null) ? new SaltedHBaseQueueStrategy(hBaseTableUtil, distributorBuckets) : new ShardedHBaseQueueStrategy(hBaseTableUtil, distributorBuckets);
                    consumers.add(queueUtil.getQueueConsumer(cConf, hTable, queueName, state, admin.getConsumerStateStore(queueName), strategy));
                }
                return consumers;
            }
        };
        return new SmartQueueConsumer(queueName, consumerConfig, consumerCreator);
    } catch (Exception e) {
        // If there is exception, nothing much can be done here besides propagating
        Throwables.propagateIfPossible(e);
        throw new IOException(e);
    }
}

Example 9 with QueueConsumer

use of co.cask.cdap.data2.queue.QueueConsumer in project cdap by caskdata.

the class LineageWriterDataFabricFacade method createConsumer.

@Override
public QueueConsumer createConsumer(QueueName queueName, ConsumerConfig consumerConfig, int numGroups) throws IOException {
    QueueConsumer consumer = queueClientFactory.createConsumer(queueName, consumerConfig, numGroups);
    if (consumer instanceof TransactionAware) {
        consumer = new CloseableQueueConsumer(datasetCache, consumer);
        datasetCache.addExtraTransactionAware((TransactionAware) consumer);
    }
    return consumer;
}

Example 10 with QueueConsumer

use of co.cask.cdap.data2.queue.QueueConsumer in project cdap by caskdata.

the class ConsumerSupplier method open.

/**
   * Updates number of instances for the consumer group that this instance belongs to. It'll close existing
   * consumer and create a new one with the new group size.
   *
   * @param groupSize New group size.
   */
void open(int groupSize) {
    try {
        close();
        ConsumerConfig config = consumerConfig;
        if (groupSize != config.getGroupSize()) {
            config = new ConsumerConfig(consumerConfig.getGroupId(), consumerConfig.getInstanceId(), groupSize, consumerConfig.getDequeueStrategy(), consumerConfig.getHashKey());
        }
        if (queueName.isQueue()) {
            QueueConsumer queueConsumer = dataFabricFacade.createConsumer(queueName, config, numGroups);
            consumerConfig = queueConsumer.getConfig();
            consumer = queueConsumer;
        } else {
            StreamId queueStream = queueName.toStreamId();
            for (EntityId owner : owners) {
                try {
                    runtimeUsageRegistry.register(owner, queueStream);
                } catch (Exception e) {
                    LOG.warn("Failed to register usage of {} -> {}", owner, queueStream, e);
                }
            }
            StreamConsumer streamConsumer = dataFabricFacade.createStreamConsumer(queueName.toStreamId(), config);
            consumerConfig = streamConsumer.getConsumerConfig();
            consumer = streamConsumer;
        }
    } catch (Exception e) {
        throw Throwables.propagate(e);
    }
}