content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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
   *
   *     https://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 org.apache.avro.ipc.netty;
  
  import java.io.IOException;
  import java.net.InetSocketAddress;
  import java.nio.ByteBuffer;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import java.util.concurrent.Semaphore;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.concurrent.atomic.AtomicReference;
  
  import org.apache.avro.AvroRuntimeException;
  import org.apache.avro.ipc.CallFuture;
  import org.apache.avro.ipc.Callback;
  import org.apache.avro.ipc.Responder;
  import org.apache.avro.ipc.Server;
  import org.apache.avro.ipc.Transceiver;
  import org.apache.avro.ipc.specific.SpecificRequestor;
  import org.apache.avro.ipc.specific.SpecificResponder;
  import org.apache.avro.test.Simple;
  import org.apache.avro.test.TestError;
  import org.apache.avro.test.TestRecord;
  import org.junit.jupiter.api.AfterAll;
  import org.junit.jupiter.api.BeforeAll;
  import org.junit.jupiter.api.Disabled;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.Timeout;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  import static org.junit.jupiter.api.Assertions.fail;
  
  /**
   * Tests asynchronous RPCs with Netty.
   */
  public class TestNettyServerWithCallbacks {
    private static Server server;
    private static Transceiver transceiver;
    private static Simple.Callback simpleClient;
    private static final AtomicBoolean ackFlag = new AtomicBoolean(false);
    private static final AtomicReference<CountDownLatch> ackLatch = new AtomicReference<>(new CountDownLatch(1));
    private static Simple simpleService = new SimpleImpl(ackFlag);
  
    @BeforeAll
    public static void initializeConnections() throws Exception {
      // start server
      Responder responder = new SpecificResponder(Simple.class, simpleService);
      server = new NettyServer(responder, new InetSocketAddress(0));
      server.start();
  
      int serverPort = server.getPort();
      System.out.println("server port : " + serverPort);
  
      transceiver = new NettyTransceiver(new InetSocketAddress(serverPort), TestNettyServer.CONNECT_TIMEOUT_MILLIS);
      simpleClient = SpecificRequestor.getClient(Simple.Callback.class, transceiver);
    }
  
    @AfterAll
    public static void tearDownConnections() throws Exception {
      if (transceiver != null) {
        transceiver.close();
      }
      if (server != null) {
        server.close();
      }
    }
  
    @Test
    void greeting() throws Exception {
      // Test synchronous RPC:
      assertEquals("Hello, how are you?", simpleClient.hello("how are you?"));
  
      // Test asynchronous RPC (future):
      CallFuture<String> future1 = new CallFuture<>();
      simpleClient.hello("World!", future1);
      assertEquals("Hello, World!", future1.get(2, TimeUnit.SECONDS));
      assertNull(future1.getError());
  
      // Test asynchronous RPC (callback):
      final CallFuture<String> future2 = new CallFuture<>();
      simpleClient.hello("what's up?", new Callback<String>() {
        @Override
        public void handleResult(String result) {
          future2.handleResult(result);
        }
  
        @Override
        public void handleError(Throwable error) {
          future2.handleError(error);
        }
      });
      assertEquals("Hello, what's up?", future2.get(2, TimeUnit.SECONDS));
      assertNull(future2.getError());
    }
  
    @Test
    void echo() throws Exception {
      TestRecord record = TestRecord.newBuilder()
          .setHash(new org.apache.avro.test.MD5(new byte[] { 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8 }))
          .setKind(org.apache.avro.test.Kind.FOO).setName("My Record").build();
  
      // Test synchronous RPC:
      assertEquals(record, simpleClient.echo(record));
  
      // Test asynchronous RPC (future):
      CallFuture<TestRecord> future1 = new CallFuture<>();
      simpleClient.echo(record, future1);
      assertEquals(record, future1.get(2, TimeUnit.SECONDS));
      assertNull(future1.getError());
  
      // Test asynchronous RPC (callback):
      final CallFuture<TestRecord> future2 = new CallFuture<>();
      simpleClient.echo(record, new Callback<TestRecord>() {
        @Override
        public void handleResult(TestRecord result) {
          future2.handleResult(result);
        }
  
        @Override
        public void handleError(Throwable error) {
          future2.handleError(error);
        }
      });
      assertEquals(record, future2.get(2, TimeUnit.SECONDS));
      assertNull(future2.getError());
    }
  
    @Test
    void add() throws Exception {
      // Test synchronous RPC:
      assertEquals(8, simpleClient.add(2, 6));
  
      // Test asynchronous RPC (future):
      CallFuture<Integer> future1 = new CallFuture<>();
      simpleClient.add(8, 8, future1);
      assertEquals(Integer.valueOf(16), future1.get(2, TimeUnit.SECONDS));
      assertNull(future1.getError());
  
      // Test asynchronous RPC (callback):
      final CallFuture<Integer> future2 = new CallFuture<>();
      simpleClient.add(512, 256, new Callback<Integer>() {
        @Override
        public void handleResult(Integer result) {
          future2.handleResult(result);
        }
  
        @Override
        public void handleError(Throwable error) {
          future2.handleError(error);
        }
      });
      assertEquals(Integer.valueOf(768), future2.get(2, TimeUnit.SECONDS));
      assertNull(future2.getError());
    }
  
    @Test
    void echoBytes() throws Exception {
      ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[] { 1, 2, 3, 4, 5, 6, 7, 8 });
  
      // Test synchronous RPC:
      assertEquals(byteBuffer, simpleClient.echoBytes(byteBuffer));
  
      // Test asynchronous RPC (future):
      CallFuture<ByteBuffer> future1 = new CallFuture<>();
      simpleClient.echoBytes(byteBuffer, future1);
      assertEquals(byteBuffer, future1.get(2, TimeUnit.SECONDS));
      assertNull(future1.getError());
  
      // Test asynchronous RPC (callback):
      final CallFuture<ByteBuffer> future2 = new CallFuture<>();
      simpleClient.echoBytes(byteBuffer, new Callback<ByteBuffer>() {
        @Override
        public void handleResult(ByteBuffer result) {
          future2.handleResult(result);
        }
  
        @Override
        public void handleError(Throwable error) {
          future2.handleError(error);
        }
      });
      assertEquals(byteBuffer, future2.get(2, TimeUnit.SECONDS));
      assertNull(future2.getError());
    }
  
    @Test
    void error() throws IOException, InterruptedException, TimeoutException {
      // Test synchronous RPC:
      try {
        simpleClient.error();
        fail("Expected " + TestError.class.getCanonicalName());
      } catch (TestError e) {
        // Expected
      }
  
      // Test asynchronous RPC (future):
      CallFuture<Void> future = new CallFuture<>();
      simpleClient.error(future);
      try {
        future.get(2, TimeUnit.SECONDS);
        fail("Expected " + TestError.class.getCanonicalName() + " to be thrown");
      } catch (ExecutionException e) {
        assertTrue(e.getCause() instanceof TestError, "Expected " + TestError.class.getCanonicalName());
      }
      assertNotNull(future.getError());
      assertTrue(future.getError() instanceof TestError, "Expected " + TestError.class.getCanonicalName());
      assertNull(future.getResult());
  
      // Test asynchronous RPC (callback):
      final CountDownLatch latch = new CountDownLatch(1);
      final AtomicReference<Throwable> errorRef = new AtomicReference<>();
      simpleClient.error(new Callback<Void>() {
        @Override
        public void handleResult(Void result) {
          fail("Expected " + TestError.class.getCanonicalName());
        }
  
        @Override
        public void handleError(Throwable error) {
          errorRef.set(error);
          latch.countDown();
        }
      });
      assertTrue(latch.await(2, TimeUnit.SECONDS), "Timed out waiting for error");
      assertNotNull(errorRef.get());
      assertTrue(errorRef.get() instanceof TestError);
    }
  
    @Test
    void ack() throws Exception {
      simpleClient.ack();
      ackLatch.get().await(2, TimeUnit.SECONDS);
      assertTrue(ackFlag.get(), "Expected ack flag to be set");
  
      ackLatch.set(new CountDownLatch(1));
      simpleClient.ack();
      ackLatch.get().await(2, TimeUnit.SECONDS);
      assertFalse(ackFlag.get(), "Expected ack flag to be cleared");
    }
  
    @Test
    void sendAfterChannelClose() throws Exception {
      // Start up a second server so that closing the server doesn't
      // interfere with the other unit tests:
      Server server2 = new NettyServer(new SpecificResponder(Simple.class, simpleService), new InetSocketAddress(0));
      server2.start();
      try {
        int serverPort = server2.getPort();
        System.out.println("server2 port : " + serverPort);
  
        try (Transceiver transceiver2 = new NettyTransceiver(new InetSocketAddress(serverPort),
            TestNettyServer.CONNECT_TIMEOUT_MILLIS)) {
          Simple.Callback simpleClient2 = SpecificRequestor.getClient(Simple.Callback.class, transceiver2);
  
          // Verify that connection works:
          assertEquals(3, simpleClient2.add(1, 2));
  
          // Try again with callbacks:
          CallFuture<Integer> addFuture = new CallFuture<>();
          simpleClient2.add(1, 2, addFuture);
          assertEquals(Integer.valueOf(3), addFuture.get());
  
          // Shut down server:
          server2.close();
          Thread.sleep(1000L);
  
          // Send a new RPC, and verify that it throws an Exception that
          // can be detected by the client:
          boolean ioeCaught = false;
          try {
            simpleClient2.add(1, 2);
            fail("Send after server close should have thrown Exception");
          } catch (AvroRuntimeException e) {
            ioeCaught = e.getCause() instanceof IOException;
            assertTrue(ioeCaught, "Expected IOException");
          } catch (Exception e) {
            e.printStackTrace();
            throw e;
          }
          assertTrue(ioeCaught, "Expected IOException");
  
          // Send a new RPC with callback, and verify that the correct Exception
          // is thrown:
          ioeCaught = false;
          try {
            addFuture = new CallFuture<>();
            simpleClient2.add(1, 2, addFuture);
            addFuture.get();
            fail("Send after server close should have thrown Exception");
          } catch (IOException e) {
            ioeCaught = true;
          } catch (Exception e) {
            e.printStackTrace();
            throw e;
          }
          assertTrue(ioeCaught, "Expected IOException");
        }
      } finally {
        server2.close();
      }
    }
  
    @Test
    void cancelPendingRequestsOnTransceiverClose() throws Exception {
      // Start up a second server so that closing the server doesn't
      // interfere with the other unit tests:
      BlockingSimpleImpl blockingSimpleImpl = new BlockingSimpleImpl();
      Server server2 = new NettyServer(new SpecificResponder(Simple.class, blockingSimpleImpl), new InetSocketAddress(0));
      server2.start();
      try {
        int serverPort = server2.getPort();
        System.out.println("server2 port : " + serverPort);
  
        CallFuture<Integer> addFuture = new CallFuture<>();
        try (Transceiver transceiver2 = new NettyTransceiver(new InetSocketAddress(serverPort),
            TestNettyServer.CONNECT_TIMEOUT_MILLIS)) {
          Simple.Callback simpleClient2 = SpecificRequestor.getClient(Simple.Callback.class, transceiver2);
  
          // The first call has to block for the handshake:
          assertEquals(3, simpleClient2.add(1, 2));
  
          // Now acquire the semaphore so that the server will block:
          blockingSimpleImpl.acquireRunPermit();
          simpleClient2.add(1, 2, addFuture);
        }
        // When the transceiver is closed, the CallFuture should get
        // an IOException
        boolean ioeThrown = false;
        try {
          addFuture.get();
        } catch (ExecutionException e) {
          ioeThrown = e.getCause() instanceof IOException;
          assertTrue(e.getCause() instanceof IOException);
        } catch (Exception e) {
          e.printStackTrace();
          fail("Unexpected Exception: " + e.toString());
        }
        assertTrue(ioeThrown, "Expected IOException to be thrown");
      } finally {
        blockingSimpleImpl.releaseRunPermit();
        server2.close();
      }
    }
  
    @Test
    @Timeout(20000)
    void cancelPendingRequestsAfterChannelCloseByServerShutdown() throws Throwable {
      // The purpose of this test is to verify that a client doesn't stay
      // blocked when a server is unexpectedly killed (or when for some
      // other reason the channel is suddenly closed) while the server
      // was in the process of handling a request (thus after it received
      // the request, and before it returned the response).
  
      // Start up a second server so that closing the server doesn't
      // interfere with the other unit tests:
      BlockingSimpleImpl blockingSimpleImpl = new BlockingSimpleImpl();
      final Server server2 = new NettyServer(new SpecificResponder(Simple.class, blockingSimpleImpl),
          new InetSocketAddress(0));
      server2.start();
  
      Transceiver transceiver2 = null;
  
      try {
        int serverPort = server2.getPort();
        System.out.println("server2 port : " + serverPort);
  
        transceiver2 = new NettyTransceiver(new InetSocketAddress(serverPort), TestNettyServer.CONNECT_TIMEOUT_MILLIS);
  
        final Simple.Callback simpleClient2 = SpecificRequestor.getClient(Simple.Callback.class, transceiver2);
  
        // Acquire the method-enter permit, which will be released by the
        // server method once we call it
        blockingSimpleImpl.acquireEnterPermit();
  
        // Acquire the run permit, to avoid that the server method returns immediately
        blockingSimpleImpl.acquireRunPermit();
  
        // Start client call
        Future<?> clientFuture = Executors.newSingleThreadExecutor().submit(() -> {
          try {
            simpleClient2.add(3, 4);
            fail("Expected an exception");
          } catch (Exception e) {
            e.printStackTrace();
            // expected
          }
        });
  
        // Wait until method is entered on the server side
        blockingSimpleImpl.acquireEnterPermit();
  
        // The server side method is now blocked waiting on the run permit
        // (= is busy handling the request)
  
        // Stop the server in a separate thread as it blocks the actual thread until the
        // server side
        // method is running
        new Thread(server2::close).start();
  
        // With the server gone, we expect the client to get some exception and exit
        // Wait for the client call to exit
        try {
          clientFuture.get(10, TimeUnit.SECONDS);
        } catch (ExecutionException e) {
          throw e.getCause();
        } catch (TimeoutException e) {
          fail("Client request should not be blocked on server shutdown");
        }
  
      } finally {
        blockingSimpleImpl.releaseRunPermit();
        server2.close();
        if (transceiver2 != null)
          transceiver2.close();
      }
    }
  
    @Test
    void clientReconnectAfterServerRestart() throws Exception {
      // Start up a second server so that closing the server doesn't
      // interfere with the other unit tests:
      SimpleImpl simpleImpl = new BlockingSimpleImpl();
      Server server2 = new NettyServer(new SpecificResponder(Simple.class, simpleImpl), new InetSocketAddress(0));
      try {
        server2.start();
        int serverPort = server2.getPort();
        System.out.println("server2 port : " + serverPort);
  
        // Initialize a client, and establish a connection to the server:
        Transceiver transceiver2 = new NettyTransceiver(new InetSocketAddress(serverPort),
            TestNettyServer.CONNECT_TIMEOUT_MILLIS);
        Simple.Callback simpleClient2 = SpecificRequestor.getClient(Simple.Callback.class, transceiver2);
        assertEquals(3, simpleClient2.add(1, 2));
  
        // Restart the server:
        server2.close();
        try {
          simpleClient2.add(2, -1);
          fail("Client should not be able to invoke RPCs " + "because server is no longer running");
        } catch (Exception e) {
          // Expected since server is no longer running
        }
        Thread.sleep(2000L);
        server2 = new NettyServer(new SpecificResponder(Simple.class, simpleImpl), new InetSocketAddress(serverPort));
        server2.start();
  
        // Invoke an RPC using the same client, which should reestablish the
        // connection to the server:
        assertEquals(3, simpleClient2.add(1, 2));
      } finally {
        server2.close();
      }
    }
  
    @Disabled
    @Test
    void performanceTest() throws Exception {
      final int threadCount = 8;
      final long runTimeMillis = 10 * 1000L;
      ExecutorService threadPool = Executors.newFixedThreadPool(threadCount);
  
      System.out.println("Running performance test for " + runTimeMillis + "ms...");
      final AtomicLong rpcCount = new AtomicLong(0L);
      final AtomicBoolean runFlag = new AtomicBoolean(true);
      final CountDownLatch startLatch = new CountDownLatch(threadCount);
      for (int ii = 0; ii < threadCount; ii++) {
        threadPool.submit(() -> {
          try {
            startLatch.countDown();
            startLatch.await(2, TimeUnit.SECONDS);
            while (runFlag.get()) {
              rpcCount.incrementAndGet();
              assertEquals("Hello, World!", simpleClient.hello("World!"));
            }
          } catch (Exception e) {
            e.printStackTrace();
          }
        });
      }
  
      startLatch.await(2, TimeUnit.SECONDS);
      Thread.sleep(runTimeMillis);
      runFlag.set(false);
      threadPool.shutdown();
      assertTrue(threadPool.awaitTermination(2, TimeUnit.SECONDS), "Timed out shutting down thread pool");
      System.out.println("Completed " + rpcCount.get() + " RPCs in " + runTimeMillis + "ms => "
          + (((double) rpcCount.get() / (double) runTimeMillis) * 1000) + " RPCs/sec, "
          + ((double) runTimeMillis / (double) rpcCount.get()) + " ms/RPC.");
    }
  
    /**
     * Implementation of the Simple interface.
     */
    private static class SimpleImpl implements Simple {
      private final AtomicBoolean ackFlag;
  
      /**
       * Creates a SimpleImpl.
       *
       * @param ackFlag the AtomicBoolean to toggle when ack() is called.
       */
      public SimpleImpl(final AtomicBoolean ackFlag) {
        this.ackFlag = ackFlag;
      }
  
      @Override
      public String hello(String greeting) {
        return "Hello, " + greeting;
      }
  
      @Override
      public TestRecord echo(TestRecord record) {
        return record;
      }
  
      @Override
      public int add(int arg1, int arg2) {
        return arg1 + arg2;
      }
  
      @Override
      public ByteBuffer echoBytes(ByteBuffer data) {
        return data;
      }
  
      @Override
      public void error() throws TestError {
        throw TestError.newBuilder().setMessage$("Test Message").build();
      }
  
      @Override
      synchronized public void ack() {
        ackFlag.set(!ackFlag.get());
        ackLatch.get().countDown();
      }
    }
  
    /**
     * A SimpleImpl that requires a semaphore permit before executing any method.
     */
    private static class BlockingSimpleImpl extends SimpleImpl {
      /** Semaphore that is released when the method is entered. */
      private final Semaphore enterSemaphore = new Semaphore(1);
      /** Semaphore that must be acquired for the method to run and exit. */
      private final Semaphore runSemaphore = new Semaphore(1);
  
      /**
       * Creates a BlockingSimpleImpl.
       */
      public BlockingSimpleImpl() {
        super(new AtomicBoolean());
      }
  
      @Override
      public String hello(String greeting) {
        releaseEnterPermit();
        acquireRunPermit();
        try {
          return super.hello(greeting);
        } finally {
          releaseRunPermit();
        }
      }
  
      @Override
      public TestRecord echo(TestRecord record) {
        releaseEnterPermit();
        acquireRunPermit();
        try {
          return super.echo(record);
        } finally {
          releaseRunPermit();
        }
      }
  
      @Override
      public int add(int arg1, int arg2) {
        releaseEnterPermit();
        acquireRunPermit();
        try {
          return super.add(arg1, arg2);
        } finally {
          releaseRunPermit();
        }
      }
  
      @Override
      public ByteBuffer echoBytes(ByteBuffer data) {
        releaseEnterPermit();
        acquireRunPermit();
        try {
          return super.echoBytes(data);
        } finally {
          releaseRunPermit();
        }
      }
  
      @Override
      public void error() throws TestError {
        releaseEnterPermit();
        acquireRunPermit();
        try {
          super.error();
        } finally {
          releaseRunPermit();
        }
      }
  
      @Override
      public void ack() {
        releaseEnterPermit();
        acquireRunPermit();
        try {
          super.ack();
        } finally {
          releaseRunPermit();
        }
      }
  
      /**
       * Acquires a single permit from the semaphore.
       */
      public void acquireRunPermit() {
        try {
          runSemaphore.acquire();
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          throw new RuntimeException(e);
        }
      }
  
      /**
       * Releases a single permit to the semaphore.
       */
      public void releaseRunPermit() {
        runSemaphore.release();
      }
  
      /**
       * Acquires a single permit from the semaphore.
       */
      public void acquireEnterPermit() {
        try {
          enterSemaphore.acquire();
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          throw new RuntimeException(e);
        }
      }
  
      /**
       * Releases a single permit to the semaphore.
       */
      public void releaseEnterPermit() {
        enterSemaphore.release();
      }
    }
  }
  