alluxio-master-2.x/core/server/common/src/main/java/alluxio/RpcUtils.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio;

import alluxio.conf.SensitiveConfigMask;
import alluxio.exception.AlluxioException;
import alluxio.exception.runtime.AlluxioRuntimeException;
import alluxio.exception.status.AlluxioStatusException;
import alluxio.exception.status.InternalException;
import alluxio.metrics.Metric;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.security.User;
import alluxio.security.authentication.AuthenticatedClientUser;
import alluxio.security.authentication.ClientContextServerInjector;

import com.codahale.metrics.Timer;
import io.grpc.StatusException;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;

import java.io.IOException;
import java.util.concurrent.CompletableFuture;

/**
 * Utilities for handling server RPC calls.
 *
 * There are two types of RPC calls: 1. RPCs that only throw IOException 2. Streaming RPCs
 * For each of these, if failureOk is set, non-fatal errors will only be logged at the DEBUG
 * level and failure metrics will not be recorded.
 */
public final class RpcUtils {
  private RpcUtils() {} // prevent instantiation

  public static final SensitiveConfigMask SENSITIVE_CONFIG_MASKER =
      RpcSensitiveConfigMask.CREDENTIAL_FIELD_MASKER;

  /**
   * Calls the given {@link RpcCallableThrowsIOException} and handles any exceptions thrown. If the
   * RPC fails, a warning or error will be logged.
   *
   * @param logger the logger to use for this call
   * @param callable the callable to call
   * @param methodName the name of the method, used for metrics
   * @param description the format string of the description, used for logging
   * @param responseObserver gRPC response observer
   * @param args the arguments for the description
   * @param <T> the return type of the callable
   */
  public static <T> void call(Logger logger, RpcCallableThrowsIOException<T> callable,
      String methodName, String description, StreamObserver<T> responseObserver, Object... args) {
    call(logger, callable, methodName, false, description, responseObserver, args);
  }

  /**
   * Calls the given {@link RpcCallableThrowsIOException} and handles any exceptions thrown.
   *
   * The failureOk parameter indicates whether or not AlluxioExceptions and IOExceptions are
   * expected results (for example it would be false for the exists() call). In this case, we do not
   * log the failure or increment failure metrics. When a RuntimeException is thrown, we always
   * treat it as a failure and log an error and increment metrics.
   *
   * @param logger the logger to use for this call
   * @param callable the callable to call
   * @param methodName the name of the method, used for metrics
   * @param failureOk whether failures are expected (affects logging and metrics)
   * @param description the format string of the description, used for logging
   * @param responseObserver gRPC response observer
   * @param args the arguments for the description
   * @param <T> the return type of the callable
   */
  public static <T> void call(Logger logger, RpcCallableThrowsIOException<T> callable,
      String methodName, boolean failureOk, String description, StreamObserver<T> responseObserver,
      Object... args) {
    T response;
    try {
      response = callAndReturn(logger, callable, methodName, failureOk, description, args);
    } catch (StatusException | StatusRuntimeException e) {
      responseObserver.onError(e);
      return;
    }
    responseObserver.onNext(response);
    responseObserver.onCompleted();
  }

  /**
   *
   * Calls the given {@link RpcCallableThrowsIOException} and returns its result. Exceptions are
   * logged, accounted for in metrics and then rethrown at the end.
   *
   * The failureOk parameter indicates whether or not AlluxioExceptions and IOExceptions are
   * expected results (for example it would be false for the exists() call). In this case, we do not
   * log the failure or increment failure metrics. When a RuntimeException is thrown, we always
   * treat it as a failure and log an error and increment metrics.
   *
   * @param logger the logger to use for this call
   * @param callable the callable to call
   * @param methodName the name of the method, used for metrics
   * @param failureOk whether failures are expected (affects logging and metrics)
   * @param description the format string of the description, used for logging
   * @param args the arguments for the description
   * @param <T> the return type of the callable
   * @return the result. (Null if failed)
   */
  public static <T> T callAndReturn(Logger logger, RpcCallableThrowsIOException<T> callable,
      String methodName, boolean failureOk, String description, Object... args)
      throws StatusException {
    // avoid string format for better performance if debug is off
    String debugDesc = logger.isDebugEnabled() ? String.format(description,
        processObjects(logger, args)) : null;
    try (MetricsSystem.MultiTimerContext ctx = new MetricsSystem.MultiTimerContext(
        MetricsSystem.timer(MetricKey.MASTER_TOTAL_RPCS.getName()),
        MetricsSystem.timer(getQualifiedMetricName(methodName)))) {
      MetricsSystem.counter(getQualifiedInProgressMetricName(methodName)).inc();
      logger.debug("Enter: {} from {}: {} client version: {}", methodName,
          ClientContextServerInjector.getIpAddress(),
          ClientContextServerInjector.getClientVersion(),
          debugDesc);
      T res = callable.call();
      logger.debug("Exit: {}: {}", methodName, debugDesc);
      return res;
    } catch (AlluxioRuntimeException e) {
      recordFailure(e, methodName, failureOk, description, debugDesc, logger, args);
      throw e.toGrpcStatusRuntimeException();
    } catch (AlluxioException e) {
      recordFailure(e, methodName, failureOk, description, debugDesc, logger, args);
      throw AlluxioStatusException.fromAlluxioException(e).toGrpcStatusException();
    } catch (IOException e) {
      recordFailure(e, methodName, failureOk, description, debugDesc, logger, args);
      throw AlluxioStatusException.fromIOException(e).toGrpcStatusException();
    } catch (RuntimeException | LinkageError e) {
      // Linkage error can happen when ufs libraries are improperly included or classloaded,
      // right now those simply fail silently.
      logger.error("Exit (Error): {}: {}", methodName,
          String.format(description, processObjects(logger, args)), e);
      MetricsSystem.counter(getQualifiedFailureMetricName(methodName)).inc();
      throw new InternalException(e).toGrpcStatusException();
    } finally {
      MetricsSystem.counter(getQualifiedInProgressMetricName(methodName)).dec();
    }
  }

  /**
   * Calls the given method and handled exception. Exceptions are
   * accounted for in metrics and then rethrown at the end.
   *
   * @param logger           the logger to use for this call
   * @param future           the future to call
   * @param methodName       the name of the method, used for metrics
   * @param description      the format string of the description, used for logging
   * @param responseObserver gRPC response observer
   * @param args             the arguments for the description
   * @param <T>              the return type of the method
   */
  public static <T> void invoke(Logger logger, CompletableFuture<T> future, String methodName,
      String description, StreamObserver<T> responseObserver, Object... args) {
    // avoid string format for better performance if debug is off
    String debugDesc =
        logger.isDebugEnabled() ? String.format(description, processObjects(logger, args)) : null;
    try (MetricsSystem.MultiTimerContext ctx = new MetricsSystem.MultiTimerContext(
        MetricsSystem.timer(MetricKey.MASTER_TOTAL_RPCS.getName()),
        MetricsSystem.timer(getQualifiedMetricName(methodName)))) {
      MetricsSystem.counter(getQualifiedInProgressMetricName(methodName)).inc();
      logger.debug("Enter: {}: {}", methodName, debugDesc);
      future.whenComplete((r, t) -> {
        if (t == null) {
          responseObserver.onNext(r);
          responseObserver.onCompleted();
        } else {
          MetricsSystem.counter(getQualifiedFailureMetricName(methodName)).inc();
          logger.warn("Exception when invoking : {}: {}", methodName, debugDesc, t);
          responseObserver.onError(
              AlluxioRuntimeException.from(t).toGrpcStatusRuntimeException());
        }
        logger.debug("Exit: {}: {}", methodName, debugDesc);
        MetricsSystem.counter(getQualifiedInProgressMetricName(methodName)).dec();
      });
    }
  }

  private static void recordFailure(Throwable e, String methodName, boolean failureOk,
      String description, String debugDesc, Logger logger, Object[] args) {
    logger.debug("Exit (Error): {}: {}", methodName, debugDesc, e);
    if (!failureOk) {
      MetricsSystem.counter(getQualifiedFailureMetricName(methodName)).inc();
      if (!logger.isDebugEnabled()) {
        logger.warn("Exit (Error): {}: {}, Error={}", methodName,
            String.format(description, processObjects(logger, args)),
            e.toString());
      }
    }
  }

  /**
   * Handles a streaming RPC callable with logging.
   *
   * @param logger the logger to use for this call
   * @param callable the callable to call
   * @param methodName the name of the method, used for metrics
   * @param sendResponse whether a response should send to the client
   * @param completeResponse whether onComplete should be called on the response observer
   * @param responseObserver gRPC response observer
   * @param description the format string of the description, used for logging
   * @param args the arguments for the description
   * @param <T> the return type of the callable
   */
  public static <T> void streamingRPCAndLog(Logger logger, StreamingRpcCallable<T> callable,
      String methodName, boolean sendResponse, boolean completeResponse,
      StreamObserver<T> responseObserver, String description, Object... args) {
    // avoid string format for better performance if debug is off
    String debugDesc = logger.isDebugEnabled() ? String.format(description,
        processObjects(logger, args)) : null;
    try (Timer.Context ctx = MetricsSystem.timer(getQualifiedMetricName(methodName)).time()) {
      MetricsSystem.counter(getQualifiedInProgressMetricName(methodName)).inc();
      logger.debug("Enter(stream): {}: {}", methodName, debugDesc);
      T result = callable.call();
      logger.debug("Exit(stream) (OK): {}: {}", methodName, debugDesc);
      if (sendResponse) {
        logger.debug("OnNext(stream): {}: {}", methodName, debugDesc);
        responseObserver.onNext(result);
      }
      if (completeResponse) {
        logger.debug("Completing(stream): {}: {}", methodName, debugDesc);
        responseObserver.onCompleted();
        logger.debug("Completed(stream): {}: {}", methodName, debugDesc);
      }
    } catch (Exception e) {
      logger.warn("Exit(stream) (Error): {}: {}, Error={}", methodName,
          String.format(description, args), e.toString());
      MetricsSystem.counter(getQualifiedFailureMetricName(methodName)).inc();
      callable.exceptionCaught(e);
    } finally {
      MetricsSystem.counter(getQualifiedInProgressMetricName(methodName)).dec();
    }
  }

  private static Object[] processObjects(Logger logger, Object... args) {
    return SENSITIVE_CONFIG_MASKER == null
        ? args : SENSITIVE_CONFIG_MASKER.maskObjects(logger, args);
  }

  private static String getQualifiedMetricName(String methodName) {
    return getQualifiedMetricNameInternal(methodName);
  }

  private static String getQualifiedFailureMetricName(String methodName) {
    return getQualifiedMetricNameInternal(methodName , "Failures");
  }

  private static String getQualifiedInProgressMetricName(String methodName) {
    return getQualifiedMetricNameInternal(methodName, "InProgress");
  }

  private static String getQualifiedMetricNameInternal(String ... components) {
    User user = AuthenticatedClientUser.getOrNull();
    if (user != null) {
      return Metric.getMetricNameWithUserTag(String.join("", components), user.getName());
    }
    return String.join("", components);
  }

  /**
   * An interface representing a callable which can only throw Alluxio or IO exceptions.
   *
   * @param <T> the return type of the callable
   */
  public interface RpcCallableThrowsIOException<T> {
    /**
     * The RPC implementation.
     *
     * @return the return value from the RPC
     */
    T call() throws AlluxioException, IOException;
  }

  /**
   * An interface representing a streaming RPC callable.
   *
   * @param <T> the return type of the callable
   */
  public interface StreamingRpcCallable<T> {
    /**
     * The RPC implementation.
     *
     * @return the return value from the RPC
     */
    T call() throws Exception;

    /**
     * Handles exception.
     *
     * @param throwable the exception
     */
    void exceptionCaught(Throwable throwable);
  }
}
