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
   *
   *    http://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.celeborn.common.network;
  
  import static org.apache.celeborn.common.util.JavaUtils.getLocalHost;
  import static org.junit.Assert.*;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.ArgumentMatchers.anyInt;
  import static org.mockito.ArgumentMatchers.anyString;
  import static org.mockito.Mockito.mock;
  
  import java.io.IOException;
  import java.util.*;
  import java.util.concurrent.atomic.AtomicInteger;
  
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  import org.mockito.Mockito;
  
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.network.client.TransportClient;
  import org.apache.celeborn.common.network.client.TransportClientFactory;
  import org.apache.celeborn.common.network.server.BaseMessageHandler;
  import org.apache.celeborn.common.network.server.TransportServer;
  import org.apache.celeborn.common.network.util.TransportConf;
  import org.apache.celeborn.common.network.util.TransportFrameDecoder;
  import org.apache.celeborn.common.util.JavaUtils;
  import org.apache.celeborn.common.util.ThreadUtils;
  
  public class TransportClientFactorySuiteJ {
  
    static final String TEST_MODULE = "shuffle";
  
    private TransportContext context;
    private TransportServer server1;
    private TransportServer server2;
  
    protected void doSetup(CelebornConf celebornConf) {
      TransportConf conf = new TransportConf(TEST_MODULE, celebornConf);
      BaseMessageHandler handler = new BaseMessageHandler();
      context = new TransportContext(conf, handler);
      server1 = context.createServer();
      server2 = context.createServer();
    }
  
    @Before
    public void setUp() {
      doSetup(new CelebornConf());
    }
  
    // for validation in subclasses
    TransportConf getTransportContextConf() {
      return context.getConf();
    }
  
    @After
    public void tearDown() {
      JavaUtils.closeQuietly(server1);
      JavaUtils.closeQuietly(server2);
      JavaUtils.closeQuietly(context);
    }
  
    /**
     * Request a bunch of clients to a single server to test we create up to maxConnections of
     * clients.
     *
     * <p>If concurrent is true, create multiple threads to create clients in parallel.
     */
    @SuppressWarnings("DoNotCall")
    private void testClientReuse(int maxConnections, boolean concurrent)
        throws IOException, InterruptedException {
  
      CelebornConf _conf = new CelebornConf();
      _conf.set("celeborn.shuffle.io.numConnectionsPerPeer", Integer.toString(maxConnections));
      TransportConf conf = new TransportConf(TEST_MODULE, _conf);
  
      BaseMessageHandler handler = new BaseMessageHandler();
      TransportContext context = new TransportContext(conf, handler);
      TransportClientFactory factory = context.createClientFactory();
      Set<TransportClient> clients = Collections.synchronizedSet(new HashSet<TransportClient>());
  
      AtomicInteger failed = new AtomicInteger();
      Thread[] attempts = new Thread[maxConnections * 10];
  
      // Launch a bunch of threads to create new clients.
      for (int i = 0; i < attempts.length; i++) {
        attempts[i] =
            ThreadUtils.newThread(
                () -> {
                  try {
                    TransportClient client = factory.createClient(getLocalHost(), server1.getPort());
                    assertTrue(client.isActive());
                    clients.add(client);
                  } catch (IOException e) {
                    failed.incrementAndGet();
                  } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                  }
                },
                "test-thread");
  
        if (concurrent) {
          attempts[i].start();
        } else {
          attempts[i].run();
        }
      }
  
      // Wait until all the threads complete.
      for (Thread attempt : attempts) {
        attempt.join();
      }
  
      assertEquals(0, failed.get());
      assertTrue(clients.size() <= maxConnections);
  
      for (TransportClient client : clients) {
        client.close();
      }
  
      factory.close();
      context.close();
    }
  
    @Test
    public void reuseClientsUpToConfigVariable() throws Exception {
      testClientReuse(1, false);
      testClientReuse(2, false);
      testClientReuse(3, false);
      testClientReuse(4, false);
    }
  
    @Test
    public void reuseClientsUpToConfigVariableConcurrent() throws Exception {
      testClientReuse(1, true);
      testClientReuse(2, true);
      testClientReuse(3, true);
      testClientReuse(4, true);
    }
  
    @Test
    public void returnDifferentClientsForDifferentServers() throws IOException, InterruptedException {
      TransportClientFactory factory = context.createClientFactory();
      TransportClient c1 = factory.createClient(getLocalHost(), server1.getPort());
      TransportClient c2 = factory.createClient(getLocalHost(), server2.getPort());
      assertTrue(c1.isActive());
      assertTrue(c2.isActive());
      assertNotSame(c1, c2);
      factory.close();
    }
  
    @Test
    public void neverReturnInactiveClients() throws IOException, InterruptedException {
      TransportClientFactory factory = context.createClientFactory();
      TransportClient c1 = factory.createClient(getLocalHost(), server1.getPort());
      c1.close();
  
      long start = System.currentTimeMillis();
      while (c1.isActive() && (System.currentTimeMillis() - start) < 3000) {
        Thread.sleep(10);
      }
      assertFalse(c1.isActive());
  
      TransportClient c2 = factory.createClient(getLocalHost(), server1.getPort());
      assertNotSame(c1, c2);
      assertTrue(c2.isActive());
      factory.close();
    }
  
    @Test
    public void closeBlockClientsWithFactory() throws IOException, InterruptedException {
      TransportClientFactory factory = context.createClientFactory();
      TransportClient c1 = factory.createClient(getLocalHost(), server1.getPort());
      TransportClient c2 = factory.createClient(getLocalHost(), server2.getPort());
      assertTrue(c1.isActive());
      assertTrue(c2.isActive());
      factory.close();
      assertFalse(c1.isActive());
      assertFalse(c2.isActive());
    }
  
    @Test
    public void closeIdleConnectionForRequestTimeOut() throws IOException, InterruptedException {
      CelebornConf _conf = new CelebornConf();
      _conf.set("celeborn.shuffle.io.connectionTimeout", "1s");
      TransportConf conf = new TransportConf(TEST_MODULE, _conf);
      TransportContext context = new TransportContext(conf, new BaseMessageHandler(), true);
      try (TransportClientFactory factory = context.createClientFactory()) {
        TransportClient c1 = factory.createClient(getLocalHost(), server1.getPort());
        assertTrue(c1.isActive());
        long expiredTime = System.currentTimeMillis() + 10000; // 10 seconds
        while (c1.isActive() && System.currentTimeMillis() < expiredTime) {
          Thread.sleep(10);
        }
        assertFalse(c1.isActive());
      }
      context.close();
    }
  
    @Test(expected = IOException.class)
    public void closeFactoryBeforeCreateClient() throws IOException, InterruptedException {
      TransportClientFactory factory = context.createClientFactory();
      factory.close();
      factory.createClient(getLocalHost(), server1.getPort());
    }
  
    @Test
    public void unlimitedConnectionAndCreationTimeouts() throws IOException, InterruptedException {
      CelebornConf _conf = new CelebornConf();
      _conf.set("celeborn.shuffle.io.connectTimeout", "-1");
      _conf.set("celeborn.shuffle.io.connectionTimeout", "-1");
      TransportConf conf = new TransportConf(TEST_MODULE, _conf);
      try (TransportContext ctx = new TransportContext(conf, new BaseMessageHandler(), true);
          TransportClientFactory factory = ctx.createClientFactory()) {
        TransportClient c1 = factory.createClient(getLocalHost(), server1.getPort());
        assertTrue(c1.isActive());
        long expiredTime = System.currentTimeMillis() + 5000;
        while (c1.isActive() && System.currentTimeMillis() < expiredTime) {
          Thread.sleep(10);
        }
        assertTrue(c1.isActive());
        // When connectionTimeout is unlimited, the connection shall be able to fail when the server
        // is not reachable.
        TransportServer server = ctx.createServer();
        int unreachablePort = server.getPort();
        JavaUtils.closeQuietly(server);
        IOException exception =
            assertThrows(
                IOException.class, () -> factory.createClient(getLocalHost(), unreachablePort));
        assertNotEquals(exception.getCause(), null);
      }
    }
  
    @Test
    public void testRetryCreateClient() throws IOException, InterruptedException {
      TransportClientFactory factory = Mockito.spy(context.createClientFactory());
      TransportClient client = mock(TransportClient.class);
      Mockito.doThrow(new IOException("xx"))
          .doReturn(client)
          .when(factory)
          .createClient(anyString(), anyInt(), anyInt(), any());
      TransportClient transportClient =
          factory.retryCreateClient("xxx", 10, 1, TransportFrameDecoder::new);
      Assert.assertEquals(transportClient, client);
    }
  }
  