content: /*
   * 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.graph;
  
  import com.google.api.client.json.GenericJson;
  import com.google.api.services.dataflow.model.InstructionOutput;
  import com.google.api.services.dataflow.model.ParDoInstruction;
  import com.google.api.services.dataflow.model.ParallelInstruction;
  import com.google.api.services.dataflow.model.PartialGroupByKeyInstruction;
  import com.google.api.services.dataflow.model.Source;
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.function.Function;
  import org.apache.beam.runners.dataflow.util.CloudObject;
  import org.apache.beam.runners.dataflow.util.CloudObjects;
  import org.apache.beam.runners.dataflow.util.PropertyNames;
  import org.apache.beam.runners.dataflow.util.Structs;
  import org.apache.beam.runners.dataflow.worker.graph.Edges.Edge;
  import org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction;
  import org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode;
  import org.apache.beam.runners.dataflow.worker.graph.Nodes.Node;
  import org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode;
  import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames;
  import org.apache.beam.sdk.coders.ByteArrayCoder;
  import org.apache.beam.sdk.coders.LengthPrefixCoder;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.graph.MutableNetwork;
  
  /** Utility for replacing or wrapping unknown coders with {@link LengthPrefixCoder}. */
  public class LengthPrefixUnknownCoders {
    private static final ImmutableSet<String> WELL_KNOWN_CODER_TYPES =
        ImmutableSet.of(
            "kind:global_window",
            "kind:interval_window",
            "kind:length_prefix",
            "kind:windowed_value",
            "kind:bytes",
            // Equivalent to beam:coder:varint:v1 (64 bit) which is different then
            // kind:var_int32 (32 bit)
            "kind:varint",
            // TODO: Remove TimerOrElementCoder as it is not truly a well known type.
            // This is to get WindowingWindmillReader to work correctly. It is not expected
            // that the SDK harness should ever see this type.
            "com.google.cloud.dataflow.sdk.util.TimerOrElement$TimerOrElementCoder",
            // DFE supplied kinds
            "kind:pair",
            "kind:stream",
            "kind:ism_record",
            "kind:fixed_big_endian_int32",
            "kind:fixed_big_endian_int64",
            "kind:var_int32",
            "kind:void",
            "org.apache.beam.sdk.coders.DoubleCoder",
            "org.apache.beam.sdk.coders.StringUtf8Coder");
  
    private static final String LENGTH_PREFIX_CODER_TYPE = "kind:length_prefix";
  
    @VisibleForTesting
    static final ImmutableSet<String> MERGE_WINDOWS_DO_FN =
        ImmutableSet.of("MergeBucketsDoFn", "MergeWindowsDoFn");
  
    @VisibleForTesting
    static final LengthPrefixCoder<byte[]> LENGTH_PREFIXED_BYTE_ARRAY_CODER =
        LengthPrefixCoder.of(ByteArrayCoder.of());
  
    /**
     * Replaces unknown coders on every {@link InstructionOutputNode} and {@link
     * ParallelInstructionNode} for the runner portion of a {@link MutableNetwork}, with a {@link
     * org.apache.beam.sdk.coders.LengthPrefixCoder LengthPrefixCoder&lt;T&gt;}, where {@code T} is a
     * {@link org.apache.beam.sdk.coders.ByteArrayCoder}.
     */
    public static MutableNetwork<Node, Edge> andReplaceForRunnerNetwork(
        MutableNetwork<Node, Edge> network) {
      Networks.replaceDirectedNetworkNodes(network, andReplaceForInstructionOutputNode());
      Networks.replaceDirectedNetworkNodes(network, andReplaceForParallelInstructionNode());
      return network;
    }
  
    /**
     * Replace unknown coders on the given {@link InstructionOutputNode} with {@link
     * org.apache.beam.sdk.coders.LengthPrefixCoder LengthPrefixCoder&lt;T&gt;} where {@code T} is a
     * {@link org.apache.beam.sdk.coders.ByteArrayCoder}.
     */
    private static Function<Node, Node> andReplaceForInstructionOutputNode() {
      return new TypeSafeNodeFunction<InstructionOutputNode>(InstructionOutputNode.class) {
        @Override
        public Node typedApply(InstructionOutputNode input) {
          InstructionOutput instructionOutput = input.getInstructionOutput();
          try {
            instructionOutput =
                forInstructionOutput(
                    instructionOutput, true /* replace unknown coder with byte array coder */);
          } catch (Exception e) {
            throw new RuntimeException(
                String.format(
                    "Failed to replace unknown coder with " + "LengthPrefixCoder for : {%s}",
                    input.getInstructionOutput()),
                e);
          }
          return InstructionOutputNode.create(instructionOutput, input.getPcollectionId());
        }
      };
    }
  
    /**
     * Replace unknown coders on the given {@link ParallelInstructionNode} with {@link
     * org.apache.beam.sdk.coders.LengthPrefixCoder LengthPrefixCoder&lt;T&gt;} where {@code T} is a
     * {@link org.apache.beam.sdk.coders.ByteArrayCoder}.
     */
    private static Function<Node, Node> andReplaceForParallelInstructionNode() {
      return new TypeSafeNodeFunction<ParallelInstructionNode>(ParallelInstructionNode.class) {
        @Override
        public Node typedApply(ParallelInstructionNode input) {
          ParallelInstruction instruction = input.getParallelInstruction();
          Nodes.ExecutionLocation location = input.getExecutionLocation();
          try {
            instruction =
                forParallelInstruction(
                    instruction, true /* replace unknown coder with byte array coder */);
          } catch (Exception e) {
            throw new RuntimeException(
                String.format(
                    "Failed to replace unknown coder with " + "LengthPrefixCoder for : {%s}",
                    input.getParallelInstruction()),
                e);
          }
          return ParallelInstructionNode.create(instruction, location);
        }
      };
    }
  
    /**
     * Recursively traverse the coder tree and wrap the first unknown coder in every branch with a
     * {@link LengthPrefixCoder} unless an ancestor coder is itself a {@link LengthPrefixCoder}. If
     * {@code replaceWithByteArrayCoder} is set, then replace that unknown coder with a {@link
     * ByteArrayCoder}.
     *
     * @param codec cloud representation of the {@link org.apache.beam.sdk.coders.Coder}.
     * @param replaceWithByteArrayCoder whether to replace an unknown coder with a {@link
     *     ByteArrayCoder}.
     */
    @VisibleForTesting
    static Map<String, Object> forCodec(
        Map<String, Object> codec, boolean replaceWithByteArrayCoder) {
      String coderType = (String) codec.get(PropertyNames.OBJECT_TYPE_NAME);
  
      // Handle well known coders.
      if (LENGTH_PREFIX_CODER_TYPE.equals(coderType)) {
        if (replaceWithByteArrayCoder) {
          return CloudObjects.asCloudObject(
              LENGTH_PREFIXED_BYTE_ARRAY_CODER, /*sdkComponents=*/ null);
        }
        return codec;
      } else if (WELL_KNOWN_CODER_TYPES.contains(coderType)) {
        // The runner knows about these types and can instantiate them so handle their
        // component encodings.
        Map<String, Object> prefixedCodec = new HashMap<>(codec);
  
        // Recursively replace component encodings
        if (codec.containsKey(PropertyNames.COMPONENT_ENCODINGS)) {
          List<Map<String, Object>> prefixedComponents = new ArrayList<>();
          for (Map<String, Object> component :
              (Iterable<Map<String, Object>>) codec.get(PropertyNames.COMPONENT_ENCODINGS)) {
            prefixedComponents.add(forCodec(component, replaceWithByteArrayCoder));
          }
          prefixedCodec.put(PropertyNames.COMPONENT_ENCODINGS, prefixedComponents);
        }
        return prefixedCodec;
      }
  
      // Wrap unknown coders with length prefix coder.
      if (replaceWithByteArrayCoder) {
        return CloudObjects.asCloudObject(LENGTH_PREFIXED_BYTE_ARRAY_CODER, /*sdkComponents=*/ null);
      } else {
        Map<String, Object> prefixedCodec = new HashMap<>();
        prefixedCodec.put(PropertyNames.OBJECT_TYPE_NAME, LENGTH_PREFIX_CODER_TYPE);
        prefixedCodec.put(PropertyNames.COMPONENT_ENCODINGS, ImmutableList.of(codec));
        return prefixedCodec;
      }
    }
  
    /**
     * Wrap unknown coders with a {@link LengthPrefixCoder} for the given {@link InstructionOutput}.
     */
    @VisibleForTesting
    static InstructionOutput forInstructionOutput(
        InstructionOutput input, boolean replaceWithByteArrayCoder) throws Exception {
      InstructionOutput cloudOutput = clone(input, InstructionOutput.class);
      cloudOutput.setCodec(forCodec(cloudOutput.getCodec(), replaceWithByteArrayCoder));
      return cloudOutput;
    }
  
    /**
     * Wrap unknown coders with a {@link LengthPrefixCoder} for the given {@link ParallelInstruction}.
     */
    @VisibleForTesting
    static ParallelInstruction forParallelInstruction(
        ParallelInstruction input, boolean replaceWithByteArrayCoder) throws Exception {
      try {
        ParallelInstruction instruction = clone(input, ParallelInstruction.class);
        if (instruction.getRead() != null) {
          Source cloudSource = instruction.getRead().getSource();
          cloudSource.setCodec(forCodec(cloudSource.getCodec(), replaceWithByteArrayCoder));
        } else if (instruction.getWrite() != null) {
          com.google.api.services.dataflow.model.Sink cloudSink = instruction.getWrite().getSink();
          cloudSink.setCodec(forCodec(cloudSink.getCodec(), replaceWithByteArrayCoder));
        } else if (instruction.getParDo() != null) {
          instruction.setParDo(
              forParDoInstruction(instruction.getParDo(), replaceWithByteArrayCoder));
        } else if (instruction.getPartialGroupByKey() != null) {
          PartialGroupByKeyInstruction pgbk = instruction.getPartialGroupByKey();
          pgbk.setInputElementCodec(forCodec(pgbk.getInputElementCodec(), replaceWithByteArrayCoder));
        } else if (instruction.getFlatten() != null) {
          // FlattenInstructions have no codecs to wrap.
        } else {
          throw new RuntimeException("Unknown parallel instruction: " + input);
        }
        return instruction;
      } catch (IOException e) {
        throw new RuntimeException(
            String.format(
                "Failed to replace unknown coder with " + "LengthPrefixCoder for : {%s}", input),
            e);
      }
    }
  
    /**
     * Wrap unknown coders with a {@link LengthPrefixCoder} for the given {@link ParDoInstruction}.
     */
    private static ParDoInstruction forParDoInstruction(
        ParDoInstruction input, boolean replaceWithByteArrayCoder) throws Exception {
      ParDoInstruction parDo = input;
      CloudObject userFn = CloudObject.fromSpec(parDo.getUserFn());
      String parDoFnType = userFn.getClassName();
      if (MERGE_WINDOWS_DO_FN.contains(parDoFnType)) {
        parDo = clone(input, ParDoInstruction.class);
        Map<String, Object> inputCoderObject =
            Structs.getObject(parDo.getUserFn(), WorkerPropertyNames.INPUT_CODER);
        parDo
            .getUserFn()
            .put(
                WorkerPropertyNames.INPUT_CODER,
                forCodec(inputCoderObject, replaceWithByteArrayCoder));
      } else if ("CreateIsmShardKeyAndSortKeyDoFn".equals(parDoFnType)
          || "ToIsmRecordForMultimapDoFn".equals(parDoFnType)
          || "StreamingPCollectionViewWriterDoFn".equals(parDoFnType)) {
        parDo = clone(input, ParDoInstruction.class);
        Map<String, Object> inputCoderObject =
            Structs.getObject(parDo.getUserFn(), PropertyNames.ENCODING);
        parDo
            .getUserFn()
            .put(PropertyNames.ENCODING, forCodec(inputCoderObject, replaceWithByteArrayCoder));
      }
      // TODO: Handle other types of ParDos.
      return parDo;
    }
  
    /** Clone a cloud object by converting it to json string and back. */
    @VisibleForTesting
    static <T extends GenericJson> T clone(T cloudObject, Class<T> type) throws IOException {
      return cloudObject.getFactory().fromString(cloudObject.toString(), type);
    }
  }
  