<commit_before>// Copyright 2017 The Bazel Authors. All rights reserved.
//
// 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 build.buildfarm.server;

import static build.buildfarm.instance.Utils.putBlobFuture;
import static com.google.common.util.concurrent.Futures.addCallback;
import static com.google.common.util.concurrent.Futures.allAsList;
import static com.google.common.util.concurrent.Futures.catching;
import static com.google.common.util.concurrent.Futures.transform;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
import static java.lang.String.format;
import static java.util.logging.Level.SEVERE;

import build.bazel.remote.execution.v2.BatchReadBlobsRequest;
import build.bazel.remote.execution.v2.BatchReadBlobsResponse;
import build.bazel.remote.execution.v2.BatchUpdateBlobsRequest;
import build.bazel.remote.execution.v2.BatchUpdateBlobsRequest.Request;
import build.bazel.remote.execution.v2.BatchUpdateBlobsResponse;
import build.bazel.remote.execution.v2.BatchUpdateBlobsResponse.Response;
import build.bazel.remote.execution.v2.ContentAddressableStorageGrpc;
import build.bazel.remote.execution.v2.Digest;
import build.bazel.remote.execution.v2.Directory;
import build.bazel.remote.execution.v2.FindMissingBlobsRequest;
import build.bazel.remote.execution.v2.FindMissingBlobsResponse;
import build.bazel.remote.execution.v2.GetTreeRequest;
import build.bazel.remote.execution.v2.GetTreeResponse;
import build.buildfarm.instance.Instance;
import com.google.common.base.Function;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.Status.Code;
import io.grpc.StatusException;
import io.grpc.stub.StreamObserver;
import java.io.IOException;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;

public class ContentAddressableStorageService extends ContentAddressableStorageGrpc.ContentAddressableStorageImplBase {
  private static final Logger logger = Logger.getLogger(ContentAddressableStorageService.class.getName());

  private final Instances instances;
  private final long writeDeadlineAfter;
  private final TimeUnit writeDeadlineAfterUnits;
  private final Level requestLogLevel;

  public ContentAddressableStorageService(
      Instances instances,
      long writeDeadlineAfter,
      TimeUnit writeDeadlineAfterUnits,
      Level requestLogLevel) {
    this.instances = instances;
    this.writeDeadlineAfter = writeDeadlineAfter;
    this.writeDeadlineAfterUnits = writeDeadlineAfterUnits;
    this.requestLogLevel = requestLogLevel;
  }

  @Override
  public void findMissingBlobs(
      FindMissingBlobsRequest request,
      StreamObserver<FindMissingBlobsResponse> responseObserver) {
    Instance instance;
    try {
      instance = instances.get(request.getInstanceName());
    } catch (InstanceNotFoundException e) {
      responseObserver.onError(BuildFarmInstances.toStatusException(e));
      return;
    }

    FindMissingBlobsResponse.Builder builder = FindMissingBlobsResponse.newBuilder();
    addCallback(
        transform(
            instance.findMissingBlobs(request.getBlobDigestsList(), newDirectExecutorService()),
            builder::addAllMissingBlobDigests,
            directExecutor()),
        new FutureCallback<FindMissingBlobsResponse.Builder>() {
          @Override
          public void onSuccess(FindMissingBlobsResponse.Builder builder) {
            try {
              responseObserver.onNext(builder.build());
              responseObserver.onCompleted();
              logger.log(requestLogLevel, format("FindMissingBlobs(%s) for %d blobs", instance.getName(), request.getBlobDigestsList().size()));
            } catch (Throwable t) {
              onFailure(t);
            }
          }

          @Override
          public void onFailure(Throwable t) {
            Status status = Status.fromThrowable(t);
            if (status.getCode() != Code.CANCELLED) {
              logger.log(SEVERE, format("findMissingBlobs(%s): %s: %d", instance.getName(), request.getInstanceName(), request.getBlobDigestsCount()), t);
              responseObserver.onError(t);
            }
          }
        },
        directExecutor());
  }

  private static com.google.rpc.Status statusForCode(Code code) {
    return com.google.rpc.Status.newBuilder()
        .setCode(code.value())
        .build();
  }

  private static ListenableFuture<Response> toResponseFuture(ListenableFuture<Code> codeFuture, Digest digest) {
    return transform(
        codeFuture,
        new Function<Code, Response>() {
          @Override
          public Response apply(Code code) {
            return Response.newBuilder()
                .setDigest(digest)
                .setStatus(statusForCode(code))
                .build();
          }
        },
        directExecutor());
  }

  private static Iterable<ListenableFuture<Response>> putAllBlobs(
      Instance instance,
      Iterable<Request> requests,
      long writeDeadlineAfter,
      TimeUnit writeDeadlineAfterUnits) {
    ImmutableList.Builder<ListenableFuture<Response>> responses = new ImmutableList.Builder<>();
    for (Request request : requests) {
      Digest digest = request.getDigest();
      responses.add(
          toResponseFuture(
              catching(
                  transform(
                      putBlobFuture(instance, digest, request.getData(), writeDeadlineAfter, writeDeadlineAfterUnits),
                      (d) -> Code.OK,
                      directExecutor()),
                  Throwable.class,
                  (e) -> Status.fromThrowable(e).getCode(),
                  directExecutor()),
              digest));
    }
    return responses.build();
  }

  @Override
  public void batchUpdateBlobs(
      BatchUpdateBlobsRequest batchRequest,
      StreamObserver<BatchUpdateBlobsResponse> responseObserver) {
    Instance instance;
    try {
      instance = instances.get(batchRequest.getInstanceName());
    } catch (InstanceNotFoundException e) {
      responseObserver.onError(BuildFarmInstances.toStatusException(e));
      return;
    }

    BatchUpdateBlobsResponse.Builder response = BatchUpdateBlobsResponse.newBuilder();
    ListenableFuture<BatchUpdateBlobsResponse> responseFuture = transform(
        allAsList(
            Iterables.transform(
                putAllBlobs(
                    instance,
                    batchRequest.getRequestsList(),
                    writeDeadlineAfter,
                    writeDeadlineAfterUnits),
                (future) -> transform(future, response::addResponses, directExecutor()))),
        (result) -> response.build(),
        directExecutor());

    addCallback(
        responseFuture,
        new FutureCallback<BatchUpdateBlobsResponse>() {
          @Override
          public void onSuccess(BatchUpdateBlobsResponse response) {
            responseObserver.onNext(response);
            responseObserver.onCompleted();
          }

          @Override
          public void onFailure(Throwable t) {
            responseObserver.onError(t);
          }
        },
        directExecutor());
  }

  private void getInstanceTree(
      Instance instance,
      Digest rootDigest,
      String pageToken,
      int pageSize,
      StreamObserver<GetTreeResponse> responseObserver) {
    try {
      do {
        ImmutableList.Builder<Directory> directories = new ImmutableList.Builder<>();
        String nextPageToken = instance.getTree(
            rootDigest, pageSize, pageToken, directories);

        responseObserver.onNext(GetTreeResponse.newBuilder()
            .addAllDirectories(directories.build())
            .setNextPageToken(nextPageToken)
            .build());
        pageToken = nextPageToken;
      } while (!pageToken.isEmpty());
      responseObserver.onCompleted();
    } catch (IOException e) {
      responseObserver.onError(Status.fromThrowable(e).asException());
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
    }
  }

  void batchReadBlobs(
      Instance instance,
      BatchReadBlobsRequest batchRequest,
      StreamObserver<BatchReadBlobsResponse> responseObserver) {
    BatchReadBlobsResponse.Builder response =
        BatchReadBlobsResponse.newBuilder();
    addCallback(
        transform(
            instance.getAllBlobsFuture(batchRequest.getDigestsList()),
            (responses) -> response.addAllResponses(responses).build(),
            directExecutor()),
        new FutureCallback<BatchReadBlobsResponse>() {
          @Override
          public void onSuccess(BatchReadBlobsResponse response) {
            responseObserver.onNext(response);
            responseObserver.onCompleted();
          }

          @Override
          public void onFailure(Throwable t) {
            responseObserver.onError(Status.fromThrowable(t).asException());
          }
        },
        directExecutor());
  }

  @Override
  public void batchReadBlobs(
      BatchReadBlobsRequest batchRequest,
      StreamObserver<BatchReadBlobsResponse> responseObserver) {
    try {
      batchReadBlobs(
          instances.get(batchRequest.getInstanceName()),
          batchRequest,
          responseObserver);
    } catch (InstanceNotFoundException e) {
      responseObserver.onError(BuildFarmInstances.toStatusException(e));
    }
  }

  @Override
  public void getTree(
      GetTreeRequest request,
      StreamObserver<GetTreeResponse> responseObserver) {
    Instance instance;
    try {
      instance = instances.get(request.getInstanceName());
    } catch (InstanceNotFoundException e) {
      responseObserver.onError(BuildFarmInstances.toStatusException(e));
      return;
    }

    int pageSize = request.getPageSize();
    if (pageSize < 0) {
      responseObserver.onError(Status.INVALID_ARGUMENT.asException());
      return;
    }

    getInstanceTree(
        instance,
        request.getRootDigest(),
        request.getPageToken(),
        pageSize,
        responseObserver);
  }
}
<commit_msg>Log elapsed time with findMissingBlobs<commit_after>// Copyright 2017 The Bazel Authors. All rights reserved.
//
// 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 build.buildfarm.server;

import static build.buildfarm.instance.Utils.putBlobFuture;
import static com.google.common.util.concurrent.Futures.addCallback;
import static com.google.common.util.concurrent.Futures.allAsList;
import static com.google.common.util.concurrent.Futures.catching;
import static com.google.common.util.concurrent.Futures.transform;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
import static java.lang.String.format;
import static java.util.concurrent.TimeUnit.MICROSECONDS;
import static java.util.logging.Level.SEVERE;

import build.bazel.remote.execution.v2.BatchReadBlobsRequest;
import build.bazel.remote.execution.v2.BatchReadBlobsResponse;
import build.bazel.remote.execution.v2.BatchUpdateBlobsRequest;
import build.bazel.remote.execution.v2.BatchUpdateBlobsRequest.Request;
import build.bazel.remote.execution.v2.BatchUpdateBlobsResponse;
import build.bazel.remote.execution.v2.BatchUpdateBlobsResponse.Response;
import build.bazel.remote.execution.v2.ContentAddressableStorageGrpc;
import build.bazel.remote.execution.v2.Digest;
import build.bazel.remote.execution.v2.Directory;
import build.bazel.remote.execution.v2.FindMissingBlobsRequest;
import build.bazel.remote.execution.v2.FindMissingBlobsResponse;
import build.bazel.remote.execution.v2.GetTreeRequest;
import build.bazel.remote.execution.v2.GetTreeResponse;
import build.buildfarm.instance.Instance;
import com.google.common.base.Function;
import com.google.common.base.Stopwatch;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.Status.Code;
import io.grpc.StatusException;
import io.grpc.stub.StreamObserver;
import java.io.IOException;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;

public class ContentAddressableStorageService extends ContentAddressableStorageGrpc.ContentAddressableStorageImplBase {
  private static final Logger logger = Logger.getLogger(ContentAddressableStorageService.class.getName());

  private final Instances instances;
  private final long writeDeadlineAfter;
  private final TimeUnit writeDeadlineAfterUnits;
  private final Level requestLogLevel;

  public ContentAddressableStorageService(
      Instances instances,
      long writeDeadlineAfter,
      TimeUnit writeDeadlineAfterUnits,
      Level requestLogLevel) {
    this.instances = instances;
    this.writeDeadlineAfter = writeDeadlineAfter;
    this.writeDeadlineAfterUnits = writeDeadlineAfterUnits;
    this.requestLogLevel = requestLogLevel;
  }

  @Override
  public void findMissingBlobs(
      FindMissingBlobsRequest request,
      StreamObserver<FindMissingBlobsResponse> responseObserver) {
    Stopwatch stopwatch = Stopwatch.createStarted();
    Instance instance;
    try {
      instance = instances.get(request.getInstanceName());
    } catch (InstanceNotFoundException e) {
      responseObserver.onError(BuildFarmInstances.toStatusException(e));
      return;
    }

    FindMissingBlobsResponse.Builder builder = FindMissingBlobsResponse.newBuilder();
    addCallback(
        transform(
            instance.findMissingBlobs(request.getBlobDigestsList(), newDirectExecutorService()),
            builder::addAllMissingBlobDigests,
            directExecutor()),
        new FutureCallback<FindMissingBlobsResponse.Builder>() {
          @Override
          public void onSuccess(FindMissingBlobsResponse.Builder builder) {
            try {
              responseObserver.onNext(builder.build());
              responseObserver.onCompleted();
              long elapsedMicros = stopwatch.elapsed(MICROSECONDS);
              logger.log(
                  requestLogLevel,
                  format(
                      "FindMissingBlobs(%s) for %d blobs in %gms",
                      instance.getName(),
                      request.getBlobDigestsList().size(),
                      elapsedMicros / 1000.0));
            } catch (Throwable t) {
              onFailure(t);
            }
          }

          @Override
          public void onFailure(Throwable t) {
            Status status = Status.fromThrowable(t);
            if (status.getCode() != Code.CANCELLED) {
              logger.log(SEVERE, format("findMissingBlobs(%s): %s: %d", instance.getName(), request.getInstanceName(), request.getBlobDigestsCount()), t);
              responseObserver.onError(t);
            }
          }
        },
        directExecutor());
  }

  private static com.google.rpc.Status statusForCode(Code code) {
    return com.google.rpc.Status.newBuilder()
        .setCode(code.value())
        .build();
  }

  private static ListenableFuture<Response> toResponseFuture(ListenableFuture<Code> codeFuture, Digest digest) {
    return transform(
        codeFuture,
        new Function<Code, Response>() {
          @Override
          public Response apply(Code code) {
            return Response.newBuilder()
                .setDigest(digest)
                .setStatus(statusForCode(code))
                .build();
          }
        },
        directExecutor());
  }

  private static Iterable<ListenableFuture<Response>> putAllBlobs(
      Instance instance,
      Iterable<Request> requests,
      long writeDeadlineAfter,
      TimeUnit writeDeadlineAfterUnits) {
    ImmutableList.Builder<ListenableFuture<Response>> responses = new ImmutableList.Builder<>();
    for (Request request : requests) {
      Digest digest = request.getDigest();
      responses.add(
          toResponseFuture(
              catching(
                  transform(
                      putBlobFuture(instance, digest, request.getData(), writeDeadlineAfter, writeDeadlineAfterUnits),
                      (d) -> Code.OK,
                      directExecutor()),
                  Throwable.class,
                  (e) -> Status.fromThrowable(e).getCode(),
                  directExecutor()),
              digest));
    }
    return responses.build();
  }

  @Override
  public void batchUpdateBlobs(
      BatchUpdateBlobsRequest batchRequest,
      StreamObserver<BatchUpdateBlobsResponse> responseObserver) {
    Instance instance;
    try {
      instance = instances.get(batchRequest.getInstanceName());
    } catch (InstanceNotFoundException e) {
      responseObserver.onError(BuildFarmInstances.toStatusException(e));
      return;
    }

    BatchUpdateBlobsResponse.Builder response = BatchUpdateBlobsResponse.newBuilder();
    ListenableFuture<BatchUpdateBlobsResponse> responseFuture = transform(
        allAsList(
            Iterables.transform(
                putAllBlobs(
                    instance,
                    batchRequest.getRequestsList(),
                    writeDeadlineAfter,
                    writeDeadlineAfterUnits),
                (future) -> transform(future, response::addResponses, directExecutor()))),
        (result) -> response.build(),
        directExecutor());

    addCallback(
        responseFuture,
        new FutureCallback<BatchUpdateBlobsResponse>() {
          @Override
          public void onSuccess(BatchUpdateBlobsResponse response) {
            responseObserver.onNext(response);
            responseObserver.onCompleted();
          }

          @Override
          public void onFailure(Throwable t) {
            responseObserver.onError(t);
          }
        },
        directExecutor());
  }

  private void getInstanceTree(
      Instance instance,
      Digest rootDigest,
      String pageToken,
      int pageSize,
      StreamObserver<GetTreeResponse> responseObserver) {
    try {
      do {
        ImmutableList.Builder<Directory> directories = new ImmutableList.Builder<>();
        String nextPageToken = instance.getTree(
            rootDigest, pageSize, pageToken, directories);

        responseObserver.onNext(GetTreeResponse.newBuilder()
            .addAllDirectories(directories.build())
            .setNextPageToken(nextPageToken)
            .build());
        pageToken = nextPageToken;
      } while (!pageToken.isEmpty());
      responseObserver.onCompleted();
    } catch (IOException e) {
      responseObserver.onError(Status.fromThrowable(e).asException());
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
    }
  }

  void batchReadBlobs(
      Instance instance,
      BatchReadBlobsRequest batchRequest,
      StreamObserver<BatchReadBlobsResponse> responseObserver) {
    BatchReadBlobsResponse.Builder response =
        BatchReadBlobsResponse.newBuilder();
    addCallback(
        transform(
            instance.getAllBlobsFuture(batchRequest.getDigestsList()),
            (responses) -> response.addAllResponses(responses).build(),
            directExecutor()),
        new FutureCallback<BatchReadBlobsResponse>() {
          @Override
          public void onSuccess(BatchReadBlobsResponse response) {
            responseObserver.onNext(response);
            responseObserver.onCompleted();
          }

          @Override
          public void onFailure(Throwable t) {
            responseObserver.onError(Status.fromThrowable(t).asException());
          }
        },
        directExecutor());
  }

  @Override
  public void batchReadBlobs(
      BatchReadBlobsRequest batchRequest,
      StreamObserver<BatchReadBlobsResponse> responseObserver) {
    try {
      batchReadBlobs(
          instances.get(batchRequest.getInstanceName()),
          batchRequest,
          responseObserver);
    } catch (InstanceNotFoundException e) {
      responseObserver.onError(BuildFarmInstances.toStatusException(e));
    }
  }

  @Override
  public void getTree(
      GetTreeRequest request,
      StreamObserver<GetTreeResponse> responseObserver) {
    Instance instance;
    try {
      instance = instances.get(request.getInstanceName());
    } catch (InstanceNotFoundException e) {
      responseObserver.onError(BuildFarmInstances.toStatusException(e));
      return;
    }

    int pageSize = request.getPageSize();
    if (pageSize < 0) {
      responseObserver.onError(Status.INVALID_ARGUMENT.asException());
      return;
    }

    getInstanceTree(
        instance,
        request.getRootDigest(),
        request.getPageToken(),
        pageSize,
        responseObserver);
  }
}

{'fraction_non_alphanumeric': 0.09084727322181423, 'fraction_numerical': 0.002221315562808831, 'mean_word_length': 2.976924463674058, 'pattern_counts': {'":': 0, '<': 47, '<?xml version=': 0, '>': 55, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3142690', 'n_tokens_mistral': 5774, 'n_tokens_neox': 5639, 'n_words': 1105}