runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SideInputTrackingIsmReader.java

/*
 * 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 java.io.Closeable;
import java.io.IOException;
import java.util.List;
import org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecordCoder;
import org.apache.beam.runners.dataflow.util.RandomAccessData;
import org.apache.beam.sdk.io.fs.ResourceId;

/**
 * A wrapper for an {@link IsmReaderImpl} object that provides counters to measure consumption of
 * side inputs.
 *
 * <p>An {@link IsmReaderImpl} object can be reused across threads, and work items because its cost
 * of initialization is high. To be able to provide metrics that are associated to the execution of
 * a single work item, we wrap it in a {@link SideInputTrackingIsmReader}, which is not reused, and
 * thus can know all the context for a single side input element.
 */
public class SideInputTrackingIsmReader<V> extends IsmReader<V> {
  private final IsmReader<V> delegate;
  private final SideInputReadCounter readCounter;

  public SideInputTrackingIsmReader(IsmReader<V> delegate, SideInputReadCounter readCounter) {
    this.delegate = delegate;
    this.readCounter = readCounter;
  }

  @Override
  public boolean isInitialized() {
    try (Closeable counterCloser = IsmReader.setSideInputReadContext(readCounter)) {
      return delegate.isInitialized();
    } catch (IOException e) {
      // This is a checked exception that comes from the {@link Closeable#close} interface method.
      // In our case, this method only realizes a memory operation, so it should never throw an
      // {@link IOException}. This should never happen.
      throw new RuntimeException(e);
    }
  }

  @Override
  public boolean isEmpty() throws IOException {
    try (Closeable counterCloser = IsmReader.setSideInputReadContext(readCounter)) {
      return delegate.isEmpty();
    }
  }

  /**
   * Returns the {@link ResourceId} associated to the IsmReader.
   *
   * <p>This is simply a getter method, so it does not call {@link
   * IsmReader#setSideInputReadContext}.
   */
  @Override
  public ResourceId getResourceId() {
    return delegate.getResourceId();
  }

  /**
   * Returns the {@link IsmRecordCoder} associated to the IsmReader.
   *
   * <p>This is simply a getter method, so it does not call {@link
   * IsmReader#setSideInputReadContext}.
   */
  @Override
  public IsmRecordCoder<V> getCoder() {
    return delegate.getCoder();
  }

  @Override
  public IsmPrefixReaderIterator iterator() throws IOException {
    try (Closeable counterCloser = IsmReader.setSideInputReadContext(readCounter)) {
      return delegate.iterator();
    }
  }

  @Override
  public IsmPrefixReaderIterator overKeyComponents(List<?> keyComponents) throws IOException {
    try (Closeable counterCloser = IsmReader.setSideInputReadContext(readCounter)) {
      return delegate.overKeyComponents(keyComponents);
    }
  }

  @Override
  public IsmPrefixReaderIterator overKeyComponents(
      List<?> keyComponents, int shardId, RandomAccessData keyBytes) throws IOException {
    try (Closeable counterCloser = IsmReader.setSideInputReadContext(readCounter)) {
      return delegate.overKeyComponents(keyComponents, shardId, keyBytes);
    }
  }
}