path: root/buildscripts/resmokelib/testing/hooks/tenant_migration.py

"""Test hook that runs tenant migrations continuously."""

import random
import re
import threading
import time
import uuid

import bson
import pymongo.errors

from buildscripts.resmokelib import errors
from buildscripts.resmokelib import utils
from buildscripts.resmokelib.testing.fixtures import interface as fixture_interface
from buildscripts.resmokelib.testing.fixtures import tenant_migration
from buildscripts.resmokelib.testing.hooks import interface


class ContinuousTenantMigration(interface.Hook):  # pylint: disable=too-many-instance-attributes
    """Starts a tenant migration thread at the beginning of each test."""

    DESCRIPTION = ("Continuous tenant migrations")

    def __init__(self, hook_logger, fixture, shell_options):
        """Initialize the ContinuousTenantMigration.

        Args:
            hook_logger: the logger instance for this hook.
            fixture: the target TenantMigrationFixture containing two replica sets.
            shell_options: contains the global_vars which contains TestData.tenantId to be used for
                           tenant migrations.

        """
        interface.Hook.__init__(self, hook_logger, fixture, ContinuousTenantMigration.DESCRIPTION)

        if not isinstance(fixture, tenant_migration.TenantMigrationFixture):
            raise ValueError("The ContinuousTenantMigration hook requires a TenantMigrationFixture")
        self._tenant_migration_fixture = fixture
        self._tenant_id = shell_options["global_vars"]["TestData"]["tenantId"]

        self._tenant_migration_thread = None

    def before_suite(self, test_report):
        """Before suite."""
        if not self._tenant_migration_fixture:
            raise ValueError("No TenantMigrationFixture to run migrations on")
        self.logger.info("Starting the tenant migration thread.")
        self._tenant_migration_thread = _TenantMigrationThread(
            self.logger, self._tenant_migration_fixture, self._tenant_id)
        self._tenant_migration_thread.start()

    def after_suite(self, test_report):
        """After suite."""
        self.logger.info("Stopping the tenant migration thread.")
        self._tenant_migration_thread.stop()
        self.logger.info("Stopped the tenant migration thread.")

    def before_test(self, test, test_report):
        """Before test."""
        self.logger.info("Resuming the tenant migration thread.")
        self._tenant_migration_thread.resume()

    def after_test(self, test, test_report):
        """After test."""
        self.logger.info("Pausing the tenant migration thread.")
        self._tenant_migration_thread.pause()
        self.logger.info("Paused the tenant migration thread.")


class TenantMigrationLifeCycle(object):
    """Class for managing the various states of the tenant migration thread.

    The job thread alternates between calling mark_test_started() and mark_test_finished(). The
    tenant migration thread is allowed to perform migrations at any point between these two calls.
    Note that the job thread synchronizes with the tenant migration thread outside the context of
    this object to know it isn't in the process of running a migration.
    """

    _TEST_STARTED_STATE = "start"
    _TEST_FINISHED_STATE = "finished"

    def __init__(self):
        """Initialize the MigrationLifecycle instance."""
        self.__lock = threading.Lock()
        self.__cond = threading.Condition(self.__lock)

        self.test_num = 0
        self.__test_state = self._TEST_FINISHED_STATE
        self.__should_stop = False

    def mark_test_started(self):
        """Signal to the tenant migration thread that a new test has started.

        This function should be called during before_test(). Calling it causes the
        wait_for_tenant_migration_permitted() function to no longer block and to instead return
        true.
        """
        with self.__lock:
            self.test_num += 1
            self.__test_state = self._TEST_STARTED_STATE
            self.__cond.notify_all()

    def mark_test_finished(self):
        """Signal to the tenant migration thread that the current test has finished.

        This function should be called during after_test(). Calling it causes the
        wait_for_tenant_migration_permitted() function to block until mark_test_started() is called
        again.
        """
        with self.__lock:
            self.__test_state = self._TEST_FINISHED_STATE
            self.__cond.notify_all()

    def is_test_finished(self):
        """Return true if the current test has finished."""
        with self.__lock:
            return self.__test_state == self._TEST_FINISHED_STATE

    def stop(self):
        """Signal to the tenant migration thread that it should exit.

        This function should be called during after_suite(). Calling it causes the
        wait_for_tenant_migration_permitted() function to no longer block and to instead return
        false.
        """
        with self.__lock:
            self.__should_stop = True
            self.__cond.notify_all()

    def wait_for_tenant_migration_permitted(self):
        """Block until migrations are permitted, or until stop() is called.

        Return true if migrations are permitted, and false if migrations are not permitted.
        """
        with self.__lock:
            while not self.__should_stop:
                if self.__test_state == self._TEST_STARTED_STATE:
                    return True

                self.__cond.wait()

        return False

    def wait_for_tenant_migration_interval(self, timeout):
        """Block for 'timeout' seconds, or until stop() is called."""
        with self.__lock:
            self.__cond.wait(timeout)

    def poll_for_idle_request(self):  # noqa: D205,D400
        """Return true if the tenant migration thread should continue running migrations, or false
        if it should temporarily stop running migrations.
        """
        with self.__lock:
            return self.__test_state == self._TEST_FINISHED_STATE


def get_certificate_and_private_key(pem_file_path):  # noqa: D205,D400
    """Return a dictionary containing the certificate and private key extracted from the given pem
    file.
    """
    lines = open(pem_file_path, 'rt').read()
    certificate = re.findall(
        re.compile("(-*BEGIN CERTIFICATE-*\n(.*\n)*-*END CERTIFICATE-*\n)", re.MULTILINE),
        lines)[0][0]
    private_key = re.findall(
        re.compile("(-*BEGIN PRIVATE KEY-*\n(.*\n)*-*END PRIVATE KEY-*\n)", re.MULTILINE),
        lines)[0][0]
    return {"certificate": certificate, "privateKey": private_key}


class _TenantMigrationOptions:
    def __init__(self, donor_rs, recipient_rs, tenant_id, read_preference):
        self.donor_rs = donor_rs
        self.recipient_rs = recipient_rs
        self.migration_id = uuid.uuid4()
        self.tenant_id = tenant_id
        self.read_preference = read_preference

    def get_donor_name(self):
        """Return the replica set name for the donor."""
        return self.donor_rs.replset_name

    def get_recipient_name(self):
        """Return the replica set name for the recipient."""
        return self.recipient_rs.replset_name

    def get_donor_primary(self):
        """Return a connection to the donor primary."""
        return self.donor_rs.get_primary()

    def get_recipient_primary(self):
        """Return a connection to the recipient primary."""
        return self.recipient_rs.get_primary()

    def __str__(self):
        opts = {
            "donor": self.get_donor_name(), "recipient": self.get_recipient_name(),
            "migration_id": self.migration_id, "tenant_id": self.tenant_id,
            "read_preference": self.read_preference
        }
        return str(opts)


class _TenantMigrationThread(threading.Thread):  # pylint: disable=too-many-instance-attributes
    WAIT_SECS_RANGES = [[0.1, 0.5], [1, 5], [5, 15]]
    MIGRATION_STATE_POLL_INTERVAL_SECS = 0.1
    NO_SUCH_MIGRATION_ERR_CODE = 327
    INTERNAL_ERR_CODE = 1
    ILLEGAL_OPERATION_ERR_CODE = 20
    INVALID_NS_ERR_CODE = 73

    def __init__(self, logger, tenant_migration_fixture, tenant_id):
        """Initialize _TenantMigrationThread."""
        threading.Thread.__init__(self, name="TenantMigrationThread")
        self.daemon = True
        self.logger = logger
        self._tenant_migration_fixture = tenant_migration_fixture
        self._tenant_id = tenant_id

        self.__lifecycle = TenantMigrationLifeCycle()
        # Event set when the thread has been stopped using the 'stop()' method.
        self._is_stopped_evt = threading.Event()
        # Event set when the thread is not performing migrations.
        self._is_idle_evt = threading.Event()
        self._is_idle_evt.set()

    def run(self):
        """Execute the thread."""
        if not self._tenant_migration_fixture:
            self.logger.warning("No TenantMigrationFixture to run migrations on.")
            return

        test_num = 0
        migration_num = 0
        donor_rs_index = 0

        try:
            while True:
                self._is_idle_evt.set()

                permitted = self.__lifecycle.wait_for_tenant_migration_permitted()
                if not permitted:
                    break

                self._is_idle_evt.clear()

                if self.__lifecycle.test_num > test_num:
                    # Reset donor_rs_index to 0 since the shell always routes all requests to rs0
                    # at the start of a test.
                    test_num = self.__lifecycle.test_num
                    donor_rs_index = 0

                recipient_rs_index = (
                    donor_rs_index + 1) % self._tenant_migration_fixture.get_num_replsets()
                migration_opts = self._create_migration_opts(donor_rs_index, recipient_rs_index)

                self.logger.info("Starting tenant migration: %s.", str(migration_opts))
                start_time = time.time()
                is_committed = self._run_migration(migration_opts)
                end_time = time.time()
                self.logger.info("Completed tenant migration in %0d ms: %s.",
                                 (end_time - start_time) * 1000, str(migration_opts))

                found_idle_request = self.__lifecycle.poll_for_idle_request()
                if found_idle_request:
                    continue

                wait_secs = random.uniform(
                    *self.WAIT_SECS_RANGES[migration_num % len(self.WAIT_SECS_RANGES)])
                self.__lifecycle.wait_for_tenant_migration_interval(wait_secs)

                if is_committed:
                    donor_rs_index = recipient_rs_index
                migration_num += 1
        except Exception:  # pylint: disable=W0703
            # Proactively log the exception when it happens so it will be flushed immediately.
            self.logger.exception("Tenant migration thread threw exception")
            # The event should be signaled whenever the thread is not performing migrations.
            self._is_idle_evt.set()

    def stop(self):
        """Stop the thread when the suite finishes."""
        self.__lifecycle.stop()
        self._is_stopped_evt.set()
        # Unpause to allow the thread to finish.
        self.resume()
        self.join()

    def pause(self):
        """Pause the thread after test."""
        self.__lifecycle.mark_test_finished()

        # Wait until we are no longer executing migrations.
        self._is_idle_evt.wait()
        # Check if the thread is alive in case it has thrown an exception while running.
        self._check_thread()

        # Check that the fixture is still running.
        if not self._tenant_migration_fixture.is_running():
            raise errors.ServerFailure(
                "TenantMigrationFixture with pids {} expected to be running in"
                " ContinuousTenantMigration, but wasn't".format(
                    self._tenant_migration_fixture.pids()))

    def resume(self):
        """Resume the thread before test."""
        self.__lifecycle.mark_test_started()

    def _wait(self, timeout):
        """Wait until stop or timeout."""
        self._is_stopped_evt.wait(timeout)

    def _check_thread(self):
        """Throw an error if the thread is not running."""
        if not self.is_alive():
            msg = "Tenant migration thread is not running."
            self.logger.error(msg)
            raise errors.ServerFailure(msg)

    def _is_fail_point_abort_reason(self, abort_reason):
        return abort_reason["code"] == self.INTERNAL_ERR_CODE and abort_reason[
            "errmsg"] == "simulate a tenant migration error"

    def _create_migration_opts(self, donor_rs_index, recipient_rs_index):
        donor_rs = self._tenant_migration_fixture.get_replset(donor_rs_index)
        recipient_rs = self._tenant_migration_fixture.get_replset(recipient_rs_index)
        read_preference = {"mode": "primary"} if random.randint(0, 1) else {"mode": "secondary"}
        return _TenantMigrationOptions(donor_rs, recipient_rs, self._tenant_id, read_preference)

    def _run_migration(self, migration_opts):  # noqa: D205,D400
        """Run donorStartMigration to start a tenant migration based on 'migration_opts', wait for
        the migration decision. Returns true if the migration commits and false otherwise.
        """
        donor_primary = migration_opts.get_donor_primary()
        donor_primary_client = donor_primary.mongo_client()

        self.logger.info(
            "Starting tenant migration with donor primary on port %d of replica set '%s'.",
            donor_primary.port, migration_opts.get_donor_name())

        cmd_obj = {
            "donorStartMigration":
                1,
            "migrationId":
                bson.Binary(migration_opts.migration_id.bytes, 4),
            "recipientConnectionString":
                migration_opts.recipient_rs.get_driver_connection_url(),
            "tenantId":
                migration_opts.tenant_id,
            "readPreference":
                migration_opts.read_preference,
            "donorCertificateForRecipient":
                get_certificate_and_private_key("jstests/libs/tenant_migration_donor.pem"),
            "recipientCertificateForDonor":
                get_certificate_and_private_key("jstests/libs/tenant_migration_recipient.pem"),
        }
        is_committed = False

        try:
            # Clean up any orphaned tenant databases on the recipient allow next migration to start.
            self._drop_tenant_databases(migration_opts.recipient_rs)

            while True:
                # Keep polling the migration state until the migration completes.
                res = donor_primary_client.admin.command(
                    cmd_obj,
                    bson.codec_options.CodecOptions(uuid_representation=bson.binary.UUID_SUBTYPE))

                if res["state"] == "committed":
                    self.logger.info("Tenant migration with donor primary on port " +
                                     str(donor_primary.port) + " of replica set '" +
                                     migration_opts.get_donor_name() + "' has committed.")
                    is_committed = True
                    break
                elif res["state"] == "aborted":
                    abort_reason = res["abortReason"]
                    if self._is_fail_point_abort_reason(abort_reason):
                        self.logger.info("Tenant migration with donor primary on port " +
                                         str(donor_primary.port) + " of replica set '" +
                                         migration_opts.get_donor_name() +
                                         "' has aborted due to failpoint: " + str(res))
                        break
                    else:
                        raise errors.ServerFailure("Tenant migration with donor primary on port " +
                                                   str(donor_primary.port) + " of replica set '" +
                                                   migration_opts.get_donor_name() +
                                                   "' has aborted due to an error: " + str(res))
                elif not res["ok"]:
                    raise errors.ServerFailure("Tenant migration with donor primary on port " +
                                               str(donor_primary.port) + " of replica set '" +
                                               migration_opts.get_donor_name() + "' has failed: " +
                                               str(res))

                time.sleep(self.MIGRATION_STATE_POLL_INTERVAL_SECS)

            # Garbage collect the migration.
            if is_committed:
                # If the migration committed, to avoid routing commands incorrectly, wait for the
                # donor/proxy to reroute at least one command before doing garbage collection. Stop
                # waiting when the test finishes.
                self._wait_for_reroute_or_test_completion(migration_opts)
            self._forget_migration(migration_opts)
            self._wait_for_migration_garbage_collection(migration_opts)

            return is_committed
        except pymongo.errors.PyMongoError:
            self.logger.exception(
                "Error running tenant migration with donor primary on port %d of replica set '%s'.",
                donor_primary.port, migration_opts.get_donor_name())
            raise

    def _forget_migration(self, migration_opts):
        """Run donorForgetMigration to garbage collection the tenant migration denoted by migration_opts'."""
        self.logger.info("Forgetting tenant migration: %s.", str(migration_opts))

        donor_primary = migration_opts.get_donor_primary()
        donor_primary_client = donor_primary.mongo_client()

        try:
            donor_primary_client.admin.command({
                "donorForgetMigration": 1, "migrationId": bson.Binary(
                    migration_opts.migration_id.bytes, 4)
            }, bson.codec_options.CodecOptions(uuid_representation=bson.binary.UUID_SUBTYPE))
        except pymongo.errors.OperationFailure as err:
            if err.code != self.NO_SUCH_MIGRATION_ERR_CODE:
                raise
            # The fixture was restarted.
            self.logger.info(
                "Could not find tenant migration '%s' on donor primary on" +
                " port %d of replica set '%s'.", migration_opts.migration_id, donor_primary.port,
                migration_opts.get_donor_name())
        except pymongo.errors.PyMongoError:
            self.logger.exception(
                "Error forgetting tenant migration '%s' on donor primary on" +
                " port %d of replica set '%s'.", migration_opts.migration_id, donor_primary.port,
                migration_opts.get_donor_name())
            raise

    def _wait_for_migration_garbage_collection(self, migration_opts):  # noqa: D205,D400
        """Wait until the persisted state for migration denoted by 'migration_opts' has been
        garbage collected on both the donor and recipient.
        """
        donor_primary = migration_opts.get_donor_primary()
        donor_primary_client = donor_primary.mongo_client()

        recipient_primary = migration_opts.get_recipient_primary()
        recipient_primary_client = recipient_primary.mongo_client()

        try:
            while True:
                res = donor_primary_client.config.command({
                    "count": "tenantMigrationDonors",
                    "query": {"tenantId": migration_opts.tenant_id}
                })
                if res["n"] == 0:
                    break
            while True:
                res = recipient_primary_client.config.command({
                    "count": "tenantMigrationRecipients",
                    "query": {"tenantId": migration_opts.tenant_id}
                })
                if res["n"] == 0:
                    break
        except pymongo.errors.PyMongoError:
            self.logger.exception(
                "Error waiting for tenant migration for tenant id '%s' from donor replica set '%s" +
                " to recipient replica set '%s' to be garbage collection.", self._tenant_id,
                migration_opts.get_donor_name(), migration_opts.get_recipient_name())
            raise

    def _wait_for_reroute_or_test_completion(self, migration_opts):
        self.logger.info(
            "Waiting for donor replica set '%s' for migration '%s' to reroute at least one " +
            "conflicting command. Stop waiting when the test finishes.",
            migration_opts.get_donor_name(), migration_opts.migration_id)

        start_time = time.time()

        while not self.__lifecycle.is_test_finished():
            donor_primary = migration_opts.get_donor_primary()
            donor_primary_client = donor_primary.mongo_client()

            try:
                doc = donor_primary_client["testTenantMigration"]["rerouted"].find_one(
                    {"_id": bson.Binary(migration_opts.migration_id.bytes, 4)})
                if doc is not None:
                    return
            except pymongo.errors.PyMongoError:
                end_time = time.time()
                self.logger.exception(
                    "Error running find command on donor primary on port %d of replica set '%s' " +
                    "after waiting for reroute for %0d ms", donor_primary.port,
                    migration_opts.get_donor_name(), (end_time - start_time) * 1000)
                raise

            time.sleep(self.MIGRATION_STATE_POLL_INTERVAL_SECS)

    def _drop_tenant_databases(self, rs):
        self.logger.info("Dropping tenant databases from replica set '%s'.", rs.replset_name)
        primary = rs.get_primary()
        primary_client = primary.mongo_client()

        try:
            res = primary_client.admin.command({"listDatabases": 1})
            for database in res["databases"]:
                db_name = database["name"]
                if db_name.startswith(self._tenant_id + "_"):
                    primary_client.drop_database(db_name)
        except pymongo.errors.PyMongoError:
            self.logger.exception(
                "Error dropping databases for tenant id '%s' on primary on" +
                " port %d of replica set '%s' to be garbage collection.", self._tenant_id,
                primary.port, rs.replset_name)
            raise