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.distributed.upgrade;
  
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.LinkedHashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.NavigableSet;
  import java.util.Objects;
  import java.util.Set;
  import java.util.function.Consumer;
  import java.util.stream.Collectors;
  
  import org.junit.After;
  import org.junit.Assume;
  import org.junit.BeforeClass;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import com.vdurmont.semver4j.Semver;
  import com.vdurmont.semver4j.Semver.SemverType;
  import org.apache.cassandra.db.DecoratedKey;
  import org.apache.cassandra.dht.Murmur3Partitioner;
  import org.apache.cassandra.distributed.UpgradeableCluster;
  import org.apache.cassandra.distributed.api.Feature;
  import org.apache.cassandra.distributed.api.ICluster;
  import org.apache.cassandra.distributed.api.IInstanceConfig;
  import org.apache.cassandra.distributed.api.TokenSupplier;
  import org.apache.cassandra.distributed.shared.DistributedTestBase;
  import org.apache.cassandra.distributed.shared.NetworkTopology;
  import org.apache.cassandra.distributed.shared.ThrowingRunnable;
  import org.apache.cassandra.distributed.shared.Versions;
  import org.apache.cassandra.utils.ByteBufferUtil;
  import org.apache.cassandra.utils.SimpleGraph;
  
  import static org.apache.cassandra.config.CassandraRelevantProperties.DTEST_ACCORD_ENABLED;
  import static org.apache.cassandra.config.CassandraRelevantProperties.SKIP_GC_INSPECTOR;
  import static org.apache.cassandra.distributed.shared.Versions.Version;
  import static org.apache.cassandra.distributed.shared.Versions.find;
  import static org.apache.cassandra.utils.SimpleGraph.sortedVertices;
  
  // checkstyle: suppress below 'blockSystemPropertyUsage'
  public class UpgradeTestBase extends DistributedTestBase
  {
      private static final Logger logger = LoggerFactory.getLogger(UpgradeTestBase.class);
  
      @After
      public void afterEach()
      {
          triggerGC();
      }
  
      @BeforeClass
      public static void beforeClass() throws Throwable
      {
          ICluster.setup();
          SKIP_GC_INSPECTOR.setBoolean(true);
          DTEST_ACCORD_ENABLED.setBoolean(false);
      }
  
  
      public UpgradeableCluster.Builder builder()
      {
          return UpgradeableCluster.build();
      }
  
      public static interface RunOnCluster
      {
          public void run(UpgradeableCluster cluster) throws Throwable;
      }
  
      public static interface RunOnClusterAndNode
      {
          public void run(UpgradeableCluster cluster, int node) throws Throwable;
      }
  
      public static final Semver v40 = new Semver("4.0-alpha1", SemverType.LOOSE);
      public static final Semver v41 = new Semver("4.1-alpha1", SemverType.LOOSE);
      public static final Semver v42 = new Semver("4.2-alpha1", SemverType.LOOSE);
      public static final Semver v50 = new Semver("5.0-alpha1", SemverType.LOOSE);
      public static final Semver v51 = new Semver("5.1-alpha1", SemverType.LOOSE);
  
      protected static final SimpleGraph<Semver> SUPPORTED_UPGRADE_PATHS = new SimpleGraph.Builder<Semver>()
                                                                           .addEdge(v40, v41)
                                                                           .addEdge(v40, v50)
                                                                           .addEdge(v40, v51)
                                                                           .addEdge(v41, v50)
                                                                           .addEdge(v41, v51)
                                                                           .addEdge(v50, v51)
                                                                           .build();
  
      // the last is always the current
      public static final Semver CURRENT = SimpleGraph.max(SUPPORTED_UPGRADE_PATHS);
      public static final Semver OLDEST = SimpleGraph.min(SUPPORTED_UPGRADE_PATHS);
  
      public static class TestVersions
      {
          final Version initial;
          final List<Version> upgrade;
          final List<Semver> upgradeVersions;
  
          public TestVersions(Version initial, List<Version> upgrade)
          {
              this.initial = initial;
              this.upgrade = upgrade;
              this.upgradeVersions = upgrade.stream().map(v -> v.version).collect(Collectors.toList());
          }
  
          @Override
          public boolean equals(Object o)
          {
              if (this == o) return true;
              if (o == null || getClass() != o.getClass()) return false;
              TestVersions that = (TestVersions) o;
              return Objects.equals(initial.version, that.initial.version) && Objects.equals(upgradeVersions, that.upgradeVersions);
          }
  
          @Override
          public int hashCode()
          {
              return Objects.hash(initial.version, upgradeVersions);
          }
  
          @Override
          public String toString()
          {
              StringBuilder sb = new StringBuilder();
              sb.append(initial.version).append(" -> ");
              sb.append(upgradeVersions);
              return sb.toString();
          }
      }
  
      public static class TestCase implements ThrowingRunnable
      {
          private final Versions versions;
          private final List<TestVersions> upgrade = new ArrayList<>();
          private int nodeCount = 3;
          private RunOnCluster setup;
          private RunOnClusterAndNode runBeforeNodeRestart;
          private RunOnClusterAndNode runAfterNodeUpgrade;
          private RunOnCluster runBeforeClusterUpgrade;
          private RunOnCluster runAfterClusterUpgrade;
          private final Set<Integer> nodesToUpgrade = new LinkedHashSet<>();
          private Consumer<IInstanceConfig> configConsumer;
          private Consumer<UpgradeableCluster.Builder> builderConsumer;
          private TokenSupplier tokenSupplier;
          private Map<Integer, NetworkTopology.DcAndRack> nodeIdTopology;
          private UpgradeListener upgradeListener = new UpgradeListener()
          {
              @Override
              public void shutdown(int i)
              {
              }
  
              @Override
              public void startup(int i)
              {
              }
          };
  
          public TestCase()
          {
              this(find());
          }
  
          public TestCase(Versions versions)
          {
              this.versions = versions;
          }
  
          public TestCase nodes(int nodeCount)
          {
              this.nodeCount = nodeCount;
              return this;
          }
  
          /** performs all supported upgrade paths that exist in between lowerBound and end on CURRENT (inclusive)
           * {@code upgradesToCurrentFrom(3.0); // produces: 3.0 -> CURRENT, 3.11 -> CURRENT, …}
           **/
          public TestCase upgradesToCurrentFrom(Semver lowerBound)
          {
              return upgradesTo(lowerBound, CURRENT);
          }
  
          /**
           * performs all supported upgrade paths to the "to" target; example
           * {@code upgradesTo(3.0, 4.0); // produces: 3.0 -> 4.0, 3.11 -> 4.0}
           */
          public TestCase upgradesTo(Semver lowerBound, Semver to)
          {
              List<TestVersions> upgrade = new ArrayList<>();
              NavigableSet<Semver> vertices = sortedVertices(SUPPORTED_UPGRADE_PATHS);
              for (Semver start : vertices.subSet(lowerBound, true, to, false))
              {
                  // only include pairs that are allowed, and start or end on CURRENT
                  if (SUPPORTED_UPGRADE_PATHS.hasEdge(start, to) && edgeTouchesTarget(start, to, CURRENT))
                      upgrade.add(new TestVersions(versions.getLatest(start), Collections.singletonList(versions.getLatest(to))));
              }
              logger.info("Adding upgrades of\n{}", upgrade.stream().map(TestVersions::toString).collect(Collectors.joining("\n")));
              this.upgrade.addAll(upgrade);
              return this;
          }
  
          /**
           * performs all supported upgrade paths from the "from" target; example
           * {@code upgradesFrom(4.0, 4.2); // produces: 4.0 -> 4.1, 4.0 -> 4.2}
           */
          public TestCase upgradesFrom(Semver from, Semver upperBound)
          {
              List<TestVersions> upgrade = new ArrayList<>();
              NavigableSet<Semver> vertices = sortedVertices(SUPPORTED_UPGRADE_PATHS);
              for (Semver end : vertices.subSet(from, false, upperBound, true))
              {
                  // only include pairs that are allowed, and start or end on CURRENT
                  if (SUPPORTED_UPGRADE_PATHS.hasEdge(from, end) && edgeTouchesTarget(from, end, CURRENT))
                      upgrade.add(new TestVersions(versions.getLatest(from), Collections.singletonList(versions.getLatest(end))));
              }
              logger.info("Adding upgrades of\n{}", upgrade.stream().map(TestVersions::toString).collect(Collectors.joining("\n")));
              this.upgrade.addAll(upgrade);
              return this;
          }
  
          /**
           * performs all supported upgrade paths that exist in between from and to that include the current version.
           * This call is equivalent to calling {@code upgradesTo(from, CURRENT).upgradesFrom(CURRENT, to)}.
           **/
          public TestCase upgrades(Semver lowerBound, Semver upperBound)
          {
              Assume.assumeTrue("Unable to do any upgrades in bounds (" + lowerBound + ", " + upperBound + "); does not cover CURRENT=" + CURRENT, rangeCoversTarget(lowerBound, upperBound, CURRENT));
              if (lowerBound.compareTo(CURRENT) < 0)
                  upgradesTo(lowerBound, CURRENT);
              if (CURRENT.compareTo(upperBound) < 0)
                  upgradesFrom(CURRENT, upperBound);
              return this;
          }
  
          private static boolean rangeCoversTarget(Semver lowerBound, Semver upperBound, Semver target)
          {
              // target >= from && target <= to
              return target.isGreaterThanOrEqualTo(lowerBound) && target.isLowerThanOrEqualTo(upperBound);
          }
  
          /** returns true if the target version has the same major and minor as either the from or the to version **/
          private static boolean edgeTouchesTarget(Semver from, Semver to, Semver target)
          {
              switch (from.diff(target))
              {
                  default:
                      return true;
                  case MAJOR:
                  case MINOR:
                      // fall through
              }
              switch (to.diff(target))
              {
                  default:
                      return true;
                  case MAJOR:
                  case MINOR:
                      return false;
              }
          }
  
          /** Will test this specific upgrade path **/
          public TestCase singleUpgradeToCurrentFrom(Semver from)
          {
              if (!SUPPORTED_UPGRADE_PATHS.hasEdge(from, CURRENT))
                  throw new AssertionError("Upgrading from " + from + " to " + CURRENT + " isn't directly supported and must go through other versions first; supported paths: " + SUPPORTED_UPGRADE_PATHS.findPaths(from, CURRENT));
              TestVersions tests = new TestVersions(this.versions.getLatest(from), Arrays.asList(this.versions.getLatest(CURRENT)));
              logger.info("Adding upgrade of {}", tests);
              this.upgrade.add(tests);
              return this;
          }
  
          public TestCase setup(RunOnCluster setup)
          {
              this.setup = setup;
              return this;
          }
  
          public TestCase runBeforeNodeRestart(RunOnClusterAndNode runBeforeNodeRestart)
          {
              this.runBeforeNodeRestart = runBeforeNodeRestart;
              return this;
          }
  
          public TestCase runAfterNodeUpgrade(RunOnClusterAndNode runAfterNodeUpgrade)
          {
              this.runAfterNodeUpgrade = runAfterNodeUpgrade;
              return this;
          }
  
          public TestCase runBeforeClusterUpgrade(RunOnCluster runBeforeClusterUpgrade)
          {
              this.runBeforeClusterUpgrade = runBeforeClusterUpgrade;
              return this;
          }
  
          public TestCase runAfterClusterUpgrade(RunOnCluster runAfterClusterUpgrade)
          {
              this.runAfterClusterUpgrade = runAfterClusterUpgrade;
              return this;
          }
  
          public TestCase withConfig(Consumer<IInstanceConfig> config)
          {
              this.configConsumer = config;
              return this;
          }
  
          public TestCase withBuilder(Consumer<UpgradeableCluster.Builder> builder)
          {
              this.builderConsumer = builder;
              return this;
          }
  
          public TestCase withUpgradeListener(UpgradeListener listener)
          {
              this.upgradeListener = listener;
              return this;
          }
  
          public TestCase withTokenSupplier(TokenSupplier tokenSupplier)
          {
              this.tokenSupplier = tokenSupplier;
              return this;
          }
  
          public TestCase withNodeIdTopology(Map<Integer, NetworkTopology.DcAndRack> nodeIdTopology)
          {
              this.nodeIdTopology = nodeIdTopology;
              return this;
          }
  
          public void run() throws Throwable
          {
              if (setup == null)
                  throw new AssertionError();
              if (upgrade.isEmpty())
                  throw new AssertionError("no upgrade paths have been specified (or exist)");
              if (runAfterClusterUpgrade == null && runAfterNodeUpgrade == null)
                  throw new AssertionError();
              if (runBeforeNodeRestart == null)
                  runBeforeNodeRestart = (c, n) -> {};
              if (runBeforeClusterUpgrade == null)
                  runBeforeClusterUpgrade = (c) -> {};
              if (runAfterClusterUpgrade == null)
                  runAfterClusterUpgrade = (c) -> {};
              if (runAfterNodeUpgrade == null)
                  runAfterNodeUpgrade = (c, n) -> {};
              if (nodesToUpgrade.isEmpty())
                  for (int n = 1; n <= nodeCount; n++)
                      nodesToUpgrade.add(n);
  
              int offset = 0;
              for (TestVersions upgrade : this.upgrade)
              {
                  logger.info("testing upgrade from {} to {}", upgrade.initial.version, upgrade.upgradeVersions);
                  try (UpgradeableCluster cluster = init(UpgradeableCluster.create(nodeCount, upgrade.initial, configConsumer, builderConsumer, tokenSupplier, nodeIdTopology)))
                  {
                      setup.run(cluster);
  
                      for (Version nextVersion : upgrade.upgrade)
                      {
                          try
                          {
                              runBeforeClusterUpgrade.run(cluster);
  
                              for (int n : nodesToUpgrade)
                              {
                                  upgradeListener.shutdown(n);
                                  cluster.get(n).shutdown().get();
                                  triggerGC();
                                  cluster.get(n).setVersion(nextVersion);
                                  runBeforeNodeRestart.run(cluster, n);
                                  cluster.get(n).startup();
                                  upgradeListener.startup(n);
                                  runAfterNodeUpgrade.run(cluster, n);
                              }
  
                              runAfterClusterUpgrade.run(cluster);
  
                              cluster.checkAndResetUncaughtExceptions();
                          }
                          catch (Throwable t)
                          {
                              throw new AssertionError(String.format("Error in test '%s' while upgrading to '%s'; successful upgrades %s", upgrade, nextVersion.version, this.upgrade.stream().limit(offset).collect(Collectors.toList())), t);
                          }
                      }
                  }
                  offset++;
              }
          }
  
          public TestCase nodesToUpgrade(int ... nodes)
          {
              Set<Integer> set = new HashSet<>(nodes.length);
              for (int n : nodes)
              {
                  set.add(n);
              }
              nodesToUpgrade.addAll(set);
              return this;
          }
  
          public TestCase nodesToUpgradeOrdered(int ... nodes)
          {
              for (int n : nodes)
              {
                  nodesToUpgrade.add(n);
              }
              return this;
          }
       }
  
      private static void triggerGC()
      {
          System.runFinalization();
          System.gc();
      }
  
      protected TestCase allUpgrades(int nodes, int... toUpgrade)
      {
          return new TestCase().nodes(nodes)
                               .upgradesToCurrentFrom(OLDEST)
                               .withConfig(c -> c.with(Feature.GOSSIP))
                               .nodesToUpgrade(toUpgrade);
      }
  
      protected static int primaryReplica(List<Long> initialTokens, Long token)
      {
          int primary = 1;
  
          for (Long initialToken : initialTokens)
          {
              if (token <= initialToken)
              {
                  break;
              }
  
              primary++;
          }
  
          return primary;
      }
  
      protected static Long tokenFrom(int key)
      {
          DecoratedKey dk = Murmur3Partitioner.instance.decorateKey(ByteBufferUtil.bytes(key));
          return (Long) dk.getToken().getTokenValue();
      }
  
      protected static int nextNode(int current, int numNodes)
      {
          return current == numNodes ? 1 : current + 1;
      }
  
      public interface UpgradeListener
      {
          void shutdown(int i);
          void startup(int i);
      }
  }
  