@Override
public void getJsonSchema(
    MetaStoreP.GetJsonSchemaRequest request,
    StreamObserver<MetaStoreP.GetJsonSchemaResponse> responseObserver) {

  try {
    AbstractRegistry registry = metaStore.registries.get(request.getRegistryName());
    ProtoDomain pContainer = registry.get();
    String jsonSchema = ProtoToJsonSchema.convert(pContainer, request.getMessageName());

    responseObserver.onNext(
        MetaStoreP.GetJsonSchemaResponse.newBuilder().setSchema(jsonSchema).build());
    responseObserver.onCompleted();
  } catch (StatusException | StatusRuntimeException e) {
    responseObserver.onError(e);
  }
}
 

@Override
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
        final MethodDescriptor<ReqT, RespT> method, CallOptions callOptions,
        final Channel next) {
    return new ClientInterceptors.CheckedForwardingClientCall<ReqT, RespT>(
            next.newCall(method, callOptions)) {
        @Override
        protected void checkedStart(Listener<RespT> responseListener, Metadata headers)
                throws StatusException {
            Metadata cachedSaved;
            URI uri = serviceUri(next, method);
            synchronized (this) {
                Map<String, List<String>> latestMetadata = getRequestMetadata(uri);
                if (mLastMetadata == null || mLastMetadata != latestMetadata) {
                    mLastMetadata = latestMetadata;
                    mCached = toHeaders(mLastMetadata);
                }
                cachedSaved = mCached;
            }
            headers.merge(cachedSaved);
            delegate().start(responseListener, headers);
        }
    };
}
 

@Test
public void ping_failsWhenTransportShutdown() throws Exception {
  initTransport();
  PingCallbackImpl callback = new PingCallbackImpl();
  clientTransport.ping(callback, MoreExecutors.directExecutor());
  assertEquals(1, getTransportStats(clientTransport).keepAlivesSent);
  assertEquals(0, callback.invocationCount);

  clientTransport.shutdown(SHUTDOWN_REASON);
  // ping failed on channel shutdown
  assertEquals(1, callback.invocationCount);
  assertTrue(callback.failureCause instanceof StatusException);
  assertSame(SHUTDOWN_REASON, ((StatusException) callback.failureCause).getStatus());

  // now that handler is in terminal state, all future pings fail immediately
  callback = new PingCallbackImpl();
  clientTransport.ping(callback, MoreExecutors.directExecutor());
  assertEquals(1, getTransportStats(clientTransport).keepAlivesSent);
  assertEquals(1, callback.invocationCount);
  assertTrue(callback.failureCause instanceof StatusException);
  assertSame(SHUTDOWN_REASON, ((StatusException) callback.failureCause).getStatus());
  shutdownAndVerify();
}
 

private void checkRequestLog(RequestLogChecker checker) throws Exception {
    final RequestLog log = requestLogQueue.take();
    assertThat(log.isComplete()).isTrue();

    final RpcRequest rpcReq = (RpcRequest) log.requestContent();
    final RpcResponse rpcRes = (RpcResponse) log.responseContent();
    assertThat(rpcReq).isNotNull();
    assertThat((Object) rpcRes).isNotNull();
    assertThat(rpcReq.serviceType()).isEqualTo(GrpcLogUtil.class);

    final Status grpcStatus;
    if (rpcRes.cause() != null) {
        grpcStatus = ((StatusException) rpcRes.cause()).getStatus();
    } else {
        grpcStatus = null;
    }

    checker.check(rpcReq, rpcRes, grpcStatus);
}
 

@Test
public void maxHeaderListSizeShouldBeEnforcedOnClient() throws Exception {
  startServer();

  NettyClientTransport transport =
      newTransport(newNegotiator(), DEFAULT_MAX_MESSAGE_SIZE, 1, null, true);
  callMeMaybe(transport.start(clientTransportListener));

  try {
    // Send a single RPC and wait for the response.
    new Rpc(transport, new Metadata()).halfClose().waitForResponse();
    fail("The stream should have been failed due to client received header exceeds header list"
        + " size limit!");
  } catch (Exception e) {
    Throwable rootCause = getRootCause(e);
    Status status = ((StatusException) rootCause).getStatus();
    assertEquals(Status.Code.INTERNAL, status.getCode());
    assertEquals("HTTP/2 error code: PROTOCOL_ERROR\nReceived Rst Stream",
        status.getDescription());
  }
}
 

@Override
public StreamObserver<PCmdActiveThreadCountRes> commandStreamActiveThreadCount(StreamObserver<Empty> streamConnectionManagerObserver) {
    final Long transportId = getTransportId();
    PinpointGrpcServer pinpointGrpcServer = grpcServerRepository.get(transportId);
    if (pinpointGrpcServer == null) {
        logger.info("{} => local. Can't find PinpointGrpcServer(transportId={})", getAgentInfo().getAgentKey(), transportId);
        streamConnectionManagerObserver.onError(new StatusException(Status.NOT_FOUND));
        return DisabledStreamObserver.DISABLED_INSTANCE;
    }

    try {
        return activeThreadCountService.handle(pinpointGrpcServer, streamConnectionManagerObserver);
    } catch (IllegalArgumentException e) {
        logger.warn("Failed to handle activeThreadCountService. agentKey={}, transportId={}", getAgentInfo().getAgentKey(), transportId, e);
        streamConnectionManagerObserver.onError(Status.INTERNAL.withDescription("Internal Server Error").asException());
        return DisabledStreamObserver.DISABLED_INSTANCE;
    }
}
 

public Book getBook(long shelfId, long bookId) throws StatusException {
  synchronized (lock) {
    @Nullable ShelfInfo shelfInfo = shelves.get(shelfId);
    if (shelfInfo == null) {
      throw Status.NOT_FOUND
          .withDescription("Unknown shelf ID")
          .asException();
    }
    @Nullable Book book = shelfInfo.books.get(bookId);
    if (book == null) {
      throw Status.NOT_FOUND
          .withDescription("Unknown book ID")
          .asException();
    }
    return book;
  }
}
 

private static Throwable toStatusFailure(Throwable throwable) {
    if (throwable instanceof StatusException || throwable instanceof StatusRuntimeException) {
        return throwable;
    } else {
        String desc = throwable.getClass().getName();
        if (throwable.getMessage() != null) {
            desc += " - " + throwable.getMessage();
        }
        if (throwable instanceof IllegalArgumentException) {
            return Status.INVALID_ARGUMENT.withDescription(desc).asException();
        }
        return Status.fromThrowable(throwable)
                .withDescription(desc)
                .asException();
    }
}
 

@Override
public void onNext(PCmdActiveThreadCountRes response) {
    if (streamChannelId == -1) {
        streamChannelId = response.getCommonStreamResponse().getResponseId();
    }

    PCmdStreamResponse headerResponse = response.getCommonStreamResponse();
    int sequenceId = headerResponse.getSequenceId();
    if (sequenceId == 1) {
        boolean success = pinpointGrpcServer.handleStreamCreateMessage(streamChannelId, connectionObserver);
        if (!success) {
            connectionObserver.onError(new StatusException(Status.NOT_FOUND));
            return;
        }
    }

    try {
        pinpointGrpcServer.handleStreamMessage(streamChannelId, response);
    } catch (StreamException e) {
        logger.warn("Failed to handle streamMessage. message:{}", e.getMessage(), e);
        connectionObserver.onError(new StatusException(Status.INTERNAL.withDescription(e.getMessage())));
    }
}
 

public Book createBook(long shelfId, Book book) throws StatusException {
  synchronized (lock) {
    @Nullable ShelfInfo shelfInfo = shelves.get(shelfId);
    if (shelfInfo == null) {
      throw Status.NOT_FOUND
          .withDescription("Unknown shelf ID")
          .asException();
    }
    shelfInfo.lastBookId++;
    book = book.toBuilder()
        .setId(shelfInfo.lastBookId)
        .build();
    shelfInfo.books.put(shelfInfo.lastBookId, book);
  }
  return book;
}
 

@Override
protected void validateAction(
    String operationName,
    Action action,
    PreconditionFailure.Builder preconditionFailure,
    RequestMetadata requestMetadata)
    throws InterruptedException, StatusException {
  if (action.hasTimeout() && config.hasMaximumActionTimeout()) {
    Duration timeout = action.getTimeout();
    Duration maximum = config.getMaximumActionTimeout();
    if (timeout.getSeconds() > maximum.getSeconds()
        || (timeout.getSeconds() == maximum.getSeconds()
            && timeout.getNanos() > maximum.getNanos())) {
      preconditionFailure
          .addViolationsBuilder()
          .setType(VIOLATION_TYPE_INVALID)
          .setSubject(Durations.toString(timeout) + " > " + Durations.toString(maximum))
          .setDescription(TIMEOUT_OUT_OF_BOUNDS);
    }
  }

  super.validateAction(operationName, action, preconditionFailure, requestMetadata);
}
 

private ByteBuf makeResponseBuffer(ByteBuffer requestBuffer)
    throws InvalidProtocolBufferException, StatusException {
  GrpcRequest grpcRequest = requestParser.parse(requestBuffer);
  GrpcResponse grpcResponse = appLogic.apply(grpcRequest);

  byte[] dataBytes = grpcResponse.toByteArray();
  int length = dataBytes.length;
  byte[] lengthByteBuffer = ByteBuffer.allocate(4).putInt(length).array();
  byte[] compressedByteBuffer = ByteBuffer.allocate(1).put((byte) 0).array();

  ByteBuf responseBuffer =
      UnpooledByteBufAllocator.DEFAULT.buffer(length + METADATA_SIZE, length + METADATA_SIZE);

  responseBuffer.writeBytes(compressedByteBuffer);
  responseBuffer.writeBytes(lengthByteBuffer);
  responseBuffer.writeBytes(dataBytes);

  return responseBuffer;
}
 

protected void validateAction(
    String operationName,
    Action action,
    PreconditionFailure.Builder preconditionFailure,
    RequestMetadata requestMetadata)
    throws InterruptedException, StatusException {
  ExecutorService service = newDirectExecutorService();
  ImmutableSet.Builder<Digest> inputDigestsBuilder = ImmutableSet.builder();
  Tree tree =
      getUnchecked(
          getTreeFuture(operationName, action.getInputRootDigest(), service, requestMetadata));
  validateAction(
      action,
      getUnchecked(expect(action.getCommandDigest(), Command.parser(), service, requestMetadata)),
      DigestUtil.proxyDirectoriesIndex(tree.getDirectories()),
      inputDigestsBuilder::add,
      preconditionFailure);
  validateInputs(inputDigestsBuilder.build(), preconditionFailure, service, requestMetadata);
}
 

@Override
public <T> ServerRequest<T> newServerRequest(Message<T> message) throws StatusException {
    final Context current = Context.current();
    final Header header = ServerContext.getAgentInfo(current);
    if (header == null) {
        throw Status.INTERNAL.withDescription("Not found request header").asException();
    }

    final TransportMetadata transportMetadata = ServerContext.getTransportMetadata(current);
    if (transportMetadata == null) {
        throw Status.INTERNAL.withDescription("Not found transportMetadata").asException();
    }

    InetSocketAddress inetSocketAddress = transportMetadata.getRemoteAddress();
    ServerRequest<T> request = new DefaultServerRequest<>(message, inetSocketAddress.getHostString(), inetSocketAddress.getPort());
    return request;
}
 

@Test
public void uploadOutputsUploadsEmptyOutputDirectories()
    throws IOException, StatusException, InterruptedException {
  Files.createDirectory(root.resolve("foo"));
  // maybe make some files...
  uploadOutputs(ImmutableList.<String>of(), ImmutableList.<String>of("foo"));
  Tree emptyTree = Tree.newBuilder().setRoot(Directory.getDefaultInstance()).build();
  ByteString emptyTreeBlob = emptyTree.toByteString();
  ArgumentCaptor<Map<HashCode, Chunker>> uploadCaptor = ArgumentCaptor.forClass(Map.class);
  verify(mockUploader).uploadBlobs(uploadCaptor.capture());
  Map<HashCode, Chunker> upload = uploadCaptor.getValue();
  Chunker chunker = upload.get(DIGEST_UTIL.computeHash(emptyTreeBlob));
  assertThat(chunker.next().getData()).isEqualTo(emptyTreeBlob);
  assertThat(resultBuilder.getOutputDirectoriesList())
      .containsExactly(
          OutputDirectory.newBuilder()
              .setPath("foo")
              .setTreeDigest(DIGEST_UTIL.compute(emptyTree))
              .build());
}
 

@Test
public void maxHeaderListSizeShouldBeEnforcedOnClient() throws Exception {
  startServer();

  NettyClientTransport transport =
      newTransport(newNegotiator(), DEFAULT_MAX_MESSAGE_SIZE, 1, null, true);
  callMeMaybe(transport.start(clientTransportListener));

  try {
    // Send a single RPC and wait for the response.
    new Rpc(transport, new Metadata()).halfClose().waitForResponse();
    fail("The stream should have been failed due to client received header exceeds header list"
        + " size limit!");
  } catch (Exception e) {
    Throwable rootCause = getRootCause(e);
    Status status = ((StatusException) rootCause).getStatus();
    assertEquals(Status.Code.INTERNAL, status.getCode());
    assertEquals("HTTP/2 error code: PROTOCOL_ERROR\nReceived Rst Stream",
        status.getDescription());
  }
}
 

private String validatePackage(String packageName) throws StatusException {
  if (packageName.contains("/")) {
    throw new StatusException(
        Status.fromCode(Status.Code.INVALID_ARGUMENT)
            .withDescription("Package name contains invalid /"));
  }
  if (packageName.endsWith(".")) {
    throw new StatusException(
        Status.fromCode(Status.Code.INVALID_ARGUMENT)
            .withDescription("Package name should not end with ."));
  }
  return packageName;
}
 

public static void logConnectionError(Logger log, Throwable t) {
  if (t instanceof StatusRuntimeException) {
    log.info("Connection status: " + niceToString(((StatusRuntimeException) t).getStatus()));
  } else if (t instanceof StatusException) {
    log.info("Connection status: " + niceToString(((StatusException) t).getStatus()));
  } else {
    log.error("Error from stream: " + t.getMessage(), t);
  }
}
 

@Override
public void updateResourceBinding(
    RegistryP.UpdateResourceBindingRequest request,
    StreamObserver<RegistryP.UpdateResourceBindingResponse> responseObserver) {
  try {
    AbstractRegistry registry = metaStore.registries.get(request.getRegistryName());
    RegistryP.ResourceBinding resourceBinding = request.getBinding();
    registry.updateResourceBinding(resourceBinding, false);
    responseObserver.onNext(RegistryP.UpdateResourceBindingResponse.newBuilder().build());
    responseObserver.onCompleted();
  } catch (StatusException e) {
    responseObserver.onError(e);
  }
}
 

@Override
public void deleteResourceBinding(
    RegistryP.DeleteResourceBindingRequest request,
    StreamObserver<RegistryP.DeleteResourceBindingResponse> responseObserver) {
  try {
    AbstractRegistry registry = metaStore.registries.get(request.getRegistryName());
    registry.deleteResourceBinding(request.getLinkedResource());
    responseObserver.onNext(RegistryP.DeleteResourceBindingResponse.newBuilder().build());
    responseObserver.onCompleted();
  } catch (StatusException e) {
    responseObserver.onError(e);
  }
}
 

/** Sends a {@link PublishLifecycleEventRequest} to the BES backend. */
private void publishLifecycleEvent(PublishLifecycleEventRequest request)
    throws DetailedStatusException, InterruptedException {
  int retryAttempt = 0;
  StatusException cause = null;
  while (retryAttempt <= MAX_NUM_RETRIES) {
    try {
      besClient.publish(request);
      return;
    } catch (StatusException e) {
      if (!shouldRetryStatus(e.getStatus())) {
        String message =
            String.format("Not retrying publishLifecycleEvent: status='%s'", e.getStatus());
        logger.atInfo().log(message);
        throw withFailureDetail(e, BuildProgress.Code.BES_STREAM_NOT_RETRYING_FAILURE, message);
      }

      cause = e;

      long sleepMillis = retrySleepMillis(retryAttempt);
      logger.atInfo().log(
          "Retrying publishLifecycleEvent: status='%s', sleepMillis=%d",
          e.getStatus(), sleepMillis);
      sleeper.sleepMillis(sleepMillis);
      retryAttempt++;
    }
  }

  // All retry attempts failed
  throw withFailureDetail(
      cause,
      BuildProgress.Code.BES_UPLOAD_RETRY_LIMIT_EXCEEDED_FAILURE,
      "All retry attempts failed.");
}
 

private void updateShadowCache() {
  ProtoDomain original = null;
  try {
    original = registries.get(shadowOf).get();
  } catch (StatusException e) {
    throw new RuntimeException("Unable to find registry with name " + shadowOf);
  }
  protoContainer = new ShadowApply().applyDelta(original, this.delta);
  protoContainer.registerOptions();
}
 

public RegistryP.ResourceBinding getResourceBinding(String linkedResource)
    throws StatusException {
  BindResult bindResult = this.bindProviders.get(0).getResourceBinding(linkedResource);
  if (bindResult == null) {
    throw Status.NOT_FOUND
        .withDescription("No binding for the linked_resource is found.")
        .asException();
  }
  return toResourceBinding(bindResult);
}
 

public static boolean causedByStatus(Throwable e, Status.Code expected) {
  if (e instanceof StatusRuntimeException) {
    return ((StatusRuntimeException) e).getStatus().getCode() == expected;
  } else if (e instanceof StatusException) {
    return ((StatusException) e).getStatus().getCode() == expected;
  } else if (e.getCause() != null) {
    return causedByStatus(e.getCause(), expected);
  }
  return false;
}
 

@Test
@SuppressWarnings("unchecked")
public void testImmediateFailure() throws Exception {
  Throwable failure = new StatusException(Status.INTERNAL);
  when(fakeServer.publishBuildToolEventStream(any()))
      .thenAnswer(
          invocation -> {
            StreamObserver<PublishBuildToolEventStreamResponse> responseObserver =
                (StreamObserver<PublishBuildToolEventStreamResponse>)
                    invocation.getArguments()[0];
            responseObserver.onError(failure);
            return NULL_OBSERVER;
          });
  assertThat(grpcClient.openStream(ack -> {}).getStatus().get()).isEqualTo(Status.INTERNAL);
}
 

private URI removePort(URI uri) throws StatusException {
    try {
        return new URI(uri.getScheme(), uri.getUserInfo(), uri.getHost(), -1 /* port */,
                uri.getPath(), uri.getQuery(), uri.getFragment());
    } catch (URISyntaxException e) {
        throw Status.UNAUTHENTICATED
                .withDescription("Unable to construct service URI after removing port")
                .withCause(e).asException();
    }
}
 

@Test
public void toStatusException() throws Exception {
  StatusException se = StatusProto.toStatusException(STATUS_PROTO);
  com.google.rpc.Status extractedStatusProto = StatusProto.fromThrowable(se);

  assertEquals(STATUS_PROTO.getCode(), se.getStatus().getCode().value());
  assertEquals(STATUS_PROTO.getMessage(), se.getStatus().getDescription());
  assertEquals(STATUS_PROTO, extractedStatusProto);
}
 

@Test
public void toStatusExceptionWithMetadata_shouldIncludeMetadata() throws Exception {
  StatusException se = StatusProto.toStatusException(STATUS_PROTO, metadata);
  com.google.rpc.Status extractedStatusProto = StatusProto.fromThrowable(se);

  assertEquals(STATUS_PROTO.getCode(), se.getStatus().getCode().value());
  assertEquals(STATUS_PROTO.getMessage(), se.getStatus().getDescription());
  assertEquals(STATUS_PROTO, extractedStatusProto);
  assertNotNull(se.getTrailers());
  assertEquals(METADATA_VALUE, se.getTrailers().get(METADATA_KEY));
}
 

@Test
public void fromThrowableWithNestedStatusException() {
  StatusException se = StatusProto.toStatusException(STATUS_PROTO);
  Throwable nestedSe = new Throwable(se);

  com.google.rpc.Status extractedStatusProto = StatusProto.fromThrowable(se);
  com.google.rpc.Status extractedStatusProtoFromNestedSe = StatusProto.fromThrowable(nestedSe);

  assertEquals(extractedStatusProto, extractedStatusProtoFromNestedSe);
}
 

@Override
public void getGauge(Metrics.GaugeRequest request,
    StreamObserver<Metrics.GaugeResponse> responseObserver) {
  String gaugeName = request.getName();
  Metrics.GaugeResponse gauge = gauges.get(gaugeName);
  if (gauge != null) {
    responseObserver.onNext(gauge);
    responseObserver.onCompleted();
  } else {
    responseObserver.onError(new StatusException(Status.NOT_FOUND));
  }
}