perfkitbenchmarker/providers/gcp/gcp_cloud_valkey.py

# Copyright 2024 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.
"""Module containing class for GCP's Cloud Valkey instances."""
import json
import logging
import os
from typing import Any

from absl import flags
from perfkitbenchmarker import errors
from perfkitbenchmarker import managed_memory_store
from perfkitbenchmarker import provider_info
from perfkitbenchmarker import vm_util
from perfkitbenchmarker.providers.gcp import flags as gcp_flags
from perfkitbenchmarker.providers.gcp import util

FLAGS = flags.FLAGS

COMMAND_TIMEOUT = 3600

_DEFAULT_VERSION = 'VALKEY_8_0'
_DEFAULT_ZONE = 'us-central1-c'


class CloudValkey(managed_memory_store.BaseManagedMemoryStore):
  """Object representing a GCP Cloud Memorystore (Valkey) instance."""

  CLOUD = provider_info.GCP
  SERVICE_TYPE = 'memorystore'
  MEMORY_STORE = managed_memory_store.VALKEY

  def __init__(self, spec):
    super().__init__(spec)
    self.project = FLAGS.project
    self.node_type = ''
    self.location = FLAGS.cloud_redis_region
    self.zone_distribution = gcp_flags.REDIS_ZONE_DISTRIBUTION.value
    self.node_type = gcp_flags.REDIS_NODE_TYPE.value
    if self.zone_distribution == 'single-zone':
      zone = FLAGS.zone[0] if FLAGS.zone else _DEFAULT_ZONE
      self.zones = [zone]
    self.version = spec.version or _DEFAULT_VERSION
    self.network = (
        'default'
        if not gcp_flags.GCE_NETWORK_NAMES.value
        else gcp_flags.GCE_NETWORK_NAMES.value[0]
    )
    self.subnet = (
        'default'
        if not gcp_flags.GCE_SUBNET_NAMES.value
        else gcp_flags.GCE_SUBNET_NAMES.value[0]
    )
    # Update the environment for gcloud commands:
    os.environ['CLOUDSDK_API_ENDPOINT_OVERRIDES_MEMORYSTORE'] = (
        gcp_flags.CLOUD_VALKEY_API_OVERRIDE.value
    )

  def GetResourceMetadata(self) -> dict[str, Any]:
    """Returns a dict containing metadata about the instance.

    Returns:
      dict mapping string property key to value.
    """
    # Metadata uses redis terminology for backwards compatibility.
    self.metadata.update({
        'cloud_redis_region': self.location,
        'cloud_redis_version': self.GetReadableVersion(),
        'cloud_redis_node_type': self.node_type,
        'cloud_redis_zone_distribution': self.zone_distribution,
    })
    return self.metadata

  def GetReadableVersion(self):
    """Parses Valkey major and minor version number."""
    if self.version.count('_') < 2:
      logging.info(
          (
              'Could not parse version string correctly, '
              'full Valkey version returned: %s'
          ),
          self.version,
      )
      return self.version
    return '.'.join(self.version.split('_')[1:])

  def _CreateServiceConnectionPolicy(self) -> None:
    """Creates a service connection policy for the VPC."""
    cmd = util.GcloudCommand(
        self,
        'network-connectivity',
        'service-connection-policies',
        'create',
        f'pkb-{FLAGS.run_uri}-policy',
    )
    cmd.flags['service-class'] = 'gcp-memorystore'
    cmd.flags['network'] = self.network
    cmd.flags['region'] = self.location
    cmd.flags['subnets'] = (
        'https://www.googleapis.com/compute/v1'
        f'/projects/{self.project}'
        f'/regions/{self.location}/subnetworks/{self.subnet}'
    )
    cmd.Issue(raise_on_failure=False)

  def _GetClusterCreateCommand(self) -> util.GcloudCommand:
    """Returns the command used to create the cluster."""
    cmd = util.GcloudCommand(
        self, 'memorystore', 'instances', 'create', self.name
    )
    cmd.flags['location'] = self.location

    network_name = f'projects/{self.project}/global/networks/{self.network}'
    cmd.flags['psc-auto-connections'] = (
        f'network={network_name},projectId={self.project}'
    )

    cmd.flags['shard-count'] = self.shard_count
    cmd.flags['replica-count'] = self.replicas_per_shard
    cmd.flags['node-type'] = self.node_type
    cmd.flags['engine-version'] = self.version

    cmd.flags['zone-distribution-config-mode'] = self.zone_distribution
    if self.zone_distribution == 'single-zone':
      cmd.flags['zone-distribution-config'] = self.zones[0]

    if self.enable_tls:
      cmd.flags['transit-encryption-mode'] = 'server-authentication'

    cmd.flags['labels'] = util.MakeFormattedDefaultTags()
    cmd.use_beta_gcloud = True

    return cmd

  def _Create(self):
    """Creates the instance."""
    self._CreateServiceConnectionPolicy()
    cmd = self._GetClusterCreateCommand()
    cmd.Issue(timeout=COMMAND_TIMEOUT)

  def _IsReady(self):
    """Returns whether cluster is ready."""
    instance_details, _, _ = self.DescribeInstance()
    return json.loads(instance_details).get('state') == 'ACTIVE'

  def _Delete(self):
    """Deletes the instance."""
    cmd = util.GcloudCommand(
        self, 'memorystore', 'instances', 'delete', self.name
    )
    cmd.flags['location'] = self.location
    cmd.use_beta_gcloud = True
    cmd.Issue(timeout=COMMAND_TIMEOUT, raise_on_failure=False)

  def _Exists(self):
    """Returns true if the instance exists."""
    _, _, retcode = self.DescribeInstance()
    return retcode == 0

  def DescribeInstance(self):
    """Calls describe instance using the gcloud tool.

    Returns:
      stdout, stderr, and retcode.
    """
    cmd = util.GcloudCommand(
        self, 'memorystore', 'instances', 'describe', self.name
    )
    cmd.flags['location'] = self.location
    cmd.use_beta_gcloud = True
    stdout, stderr, retcode = cmd.Issue(raise_on_failure=False)
    if retcode != 0:
      logging.info('Could not find memorystore instance %s', self.name)
    return stdout, stderr, retcode

  @vm_util.Retry(max_retries=5)
  def _PopulateEndpoint(self):
    """Populates endpoint information about the instance.

    Raises:
      errors.Resource.RetryableGetError:
      Failed to retrieve information on instance
    """
    stdout, _, retcode = self.DescribeInstance()
    if retcode != 0:
      raise errors.Resource.RetryableGetError(
          'Failed to retrieve information on {}'.format(self.name)
      )
    if self.clustered:
      self._ip = json.loads(stdout)['discoveryEndpoints'][0]['address']
      self._port = 6379
      return
    self._ip = json.loads(stdout)['host']
    self._port = json.loads(stdout)['port']