cassandra.py

# Copyright 2015 PerfKitBenchmarker Authors. All rights reserved.
#
# Licensed 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.

"""Installs/Configures Cassandra.

See 'perfkitbenchmarker/data/cassandra/' for configuration files used.

Cassandra homepage: http://cassandra.apache.org
"""


import logging
import os
import posixpath
import re
import time
from perfkitbenchmarker import background_tasks
from perfkitbenchmarker import data
from perfkitbenchmarker import errors
from perfkitbenchmarker import linux_packages


JNA_JAR_URL = (
    'https://maven.java.net/content/repositories/releases/'
    'net/java/dev/jna/jna/4.1.0/jna-4.1.0.jar'
)
CASSANDRA_GIT_REPRO = 'https://github.com/apache/cassandra.git'
CASSANDRA_YAML_TEMPLATE = 'cassandra/cassandra.yaml.j2'
CASSANDRA_RACKDC_TEMPLATE = 'cassandra/cassandra-rackdc.properties.j2'
CASSANDRA_KEYSPACE_TEMPLATE = (
    'cassandra/create-keyspace-cassandra-stress.cql.j2'
)
CASSANDRA_ROW_CACHE_TEMPLATE = 'cassandra/enable-row-caching.cql.j2'
CASSANDRA_VERSION = 'apache-cassandra-5.0.0'
CASSANDRA_DIR = posixpath.join(linux_packages.INSTALL_DIR, CASSANDRA_VERSION)
CASSANDRA_PID = posixpath.join(CASSANDRA_DIR, 'cassandra.pid')
CASSANDRA_OUT = posixpath.join(CASSANDRA_DIR, 'cassandra.out')
CASSANDRA_ERR = posixpath.join(CASSANDRA_DIR, 'cassandra.err')


# Number of times to attempt to start the cluster.
CLUSTER_START_TRIES = 10
CLUSTER_START_SLEEP = 30
# Time, in seconds, to sleep between node starts.
NODE_START_SLEEP = 30


def CheckPrerequisites():
  """Verifies that the required resources are present.

  Raises:
    perfkitbenchmarker.data.ResourceNotFound: On missing resource.
  """
  for resource in (CASSANDRA_YAML_TEMPLATE, CASSANDRA_RACKDC_TEMPLATE):
    data.ResourcePath(resource)


def _Install(vm):
  """Installs Cassandra as a debian package.

  Args:
    vm: VirtualMachine. The VM to install Cassandra on.
  """
  vm.Install('openjdk')
  vm.Install('curl')
  vm.RemoteCommand(
      'curl -o /opt/pkb/cassandra.tar.gz'
      f' https://archive.apache.org/dist/cassandra/{CASSANDRA_VERSION.split("-")[-1]}/{CASSANDRA_VERSION}-bin.tar.gz'
  )
  vm.RemoteCommand('tar xzvf /opt/pkb/cassandra.tar.gz --directory /opt/pkb')


def GetCassandraVersion(vm) -> str:
  """Returns the Cassandra version installed on the VM."""
  stdout, _ = vm.RemoteCommand(f'{GetCassandraPath()} -v')
  return stdout


def YumInstall(vm):
  """Installs Cassandra on the VM."""
  _Install(vm)


def AptInstall(vm):
  """Installs Cassandra on the VM."""
  _Install(vm)


def JujuInstall(vm, vm_group_name):
  """Installs the Cassandra charm on the VM."""
  vm.JujuDeploy('cs:trusty/cassandra', vm_group_name)

  # The charm defaults to Cassandra 2.2.x, which has deprecated
  # cassandra-cli. Specify the sources to downgrade to Cassandra 2.1.x
  # to match the cassandra benchmark(s) expectations.
  sources = [
      'deb https://www.apache.org/dist/cassandra/debian 21x main',
      'ppa:openjdk-r/ppa',
      'ppa:stub/cassandra',
  ]

  keys = ['F758CE318D77295D', 'null', 'null']

  vm.JujuSet(
      'cassandra',
      [
          # Allow authentication from all units
          'authenticator=AllowAllAuthenticator',
          'install_sources="[%s]"'
          % ', '.join(["'" + x + "'" for x in sources]),
          'install_keys="[%s]"' % ', '.join(keys),
      ],
  )

  # Wait for cassandra to be installed and configured
  vm.JujuWait()

  for unit in vm.units:
    # Make sure the cassandra/conf dir is created, since we're skipping
    # the manual installation to /opt/pkb.
    remote_path = posixpath.join(CASSANDRA_DIR, 'conf')
    unit.RemoteCommand('mkdir -p %s' % remote_path)


def Configure(vm, seed_vms, custom_cassandra_conf=None):
  """Configure Cassandra on 'vm'.

  Args:
    vm: VirtualMachine. The VM to configure.
    seed_vms: List of VirtualMachine. The seed virtual machine(s).
    custom_cassandra_conf: Dict. Custom configurations for Cassandra.
  """
  if not custom_cassandra_conf:
    custom_cassandra_conf = {}
  context = {
      'ip_address': vm.internal_ip,
      'data_path': posixpath.join(vm.GetScratchDir(), 'cassandra'),
      'seeds': ','.join(f'{vm.internal_ip}:7000' for vm in seed_vms),
      'num_cpus': vm.NumCpusForBenchmark(),
      'cluster_name': 'Test cluster',
      'datacenter': 'datacenter',
      'rack': vm.name,
  }
  context.update(custom_cassandra_conf)
  logging.info('cassandra yaml context: %s', context)
  for template in [CASSANDRA_YAML_TEMPLATE, CASSANDRA_RACKDC_TEMPLATE]:
    local_path = data.ResourcePath(template)
    cassandra_conf_path = posixpath.join(CASSANDRA_DIR, 'conf')
    remote_path = posixpath.join(
        '~', os.path.splitext(os.path.basename(template))[0]
    )
    vm.RenderTemplate(local_path, remote_path, context=context)
    vm.RemoteCommand(f'sudo cp {remote_path} {cassandra_conf_path}')
  vm.RemoteCommand(f'mkdir {vm.GetScratchDir()}/cassandra')
  vm.RemoteCommand(f'sudo chmod -R 755 {vm.GetScratchDir()}/cassandra')
  _SetupPy311(vm)


def _SetupPy311(vm):
  """Setup Python 3.11 on a VM.

  CQL has known issues with Python 3.12. Using Python 3.11 for cassandra
  benchmark. CQL fails with error :
  "unsupported version of Python, required 3.6-3.11 but found 3.12"

  Args:
    vm: The target vm.
  """
  vm.RemoteCommand('sudo apt update')
  vm.RemoteCommand('sudo apt install software-properties-common -y')
  vm.RemoteCommand('sudo add-apt-repository ppa:deadsnakes/ppa -y')
  vm.RemoteCommand(
      'sudo apt install python3.11 python3.11-dev python3.11-venv -y'
  )
  vm.RemoteCommand('python3.11 -m venv ~/cqlsh_venv')
  vm.RemoteCommand('~/cqlsh_venv/bin/pip install cassandra-driver')


def Start(vm):
  """Start Cassandra on a VM.

  Args:
    vm: The target vm. Should already be configured via 'Configure'.
  """
  vm.RemoteCommand(f'{GetCassandraPath()}')


def CreateKeyspace(vm, replication_factor, is_row_cache_enabled):
  """Create a keyspace on a VM."""
  RunCql(vm, CASSANDRA_KEYSPACE_TEMPLATE, replication_factor)
  if is_row_cache_enabled:
    RunCql(vm, CASSANDRA_ROW_CACHE_TEMPLATE, replication_factor)


def RunCql(vm, template, replication_factor):
  """Run a CQL file on a VM."""
  cassandra_conf_path = posixpath.join(CASSANDRA_DIR, 'conf')
  template_path = data.ResourcePath(template)
  file_name = os.path.basename(os.path.splitext(template_path)[0])
  remote_path = os.path.join(
      '~',
      file_name,
  )
  vm.RenderTemplate(
      template_path,
      remote_path,
      context={
          'keyspace': 'keyspace1',
          'replication_factor': replication_factor,
      },
  )
  vm.RemoteCommand(f'sudo cp {remote_path} {cassandra_conf_path}')
  vm.RemoteCommand(
      'PYTHONPATH=~/cqlsh_venv/lib/python3.11/site-packages sudo -E'
      f' ~/cqlsh_venv/bin/python {CASSANDRA_DIR}/bin/cqlsh.py -f'
      f' {cassandra_conf_path}/{file_name}'
  )


def Stop(vm):
  """Stops Cassandra on 'vm'."""
  vm.RemoteCommand('kill $(cat {})'.format(CASSANDRA_PID), ignore_failure=True)


def IsRunning(vm):
  """Returns a boolean indicating whether Cassandra is running on 'vm'."""
  try:
    _, stderr = vm.RemoteCommand(f'{GetNodetoolPath()} status')
    if stderr:
      return False
    return True
  except errors.VirtualMachine.RemoteCommandError as ex:
    logging.warning('Exception: %s', ex)
    return False


def GetCompactionStats(vm):
  """Returns compaction stats for the given VM.

  Args:
    vm: VirtualMachine. The VM to get compaction stats from.

  Sample Output of compaction stats:
  pending tasks: 5
  - keyspace1.standard1: 5

  id  compaction type keyspace  table     completed  total      unit  progress
  c69  Compaction    keyspace1 standard1 437886277  20868111340 bytes 2.10%
  Active compaction remaining time :   0h05m33s
  """
  stdout, _ = vm.RemoteCommand(f'{GetNodetoolPath()} compactionstats')
  return stdout


def UpdateCassandraConcurrentCompactorCount(vm, concurrent_compactors_count):
  """Updates the cassandra concurrent compactor count."""
  vm.RemoteCommand(
      f'{GetNodetoolPath()} setconcurrentcompactors'
      f' {concurrent_compactors_count}'
  )


def UpdateCassandraCompactionThroughput(vm, compaction_throughput_mb_per_sec):
  """Updates the cassandra compaction throughput."""
  vm.RemoteCommand(
      f'{GetNodetoolPath()} setcompactionthroughput'
      f' {compaction_throughput_mb_per_sec}'
  )


def GetPendingTaskCountFromCompactionStats(cassandra_vms):
  """Parses the compaction stats for the given VMs and returns the pending task count.

  Args:
    cassandra_vms: List of VirtualMachine. The Cassandra VMs to get compaction
      stats from.

  Returns:
    List of int. The pending task count for each VM.
  """
  compaction_stats = background_tasks.RunThreaded(
      GetCompactionStats, cassandra_vms
  )
  pending_tasks = []
  for stats in compaction_stats:
    line = re.search(r'pending tasks:? *(\d*)', stats)
    if line:
      value = line.group(1)
      pending_tasks.append(int(value))
  return pending_tasks


def CleanNode(vm):
  """Remove Cassandra data from 'vm'.

  Args:
    vm: VirtualMachine. VM to clean.
  """
  data_path = posixpath.join(vm.GetScratchDir(), 'cassandra')
  vm.RemoteCommand('rm -rf {}'.format(data_path))


def _StartCassandraIfNotRunning(vm):
  """Starts Cassandra on 'vm' if not currently running."""
  if not IsRunning(vm):
    logging.info('Retrying starting cassandra on %s', vm)
    Start(vm)


def GetCassandraCliPath(_):
  return posixpath.join(CASSANDRA_DIR, 'bin', 'cassandra-cli')


def GetCassandraPath():
  return posixpath.join(CASSANDRA_DIR, 'bin', 'cassandra')


def GetNodetoolPath():
  return posixpath.join(CASSANDRA_DIR, 'bin', 'nodetool')


def GetCassandraStressPath(_):
  return posixpath.join(CASSANDRA_DIR, 'tools', 'bin', 'cassandra-stress')


def GetNumberOfNodesUp(vm):
  """Gets the number of VMs which are up in a Cassandra cluster.

  Args:
    vm: VirtualMachine. The VM to use to check the cluster status.

  Returns:
    int. The number of VMs which are up in a Cassandra cluster.
  """
  vms_up = vm.RemoteCommand(f'{GetNodetoolPath()} status | grep -c "^UN"')[
      0
  ].strip()
  return int(vms_up)


def GetNodetoolStatus(vms):
  for vm in vms:
    vm.RemoteCommand(f'{GetNodetoolPath()} status')


def StartCluster(seed_vm, vms):
  """Starts a Cassandra cluster.

  Starts a Cassandra cluster, first starting 'seed_vm', then remaining VMs in
  'vms'.

  Args:
    seed_vm: VirtualMachine. Machine which will function as the sole seed. It
      will be started before all other VMs.
    vms: list of VirtualMachines. VMs *other than* seed_vm which should be
      started.

  Raises:
    OSError: if cluster startup fails.
  """
  vm_count = len(vms) + 1

  # Cassandra setup
  logging.info('Starting seed VM %s', seed_vm)
  Start(seed_vm)
  logging.info('Waiting %ds for seed to start', NODE_START_SLEEP)
  time.sleep(NODE_START_SLEEP)
  for i in range(5):
    if not IsRunning(seed_vm):
      logging.warn(
          'Seed %s: Cassandra not running yet (try %d). Waiting %ds.',
          seed_vm,
          i,
          NODE_START_SLEEP,
      )
      time.sleep(NODE_START_SLEEP)
    else:
      break
  else:
    raise ValueError('Cassandra failed to start on seed.')

  if vms:
    logging.info('Starting remaining %d nodes', len(vms))
    # Start the VMs with a small pause in between each, to allow the node to
    # join.
    # Starting Cassandra nodes fails when multiple nodes attempt to join the
    # cluster concurrently.
    for i, vm in enumerate(vms):
      logging.info('Starting non-seed VM %d/%d.', i + 1, len(vms))
      Start(vm)
      time.sleep(NODE_START_SLEEP)
    logging.info('Waiting %ds for nodes to join', CLUSTER_START_SLEEP)
    time.sleep(CLUSTER_START_SLEEP)
  for i in range(CLUSTER_START_TRIES):
    vms_up = GetNumberOfNodesUp(seed_vm)
    if vms_up == vm_count:
      logging.info('All %d nodes up!', vm_count)
      break

    logging.warn(
        'Try %d: only %s of %s up. Restarting and sleeping %ds',
        i,
        vms_up,
        vm_count,
        NODE_START_SLEEP,
    )
    background_tasks.RunThreaded(_StartCassandraIfNotRunning, vms)
    time.sleep(NODE_START_SLEEP)
  else:
    raise OSError('Failed to start Cassandra cluster.')