| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.runners.dataflow.worker; |
| |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState; |
| |
| import com.google.auto.value.AutoValue; |
| import java.io.IOException; |
| import java.util.concurrent.CompletionStage; |
| import java.util.concurrent.ConcurrentLinkedQueue; |
| import java.util.concurrent.ExecutionException; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.RegisterRequest; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort; |
| import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor; |
| import org.apache.beam.model.pipeline.v1.RunnerApi; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.SdkFunctionSpec; |
| import org.apache.beam.runners.core.construction.CoderTranslation; |
| import org.apache.beam.runners.core.construction.RehydratedComponents; |
| import org.apache.beam.runners.core.construction.SdkComponents; |
| import org.apache.beam.runners.fnexecution.control.InstructionRequestHandler; |
| import org.apache.beam.runners.fnexecution.data.FnDataService; |
| import org.apache.beam.sdk.coders.ByteArrayCoder; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.KvCoder; |
| import org.apache.beam.sdk.fn.IdGenerator; |
| import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver; |
| import org.apache.beam.sdk.fn.data.InboundDataClient; |
| import org.apache.beam.sdk.fn.data.LogicalEndpoint; |
| import org.apache.beam.sdk.fn.data.RemoteGrpcPortRead; |
| import org.apache.beam.sdk.fn.data.RemoteGrpcPortWrite; |
| import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| import org.apache.beam.sdk.transforms.windowing.GlobalWindow; |
| import org.apache.beam.sdk.transforms.windowing.WindowMappingFn; |
| import org.apache.beam.sdk.util.MoreFutures; |
| import org.apache.beam.sdk.util.WindowedValue; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.WindowingStrategy; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.Cache; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** |
| * Maps main input windows onto side input windows. |
| * |
| * <p>Note that this {@link WindowMappingFn} performs blocking calls over the data plane. The thread |
| * which maps side input windows should not be on the same thread that is reading from the same data |
| * plane channel. |
| * |
| * <p>TODO: Swap this with an implementation which streams all required window mappings per bundle |
| * instead of per mapping request. This requires rewriting the {@link StreamingSideInputFetcher} to |
| * not be inline calls and process elements over a stream. |
| */ |
| class FnApiWindowMappingFn<TargetWindowT extends BoundedWindow> |
| extends WindowMappingFn<TargetWindowT> { |
| private static final Logger LOG = LoggerFactory.getLogger(FnApiWindowMappingFn.class); |
| private static final byte[] EMPTY_ARRAY = new byte[0]; |
| |
| @AutoValue |
| public abstract static class CacheKey { |
| public static CacheKey create(SdkFunctionSpec windowMappingFn, BoundedWindow mainWindow) { |
| return new AutoValue_FnApiWindowMappingFn_CacheKey(windowMappingFn, mainWindow); |
| } |
| |
| public abstract SdkFunctionSpec getWindowMappingFn(); |
| |
| public abstract BoundedWindow getMainWindow(); |
| } |
| |
| private static final Cache<CacheKey, BoundedWindow> sideInputMappingCache = |
| CacheBuilder.newBuilder().maximumSize(1000).build(); |
| |
| private final IdGenerator idGenerator; |
| private final FnDataService beamFnDataService; |
| private final InstructionRequestHandler instructionRequestHandler; |
| private final SdkFunctionSpec windowMappingFn; |
| private final Coder<WindowedValue<KV<byte[], BoundedWindow>>> outboundCoder; |
| private final Coder<WindowedValue<KV<byte[], TargetWindowT>>> inboundCoder; |
| private final ProcessBundleDescriptor processBundleDescriptor; |
| |
| FnApiWindowMappingFn( |
| IdGenerator idGenerator, |
| InstructionRequestHandler instructionRequestHandler, |
| ApiServiceDescriptor dataServiceApiServiceDescriptor, |
| FnDataService beamFnDataService, |
| SdkFunctionSpec windowMappingFn, |
| Coder<BoundedWindow> mainInputWindowCoder, |
| Coder<TargetWindowT> sideInputWindowCoder) { |
| this.idGenerator = idGenerator; |
| this.instructionRequestHandler = instructionRequestHandler; |
| this.beamFnDataService = beamFnDataService; |
| this.windowMappingFn = windowMappingFn; |
| |
| SdkComponents sdkComponents = SdkComponents.create(); |
| sdkComponents.registerEnvironment(RunnerApi.Environment.getDefaultInstance()); |
| String mainInputWindowCoderId; |
| String sideInputWindowCoderId; |
| String windowingStrategyId; |
| RunnerApi.Components components; |
| try { |
| mainInputWindowCoderId = |
| sdkComponents.registerCoder( |
| WindowedValue.getFullCoder( |
| KvCoder.of(ByteArrayCoder.of(), mainInputWindowCoder), |
| GlobalWindow.Coder.INSTANCE)); |
| sideInputWindowCoderId = |
| sdkComponents.registerCoder( |
| WindowedValue.getFullCoder( |
| KvCoder.of(ByteArrayCoder.of(), sideInputWindowCoder), |
| GlobalWindow.Coder.INSTANCE)); |
| windowingStrategyId = |
| sdkComponents.registerWindowingStrategy(WindowingStrategy.globalDefault()); |
| RehydratedComponents rehydratedComponents = |
| RehydratedComponents.forComponents(sdkComponents.toComponents()); |
| components = sdkComponents.toComponents(); |
| outboundCoder = |
| (Coder) |
| CoderTranslation.fromProto( |
| components.getCodersOrThrow(mainInputWindowCoderId), rehydratedComponents); |
| inboundCoder = |
| (Coder) |
| CoderTranslation.fromProto( |
| components.getCodersOrThrow(sideInputWindowCoderId), rehydratedComponents); |
| } catch (IOException e) { |
| throw new IllegalStateException( |
| "Unable to create side input window mapping process bundle specification.", e); |
| } |
| |
| processBundleDescriptor = |
| ProcessBundleDescriptor.newBuilder() |
| .putAllCoders(components.getCodersMap()) |
| .putPcollections( |
| "inPC", |
| PCollection.newBuilder() |
| .setCoderId(mainInputWindowCoderId) |
| .setWindowingStrategyId(windowingStrategyId) |
| .build()) |
| .putPcollections( |
| "outPC", |
| PCollection.newBuilder() |
| .setCoderId(sideInputWindowCoderId) |
| .setWindowingStrategyId(windowingStrategyId) |
| .build()) |
| .putTransforms( |
| "read", |
| RunnerApi.PTransform.newBuilder() |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(RemoteGrpcPortRead.URN) |
| .setPayload( |
| RemoteGrpcPort.newBuilder() |
| .setApiServiceDescriptor(dataServiceApiServiceDescriptor) |
| .setCoderId(mainInputWindowCoderId) |
| .build() |
| .toByteString())) |
| .putOutputs("in", "inPC") |
| .build()) |
| .putTransforms( |
| "map", |
| RunnerApi.PTransform.newBuilder() |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn("beam:transform:map_windows:v1") |
| .setPayload(windowMappingFn.toByteString())) |
| .putInputs("in", "inPC") |
| .putOutputs("out", "outPC") |
| .build()) |
| .putTransforms( |
| "write", |
| RunnerApi.PTransform.newBuilder() |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(RemoteGrpcPortWrite.URN) |
| .setPayload( |
| RemoteGrpcPort.newBuilder() |
| .setApiServiceDescriptor(dataServiceApiServiceDescriptor) |
| .setCoderId(sideInputWindowCoderId) |
| .build() |
| .toByteString())) |
| .putInputs("out", "outPC") |
| .build()) |
| .putAllWindowingStrategies(components.getWindowingStrategiesMap()) |
| .build(); |
| } |
| |
| @Override |
| public TargetWindowT getSideInputWindow(BoundedWindow mainWindow) { |
| try { |
| TargetWindowT rval = |
| (TargetWindowT) |
| sideInputMappingCache.get( |
| CacheKey.create(windowMappingFn, mainWindow), |
| () -> loadIfNeeded(windowMappingFn, mainWindow)); |
| return rval; |
| } catch (ExecutionException e) { |
| throw new IllegalStateException( |
| String.format("Unable to load side input window for %s", mainWindow), e); |
| } |
| } |
| |
| private TargetWindowT loadIfNeeded(SdkFunctionSpec windowMappingFn, BoundedWindow mainWindow) { |
| try { |
| String processRequestInstructionId = idGenerator.getId(); |
| InstructionRequest processRequest = |
| InstructionRequest.newBuilder() |
| .setInstructionId(processRequestInstructionId) |
| .setProcessBundle( |
| ProcessBundleRequest.newBuilder() |
| .setProcessBundleDescriptorId(registerIfRequired())) |
| .build(); |
| |
| ConcurrentLinkedQueue<WindowedValue<KV<byte[], TargetWindowT>>> outputValue = |
| new ConcurrentLinkedQueue<>(); |
| |
| // Open the inbound consumer |
| InboundDataClient waitForInboundTermination = |
| beamFnDataService.receive( |
| LogicalEndpoint.of(processRequestInstructionId, "write"), |
| inboundCoder, |
| outputValue::add); |
| |
| CompletionStage<InstructionResponse> processResponse = |
| instructionRequestHandler.handle(processRequest); |
| |
| // Open the outbound consumer |
| try (CloseableFnDataReceiver<WindowedValue<KV<byte[], BoundedWindow>>> outboundConsumer = |
| beamFnDataService.send( |
| LogicalEndpoint.of(processRequestInstructionId, "read"), outboundCoder)) { |
| |
| outboundConsumer.accept(WindowedValue.valueInGlobalWindow(KV.of(EMPTY_ARRAY, mainWindow))); |
| } |
| |
| // Check to see if processing the request failed. |
| throwIfFailure(processResponse); |
| |
| waitForInboundTermination.awaitCompletion(); |
| WindowedValue<KV<byte[], TargetWindowT>> sideInputWindow = outputValue.poll(); |
| checkState( |
| sideInputWindow != null, |
| "Expected side input window to have been emitted by SDK harness."); |
| checkState( |
| sideInputWindow.getValue() != null, |
| "Side input window emitted by SDK harness was a WindowedValue with no value in it."); |
| checkState( |
| sideInputWindow.getValue().getValue() != null, |
| "Side input window emitted by SDK harness was a WindowedValue<KV<...>> with a null V."); |
| checkState( |
| outputValue.isEmpty(), |
| "Expected only a single side input window to have been emitted by " |
| + "the SDK harness but also received %s", |
| outputValue); |
| return sideInputWindow.getValue().getValue(); |
| } catch (Throwable e) { |
| LOG.error("Unable to map main input window {} to side input window.", mainWindow, e); |
| throw new IllegalStateException(e); |
| } |
| } |
| |
| /** Should only be accessed from within {@link #registerIfRequired}. */ |
| private String processBundleDescriptorId; |
| |
| /** |
| * Register a process bundle descriptor for this remote window mapping fn. |
| * |
| * <p>Caches the result after the first registration. |
| */ |
| private synchronized String registerIfRequired() throws ExecutionException, InterruptedException { |
| if (processBundleDescriptorId == null) { |
| String descriptorId = idGenerator.getId(); |
| |
| CompletionStage<InstructionResponse> response = |
| instructionRequestHandler.handle( |
| InstructionRequest.newBuilder() |
| .setInstructionId(idGenerator.getId()) |
| .setRegister( |
| RegisterRequest.newBuilder() |
| .addProcessBundleDescriptor( |
| processBundleDescriptor.toBuilder().setId(descriptorId).build()) |
| .build()) |
| .build()); |
| throwIfFailure(response); |
| processBundleDescriptorId = descriptorId; |
| } |
| return processBundleDescriptorId; |
| } |
| |
| private static InstructionResponse throwIfFailure( |
| CompletionStage<InstructionResponse> responseFuture) |
| throws ExecutionException, InterruptedException { |
| InstructionResponse response = MoreFutures.get(responseFuture); |
| if (!Strings.isNullOrEmpty(response.getError())) { |
| throw new IllegalStateException( |
| String.format( |
| "Client failed to process %s with error [%s].", |
| response.getInstructionId(), response.getError())); |
| } |
| return response; |
| } |
| } |
