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.windmill.client.grpc;
  
  import com.google.auto.value.AutoValue;
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Optional;
  import javax.annotation.Nullable;
  import javax.annotation.concurrent.NotThreadSafe;
  import org.apache.beam.runners.dataflow.worker.WindmillTimeUtils;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem;
  import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
  import org.joda.time.Instant;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * {@link ByteString} buffer of {@link
   * org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingGetWorkResponseChunk}(s).
   *
   * <p>Once all serialized chunks of an {@link WorkItem} have been received flushes (deserializes)
   * the chunk of bytes and metadata into an {@link AssembledWorkItem}.
   *
   * @implNote This class is not thread safe, and provides no synchronization underneath.
   */
  @NotThreadSafe
  final class GetWorkResponseChunkAssembler {
  
    private static final Logger LOG = LoggerFactory.getLogger(GetWorkResponseChunkAssembler.class);
  
    private final GetWorkTimingInfosTracker workTimingInfosTracker;
    private @Nullable ComputationMetadata metadata;
    private final WorkItem.Builder workItemBuilder; // Reused to reduce GC overhead.
    private ByteString data;
    private long bufferedSize;
  
    GetWorkResponseChunkAssembler() {
      workTimingInfosTracker = new GetWorkTimingInfosTracker(System::currentTimeMillis);
      data = ByteString.EMPTY;
      bufferedSize = 0;
      metadata = null;
      workItemBuilder = WorkItem.newBuilder();
    }
  
    /**
     * Appends the response chunk bytes to the {@link #data }byte buffer. Return the assembled
     * WorkItem if all response chunks for a WorkItem have been received.
     */
    List<AssembledWorkItem> append(Windmill.StreamingGetWorkResponseChunk chunk) {
      if (chunk.hasComputationMetadata()) {
        metadata = ComputationMetadata.fromProto(chunk.getComputationMetadata());
      }
      workTimingInfosTracker.addTimingInfo(chunk.getPerWorkItemTimingInfosList());
  
      List<AssembledWorkItem> response = new ArrayList<>();
      for (int i = 0; i < chunk.getSerializedWorkItemList().size(); i++) {
        ByteString serializedWorkItem = chunk.getSerializedWorkItemList().get(i);
        data = data.concat(serializedWorkItem);
        bufferedSize += serializedWorkItem.size();
        long remainingSize = 0;
        if (i == chunk.getSerializedWorkItemList().size() - 1) {
          remainingSize = chunk.getRemainingBytesForWorkItem();
        }
        if (remainingSize == 0) {
          flushToWorkItem().ifPresent(response::add);
        }
      }
      return response;
    }
  
    /**
     * Attempt to flush the {@link #data} bytes into a {@link WorkItem} w/ it's metadata. Resets the
     * data byte string and tracking metadata afterwards, whether the {@link WorkItem} deserialization
     * was successful or not.
     */
    private Optional<AssembledWorkItem> flushToWorkItem() {
      try {
        workItemBuilder.mergeFrom(data);
        return Optional.of(
            AssembledWorkItem.create(
                workItemBuilder.build(),
                Preconditions.checkNotNull(metadata),
                workTimingInfosTracker.getLatencyAttributions(),
                bufferedSize));
      } catch (IOException e) {
        LOG.error("Failed to parse work item from stream: ", e);
      } finally {
        workItemBuilder.clear();
        workTimingInfosTracker.reset();
        data = ByteString.EMPTY;
        bufferedSize = 0;
      }
  
      return Optional.empty();
    }
  
    @AutoValue
    abstract static class ComputationMetadata {
  
      private static ComputationMetadata fromProto(
          Windmill.ComputationWorkItemMetadata metadataProto) {
        return new AutoValue_GetWorkResponseChunkAssembler_ComputationMetadata(
            metadataProto.getComputationId(),
            WindmillTimeUtils.windmillToHarnessWatermark(metadataProto.getInputDataWatermark()),
            WindmillTimeUtils.windmillToHarnessWatermark(
                metadataProto.getDependentRealtimeInputWatermark()));
      }
  
      abstract String computationId();
  
      abstract @Nullable Instant inputDataWatermark();
  
      abstract @Nullable Instant synchronizedProcessingTime();
    }
  
    @AutoValue
    abstract static class AssembledWorkItem {
  
      private static AssembledWorkItem create(
          WorkItem workItem,
          ComputationMetadata computationMetadata,
          ImmutableList<LatencyAttribution> latencyAttributions,
          long size) {
        return new AutoValue_GetWorkResponseChunkAssembler_AssembledWorkItem(
            workItem, computationMetadata, latencyAttributions, size);
      }
  
      abstract WorkItem workItem();
  
      abstract ComputationMetadata computationMetadata();
  
      abstract ImmutableList<LatencyAttribution> latencyAttributions();
  
      abstract long bufferedSize();
    }
  }
  