alluxio-master-2.x/core/server/common/src/test/java/alluxio/master/transport/GrpcMessagingTransportTest.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 alluxio.conf.Configuration;
import alluxio.security.user.ServerUserState;

import io.atomix.catalyst.buffer.BufferInput;
import io.atomix.catalyst.buffer.BufferOutput;
import io.atomix.catalyst.serializer.CatalystSerializable;
import io.atomix.catalyst.serializer.Serializer;
import io.grpc.StatusRuntimeException;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Supplier;

/**
 * Units tests for {@link GrpcMessagingTransport}.
 */
public class GrpcMessagingTransportTest {

  private GrpcMessagingTransport mTransport;

  @Before
  public void before() {
    mTransport = new GrpcMessagingTransport(
        Configuration.global(), ServerUserState.global(), "TestClient");
  }

  @After
  public void after() throws Exception {
    mTransport.close();
  }

  @Test
  public void testEstablishConnection() throws Exception {
    // Set by server connection listener, when new connection is opened to server.
    final AtomicBoolean connectionEstablished = new AtomicBoolean(false);
    // Server connection lister that validates connection establishment.
    Consumer<GrpcMessagingConnection> connectionListener
        = new MessagingTransportTestListener((connection) -> {
          // Set connection as established.
          connectionEstablished.set(true);
        });

    // Catalyst thread context for managing client/server.
    GrpcMessagingContext connectionContext
        = createSingleThreadContext("ClientServerCtx");

    // Create and bind transport server.
    InetSocketAddress address
        = bindServer(connectionContext, mTransport.server(), connectionListener);

    // Open a client connection to server.
    connectClient(connectionContext, mTransport.client(), address);

    // Assert server has established the connection.
    Assert.assertTrue(connectionEstablished.get());
  }

  @Test
  public void testConnectionIsolation() throws Exception {

    // Catalyst thread context for managing client/server.
    GrpcMessagingContext connectionContext = createSingleThreadContext("ClientServerCtx");

    // Create and bind transport server.
    InetSocketAddress address =
        bindServer(connectionContext, mTransport.server(), new MessagingTransportTestListener());

    GrpcMessagingClient transportClient = mTransport.client();
    // Open 2 client connections to server.
    GrpcMessagingConnection clientConnection1
        = connectClient(connectionContext, transportClient, address);
    GrpcMessagingConnection clientConnection2
        = connectClient(connectionContext, transportClient, address);

    // Close connection-1.
    clientConnection1.close().get();

    // Sent request over connection-2. Assert response present and null.
    Assert.assertNull(sendRequest(clientConnection2, new DummyRequest("dummy")).get());
  }

  @Test
  public void testConnectionClosed() throws Exception {

    // Catalyst thread context for managing client/server.
    GrpcMessagingContext connectionContext = createSingleThreadContext("ClientServerCtx");

    // Create and bind transport server.
    InetSocketAddress address =
        bindServer(connectionContext, mTransport.server(), new MessagingTransportTestListener());

    // Open a client connection to server.
    GrpcMessagingConnection clientConnection
        = connectClient(connectionContext, mTransport.client(), address);

    // Close connection.
    clientConnection.close().get();

    // Sent request over connection. Assert request can't be sent over closed connection.
    boolean failed = false;
    try {
      sendRequest(clientConnection, new DummyRequest("dummy")).get();
    } catch (ExecutionException e) {
      Assert.assertTrue(e.getCause() instanceof IllegalStateException);
      failed = true;
    }
    Assert.assertTrue(failed);
  }

  @Test
  public void testServerClosed() throws Exception {
    // Catalyst thread context for managing client/server.
    GrpcMessagingContext connectionContext = createSingleThreadContext("ClientServerCtx");

    // Create transport server.
    GrpcMessagingServer server = mTransport.server();
    // Bind transport server.
    InetSocketAddress address
        = bindServer(connectionContext, server, new MessagingTransportTestListener());

    // Open a client connection to server.
    GrpcMessagingConnection clientConnection
        = connectClient(connectionContext, mTransport.client(), address);

    // Close server.
    server.close().get();

    // Sent request over connection. Assert request can't be sent over closed connection.
    boolean failed = false;
    try {
      sendRequest(clientConnection, new DummyRequest("dummy")).get();
    } catch (ExecutionException e) {
      Assert.assertTrue(e.getCause() instanceof IllegalStateException
          || e.getCause() instanceof StatusRuntimeException);
      failed = true;
    }
    Assert.assertTrue(failed);
  }

  /**
   * Creates and binds transport server.
   *
   * @param context catalyst context
   * @param listener listener for new connections
   * @return address to which the server is bound
   * @throws Exception
   */
  private InetSocketAddress bindServer(GrpcMessagingContext context,
      GrpcMessagingServer server, Consumer<GrpcMessagingConnection> listener)
      throws Exception {
    ServerSocket autoBindSocket = new ServerSocket(0);
    InetSocketAddress serverAddress
        = new InetSocketAddress("localhost", autoBindSocket.getLocalPort());
    autoBindSocket.close();

    context.execute(() -> {
      try {
        // Bind server.
        return server.listen(serverAddress, listener);
      } catch (Exception e) {
        throw new RuntimeException(e);
      }
    }).get().get();

    return serverAddress;
  }

  /**
   * Opens a client connection to server.
   *
   * @param context catalyst context
   * @param client transport client
   * @param serverAddress server address
   * @return client connection
   * @throws Exception
   */
  private GrpcMessagingConnection connectClient(GrpcMessagingContext context,
      GrpcMessagingClient client, InetSocketAddress serverAddress)
      throws Exception {
    Supplier<CompletableFuture<GrpcMessagingConnection>> connectionSupplier = () -> {
      try {
        // Create client connection to server.
        return client.connect(serverAddress);
      } catch (Exception e) {
        throw new RuntimeException(e);
      }
    };
    // Run supplier on given context.
    GrpcMessagingConnection clientConnection = context.execute(connectionSupplier).get().get();

    /*
     * gRPC won't establish stream until message is sent over. Explicitly send a command to cause
     * underlying gRPC stream to be established.
     *
     */

    // Validate connection.
    Assert.assertNull(sendRequest(clientConnection, new DummyRequest("dummy")).get());

    return clientConnection;
  }

  /**
   * Sends a request over connection.
   *
   * @param connection connection for sending
   * @param request request to send
   * @return response to request
   * @throws Exception
   */
  private CompletableFuture<Object> sendRequest(GrpcMessagingConnection connection, Object request)
      throws Exception {
    // Future for receiving command completion.
    CompletableFuture<Object> commandFuture = new CompletableFuture<>();
    // Sent a dummy client request.
    createSingleThreadContext("CommandCtx").execute(() -> {
      connection.sendAndReceive(request).whenComplete((result, error) -> {
        if (error != null) {
          commandFuture.completeExceptionally(error);
        } else {
          commandFuture.complete(result);
        }
      });
    }).get();
    return commandFuture;
  }

  /**
   * Creates test serializer.
   *
   * @return the serializer
   */
  private Serializer createTestSerializer() {
    Serializer serializer = new Serializer();

    // Register dummy test command.
    serializer.register(DummyRequest.class);

    return serializer;
  }

  /**
   * Create a single threaded catalyst context with test serializer.
   *
   * @param contextName context name
   * @return thread context
   */
  private GrpcMessagingContext createSingleThreadContext(String contextName) {
    return new GrpcMessagingContext(contextName, createTestSerializer());
  }

  /**
   * Dummy request class that keeps a single string value.
   *
   * Note: Defined as "public static" for allowing catalyst to see empty constructor.
   */
  public static class DummyRequest implements CatalystSerializable {
    private String mContent;

    /**
     * Required by catalyst.
     */
    public DummyRequest() {}

    /**
     * Creates dummy request with given content.
     *
     * @param content request content
     */
    public DummyRequest(String content) {
      mContent = content;
    }

    @Override
    public void writeObject(BufferOutput<?> buffer, Serializer serializer) {
      byte[] contentBytes = mContent.getBytes();
      buffer.writeInt(contentBytes.length);
      buffer.writeBytes(contentBytes);
    }

    @Override
    public void readObject(BufferInput<?> buffer, Serializer serializer) {
      mContent = new String(buffer.readBytes(buffer.readInt()));
    }
  }

  /**
   * Test listener that is used to register test request types on server connection.
   */
  class MessagingTransportTestListener implements Consumer<GrpcMessagingConnection> {
    Consumer<GrpcMessagingConnection> mNestedListener;

    /**
     * Creates test listener.
     */
    public MessagingTransportTestListener() {
      this(null);
    }

    /**
     * Creates test listener with nested listener.
     *
     * @param nestedListener nested listener
     */
    public MessagingTransportTestListener(Consumer<GrpcMessagingConnection> nestedListener) {
      mNestedListener = nestedListener;
    }

    @Override
    public void accept(GrpcMessagingConnection connection) {
      // Register request handler for 'DummyRequest'.
      connection.handler(DummyRequest.class, (command) -> {
        return CompletableFuture.completedFuture(null);
      });
      if (mNestedListener != null) {
        mNestedListener.accept(connection);
      }
    }
  }
}
