Examples with CloseableIterator io.cdap.cdap.api.dataset.lib.CloseableIterator used on opensource projects

Example 11 with CloseableIterator

use of io.cdap.cdap.api.dataset.lib.CloseableIterator in project cdap by cdapio.

the class SparkTest method testTransaction.

@Test
public void testTransaction() throws Exception {
    ApplicationManager applicationManager = deploy(TestSparkApp.class);
    // Write some data to a local file
    File inputFile = TEMP_FOLDER.newFile();
    try (PrintWriter writer = new PrintWriter(Files.newBufferedWriter(inputFile.toPath(), StandardCharsets.UTF_8))) {
        writer.println("red fox");
        writer.println("brown fox");
        writer.println("grey fox");
        writer.println("brown bear");
        writer.println("black bear");
    }
    // Run the spark program
    SparkManager sparkManager = applicationManager.getSparkManager(TransactionSpark.class.getSimpleName());
    sparkManager.start(ImmutableMap.of("input.file", inputFile.getAbsolutePath(), "keyvalue.table", "KeyValueTable", "result.all.dataset", "SparkResult", "result.threshold", "2", "result.threshold.dataset", "SparkThresholdResult"));
    // Verify result from dataset before the Spark program terminates
    final DataSetManager<KeyValueTable> resultManager = getDataset("SparkThresholdResult");
    final KeyValueTable resultTable = resultManager.get();
    // Expect the threshold result dataset, with threshold >=2, contains [brown, fox, bear]
    Tasks.waitFor(ImmutableSet.of("brown", "fox", "bear"), () -> {
        // This is to start a new TX
        resultManager.flush();
        LOG.info("Reading from threshold result");
        try (CloseableIterator<KeyValue<byte[], byte[]>> itor = resultTable.scan(null, null)) {
            return ImmutableSet.copyOf(Iterators.transform(itor, input -> {
                String word = Bytes.toString(input.getKey());
                LOG.info("{}, {}", word, Bytes.toInt(input.getValue()));
                return word;
            }));
        }
    }, 3, TimeUnit.MINUTES, 1, TimeUnit.SECONDS);
    sparkManager.stop();
    sparkManager.waitForRun(ProgramRunStatus.KILLED, 60, TimeUnit.SECONDS);
}

Example 12 with CloseableIterator

use of io.cdap.cdap.api.dataset.lib.CloseableIterator in project cdap by cdapio.

the class LeaderElectionMessagingServiceTest method testTransition.

@Test
public void testTransition() throws Throwable {
    final TopicId topicId = NamespaceId.SYSTEM.topic("topic");
    Injector injector1 = createInjector(0);
    Injector injector2 = createInjector(1);
    // Start a messaging service, which would becomes leader
    ZKClientService zkClient1 = injector1.getInstance(ZKClientService.class);
    zkClient1.startAndWait();
    final MessagingService firstService = injector1.getInstance(MessagingService.class);
    if (firstService instanceof Service) {
        ((Service) firstService).startAndWait();
    }
    // Publish a message with the leader
    firstService.publish(StoreRequestBuilder.of(topicId).addPayload("Testing1").build());
    // Start another messaging service, this one would be follower
    ZKClientService zkClient2 = injector2.getInstance(ZKClientService.class);
    zkClient2.startAndWait();
    final MessagingService secondService = injector2.getInstance(MessagingService.class);
    if (secondService instanceof Service) {
        ((Service) secondService).startAndWait();
    }
    // Try to call the follower, should get service unavailable.
    try {
        secondService.listTopics(NamespaceId.SYSTEM);
        Assert.fail("Expected service unavailable");
    } catch (ServiceUnavailableException e) {
    // Expected
    }
    // Make the ZK session timeout for the leader service. The second one should pickup.
    KillZKSession.kill(zkClient1.getZooKeeperSupplier().get(), zkClient1.getConnectString(), 10000);
    // Publish one more message and then fetch from the current leader
    List<String> messages = Retries.callWithRetries(new Retries.Callable<List<String>, Throwable>() {

        @Override
        public List<String> call() throws Throwable {
            secondService.publish(StoreRequestBuilder.of(topicId).addPayload("Testing2").build());
            List<String> messages = new ArrayList<>();
            try (CloseableIterator<RawMessage> iterator = secondService.prepareFetch(topicId).fetch()) {
                while (iterator.hasNext()) {
                    messages.add(new String(iterator.next().getPayload(), "UTF-8"));
                }
            }
            return messages;
        }
    }, RetryStrategies.timeLimit(10, TimeUnit.SECONDS, RetryStrategies.fixDelay(1, TimeUnit.SECONDS)));
    Assert.assertEquals(Arrays.asList("Testing1", "Testing2"), messages);
    // Shutdown the current leader. The session timeout one should becomes leader again.
    if (secondService instanceof Service) {
        ((Service) secondService).stopAndWait();
    }
    // Try to fetch message from the current leader again.
    // Should see two messages (because the cache is cleared and fetch is from the backing store).
    messages = Retries.callWithRetries(new Retries.Callable<List<String>, Throwable>() {

        @Override
        public List<String> call() throws Throwable {
            List<String> messages = new ArrayList<>();
            try (CloseableIterator<RawMessage> iterator = firstService.prepareFetch(topicId).fetch()) {
                while (iterator.hasNext()) {
                    messages.add(new String(iterator.next().getPayload(), "UTF-8"));
                }
            }
            return messages;
        }
    }, RetryStrategies.timeLimit(10, TimeUnit.SECONDS, RetryStrategies.fixDelay(1, TimeUnit.SECONDS)));
    Assert.assertEquals(Arrays.asList("Testing1", "Testing2"), messages);
    zkClient1.stopAndWait();
    zkClient2.stopAndWait();
}

Example 13 with CloseableIterator

use of io.cdap.cdap.api.dataset.lib.CloseableIterator in project cdap by cdapio.

the class AppMetadataStore method scanApplications.

/**
 * Scans applications. Allows to optionally set namespace / filters and implement pagination. For pagination
 * set {@link ScanApplicationsRequest#getScanFrom()} to the last application id of the previous page.
 *
 * @param request parameters defining filters and sorting
 * @param func a {@link Function} to consume application metadata entries generated by the scan. The boolean
 *             value returned is {@code true}, the scan will continue; otherwise the scan will stop and return.
 *             Note that the parameter is a {@link Map.Entry} to allow lazy deserialization of
 *             {@link ApplicationMeta} and it should not be replaced with {@link BiFunction}.
 * @see ScanApplicationsRequest#builder(ScanApplicationsRequest) to create a next page / batch request
 * @throws IOException if failed to scan the storage
 */
public void scanApplications(ScanApplicationsRequest request, Function<Map.Entry<ApplicationId, ApplicationMeta>, Boolean> func) throws IOException {
    Range.Bound startBound = Range.Bound.INCLUSIVE;
    Collection<Field<?>> startFields = request.getNamespaceId() == null ? Collections.emptyList() : Collections.singletonList(Fields.stringField(StoreDefinition.AppMetadataStore.NAMESPACE_FIELD, request.getNamespaceId().getNamespace()));
    Range.Bound endBound = Range.Bound.INCLUSIVE;
    Collection<Field<?>> endFields = startFields;
    if (request.getScanFrom() != null) {
        if (request.getNamespaceId() != null && !request.getNamespaceId().equals(request.getScanFrom().getNamespaceId())) {
            throw new IllegalArgumentException("Requested to start scan from application " + request.getScanFrom() + " that is outside of scan namespace " + request.getNamespaceId());
        }
        startBound = Range.Bound.EXCLUSIVE;
        startFields = getApplicationPrimaryKeys(request.getScanFrom());
    }
    if (request.getScanTo() != null) {
        if (request.getNamespaceId() != null && !request.getNamespaceId().equals(request.getScanTo().getNamespaceId())) {
            throw new IllegalArgumentException("Requested to finish scan at application " + request.getScanTo() + " that is outside of scan namespace " + request.getNamespaceId());
        }
        endBound = Range.Bound.EXCLUSIVE;
        endFields = getApplicationPrimaryKeys(request.getScanTo());
    }
    Range range;
    if (request.getSortOrder() == SortOrder.ASC) {
        range = Range.create(startFields, startBound, endFields, endBound);
    } else {
        range = Range.create(endFields, endBound, startFields, startBound);
    }
    // As of now this is where we push filter to. it does not go to the StructuredTable,
    // but we don't deserialize ApplicationMeta unless needed
    Predicate<AppScanEntry> scanEntryPredicate = e -> true;
    for (ApplicationFilter filter : request.getFilters()) {
        if (filter instanceof ApplicationFilter.ApplicationIdFilter) {
            scanEntryPredicate = scanEntryPredicate.and(e -> ((ApplicationFilter.ApplicationIdFilter) filter).test(e.getKey()));
        } else if (filter instanceof ApplicationFilter.ArtifactIdFilter) {
            scanEntryPredicate = scanEntryPredicate.and(e -> ((ApplicationFilter.ArtifactIdFilter) filter).test(e.getValue().getSpec().getArtifactId()));
        } else {
            throw new UnsupportedOperationException("Application filter " + filter + " is not supported");
        }
    }
    StructuredTable table = getApplicationSpecificationTable();
    int limit = request.getLimit();
    try (CloseableIterator<StructuredRow> iterator = table.scan(range, Integer.MAX_VALUE, request.getSortOrder())) {
        boolean keepScanning = true;
        while (iterator.hasNext() && keepScanning && limit > 0) {
            StructuredRow row = iterator.next();
            AppScanEntry scanEntry = new AppScanEntry(row);
            if (scanEntryPredicate.test(scanEntry)) {
                keepScanning = func.apply(scanEntry);
                limit--;
            }
        }
    }
}

Example 14 with CloseableIterator

use of io.cdap.cdap.api.dataset.lib.CloseableIterator in project cdap by caskdata.

the class AppMetadataStore method scanApplications.

/**
 * Scans applications. Allows to optionally set namespace / filters and implement pagination. For pagination
 * set {@link ScanApplicationsRequest#getScanFrom()} to the last application id of the previous page.
 *
 * @param request parameters defining filters and sorting
 * @param func a {@link Function} to consume application metadata entries generated by the scan. The boolean
 *             value returned is {@code true}, the scan will continue; otherwise the scan will stop and return.
 *             Note that the parameter is a {@link Map.Entry} to allow lazy deserialization of
 *             {@link ApplicationMeta} and it should not be replaced with {@link BiFunction}.
 * @see ScanApplicationsRequest#builder(ScanApplicationsRequest) to create a next page / batch request
 * @throws IOException if failed to scan the storage
 */
public void scanApplications(ScanApplicationsRequest request, Function<Map.Entry<ApplicationId, ApplicationMeta>, Boolean> func) throws IOException {
    Range.Bound startBound = Range.Bound.INCLUSIVE;
    Collection<Field<?>> startFields = request.getNamespaceId() == null ? Collections.emptyList() : Collections.singletonList(Fields.stringField(StoreDefinition.AppMetadataStore.NAMESPACE_FIELD, request.getNamespaceId().getNamespace()));
    Range.Bound endBound = Range.Bound.INCLUSIVE;
    Collection<Field<?>> endFields = startFields;
    if (request.getScanFrom() != null) {
        if (request.getNamespaceId() != null && !request.getNamespaceId().equals(request.getScanFrom().getNamespaceId())) {
            throw new IllegalArgumentException("Requested to start scan from application " + request.getScanFrom() + " that is outside of scan namespace " + request.getNamespaceId());
        }
        startBound = Range.Bound.EXCLUSIVE;
        startFields = getApplicationPrimaryKeys(request.getScanFrom());
    }
    if (request.getScanTo() != null) {
        if (request.getNamespaceId() != null && !request.getNamespaceId().equals(request.getScanTo().getNamespaceId())) {
            throw new IllegalArgumentException("Requested to finish scan at application " + request.getScanTo() + " that is outside of scan namespace " + request.getNamespaceId());
        }
        endBound = Range.Bound.EXCLUSIVE;
        endFields = getApplicationPrimaryKeys(request.getScanTo());
    }
    Range range;
    if (request.getSortOrder() == SortOrder.ASC) {
        range = Range.create(startFields, startBound, endFields, endBound);
    } else {
        range = Range.create(endFields, endBound, startFields, startBound);
    }
    // As of now this is where we push filter to. it does not go to the StructuredTable,
    // but we don't deserialize ApplicationMeta unless needed
    Predicate<AppScanEntry> scanEntryPredicate = e -> true;
    for (ApplicationFilter filter : request.getFilters()) {
        if (filter instanceof ApplicationFilter.ApplicationIdFilter) {
            scanEntryPredicate = scanEntryPredicate.and(e -> ((ApplicationFilter.ApplicationIdFilter) filter).test(e.getKey()));
        } else if (filter instanceof ApplicationFilter.ArtifactIdFilter) {
            scanEntryPredicate = scanEntryPredicate.and(e -> ((ApplicationFilter.ArtifactIdFilter) filter).test(e.getValue().getSpec().getArtifactId()));
        } else {
            throw new UnsupportedOperationException("Application filter " + filter + " is not supported");
        }
    }
    StructuredTable table = getApplicationSpecificationTable();
    int limit = request.getLimit();
    try (CloseableIterator<StructuredRow> iterator = table.scan(range, Integer.MAX_VALUE, request.getSortOrder())) {
        boolean keepScanning = true;
        while (iterator.hasNext() && keepScanning && limit > 0) {
            StructuredRow row = iterator.next();
            AppScanEntry scanEntry = new AppScanEntry(row);
            if (scanEntryPredicate.test(scanEntry)) {
                keepScanning = func.apply(scanEntry);
                limit--;
            }
        }
    }
}

Example 15 with CloseableIterator

use of io.cdap.cdap.api.dataset.lib.CloseableIterator in project cdap by caskdata.

the class MessagingAppTestRun method testWithWorker.

@Test
public void testWithWorker() throws Exception {
    ApplicationManager appManager = deployWithArtifact(NAMESPACE, MessagingApp.class, artifactJar);
    final WorkerManager workerManager = appManager.getWorkerManager(MessagingApp.MessagingWorker.class.getSimpleName()).start();
    MessagingContext messagingContext = getMessagingContext();
    final MessagingAdmin messagingAdmin = getMessagingAdmin(NAMESPACE);
    // Wait for the worker to create the topic
    Tasks.waitFor(true, new Callable<Boolean>() {

        @Override
        public Boolean call() throws Exception {
            try {
                messagingAdmin.getTopicProperties(MessagingApp.TOPIC);
                return true;
            } catch (TopicNotFoundException e) {
                return false;
            }
        }
    }, 5, TimeUnit.SECONDS, 100, TimeUnit.MILLISECONDS);
    // Publish a message
    String message = "message";
    MessagePublisher messagePublisher = messagingContext.getMessagePublisher();
    messagePublisher.publish(NAMESPACE.getNamespace(), MessagingApp.TOPIC, message);
    // The worker will publish back a message with payload as concat(message, message)
    final MessageFetcher messageFetcher = messagingContext.getMessageFetcher();
    Tasks.waitFor(message + message, new Callable<String>() {

        @Override
        public String call() throws Exception {
            try (CloseableIterator<Message> iterator = messageFetcher.fetch(NAMESPACE.getNamespace(), MessagingApp.TOPIC, Integer.MAX_VALUE, 0L)) {
                Message message = Iterators.getLast(iterator, null);
                return message == null ? null : message.getPayloadAsString();
            }
        }
    }, 5, TimeUnit.SECONDS, 100, TimeUnit.MILLISECONDS);
    // Publish concat(message + message) to the app
    messagePublisher.publish(NAMESPACE.getNamespace(), MessagingApp.TOPIC, message + message);
    // timeout.
    try {
        Tasks.waitFor(message + message + message + message, new Callable<String>() {

            @Override
            public String call() throws Exception {
                try (CloseableIterator<Message> iterator = messageFetcher.fetch(NAMESPACE.getNamespace(), MessagingApp.TOPIC, Integer.MAX_VALUE, 0L)) {
                    Message message = Iterators.getLast(iterator, null);
                    return message == null ? null : message.getPayloadAsString();
                }
            }
        }, 2, TimeUnit.SECONDS, 100, TimeUnit.MILLISECONDS);
        Assert.fail("Expected timeout exception");
    } catch (TimeoutException e) {
    // expected
    }
    // Now publish a message to the control topic, to unblock the transaction block.
    messagePublisher.publish(NAMESPACE.getNamespace(), MessagingApp.CONTROL_TOPIC, message);
    // Should expect a new message as concat(message, message, message, message)
    Tasks.waitFor(message + message + message + message, new Callable<String>() {

        @Override
        public String call() throws Exception {
            try (CloseableIterator<Message> iterator = messageFetcher.fetch(NAMESPACE.getNamespace(), MessagingApp.TOPIC, Integer.MAX_VALUE, 0L)) {
                Message message = Iterators.getLast(iterator, null);
                return message == null ? null : message.getPayloadAsString();
            }
        }
    }, 5, TimeUnit.SECONDS, 100, TimeUnit.MILLISECONDS);
    // Wait for the worker to finish and verify that it completes successfully.
    workerManager.waitForRun(ProgramRunStatus.COMPLETED, 5, TimeUnit.SECONDS);
}