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.cassandra.sidecar.cluster.driver;
  
  import java.net.InetSocketAddress;
  import java.util.Collections;
  import java.util.List;
  import java.util.Set;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.TimeUnit;
  import java.util.stream.Collectors;
  
  import com.google.common.util.concurrent.Uninterruptibles;
  import org.junit.jupiter.api.Assertions;
  
  import com.datastax.driver.core.DriverUtils;
  import com.datastax.driver.core.Host;
  import org.apache.cassandra.distributed.api.IInstance;
  import org.apache.cassandra.sidecar.cluster.instance.InstanceMetadata;
  import org.apache.cassandra.sidecar.testing.IntegrationTestBase;
  import org.apache.cassandra.testing.CassandraIntegrationTest;
  import org.apache.cassandra.testing.IClusterExtension;
  
  import static org.assertj.core.api.Assertions.assertThat;
  
  /**
   * A test for the SidecarLoadBalancingPolicy
   */
  public class SidecarLoadBalancingPolicyTest extends IntegrationTestBase
  {
  
      public static final int SIDECAR_MANAGED_INSTANCES = 2;
  
      private static List<Host> getConnectedHosts(Set<Host> hosts)
      {
          return hosts.stream()
                      .filter(DriverUtils::hasActiveConnections)
                      .collect(Collectors.toList());
      }
  
      @Override
      protected int[] getInstancesToManage(int clusterSize)
      {
          return new int[]{ 1, 2 }; // we only want to manage the first 2 instances in the "cluster"
      }
  
      @CassandraIntegrationTest(nodesPerDc = 6)
      public void shouldMaintainMinimumConnections() throws ExecutionException, InterruptedException
      {
          Set<Host> hosts = sidecarTestContext.session().getCluster().getMetadata().getAllHosts();
          List<Host> connectedHosts = getConnectedHosts(hosts);
          // We manage 2 hosts, and ask for an additional 2 (the default) for connections.
          // Therefore, we expect 4 hosts to have connections at startup.
          int expectedConnections = SIDECAR_MANAGED_INSTANCES + SidecarLoadBalancingPolicy.MIN_NON_LOCAL_CONNECTIONS;
          assertThat(connectedHosts.size()).isEqualTo(expectedConnections);
          // Now, shut down one of the hosts and make sure that we connect to a different node
          IClusterExtension<? extends IInstance> cluster = sidecarTestContext.cluster();
          IInstance inst = instanceToShutdown(cluster, sidecarTestContext.instancesMetadata().instances());
          InetSocketAddress downInstanceAddress = new InetSocketAddress(inst.broadcastAddress().getAddress(),
                                                                        inst.config().getInt("native_transport_port"));
          inst.shutdown(true).get();
          assertThat(inst.isShutdown()).isTrue();
          assertConnectionsWithRetry(downInstanceAddress, expectedConnections);
      }
  
      private void assertConnectionsWithRetry(InetSocketAddress downInstanceAddress, int expectedConnections)
      {
          List<Host> connectedHosts = Collections.emptyList();
          int attempts = 0;
          // Retry for up to 2 minutes, but passes much more quickly most of the time, so this should be safe.
          while (attempts <= 24)
          {
              Set<Host> hosts = sidecarTestContext.session().getCluster().getMetadata().getAllHosts();
              connectedHosts = getConnectedHosts(hosts);
              List<InetSocketAddress> connectedAddresses = getAddresses(connectedHosts);
              if (connectedHosts.size() == expectedConnections && !connectedAddresses.contains(downInstanceAddress))
              {
                  return;
              }
              attempts++;
              Uninterruptibles.sleepUninterruptibly(5, TimeUnit.SECONDS);
          }
          List<InetSocketAddress> connectedAddresses = getAddresses(connectedHosts);
          assertThat(connectedAddresses).doesNotContain(downInstanceAddress);
          if (connectedHosts.size() == expectedConnections)
          {
              return;
          }
          String message =
          String.format("Waited 2 minutes for connected hosts (%d) to be the expected number (%d) but failed",
                        connectedHosts.size(), expectedConnections);
          Assertions.fail(message);
      }
  
      private List<InetSocketAddress> getAddresses(List<Host> connectedHosts)
      {
          return connectedHosts.stream()
                               .map(h -> h.getEndPoint().resolve())
                               .collect(Collectors.toList());
      }
  
      private IInstance instanceToShutdown(IClusterExtension<? extends IInstance> cluster,
                                           List<InstanceMetadata> instances)
      {
          Set<InetSocketAddress> localInstances = instances.stream().map(i -> new InetSocketAddress(i.host(), i.port()))
                                                           .collect(Collectors.toSet());
          for (IInstance inst : cluster)
          {
              InetSocketAddress nativeAddress = new InetSocketAddress(inst.config().broadcastAddress().getAddress(),
                                                                      inst.config().getInt("native_transport_port"));
              if (localInstances.contains(nativeAddress))
              {
                  continue;
              }
              return inst;
          }
          throw new RuntimeException("Could not find instance to shut down");
      }
  }
  