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.client;
  
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertThrows;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  
  import java.io.IOException;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.function.Supplier;
  
  import scala.concurrent.Future;
  import scala.concurrent.Future$;
  
  import com.google.common.util.concurrent.SettableFuture;
  import org.junit.Before;
  import org.junit.Test;
  import org.mockito.Mockito;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.exception.CelebornException;
  import org.apache.celeborn.common.protocol.message.ControlMessages.HeartbeatFromApplication;
  import org.apache.celeborn.common.protocol.message.ControlMessages.HeartbeatFromWorker;
  import org.apache.celeborn.common.protocol.message.ControlMessages.HeartbeatFromWorkerResponse;
  import org.apache.celeborn.common.protocol.message.ControlMessages.OneWayMessageResponse$;
  import org.apache.celeborn.common.rpc.RpcAddress;
  import org.apache.celeborn.common.rpc.RpcEndpointRef;
  import org.apache.celeborn.common.rpc.RpcEnv;
  import org.apache.celeborn.common.rpc.RpcTimeoutException;
  
  public class MasterClientSuiteJ {
    private static final Logger LOG = LoggerFactory.getLogger(MasterClientSuiteJ.class);
  
    private final String masterHost = "localhost";
    private final int masterPort = 9097;
    private final CelebornConf conf = new CelebornConf(false);
    private final OneWayMessageResponse$ response = OneWayMessageResponse$.MODULE$;
    private final HeartbeatFromWorkerResponse mockResponse =
        Mockito.mock(HeartbeatFromWorkerResponse.class);
  
    private RpcEnv rpcEnv = null;
    private RpcEndpointRef endpointRef = null;
  
    @Before
    public void beforeEach() {
      conf.set(CelebornConf.RPC_ASK_TIMEOUT().key(), "5s");
      conf.set(CelebornConf.NETWORK_TIMEOUT().key(), "5s");
      rpcEnv = Mockito.mock(RpcEnv.class);
      endpointRef = Mockito.mock(RpcEndpointRef.class);
    }
  
    @Test
    public void testSendOneWayMessageWithoutHA() throws Exception {
      final SettableFuture<Boolean> success = SettableFuture.create();
      final CelebornConf conf = prepareForCelebornConfWithoutHA();
  
      prepareForEndpointRefWithoutRetry(
          () -> {
            success.set(true);
            return Future$.MODULE$.successful(response);
          });
      prepareForRpcEnvWithoutHA();
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromApplication message = Mockito.mock(HeartbeatFromApplication.class);
  
      try {
        client.send(message);
      } catch (Throwable t) {
        LOG.error("It should be no exceptions when sending one-way message.", t);
        fail("It should be no exceptions when sending one-way message.");
      }
  
      assertTrue(success.get(5, TimeUnit.SECONDS));
    }
  
    @Test
    public void testSendOneWayMessageWithoutHAWithRetry() throws Exception {
      final AtomicInteger numTries = new AtomicInteger(0);
      final SettableFuture<Boolean> success = SettableFuture.create();
      final CelebornConf conf = prepareForCelebornConfWithoutHA();
  
      prepareForEndpointRefWithRetry(
          numTries,
          () -> {
            success.set(true);
            return Future$.MODULE$.successful(response);
          });
      prepareForRpcEnvWithoutHA();
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromApplication message = Mockito.mock(HeartbeatFromApplication.class);
  
      try {
        client.send(message);
      } catch (Throwable t) {
        LOG.error("It should be no exceptions when sending one-way message.", t);
        fail("It should be no exceptions when sending one-way message.");
      }
  
      assertTrue(success.get(5, TimeUnit.SECONDS));
      assertEquals(3, numTries.get());
    }
  
    @Test
    public void testSendOneWayMessageWithHA() throws Exception {
      final CelebornConf conf = prepareForCelebornConfWithHA();
  
      final SettableFuture<Boolean> success = SettableFuture.create();
  
      prepareForRpcEnvWithHA(
          () -> {
            success.set(true);
            return Future$.MODULE$.successful(response);
          });
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromApplication message = Mockito.mock(HeartbeatFromApplication.class);
  
      try {
        client.send(message);
      } catch (Throwable t) {
        LOG.error("It should be no exceptions when sending one-way message.", t);
        fail("It should be no exceptions when sending one-way message.");
      }
  
      assertTrue(success.get(5, TimeUnit.SECONDS));
    }
  
    @Test
    public void testSendMessageWithoutHA() {
      final CelebornConf conf = prepareForCelebornConfWithoutHA();
  
      prepareForEndpointRefWithoutRetry(() -> Future$.MODULE$.successful(mockResponse));
      prepareForRpcEnvWithoutHA();
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
  
      HeartbeatFromWorkerResponse response = null;
      try {
        response = client.askSync(message, HeartbeatFromWorkerResponse.class);
      } catch (Throwable t) {
        LOG.error("It should be no exceptions when sending one-way message.", t);
        fail("It should be no exceptions when sending one-way message.");
      }
  
      assertEquals(mockResponse, response);
    }
  
    @Test
    public void testSendMessageWithoutHAWithRetry() {
      final AtomicInteger numTries = new AtomicInteger(0);
      final CelebornConf conf = prepareForCelebornConfWithoutHA();
  
      prepareForEndpointRefWithRetry(numTries, () -> Future$.MODULE$.successful(mockResponse));
      prepareForRpcEnvWithoutHA();
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
  
      HeartbeatFromWorkerResponse response = null;
      try {
        response = client.askSync(message, HeartbeatFromWorkerResponse.class);
      } catch (Throwable t) {
        LOG.error("It should be no exceptions when sending one-way message.", t);
        fail("It should be no exceptions when sending one-way message.");
      }
  
      assertEquals(mockResponse, response);
    }
  
    @Test
    public void testSendMessageWithoutHAWithoutRetry() {
      final AtomicInteger numTries = new AtomicInteger(0);
      final CelebornConf conf =
          prepareForCelebornConfWithoutHA().set(CelebornConf.MASTER_CLIENT_MAX_RETRIES(), 0);
  
      prepareForEndpointRefWithRetry(numTries, () -> Future$.MODULE$.failed(new IOException()));
      prepareForRpcEnvWithoutHA();
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
  
      Throwable throwable =
          assertThrows(
              Throwable.class, () -> client.askSync(message, HeartbeatFromWorkerResponse.class));
      assertTrue(throwable.getCause() instanceof IOException);
      assertEquals(1, client.getMaxRetries());
    }
  
    @Test
    public void testSendMessageWithHA() {
      final CelebornConf conf = prepareForCelebornConfWithHA();
  
      prepareForRpcEnvWithHA(() -> Future$.MODULE$.successful(mockResponse));
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
  
      HeartbeatFromWorkerResponse response = null;
      try {
        response = client.askSync(message, HeartbeatFromWorkerResponse.class);
      } catch (Throwable t) {
        LOG.error("It should be no exceptions when sending one-way message.", t);
        fail("It should be no exceptions when sending one-way message.");
      }
  
      assertEquals(mockResponse, response);
    }
  
    @Test
    public void testOneMasterDownCausedByIOExceptionInHA() {
      checkOneMasterDownInHA(new IOException("test"));
    }
  
    @Test
    public void testOneMasterDownCausedByRuntimeExceptionInHA() {
      checkOneMasterDownInHA(new RuntimeException("test"));
    }
  
    @Test
    public void testOneMasterTimeoutInHA() {
      checkOneMasterAskFailedInHA(new RpcTimeoutException("test", new TimeoutException("test")));
    }
  
    private void checkOneMasterDownInHA(Exception causedByException) {
      final CelebornConf conf = prepareForCelebornConfWithHA();
  
      final RpcEndpointRef master1 = Mockito.mock(RpcEndpointRef.class);
      final RpcEndpointRef master3 = Mockito.mock(RpcEndpointRef.class);
  
      // master leader switch to host2
      Mockito.doReturn(
              Future$.MODULE$.failed(new MasterNotLeaderException("host1:9097", "host2:9097", null)))
          .when(master1)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
  
      Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
          .when(master3)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
  
      Mockito.doAnswer(
              (invocation) -> {
                RpcAddress address = invocation.getArgument(0, RpcAddress.class);
                switch (address.host()) {
                  case "host1":
                    return master1;
                  case "host2":
                    throw new CelebornException("test", causedByException);
                  case "host3":
                    return master3;
                  default:
                    fail(
                        "Should use master host1/host2/host3:" + masterPort + ", but use " + address);
                }
                return null;
              })
          .when(rpcEnv)
          .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
  
      HeartbeatFromWorkerResponse response = null;
      try {
        response = client.askSync(message, HeartbeatFromWorkerResponse.class);
      } catch (Throwable t) {
        LOG.error("It should be no exceptions when sending one-way message.", t);
        fail("It should be no exceptions when sending one-way message.");
      }
  
      assertEquals(mockResponse, response);
    }
  
    private void checkOneMasterAskFailedInHA(Exception exception) {
      final CelebornConf conf = prepareForCelebornConfWithHA();
  
      final RpcEndpointRef master1 = Mockito.mock(RpcEndpointRef.class);
      final RpcEndpointRef master2 = Mockito.mock(RpcEndpointRef.class);
      final RpcEndpointRef master3 = Mockito.mock(RpcEndpointRef.class);
  
      // master leader switch to host2
      Mockito.doReturn(
              Future$.MODULE$.failed(new MasterNotLeaderException("host1:9097", "host2:9097", null)))
          .when(master1)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
  
      // Assume master2 get exception.
      Mockito.doReturn(Future$.MODULE$.failed(exception))
          .when(master2)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
  
      Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
          .when(master3)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
  
      Mockito.doAnswer(
              (invocation) -> {
                RpcAddress address = invocation.getArgument(0, RpcAddress.class);
                switch (address.host()) {
                  case "host1":
                    return master1;
                  case "host2":
                    return master2;
                  case "host3":
                    return master3;
                  default:
                    fail(
                        "Should use master host1/host2/host3:" + masterPort + ", but use " + address);
                }
                return null;
              })
          .when(rpcEnv)
          .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
  
      MasterClient client = new MasterClient(rpcEnv, conf, false);
      HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
  
      HeartbeatFromWorkerResponse response = null;
      try {
        response = client.askSync(message, HeartbeatFromWorkerResponse.class);
      } catch (Throwable t) {
        LOG.error("It should be no exceptions when sending one-way message.", t);
        fail("It should be no exceptions when sending one-way message.");
      }
  
      assertEquals(mockResponse, response);
    }
  
    private void prepareForRpcEnvWithHA(final Supplier<Future<?>> supplier) {
      final RpcEndpointRef ref1 = Mockito.mock(RpcEndpointRef.class);
      final RpcEndpointRef ref2 = Mockito.mock(RpcEndpointRef.class);
      final RpcEndpointRef ref3 = Mockito.mock(RpcEndpointRef.class);
  
      // master leader switch to host2
      Mockito.doReturn(
              Future$.MODULE$.failed(new MasterNotLeaderException("host1:9097", "host2:9097", null)))
          .when(ref1)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
      // Assume host2 down.
      Mockito.doReturn(Future$.MODULE$.failed(new IOException("Test IOException")))
          .when(ref2)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
      // master leader switch to host3 after host2 down.
      Mockito.doReturn(supplier.get()).when(ref3).ask(Mockito.any(), Mockito.any(), Mockito.any());
  
      Mockito.doAnswer(
              invocation -> {
                RpcAddress address = invocation.getArgument(0, RpcAddress.class);
                if (address.port() == masterPort) {
                  switch (address.host()) {
                    case "host1":
                      return ref1;
                    case "host2":
                      return ref2;
                    case "host3":
                      return ref3;
                    default:
                      fail(
                          "Should use master host1/host2/host3:"
                              + masterPort
                              + ", but use "
                              + address);
                  }
                } else {
                  fail("Should use master host1/host2/host3:" + masterPort + ", but use " + address);
                }
                return null;
              })
          .when(rpcEnv)
          .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
    }
  
    private void prepareForRpcEnvWithoutHA() {
      Mockito.doAnswer(
              (invocationOnMock) -> {
                RpcAddress address = invocationOnMock.getArgument(0, RpcAddress.class);
                if (address.host().equals(masterHost) && address.port() == masterPort) {
                  return endpointRef;
                } else {
                  fail(
                      "Should only use master + "
                          + masterHost
                          + ":"
                          + masterPort
                          + ", but use "
                          + address);
                  return null;
                }
              })
          .when(rpcEnv)
          .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
    }
  
    private void prepareForEndpointRefWithRetry(
        final AtomicInteger numTries, Supplier<Future<?>> supplier) {
      Mockito.doAnswer(
              invocation -> {
                switch (numTries.getAndIncrement()) {
                  case 0:
                    return Future$.MODULE$.failed(new IOException("Test 1"));
                  case 1:
                    return Future$.MODULE$.failed(new IOException("Test 2"));
                  case 2:
                    return supplier.get();
                  default:
                    return Future$.MODULE$.failed(new IllegalStateException("too many tries."));
                }
              })
          .when(endpointRef)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
    }
  
    private void prepareForEndpointRefWithoutRetry(Supplier<Future<?>> supplier) {
      Mockito.doAnswer(invocation -> supplier.get())
          .when(endpointRef)
          .ask(Mockito.any(), Mockito.any(), Mockito.any());
    }
  
    private CelebornConf prepareForCelebornConfWithoutHA() {
      return conf.clone()
          .set(CelebornConf.HA_ENABLED().key(), "false")
          .set(CelebornConf.MASTER_ENDPOINTS().key(), masterHost + ":" + masterPort);
    }
  
    private CelebornConf prepareForCelebornConfWithHA() {
      return conf.clone()
          .set(CelebornConf.HA_ENABLED().key(), "true")
          .set(CelebornConf.MASTER_ENDPOINTS().key(), "host1:9097,host2:9097,host3:9097")
          .set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
    }
  }
  