runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/profiler/ScopedProfiler.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.profiler;

import org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** A wrapper around {@link Profiler} to support more idiomatic usage within Java. */
public class ScopedProfiler {

  /** A thin wrapper around {@link Profiler} to allow mocking in tests. */
  @VisibleForTesting
  static class ProfilerWrapper {
    /**
     * Registers a nonzero index to be used in a subsequent call to {@link #setAttribute}. The value
     * is permanently registered for the lifetime of the JVM. Registering the same string multiple
     * times will return the same value.
     *
     * @return a nonzero index representing the string, to be used on calls to setAttribute.
     * @throws java.lang.UnsatisfiedLinkError if the agent hasn't been loaded.
     */
    public int registerAttribute(String value) {
      return Profiler.registerAttribute(value);
    }

    /**
     * Sets attribute for the current thread. Can be called before profiling is enabled or while
     * profiling is active. Attributes must be either 0 (to clear attributes), or a value returned
     * by {@link #registerAttribute}.
     *
     * <p>Samples collected by this thread will be annotated with the string associated to this
     * value.
     *
     * @return the previous value of the thread attribute.
     * @throws java.lang.UnsatisfiedLinkError if the agent hasn't been loaded.
     */
    public int setAttribute(int value) {
      return Profiler.setAttribute(value);
    }

    /**
     * Gets attribute for the current thread. Can be called before profiling is enabled or while
     * profiling is active. Will return the value passed to the most recent call to {@link
     * #setAttribute} from this thread.
     *
     * @return the previous value of the thread attribute.
     * @throws java.lang.UnsatisfiedLinkError if the agent hasn't been loaded.
     */
    public int getAttribute() {
      return Profiler.getAttribute();
    }
  }

  /** Interface for a specific scope in the profile. */
  public interface ProfileScope {

    /** Sets the current thread to the given scope. */
    void activate();
  }

  private static final Logger LOG = LoggerFactory.getLogger(ScopedProfiler.class);

  private enum ProfilingState {
    PROFILING_PRESENT {
      @Override
      public ProfileScope createAttribute(ProfilerWrapper profiler, String attributeTag) {
        return new TaggedProfileScope(profiler, attributeTag);
      }

      @Override
      public ProfileScope currentScope(ProfilerWrapper profiler) {
        return new TaggedProfileScope(profiler, profiler.getAttribute());
      }

      @Override
      public ProfileScope emptyScope(ProfilerWrapper profiler) {
        return new TaggedProfileScope(profiler, 0);
      }
    },
    PROFILING_ABSENT {
      @Override
      public ProfileScope createAttribute(ProfilerWrapper profiler, String attributeTag) {
        return NoopProfileScope.NOOP;
      }

      @Override
      public ProfileScope currentScope(ProfilerWrapper profiler) {
        return NoopProfileScope.NOOP;
      }

      @Override
      public ProfileScope emptyScope(ProfilerWrapper profiler) {
        return NoopProfileScope.NOOP;
      }
    };

    public abstract ProfileScope createAttribute(ProfilerWrapper profiler, String attributeTag);

    public abstract ProfileScope currentScope(ProfilerWrapper profiler);

    public abstract ProfileScope emptyScope(ProfilerWrapper profiler);
  }

  private final ProfilerWrapper profiler;
  private final ProfilingState profilingState;

  public static final ScopedProfiler INSTANCE = new ScopedProfiler(new ProfilerWrapper());

  @VisibleForTesting
  ScopedProfiler(ProfilerWrapper wrapper) {
    profiler = wrapper;
    profilingState = checkProfilingState();
  }

  private ProfilingState checkProfilingState() {
    try {
      // Call a method from the profiler to see if it is available
      profiler.getAttribute();

      // If we make it here, then we successfully invoked the above method, which means the profiler
      // is available.
      LOG.info("Profiling Agent found. Per-step profiling is enabled.");
      return ProfilingState.PROFILING_PRESENT;
    } catch (UnsatisfiedLinkError e) {
      // If we make it here, then the profiling agent wasn't linked in.
      LOG.info("Profiling Agent not found. Profiles will not be available from this worker.");
      return ProfilingState.PROFILING_ABSENT;
    }
  }

  /**
   * Return an {@link ProfileScope} object to use for managing blocks associated with the given tag.
   *
   * <p>Registering a new scope should not happen within hot-loops since it is potentially
   * expensive. Instead, register a scope with a name outside the loop, and then enter the scope as
   * many times as necessary.
   */
  public ProfileScope registerScope(String scopeName) {
    return profilingState.createAttribute(profiler, scopeName);
  }

  public ProfileScope emptyScope() {
    return profilingState.emptyScope(profiler);
  }

  public ProfileScope currentScope() {
    return profilingState.currentScope(profiler);
  }

  /** The implementation of {@link ProfileScope} used when the profiling agent is absent. */
  public enum NoopProfileScope implements ProfileScope {
    NOOP {
      @Override
      public void activate() {
        // Do nothing
      }
    }
  }

  /** The implementation of {@link ProfileScope} used when the profiling agent is present. */
  private static class TaggedProfileScope implements ProfileScope {
    private final int attribute;
    private final ProfilerWrapper profiler;

    private TaggedProfileScope(ProfilerWrapper profiler, String attributeTag) {
      this.profiler = profiler;
      this.attribute = profiler.registerAttribute(attributeTag);
    }

    private TaggedProfileScope(ProfilerWrapper profiler, int attribute) {
      this.profiler = profiler;
      this.attribute = attribute;
    }

    @Override
    public void activate() {
      profiler.setAttribute(attribute);
    }
  }
}