<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed 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 com.facebook.buck.core.build.engine.impl;

import com.facebook.buck.artifact_cache.ArtifactCache;
import com.facebook.buck.artifact_cache.CacheResult;
import com.facebook.buck.core.build.action.resolver.BuildEngineActionToBuildRuleResolver;
import com.facebook.buck.core.build.buildable.context.BuildableContext;
import com.facebook.buck.core.build.engine.BuildEngine;
import com.facebook.buck.core.build.engine.BuildEngineBuildContext;
import com.facebook.buck.core.build.engine.BuildResult;
import com.facebook.buck.core.build.engine.BuildRuleStatus;
import com.facebook.buck.core.build.engine.BuildRuleSuccessType;
import com.facebook.buck.core.build.engine.RuleDepsCache;
import com.facebook.buck.core.build.engine.buildinfo.BuildInfo;
import com.facebook.buck.core.build.engine.buildinfo.BuildInfoRecorder;
import com.facebook.buck.core.build.engine.buildinfo.BuildInfoStore;
import com.facebook.buck.core.build.engine.buildinfo.OnDiskBuildInfo;
import com.facebook.buck.core.build.engine.cache.manager.BuildInfoStoreManager;
import com.facebook.buck.core.build.engine.config.ResourceAwareSchedulingInfo;
import com.facebook.buck.core.build.engine.delegate.CachingBuildEngineDelegate;
import com.facebook.buck.core.build.engine.type.BuildType;
import com.facebook.buck.core.build.engine.type.DepFiles;
import com.facebook.buck.core.build.event.BuildRuleEvent;
import com.facebook.buck.core.build.execution.context.ExecutionContext;
import com.facebook.buck.core.build.stats.BuildRuleDurationTracker;
import com.facebook.buck.core.model.BuildTarget;
import com.facebook.buck.core.model.TargetConfigurationSerializer;
import com.facebook.buck.core.rulekey.RuleKey;
import com.facebook.buck.core.rulekey.calculator.ParallelRuleKeyCalculator;
import com.facebook.buck.core.rules.BuildRule;
import com.facebook.buck.core.rules.BuildRuleResolver;
import com.facebook.buck.core.rules.attr.HasRuntimeDeps;
import com.facebook.buck.core.rules.build.strategy.BuildRuleStrategy;
import com.facebook.buck.event.BuckEventBus;
import com.facebook.buck.event.LeafEvents;
import com.facebook.buck.event.TopLevelRuleKeyCalculatedEvent;
import com.facebook.buck.rules.keys.RuleKeyDiagnostics;
import com.facebook.buck.rules.keys.RuleKeyFactories;
import com.facebook.buck.rules.keys.hasher.StringRuleKeyHasher;
import com.facebook.buck.util.Scope;
import com.facebook.buck.util.cache.FileHashCache;
import com.facebook.buck.util.collect.SortedSets;
import com.facebook.buck.util.concurrent.MoreFutures;
import com.facebook.buck.util.concurrent.ResourceAmounts;
import com.facebook.buck.util.concurrent.WeightedListeningExecutorService;
import com.facebook.buck.util.types.Unit;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.SettableFuture;
import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import javax.annotation.Nullable;

/**
 * A build engine used to build a {@link BuildRule} which also caches the results. If the current
 * {@link RuleKey} of the build rules matches the one on disk, it does not do any work. It also
 * tries to fetch its output from an {@link ArtifactCache} to avoid doing any computation.
 */
public class CachingBuildEngine implements BuildEngine, Closeable {
  public static final ResourceAmounts CACHE_CHECK_RESOURCE_AMOUNTS = ResourceAmounts.of(0, 0, 1, 1);

  public static final ResourceAmounts RULE_KEY_COMPUTATION_RESOURCE_AMOUNTS =
      ResourceAmounts.of(0, 0, 1, 0);
  public static final ResourceAmounts SCHEDULING_MORE_WORK_RESOURCE_AMOUNTS =
      ResourceAmounts.zero();

  static final String BUILD_RULE_TYPE_CONTEXT_KEY = "build_rule_type";
  static final String STEP_TYPE_CONTEXT_KEY = "step_type";
  private final ConcurrentLinkedQueue<ListenableFuture<Unit>> asyncCallbacks;

  enum StepType {
    BUILD_STEP,
    POST_BUILD_STEP,
    ;
  }

  /**
   * These are the values returned by {@link BuildEngine#build(BuildEngineBuildContext,
   * ExecutionContext, BuildRule)}. This must always return the same value for the build of each
   * target.
   */
  private final ConcurrentMap<BuildTarget, ListenableFuture<BuildResult>> results =
      Maps.newConcurrentMap();

  private final AtomicReference<Throwable> firstFailure = new AtomicReference<>();

  private final ConcurrentHashMap<BuildTarget, CachingBuildRuleBuilder> liveRuleBuilders =
      new ConcurrentHashMap<>();

  private final CachingBuildEngineDelegate cachingBuildEngineDelegate;

  private final WeightedListeningExecutorService service;
  private final BuildType buildMode;
  private final DepFiles depFiles;
  private final long maxDepFileCacheEntries;
  private final BuildRuleResolver resolver;
  private final TargetConfigurationSerializer targetConfigurationSerializer;
  private final Optional<Long> artifactCacheSizeLimit;
  private final Optional<Long> defaultOutputHashSizeLimit;
  private final ImmutableMap<String, Long> ruleTypeOutputHashSizeLimit;
  private final FileHashCache fileHashCache;
  @VisibleForTesting final RuleKeyFactories ruleKeyFactories;
  private final ResourceAwareSchedulingInfo resourceAwareSchedulingInfo;

  private final RuleDepsCache ruleDeps;
  private final Optional<UnskippedBuildEngineActionTracker> unskippedRulesTracker;
  private final BuildRuleDurationTracker buildRuleDurationTracker = new BuildRuleDurationTracker();
  private final RuleKeyDiagnostics<RuleKey, String> defaultRuleKeyDiagnostics;
  private final BuildRulePipelinesRunner pipelinesRunner = new BuildRulePipelinesRunner();
  private final ParallelRuleKeyCalculator<RuleKey> ruleKeyCalculator;

  private final BuildInfoStoreManager buildInfoStoreManager;

  private final boolean consoleLogBuildFailuresInline;

  private final Optional<BuildRuleStrategy> customBuildRuleStrategy;

  public CachingBuildEngine(
      CachingBuildEngineDelegate cachingBuildEngineDelegate,
      Optional<BuildRuleStrategy> customBuildRuleStrategy,
      WeightedListeningExecutorService service,
      BuildType buildMode,
      DepFiles depFiles,
      long maxDepFileCacheEntries,
      Optional<Long> artifactCacheSizeLimit,
      Optional<Long> defaultOutputHashSizeLimit,
      ImmutableMap<String, Long> ruleTypeOutputHashSizeLimit,
      BuildRuleResolver resolver,
      BuildEngineActionToBuildRuleResolver actionToBuildRuleResolver,
      TargetConfigurationSerializer targetConfigurationSerializer,
      BuildInfoStoreManager buildInfoStoreManager,
      ResourceAwareSchedulingInfo resourceAwareSchedulingInfo,
      boolean consoleLogBuildFailuresInline,
      RuleKeyFactories ruleKeyFactories) {
    this(
        cachingBuildEngineDelegate,
        customBuildRuleStrategy,
        service,
        buildMode,
        depFiles,
        maxDepFileCacheEntries,
        artifactCacheSizeLimit,
        defaultOutputHashSizeLimit,
        ruleTypeOutputHashSizeLimit,
        resolver,
        buildInfoStoreManager,
        actionToBuildRuleResolver,
        targetConfigurationSerializer,
        ruleKeyFactories,
        resourceAwareSchedulingInfo,
        new RuleKeyDiagnostics<>(
            rule ->
                ruleKeyFactories
                    .getDefaultRuleKeyFactory()
                    .buildForDiagnostics(rule, new StringRuleKeyHasher()),
            appendable ->
                ruleKeyFactories
                    .getDefaultRuleKeyFactory()
                    .buildForDiagnostics(appendable, new StringRuleKeyHasher())),
        consoleLogBuildFailuresInline);
  }

  /** This constructor MUST ONLY BE USED FOR TESTS. */
  @VisibleForTesting
  CachingBuildEngine(
      CachingBuildEngineDelegate cachingBuildEngineDelegate,
      Optional<BuildRuleStrategy> customBuildRuleStrategy,
      WeightedListeningExecutorService service,
      BuildType buildMode,
      DepFiles depFiles,
      long maxDepFileCacheEntries,
      Optional<Long> artifactCacheSizeLimit,
      Optional<Long> defaultOutputHashSizeLimit,
      ImmutableMap<String, Long> ruleTypeOutputHashSizeLimit,
      BuildRuleResolver resolver,
      BuildInfoStoreManager buildInfoStoreManager,
      BuildEngineActionToBuildRuleResolver actionToBuildRuleResolver,
      TargetConfigurationSerializer targetConfigurationSerializer,
      RuleKeyFactories ruleKeyFactories,
      ResourceAwareSchedulingInfo resourceAwareSchedulingInfo,
      RuleKeyDiagnostics<RuleKey, String> defaultRuleKeyDiagnostics,
      boolean consoleLogBuildFailuresInline) {
    this.cachingBuildEngineDelegate = cachingBuildEngineDelegate;
    this.customBuildRuleStrategy = customBuildRuleStrategy;

    this.service = service;
    this.buildMode = buildMode;
    this.depFiles = depFiles;
    this.maxDepFileCacheEntries = maxDepFileCacheEntries;
    this.artifactCacheSizeLimit = artifactCacheSizeLimit;
    this.defaultOutputHashSizeLimit = defaultOutputHashSizeLimit;
    this.ruleTypeOutputHashSizeLimit = ruleTypeOutputHashSizeLimit;
    this.resolver = resolver;
    this.targetConfigurationSerializer = targetConfigurationSerializer;

    this.fileHashCache = cachingBuildEngineDelegate.getFileHashCache();
    this.ruleKeyFactories = ruleKeyFactories;
    this.resourceAwareSchedulingInfo = resourceAwareSchedulingInfo;
    this.buildInfoStoreManager = buildInfoStoreManager;

    this.ruleDeps = new DefaultRuleDepsCache(resolver, actionToBuildRuleResolver);
    this.unskippedRulesTracker = createUnskippedRulesTracker(buildMode, ruleDeps, resolver);
    this.defaultRuleKeyDiagnostics = defaultRuleKeyDiagnostics;
    this.consoleLogBuildFailuresInline = consoleLogBuildFailuresInline;
    this.asyncCallbacks = new ConcurrentLinkedQueue<>();
    this.ruleKeyCalculator =
        new ParallelRuleKeyCalculator<>(
            serviceByAdjustingDefaultWeightsTo(RULE_KEY_COMPUTATION_RESOURCE_AMOUNTS),
            ruleKeyFactories.getDefaultRuleKeyFactory(),
            ruleDeps,
            (eventBus, rule) ->
                BuildRuleEvent.ruleKeyCalculationScope(
                    eventBus,
                    rule,
                    buildRuleDurationTracker,
                    ruleKeyFactories.getDefaultRuleKeyFactory()));
  }

  @Override
  public void close() {
    try {
      if (customBuildRuleStrategy.isPresent()) {
        customBuildRuleStrategy.get().close();
      }
      terminateBuildWithFailure(new CancellationException("Cancelling due to engine shutdown."));
      Futures.allAsList(asyncCallbacks).get();
    } catch (InterruptedException e) {
      e.printStackTrace();
    } catch (IOException | ExecutionException e) {
      throw new RuntimeException(e);
    }
  }

  /// We might want to share rule-key calculation with other parts of code.
  public ParallelRuleKeyCalculator<RuleKey> getRuleKeyCalculator() {
    return ruleKeyCalculator;
  }

  /**
   * We have a lot of places where tasks are submitted into a service implicitly. There is no way to
   * assign custom weights to such tasks. By creating a temporary service with adjusted weights it
   * is possible to trick the system and tweak the weights.
   */
  private WeightedListeningExecutorService serviceByAdjustingDefaultWeightsTo(
      ResourceAmounts defaultAmounts) {
    return resourceAwareSchedulingInfo.adjustServiceDefaultWeightsTo(defaultAmounts, service);
  }

  private static Optional<UnskippedBuildEngineActionTracker> createUnskippedRulesTracker(
      BuildType buildMode, RuleDepsCache ruleDeps, BuildRuleResolver resolver) {
    if (buildMode == BuildType.DEEP || buildMode == BuildType.POPULATE_FROM_REMOTE_CACHE) {
      // Those modes never skip rules, there is no need to track unskipped rules.
      return Optional.empty();
    }
    return Optional.of(new UnskippedBuildEngineActionTracker(ruleDeps, resolver));
  }

  @VisibleForTesting
  void setBuildRuleResult(
      BuildRule buildRule, BuildRuleSuccessType success, CacheResult cacheResult) {
    results.put(
        buildRule.getBuildTarget(),
        Futures.immediateFuture(BuildResult.success(buildRule, success, cacheResult)));
  }

  @Override
  public boolean isRuleBuilt(BuildTarget buildTarget) throws InterruptedException {
    ListenableFuture<BuildResult> resultFuture = results.get(buildTarget);
    return resultFuture != null && MoreFutures.isSuccess(resultFuture);
  }

  @Override
  public void terminateBuildWithFailure(Throwable failure) {
    // TODO(cjhopman): Change this to only accept specific exception types to enforce that we get
    // the information that we want.
    if (firstFailure.compareAndSet(null, failure)) {
      forEachLiveBuilder(builder -> builder.cancel(failure));
    }
  }

  private void forEachLiveBuilder(Consumer<CachingBuildRuleBuilder> action) {
    for (CachingBuildRuleBuilder builder : liveRuleBuilders.values()) {
      action.accept(builder);
    }
  }

  // Dispatch and return a future resolving to a list of all results of this rules dependencies.
  private ListenableFuture<List<BuildResult>> getDepResults(
      BuildRule rule, BuildEngineBuildContext buildContext, ExecutionContext executionContext) {
    List<ListenableFuture<BuildResult>> depResults =
        new ArrayList<>(SortedSets.sizeEstimate(rule.getBuildDeps()));
    for (BuildRule dep : shuffled(rule.getBuildDeps())) {
      depResults.add(getBuildRuleResultWithRuntimeDeps(dep, buildContext, executionContext, false));
    }
    return Futures.allAsList(depResults);
  }

  private static List<BuildRule> shuffled(Iterable<BuildRule> rules) {
    ArrayList<BuildRule> rulesList = Lists.newArrayList(rules);
    Collections.shuffle(rulesList);
    return rulesList;
  }

  private void registerTopLevelRule(BuildRule rule, BuckEventBus eventBus) {
    unskippedRulesTracker.ifPresent(
        tracker -> {
          try (Scope scope = LeafEvents.scope(eventBus, "register-top-level-rule")) {
            tracker.registerTopLevelRule(rule, eventBus);
          }
        });
  }

  private void markRuleAsUsed(BuildRule rule, BuckEventBus eventBus) {
    unskippedRulesTracker.ifPresent(tracker -> tracker.markRuleAsUsed(rule, eventBus));
  }

  // Provide a future that resolves to the result of executing this rule and its runtime
  // dependencies.
  private ListenableFuture<BuildResult> getBuildRuleResultWithRuntimeDeps(
      BuildRule rule,
      BuildEngineBuildContext buildContext,
      ExecutionContext executionContext,
      boolean isTopLevel) {

    // If the rule is already executing, return its result future from the cache.
    ListenableFuture<BuildResult> existingResult = results.get(rule.getBuildTarget());
    if (existingResult != null) {
      return existingResult;
    }

    // Create a `SettableFuture` and atomically put it in the results map.  At this point, if it's
    // not already created, we should proceed.
    SettableFuture<BuildResult> future = SettableFuture.create();
    existingResult = results.putIfAbsent(rule.getBuildTarget(), future);
    if (existingResult != null) {
      return existingResult;
    }

    // Get the future holding the result for this rule and, if we have no additional runtime deps
    // to attach, return it.
    ListenableFuture<RuleKey> ruleKey = calculateRuleKey(rule, buildContext);
    ListenableFuture<BuildResult> result =
        Futures.transformAsync(
            ruleKey,
            input -> {
              if (isTopLevel) {
                buildContext
                    .getEventBus()
                    .post(new TopLevelRuleKeyCalculatedEvent(rule.getBuildTarget(), ruleKey.get()));
              }
              return processBuildRule(rule, buildContext, executionContext);
            },
            serviceByAdjustingDefaultWeightsTo(SCHEDULING_MORE_WORK_RESOURCE_AMOUNTS));
    if (!(rule instanceof HasRuntimeDeps)) {
      future.setFuture(result);
      return future;
    }

    // Collect any runtime deps we have into a list of futures.
    List<ListenableFuture<BuildResult>> runtimeDepResults =
        ((HasRuntimeDeps) rule)
            .getRuntimeDeps(resolver)
            .map(resolver::getRule)
            .map(
                dep ->
                    getBuildRuleResultWithRuntimeDeps(dep, buildContext, executionContext, false))
            .collect(ImmutableList.toImmutableList());

    // If we don't have any runtime deps we can short circuit here
    if (runtimeDepResults.isEmpty()) {
      future.setFuture(result);
      return future;
    }

    // Create a new combined future, which runs the original rule and all the runtime deps in
    // parallel, but which propagates an error if any one of them fails.
    // It also checks that all runtime deps succeeded.
    ListenableFuture<BuildResult> chainedResult =
        Futures.transformAsync(
            Futures.allAsList(runtimeDepResults),
            results -> {
              Optional<BuildResult> cancelledResult = Optional.empty();

              for (BuildResult buildResult : results) {
                if (!buildResult.isSuccess()) {
                  if (buildResult.getStatus() == BuildRuleStatus.CANCELED) {
                    cancelledResult =
                        Optional.of(BuildResult.canceled(rule, buildResult.getFailure()));
                  } else {
                    return Futures.immediateFuture(
                        BuildResult.failure(rule, buildResult.getFailure()));
                  }
                }
              }
              return cancelledResult.isPresent()
                  ? Futures.immediateFuture(cancelledResult.get())
                  : result;
            },
            MoreExecutors.directExecutor());
    future.setFuture(chainedResult);
    return future;
  }

  public ListenableFuture<?> walkRule(BuildRule rule, Set<BuildRule> seen) {
    return Futures.transformAsync(
        Futures.immediateFuture(ruleDeps.get(rule)),
        deps -> {
          List<ListenableFuture<?>> results1 = new ArrayList<>(Iterables.size(deps));
          for (BuildRule dep : deps) {
            if (seen.add(dep)) {
              results1.add(walkRule(dep, seen));
            }
          }
          return Futures.allAsList(results1);
        },
        serviceByAdjustingDefaultWeightsTo(SCHEDULING_MORE_WORK_RESOURCE_AMOUNTS));
  }

  @Override
  public ListenableFuture<Integer> getNumRulesToBuild(Iterable<BuildRule> rules) {
    Set<BuildRule> seen = Sets.newConcurrentHashSet();
    ImmutableList.Builder<ListenableFuture<?>> results = ImmutableList.builder();
    for (BuildRule rule : rules) {
      if (seen.add(rule)) {
        results.add(walkRule(rule, seen));
      }
    }
    return Futures.transform(
        Futures.allAsList(results.build()), unused -> seen.size(), MoreExecutors.directExecutor());
  }

  private ListenableFuture<RuleKey> calculateRuleKey(
      BuildRule rule, BuildEngineBuildContext context) {
    return ruleKeyCalculator.calculate(context.getEventBus(), rule);
  }

  @Override
  public BuildEngine.BuildEngineResult build(
      BuildEngineBuildContext buildContext, ExecutionContext executionContext, BuildRule rule) {
    // Keep track of all jobs that run asynchronously with respect to the build dep chain.  We want
    // to make sure we wait for these before calling yielding the final build result.
    registerTopLevelRule(rule, buildContext.getEventBus());
    ListenableFuture<BuildResult> resultFuture =
        getBuildRuleResultWithRuntimeDeps(rule, buildContext, executionContext, true);
    return BuildEngine.BuildEngineResult.of(resultFuture);
  }

  @Nullable
  @Override
  public BuildResult getBuildRuleResult(BuildTarget buildTarget)
      throws ExecutionException, InterruptedException {
    ListenableFuture<BuildResult> result = results.get(buildTarget);
    if (result == null) {
      return null;
    }
    return result.get();
  }

  private boolean isKeepGoingEnabled(BuildEngineBuildContext buildContext) {
    return buildMode == BuildType.POPULATE_FROM_REMOTE_CACHE || buildContext.isKeepGoing();
  }
<commit_msg>CachingBuildEngine: run `registerTopLevelRule` in parallel to build<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed 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 com.facebook.buck.core.build.engine.impl;

import com.facebook.buck.artifact_cache.ArtifactCache;
import com.facebook.buck.artifact_cache.CacheResult;
import com.facebook.buck.core.build.action.resolver.BuildEngineActionToBuildRuleResolver;
import com.facebook.buck.core.build.buildable.context.BuildableContext;
import com.facebook.buck.core.build.engine.BuildEngine;
import com.facebook.buck.core.build.engine.BuildEngineBuildContext;
import com.facebook.buck.core.build.engine.BuildResult;
import com.facebook.buck.core.build.engine.BuildRuleStatus;
import com.facebook.buck.core.build.engine.BuildRuleSuccessType;
import com.facebook.buck.core.build.engine.RuleDepsCache;
import com.facebook.buck.core.build.engine.buildinfo.BuildInfo;
import com.facebook.buck.core.build.engine.buildinfo.BuildInfoRecorder;
import com.facebook.buck.core.build.engine.buildinfo.BuildInfoStore;
import com.facebook.buck.core.build.engine.buildinfo.OnDiskBuildInfo;
import com.facebook.buck.core.build.engine.cache.manager.BuildInfoStoreManager;
import com.facebook.buck.core.build.engine.config.ResourceAwareSchedulingInfo;
import com.facebook.buck.core.build.engine.delegate.CachingBuildEngineDelegate;
import com.facebook.buck.core.build.engine.type.BuildType;
import com.facebook.buck.core.build.engine.type.DepFiles;
import com.facebook.buck.core.build.event.BuildRuleEvent;
import com.facebook.buck.core.build.execution.context.ExecutionContext;
import com.facebook.buck.core.build.stats.BuildRuleDurationTracker;
import com.facebook.buck.core.model.BuildTarget;
import com.facebook.buck.core.model.TargetConfigurationSerializer;
import com.facebook.buck.core.rulekey.RuleKey;
import com.facebook.buck.core.rulekey.calculator.ParallelRuleKeyCalculator;
import com.facebook.buck.core.rules.BuildRule;
import com.facebook.buck.core.rules.BuildRuleResolver;
import com.facebook.buck.core.rules.attr.HasRuntimeDeps;
import com.facebook.buck.core.rules.build.strategy.BuildRuleStrategy;
import com.facebook.buck.event.BuckEventBus;
import com.facebook.buck.event.LeafEvents;
import com.facebook.buck.event.TopLevelRuleKeyCalculatedEvent;
import com.facebook.buck.rules.keys.RuleKeyDiagnostics;
import com.facebook.buck.rules.keys.RuleKeyFactories;
import com.facebook.buck.rules.keys.hasher.StringRuleKeyHasher;
import com.facebook.buck.util.Scope;
import com.facebook.buck.util.cache.FileHashCache;
import com.facebook.buck.util.collect.SortedSets;
import com.facebook.buck.util.concurrent.MoreFutures;
import com.facebook.buck.util.concurrent.ResourceAmounts;
import com.facebook.buck.util.concurrent.WeightedListeningExecutorService;
import com.facebook.buck.util.types.Unit;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.SettableFuture;
import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import javax.annotation.Nullable;

/**
 * A build engine used to build a {@link BuildRule} which also caches the results. If the current
 * {@link RuleKey} of the build rules matches the one on disk, it does not do any work. It also
 * tries to fetch its output from an {@link ArtifactCache} to avoid doing any computation.
 */
public class CachingBuildEngine implements BuildEngine, Closeable {
  public static final ResourceAmounts CACHE_CHECK_RESOURCE_AMOUNTS = ResourceAmounts.of(0, 0, 1, 1);

  public static final ResourceAmounts RULE_KEY_COMPUTATION_RESOURCE_AMOUNTS =
      ResourceAmounts.of(0, 0, 1, 0);
  public static final ResourceAmounts SCHEDULING_MORE_WORK_RESOURCE_AMOUNTS =
      ResourceAmounts.zero();

  static final String BUILD_RULE_TYPE_CONTEXT_KEY = "build_rule_type";
  static final String STEP_TYPE_CONTEXT_KEY = "step_type";
  private final ConcurrentLinkedQueue<ListenableFuture<Unit>> asyncCallbacks;

  enum StepType {
    BUILD_STEP,
    POST_BUILD_STEP,
    ;
  }

  /**
   * These are the values returned by {@link BuildEngine#build(BuildEngineBuildContext,
   * ExecutionContext, BuildRule)}. This must always return the same value for the build of each
   * target.
   */
  private final ConcurrentMap<BuildTarget, ListenableFuture<BuildResult>> results =
      Maps.newConcurrentMap();

  private final AtomicReference<Throwable> firstFailure = new AtomicReference<>();

  private final ConcurrentHashMap<BuildTarget, CachingBuildRuleBuilder> liveRuleBuilders =
      new ConcurrentHashMap<>();

  private final CachingBuildEngineDelegate cachingBuildEngineDelegate;

  private final WeightedListeningExecutorService service;
  private final BuildType buildMode;
  private final DepFiles depFiles;
  private final long maxDepFileCacheEntries;
  private final BuildRuleResolver resolver;
  private final TargetConfigurationSerializer targetConfigurationSerializer;
  private final Optional<Long> artifactCacheSizeLimit;
  private final Optional<Long> defaultOutputHashSizeLimit;
  private final ImmutableMap<String, Long> ruleTypeOutputHashSizeLimit;
  private final FileHashCache fileHashCache;
  @VisibleForTesting final RuleKeyFactories ruleKeyFactories;
  private final ResourceAwareSchedulingInfo resourceAwareSchedulingInfo;

  private final RuleDepsCache ruleDeps;
  private final Optional<UnskippedBuildEngineActionTracker> unskippedRulesTracker;
  private final BuildRuleDurationTracker buildRuleDurationTracker = new BuildRuleDurationTracker();
  private final RuleKeyDiagnostics<RuleKey, String> defaultRuleKeyDiagnostics;
  private final BuildRulePipelinesRunner pipelinesRunner = new BuildRulePipelinesRunner();
  private final ParallelRuleKeyCalculator<RuleKey> ruleKeyCalculator;

  private final BuildInfoStoreManager buildInfoStoreManager;

  private final boolean consoleLogBuildFailuresInline;

  private final Optional<BuildRuleStrategy> customBuildRuleStrategy;

  public CachingBuildEngine(
      CachingBuildEngineDelegate cachingBuildEngineDelegate,
      Optional<BuildRuleStrategy> customBuildRuleStrategy,
      WeightedListeningExecutorService service,
      BuildType buildMode,
      DepFiles depFiles,
      long maxDepFileCacheEntries,
      Optional<Long> artifactCacheSizeLimit,
      Optional<Long> defaultOutputHashSizeLimit,
      ImmutableMap<String, Long> ruleTypeOutputHashSizeLimit,
      BuildRuleResolver resolver,
      BuildEngineActionToBuildRuleResolver actionToBuildRuleResolver,
      TargetConfigurationSerializer targetConfigurationSerializer,
      BuildInfoStoreManager buildInfoStoreManager,
      ResourceAwareSchedulingInfo resourceAwareSchedulingInfo,
      boolean consoleLogBuildFailuresInline,
      RuleKeyFactories ruleKeyFactories) {
    this(
        cachingBuildEngineDelegate,
        customBuildRuleStrategy,
        service,
        buildMode,
        depFiles,
        maxDepFileCacheEntries,
        artifactCacheSizeLimit,
        defaultOutputHashSizeLimit,
        ruleTypeOutputHashSizeLimit,
        resolver,
        buildInfoStoreManager,
        actionToBuildRuleResolver,
        targetConfigurationSerializer,
        ruleKeyFactories,
        resourceAwareSchedulingInfo,
        new RuleKeyDiagnostics<>(
            rule ->
                ruleKeyFactories
                    .getDefaultRuleKeyFactory()
                    .buildForDiagnostics(rule, new StringRuleKeyHasher()),
            appendable ->
                ruleKeyFactories
                    .getDefaultRuleKeyFactory()
                    .buildForDiagnostics(appendable, new StringRuleKeyHasher())),
        consoleLogBuildFailuresInline);
  }

  /** This constructor MUST ONLY BE USED FOR TESTS. */
  @VisibleForTesting
  CachingBuildEngine(
      CachingBuildEngineDelegate cachingBuildEngineDelegate,
      Optional<BuildRuleStrategy> customBuildRuleStrategy,
      WeightedListeningExecutorService service,
      BuildType buildMode,
      DepFiles depFiles,
      long maxDepFileCacheEntries,
      Optional<Long> artifactCacheSizeLimit,
      Optional<Long> defaultOutputHashSizeLimit,
      ImmutableMap<String, Long> ruleTypeOutputHashSizeLimit,
      BuildRuleResolver resolver,
      BuildInfoStoreManager buildInfoStoreManager,
      BuildEngineActionToBuildRuleResolver actionToBuildRuleResolver,
      TargetConfigurationSerializer targetConfigurationSerializer,
      RuleKeyFactories ruleKeyFactories,
      ResourceAwareSchedulingInfo resourceAwareSchedulingInfo,
      RuleKeyDiagnostics<RuleKey, String> defaultRuleKeyDiagnostics,
      boolean consoleLogBuildFailuresInline) {
    this.cachingBuildEngineDelegate = cachingBuildEngineDelegate;
    this.customBuildRuleStrategy = customBuildRuleStrategy;

    this.service = service;
    this.buildMode = buildMode;
    this.depFiles = depFiles;
    this.maxDepFileCacheEntries = maxDepFileCacheEntries;
    this.artifactCacheSizeLimit = artifactCacheSizeLimit;
    this.defaultOutputHashSizeLimit = defaultOutputHashSizeLimit;
    this.ruleTypeOutputHashSizeLimit = ruleTypeOutputHashSizeLimit;
    this.resolver = resolver;
    this.targetConfigurationSerializer = targetConfigurationSerializer;

    this.fileHashCache = cachingBuildEngineDelegate.getFileHashCache();
    this.ruleKeyFactories = ruleKeyFactories;
    this.resourceAwareSchedulingInfo = resourceAwareSchedulingInfo;
    this.buildInfoStoreManager = buildInfoStoreManager;

    this.ruleDeps = new DefaultRuleDepsCache(resolver, actionToBuildRuleResolver);
    this.unskippedRulesTracker = createUnskippedRulesTracker(buildMode, ruleDeps, resolver);
    this.defaultRuleKeyDiagnostics = defaultRuleKeyDiagnostics;
    this.consoleLogBuildFailuresInline = consoleLogBuildFailuresInline;
    this.asyncCallbacks = new ConcurrentLinkedQueue<>();
    this.ruleKeyCalculator =
        new ParallelRuleKeyCalculator<>(
            serviceByAdjustingDefaultWeightsTo(RULE_KEY_COMPUTATION_RESOURCE_AMOUNTS),
            ruleKeyFactories.getDefaultRuleKeyFactory(),
            ruleDeps,
            (eventBus, rule) ->
                BuildRuleEvent.ruleKeyCalculationScope(
                    eventBus,
                    rule,
                    buildRuleDurationTracker,
                    ruleKeyFactories.getDefaultRuleKeyFactory()));
  }

  @Override
  public void close() {
    try {
      if (customBuildRuleStrategy.isPresent()) {
        customBuildRuleStrategy.get().close();
      }
      terminateBuildWithFailure(new CancellationException("Cancelling due to engine shutdown."));
      Futures.allAsList(asyncCallbacks).get();
    } catch (InterruptedException e) {
      e.printStackTrace();
    } catch (IOException | ExecutionException e) {
      throw new RuntimeException(e);
    }
  }

  /// We might want to share rule-key calculation with other parts of code.
  public ParallelRuleKeyCalculator<RuleKey> getRuleKeyCalculator() {
    return ruleKeyCalculator;
  }

  /**
   * We have a lot of places where tasks are submitted into a service implicitly. There is no way to
   * assign custom weights to such tasks. By creating a temporary service with adjusted weights it
   * is possible to trick the system and tweak the weights.
   */
  private WeightedListeningExecutorService serviceByAdjustingDefaultWeightsTo(
      ResourceAmounts defaultAmounts) {
    return resourceAwareSchedulingInfo.adjustServiceDefaultWeightsTo(defaultAmounts, service);
  }

  private static Optional<UnskippedBuildEngineActionTracker> createUnskippedRulesTracker(
      BuildType buildMode, RuleDepsCache ruleDeps, BuildRuleResolver resolver) {
    if (buildMode == BuildType.DEEP || buildMode == BuildType.POPULATE_FROM_REMOTE_CACHE) {
      // Those modes never skip rules, there is no need to track unskipped rules.
      return Optional.empty();
    }
    return Optional.of(new UnskippedBuildEngineActionTracker(ruleDeps, resolver));
  }

  @VisibleForTesting
  void setBuildRuleResult(
      BuildRule buildRule, BuildRuleSuccessType success, CacheResult cacheResult) {
    results.put(
        buildRule.getBuildTarget(),
        Futures.immediateFuture(BuildResult.success(buildRule, success, cacheResult)));
  }

  @Override
  public boolean isRuleBuilt(BuildTarget buildTarget) throws InterruptedException {
    ListenableFuture<BuildResult> resultFuture = results.get(buildTarget);
    return resultFuture != null && MoreFutures.isSuccess(resultFuture);
  }

  @Override
  public void terminateBuildWithFailure(Throwable failure) {
    // TODO(cjhopman): Change this to only accept specific exception types to enforce that we get
    // the information that we want.
    if (firstFailure.compareAndSet(null, failure)) {
      forEachLiveBuilder(builder -> builder.cancel(failure));
    }
  }

  private void forEachLiveBuilder(Consumer<CachingBuildRuleBuilder> action) {
    for (CachingBuildRuleBuilder builder : liveRuleBuilders.values()) {
      action.accept(builder);
    }
  }

  // Dispatch and return a future resolving to a list of all results of this rules dependencies.
  private ListenableFuture<List<BuildResult>> getDepResults(
      BuildRule rule, BuildEngineBuildContext buildContext, ExecutionContext executionContext) {
    List<ListenableFuture<BuildResult>> depResults =
        new ArrayList<>(SortedSets.sizeEstimate(rule.getBuildDeps()));
    for (BuildRule dep : shuffled(rule.getBuildDeps())) {
      depResults.add(getBuildRuleResultWithRuntimeDeps(dep, buildContext, executionContext, false));
    }
    return Futures.allAsList(depResults);
  }

  private static List<BuildRule> shuffled(Iterable<BuildRule> rules) {
    ArrayList<BuildRule> rulesList = Lists.newArrayList(rules);
    Collections.shuffle(rulesList);
    return rulesList;
  }

  private void registerTopLevelRule(BuildRule rule, BuckEventBus eventBus) {
    unskippedRulesTracker.ifPresent(
        tracker -> {
          try (Scope scope = LeafEvents.scope(eventBus, "register-top-level-rule")) {
            tracker.registerTopLevelRule(rule, eventBus);
          }
        });
  }

  private void markRuleAsUsed(BuildRule rule, BuckEventBus eventBus) {
    unskippedRulesTracker.ifPresent(tracker -> tracker.markRuleAsUsed(rule, eventBus));
  }

  // Provide a future that resolves to the result of executing this rule and its runtime
  // dependencies.
  private ListenableFuture<BuildResult> getBuildRuleResultWithRuntimeDeps(
      BuildRule rule,
      BuildEngineBuildContext buildContext,
      ExecutionContext executionContext,
      boolean isTopLevel) {

    // If the rule is already executing, return its result future from the cache.
    ListenableFuture<BuildResult> existingResult = results.get(rule.getBuildTarget());
    if (existingResult != null) {
      return existingResult;
    }

    // Create a `SettableFuture` and atomically put it in the results map.  At this point, if it's
    // not already created, we should proceed.
    SettableFuture<BuildResult> future = SettableFuture.create();
    existingResult = results.putIfAbsent(rule.getBuildTarget(), future);
    if (existingResult != null) {
      return existingResult;
    }

    // Get the future holding the result for this rule and, if we have no additional runtime deps
    // to attach, return it.
    ListenableFuture<RuleKey> ruleKey = calculateRuleKey(rule, buildContext);
    ListenableFuture<BuildResult> result =
        Futures.transformAsync(
            ruleKey,
            input -> {
              if (isTopLevel) {
                buildContext
                    .getEventBus()
                    .post(new TopLevelRuleKeyCalculatedEvent(rule.getBuildTarget(), ruleKey.get()));
              }
              return processBuildRule(rule, buildContext, executionContext);
            },
            serviceByAdjustingDefaultWeightsTo(SCHEDULING_MORE_WORK_RESOURCE_AMOUNTS));
    if (!(rule instanceof HasRuntimeDeps)) {
      future.setFuture(result);
      return future;
    }

    // Collect any runtime deps we have into a list of futures.
    List<ListenableFuture<BuildResult>> runtimeDepResults =
        ((HasRuntimeDeps) rule)
            .getRuntimeDeps(resolver)
            .map(resolver::getRule)
            .map(
                dep ->
                    getBuildRuleResultWithRuntimeDeps(dep, buildContext, executionContext, false))
            .collect(ImmutableList.toImmutableList());

    // If we don't have any runtime deps we can short circuit here
    if (runtimeDepResults.isEmpty()) {
      future.setFuture(result);
      return future;
    }

    // Create a new combined future, which runs the original rule and all the runtime deps in
    // parallel, but which propagates an error if any one of them fails.
    // It also checks that all runtime deps succeeded.
    ListenableFuture<BuildResult> chainedResult =
        Futures.transformAsync(
            Futures.allAsList(runtimeDepResults),
            results -> {
              Optional<BuildResult> cancelledResult = Optional.empty();

              for (BuildResult buildResult : results) {
                if (!buildResult.isSuccess()) {
                  if (buildResult.getStatus() == BuildRuleStatus.CANCELED) {
                    cancelledResult =
                        Optional.of(BuildResult.canceled(rule, buildResult.getFailure()));
                  } else {
                    return Futures.immediateFuture(
                        BuildResult.failure(rule, buildResult.getFailure()));
                  }
                }
              }
              return cancelledResult.isPresent()
                  ? Futures.immediateFuture(cancelledResult.get())
                  : result;
            },
            MoreExecutors.directExecutor());
    future.setFuture(chainedResult);
    return future;
  }

  public ListenableFuture<?> walkRule(BuildRule rule, Set<BuildRule> seen) {
    return Futures.transformAsync(
        Futures.immediateFuture(ruleDeps.get(rule)),
        deps -> {
          List<ListenableFuture<?>> results1 = new ArrayList<>(Iterables.size(deps));
          for (BuildRule dep : deps) {
            if (seen.add(dep)) {
              results1.add(walkRule(dep, seen));
            }
          }
          return Futures.allAsList(results1);
        },
        serviceByAdjustingDefaultWeightsTo(SCHEDULING_MORE_WORK_RESOURCE_AMOUNTS));
  }

  @Override
  public ListenableFuture<Integer> getNumRulesToBuild(Iterable<BuildRule> rules) {
    Set<BuildRule> seen = Sets.newConcurrentHashSet();
    ImmutableList.Builder<ListenableFuture<?>> results = ImmutableList.builder();
    for (BuildRule rule : rules) {
      if (seen.add(rule)) {
        results.add(walkRule(rule, seen));
      }
    }
    return Futures.transform(
        Futures.allAsList(results.build()), unused -> seen.size(), MoreExecutors.directExecutor());
  }

  private ListenableFuture<RuleKey> calculateRuleKey(
      BuildRule rule, BuildEngineBuildContext context) {
    return ruleKeyCalculator.calculate(context.getEventBus(), rule);
  }

  @Override
  public BuildEngine.BuildEngineResult build(
      BuildEngineBuildContext buildContext, ExecutionContext executionContext, BuildRule rule) {
    ListenableFuture<BuildResult> resultFuture =
        getBuildRuleResultWithRuntimeDeps(rule, buildContext, executionContext, true);
    return BuildEngine.BuildEngineResult.of(
        Futures.transformAsync(
            // Keep track of all jobs that run asynchronously with respect to the build dep chain.
            // We want
            // to make sure we wait for these before calling yielding the final build result.
            service.submit(() -> registerTopLevelRule(rule, buildContext.getEventBus())),
            unused -> resultFuture,
            MoreExecutors.directExecutor()));
  }

  @Nullable
  @Override
  public BuildResult getBuildRuleResult(BuildTarget buildTarget)
      throws ExecutionException, InterruptedException {
    ListenableFuture<BuildResult> result = results.get(buildTarget);
    if (result == null) {
      return null;
    }
    return result.get();
  }

  private boolean isKeepGoingEnabled(BuildEngineBuildContext buildContext) {
    return buildMode == BuildType.POPULATE_FROM_REMOTE_CACHE || buildContext.isKeepGoing();
  }

{'fraction_non_alphanumeric': 0.07775252759484277, 'fraction_numerical': 0.0006956683053520082, 'mean_word_length': 3.901125127855438, 'pattern_counts': {'":': 0, '<': 125, '<?xml version=': 0, '>': 149, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5629623', 'n_tokens_mistral': 11569, 'n_tokens_neox': 11258, 'n_words': 2569}