alluxio-master-2.x/core/server/common/src/main/java/alluxio/master/transport/GrpcMessagingContext.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.master.transport;

import com.google.common.base.Preconditions;
import io.atomix.catalyst.serializer.Serializer;
import org.apache.http.concurrent.Cancellable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;

/**
 * The context for Grpc messaging single thread.
 * This context uses a {@link ScheduledExecutorService} to schedule events on the context thread.
 */
public class GrpcMessagingContext {
  private static final Logger LOG = LoggerFactory.getLogger(GrpcMessagingContext.class);
  private final ScheduledExecutorService mExecutor;
  // TODO(lu) remove the catalyst serializer and change to gRPC proto
  private final Serializer mSerializer;
  private final Executor mWrappedExecutor = new Executor() {
    @Override
    public void execute(Runnable command) {
      try {
        mExecutor.execute(logFailure(command));
      } catch (RejectedExecutionException e) {
        // Ignore the rejected exception
      }
    }
  };

  /**
   * Constructs a new {@link GrpcMessagingContext}.
   *
   * @param nameFormat the name format used when instantiating the context thread
   * @param serializer the context serializer
   */
  public GrpcMessagingContext(String nameFormat, Serializer serializer) {
    this(new GrpcMessagingThreadFactory(nameFormat), serializer);
  }

  /**
   * Constructs a new single thread {@link GrpcMessagingContext}.
   *
   * @param factory factory to use when creating the new thread
   * @param serializer the serializer
   */
  public GrpcMessagingContext(GrpcMessagingThreadFactory factory, Serializer serializer) {
    this((ScheduledExecutorService) (new ScheduledThreadPoolExecutor(1, factory)), serializer);
  }

  /**
   * Constructs a new single thread {@link GrpcMessagingContext}.
   *
   * @param executor the single thread scheduled executor to schedule events on
   * @param serializer the serializer
   */
  public GrpcMessagingContext(ScheduledExecutorService executor, Serializer serializer) {
    this(getThread(executor), executor, serializer);
  }

  /**
   * Constructs a new {@link GrpcMessagingContext}.
   *
   * @param thread the thread
   * @param executor the executor
   * @param serializer the serializer
   */
  public GrpcMessagingContext(Thread thread,
      ScheduledExecutorService executor, Serializer serializer) {
    mExecutor = executor;
    mSerializer = serializer;
    Preconditions.checkState(thread instanceof GrpcMessagingThread,
        "not a Grpc messaging thread");
    ((GrpcMessagingThread) thread).setContext(this);
  }

  /**
   * Gets the thread from a single threaded executor service.
   */
  protected static GrpcMessagingThread getThread(ExecutorService executor) {
    final AtomicReference<GrpcMessagingThread> thread = new AtomicReference<>();
    try {
      executor.submit(() -> {
        thread.set((GrpcMessagingThread) Thread.currentThread());
      }).get();
    } catch (ExecutionException | InterruptedException e) {
      throw new IllegalStateException("Failed to initialize thread state", e);
    }
    return thread.get();
  }

  /**
   * @return the serializer of this context
   */
  public Serializer serializer() {
    return mSerializer;
  }

  /**
   * @return the executor of this context
   */
  public Executor executor() {
    return mWrappedExecutor;
  }

  /**
   * Submits a one-shot task that becomes enabled after the given delay.
   *
   * @param delay the time from now to delay execution
   * @param runnable the task to execute
   * @return task cancellable
   */
  public Cancellable schedule(Duration delay, Runnable runnable) {
    ScheduledFuture<?> future = mExecutor.schedule(logFailure(runnable),
        delay.toMillis(), TimeUnit.MILLISECONDS);
    return () -> future.cancel(false);
  }

  /**
   * Submits a one-shot task that becomes enabled after the given delay.
   *
   * @param delay the time from now to delay execution
   * @param interval the interval between successive executions
   * @param runnable the task to execute
   * @return task cancellable
   */
  public Cancellable schedule(Duration delay, Duration interval, Runnable runnable) {
    ScheduledFuture<?> future = mExecutor.scheduleAtFixedRate(logFailure(runnable),
        delay.toMillis(), interval.toMillis(), TimeUnit.MILLISECONDS);
    return () -> future.cancel(false);
  }

  /**
   * Closes the context.
   */
  public void close() {
    mExecutor.shutdownNow();
  }

  /**
   * Executes the given callback.
   *
   * @param callback the callback to execute
   * @return completable future of the callback run
   */
  public CompletableFuture<Void> execute(Runnable callback) {
    CompletableFuture<Void> future = new CompletableFuture();
    executor().execute(() -> {
      try {
        callback.run();
        future.complete(null);
      } catch (Throwable var3) {
        future.completeExceptionally(var3);
      }
    });
    return future;
  }

  /**
   * Executes the given callback.
   *
   * @param callback the callback to execute
   * @return completable future of the callback run
   * @param <T> the supplier type
   */
  public <T> CompletableFuture<T> execute(Supplier<T> callback) {
    CompletableFuture<T> future = new CompletableFuture();
    executor().execute(() -> {
      try {
        future.complete(callback.get());
      } catch (Throwable var3) {
        future.completeExceptionally(var3);
      }
    });
    return future;
  }

  /**
   * @return The current thread context or {@code null} if no context exists
   */
  static GrpcMessagingContext currentContext() {
    Thread thread = Thread.currentThread();
    return thread instanceof GrpcMessagingThread
        ? ((GrpcMessagingThread) thread).getContext() : null;
  }

  /**
   * @return The current thread context or {@code null} if no context exists
   */
  static GrpcMessagingContext currentContextOrThrow() {
    GrpcMessagingContext context = currentContext();
    Preconditions.checkNotNull(context, "not on a Grpc messaging thread");
    return context;
  }

  /**
   * Returns a wrapped runnable that logs and rethrows uncaught exceptions.
   */
  Runnable logFailure(Runnable runnable) {
    return () -> {
      try {
        runnable.run();
      } catch (Throwable t) {
        if (!(t instanceof RejectedExecutionException)) {
          LOG.error("An uncaught exception occurred", t);
        }

        throw t;
      }
    };
  }
}
