Examples with VisibleForTesting org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting used on opensource projects

Example 41 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class NestedPayloadKafkaTable method transformOutput.

// Suppress nullability warnings: ProducerRecord is supposed to accept null arguments.
@SuppressWarnings("argument.type.incompatible")
@VisibleForTesting
ProducerRecord<byte[], byte[]> transformOutput(Row row) {
    row = castRow(row, row.getSchema(), schema);
    String topic = Iterables.getOnlyElement(getTopics());
    byte[] key = null;
    byte[] payload;
    List<Header> headers = ImmutableList.of();
    Long timestampMillis = null;
    if (schema.hasField(Schemas.MESSAGE_KEY_FIELD)) {
        key = row.getBytes(Schemas.MESSAGE_KEY_FIELD);
    }
    if (schema.hasField(Schemas.EVENT_TIMESTAMP_FIELD)) {
        ReadableDateTime time = row.getDateTime(Schemas.EVENT_TIMESTAMP_FIELD);
        if (time != null) {
            timestampMillis = time.getMillis();
        }
    }
    if (schema.hasField(Schemas.HEADERS_FIELD)) {
        Collection<Row> headerRows = checkArgumentNotNull(row.getArray(Schemas.HEADERS_FIELD));
        ImmutableList.Builder<Header> headersBuilder = ImmutableList.builder();
        headerRows.forEach(entry -> {
            String headerKey = checkArgumentNotNull(entry.getString(Schemas.HEADERS_KEY_FIELD));
            Collection<byte[]> values = checkArgumentNotNull(entry.getArray(Schemas.HEADERS_VALUES_FIELD));
            values.forEach(value -> headersBuilder.add(new RecordHeader(headerKey, value)));
        });
        headers = headersBuilder.build();
    }
    if (payloadSerializer == null) {
        payload = row.getBytes(Schemas.PAYLOAD_FIELD);
    } else {
        payload = payloadSerializer.serialize(checkArgumentNotNull(row.getRow(Schemas.PAYLOAD_FIELD)));
    }
    return new ProducerRecord<>(topic, null, timestampMillis, key, payload, headers);
}

Example 42 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class BatchDataflowWorker method doWork.

/**
 * Performs the given work; returns true if successful.
 *
 * @throws IOException Only if the WorkUnitClient fails.
 */
@VisibleForTesting
boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) throws IOException {
    LOG.debug("Executing: {}", workItem);
    DataflowWorkExecutor worker = null;
    SdkWorkerHarness sdkWorkerHarness = sdkHarnessRegistry.getAvailableWorkerAndAssignWork();
    try {
        // Populate PipelineOptions with data from work unit.
        options.setProject(workItem.getProjectId());
        final String stageName;
        if (workItem.getMapTask() != null) {
            stageName = workItem.getMapTask().getStageName();
        } else if (workItem.getSourceOperationTask() != null) {
            stageName = workItem.getSourceOperationTask().getStageName();
        } else {
            throw new RuntimeException("Unknown kind of work item: " + workItem.toString());
        }
        CounterSet counterSet = new CounterSet();
        BatchModeExecutionContext executionContext = BatchModeExecutionContext.create(counterSet, sideInputDataCache, sideInputWeakReferenceCache, readerRegistry, options, stageName, String.valueOf(workItem.getId()));
        if (workItem.getMapTask() != null) {
            MutableNetwork<Node, Edge> network = mapTaskToNetwork.apply(workItem.getMapTask());
            if (LOG.isDebugEnabled()) {
                LOG.debug("Network as Graphviz .dot: {}", Networks.toDot(network));
            }
            worker = mapTaskExecutorFactory.create(sdkWorkerHarness.getControlClientHandler(), sdkWorkerHarness.getGrpcDataFnServer(), sdkHarnessRegistry.beamFnDataApiServiceDescriptor(), sdkWorkerHarness.getGrpcStateFnServer(), network, options, stageName, readerRegistry, sinkRegistry, executionContext, counterSet, idGenerator);
        } else if (workItem.getSourceOperationTask() != null) {
            worker = SourceOperationExecutorFactory.create(options, workItem.getSourceOperationTask(), counterSet, executionContext, stageName);
        } else {
            throw new IllegalStateException("Work Item was neither a MapTask nor a SourceOperation");
        }
        workItemStatusClient.setWorker(worker, executionContext);
        DataflowWorkProgressUpdater progressUpdater = new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, options);
        executeWork(worker, progressUpdater);
        workItemStatusClient.reportSuccess();
        return true;
    } catch (Throwable e) {
        workItemStatusClient.reportError(e);
        return false;
    } finally {
        if (worker != null) {
            try {
                worker.close();
            } catch (Exception exn) {
                LOG.warn("Uncaught exception while closing worker. All work has already committed or " + "been marked for retry.", exn);
            }
        }
        if (sdkWorkerHarness != null) {
            sdkHarnessRegistry.completeWork(sdkWorkerHarness);
        }
    }
}

Example 43 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class DataflowBatchWorkerHarness method processWork.

@VisibleForTesting
static void processWork(DataflowWorkerHarnessOptions pipelineOptions, final BatchDataflowWorker worker, Sleeper sleeper) throws InterruptedException {
    int numThreads = chooseNumberOfThreads(pipelineOptions);
    ExecutorService executor = pipelineOptions.getExecutorService();
    final List<Callable<Boolean>> tasks = new ArrayList<>();
    LOG.debug("Starting {} worker threads", numThreads);
    // We start the appropriate number of threads.
    for (int i = 0; i < numThreads; ++i) {
        tasks.add(new WorkerThread(worker, sleeper));
    }
    LOG.debug("Waiting for {} worker threads", numThreads);
    // We wait forever unless there is a big problem.
    executor.invokeAll(tasks);
    LOG.error("All threads died.");
}

Example 44 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class ParDoTranslation method fromProto.

@VisibleForTesting
static StateSpec<?> fromProto(RunnerApi.StateSpec stateSpec, RehydratedComponents components) throws IOException {
    switch(stateSpec.getSpecCase()) {
        case READ_MODIFY_WRITE_SPEC:
            return StateSpecs.value(components.getCoder(stateSpec.getReadModifyWriteSpec().getCoderId()));
        case BAG_SPEC:
            return StateSpecs.bag(components.getCoder(stateSpec.getBagSpec().getElementCoderId()));
        case COMBINING_SPEC:
            FunctionSpec combineFnSpec = stateSpec.getCombiningSpec().getCombineFn();
            if (!combineFnSpec.getUrn().equals(CombineTranslation.JAVA_SERIALIZED_COMBINE_FN_URN)) {
                throw new UnsupportedOperationException(String.format("Cannot create %s from non-Java %s: %s", StateSpec.class.getSimpleName(), Combine.CombineFn.class.getSimpleName(), combineFnSpec.getUrn()));
            }
            Combine.CombineFn<?, ?, ?> combineFn = (Combine.CombineFn<?, ?, ?>) SerializableUtils.deserializeFromByteArray(combineFnSpec.getPayload().toByteArray(), Combine.CombineFn.class.getSimpleName());
            // for the CombineFn, by construction
            return StateSpecs.combining((Coder) components.getCoder(stateSpec.getCombiningSpec().getAccumulatorCoderId()), combineFn);
        case MAP_SPEC:
            return StateSpecs.map(components.getCoder(stateSpec.getMapSpec().getKeyCoderId()), components.getCoder(stateSpec.getMapSpec().getValueCoderId()));
        case SET_SPEC:
            return StateSpecs.set(components.getCoder(stateSpec.getSetSpec().getElementCoderId()));
        case SPEC_NOT_SET:
        default:
            throw new IllegalArgumentException(String.format("Unknown %s: %s", RunnerApi.StateSpec.class.getName(), stateSpec));
    }
}

Example 45 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class DataflowRunner method validateWorkerSettings.

@VisibleForTesting
static void validateWorkerSettings(DataflowPipelineWorkerPoolOptions workerOptions) {
    DataflowPipelineOptions dataflowOptions = workerOptions.as(DataflowPipelineOptions.class);
    validateSdkContainerImageOptions(workerOptions);
    GcpOptions gcpOptions = workerOptions.as(GcpOptions.class);
    Preconditions.checkArgument(gcpOptions.getZone() == null || gcpOptions.getWorkerRegion() == null, "Cannot use option zone with workerRegion. Prefer either workerZone or workerRegion.");
    Preconditions.checkArgument(gcpOptions.getZone() == null || gcpOptions.getWorkerZone() == null, "Cannot use option zone with workerZone. Prefer workerZone.");
    Preconditions.checkArgument(gcpOptions.getWorkerRegion() == null || gcpOptions.getWorkerZone() == null, "workerRegion and workerZone options are mutually exclusive.");
    boolean hasExperimentWorkerRegion = false;
    if (dataflowOptions.getExperiments() != null) {
        for (String experiment : dataflowOptions.getExperiments()) {
            if (experiment.startsWith("worker_region")) {
                hasExperimentWorkerRegion = true;
                break;
            }
        }
    }
    Preconditions.checkArgument(!hasExperimentWorkerRegion || gcpOptions.getWorkerRegion() == null, "Experiment worker_region and option workerRegion are mutually exclusive.");
    Preconditions.checkArgument(!hasExperimentWorkerRegion || gcpOptions.getWorkerZone() == null, "Experiment worker_region and option workerZone are mutually exclusive.");
    if (gcpOptions.getZone() != null) {
        LOG.warn("Option --zone is deprecated. Please use --workerZone instead.");
        gcpOptions.setWorkerZone(gcpOptions.getZone());
        gcpOptions.setZone(null);
    }
}