• Home
• software.amazon.kinesis
• amazon-kinesis-client
• 3.0.2
• source code
• DynamicMigrationComponentsInitializer.java

×

Project download

Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $

You can buy this project and download/modify it how often you want.

software.amazon.kinesis.coordinator.DynamicMigrationComponentsInitializer Maven / Gradle / Ivy

/*
 * Copyright 2024 Amazon.com, Inc. or its affiliates.
 * 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.
 */
package software.amazon.kinesis.coordinator;

import java.time.Duration;
import java.util.Objects;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.BiFunction;
import java.util.function.Supplier;

import lombok.AccessLevel;
import lombok.Builder;
import lombok.Getter;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import software.amazon.awssdk.annotations.ThreadSafe;
import software.amazon.kinesis.annotations.KinesisClientInternalApi;
import software.amazon.kinesis.coordinator.MigrationAdaptiveLeaseAssignmentModeProvider.LeaseAssignmentMode;
import software.amazon.kinesis.coordinator.assignment.LeaseAssignmentManager;
import software.amazon.kinesis.coordinator.migration.ClientVersion;
import software.amazon.kinesis.leader.DynamoDBLockBasedLeaderDecider;
import software.amazon.kinesis.leader.MigrationAdaptiveLeaderDecider;
import software.amazon.kinesis.leases.LeaseManagementConfig.WorkerUtilizationAwareAssignmentConfig;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.metrics.MetricsFactory;
import software.amazon.kinesis.worker.metricstats.WorkerMetricStatsDAO;
import software.amazon.kinesis.worker.metricstats.WorkerMetricStatsManager;
import software.amazon.kinesis.worker.metricstats.WorkerMetricStatsReporter;

import static software.amazon.kinesis.coordinator.MigrationAdaptiveLeaseAssignmentModeProvider.LeaseAssignmentMode.DEFAULT_LEASE_COUNT_BASED_ASSIGNMENT;
import static software.amazon.kinesis.coordinator.MigrationAdaptiveLeaseAssignmentModeProvider.LeaseAssignmentMode.WORKER_UTILIZATION_AWARE_ASSIGNMENT;
import static software.amazon.kinesis.coordinator.assignment.LeaseAssignmentManager.DEFAULT_NO_OF_SKIP_STAT_FOR_DEAD_WORKER_THRESHOLD;

/**
 * This class is responsible for initializing the KCL components that supports
 * seamless upgrade from v2.x to v3.x.
 * During specific versions, it also dynamically switches the functionality
 * to be either vanilla 3.x or 2.x compatible.
 *
 * It is responsible for creating:
 * 1. LeaderDecider
 * 2. LAM
 * 3. WorkerMetricStatsReporter
 *
 * It manages initializing the following components at initialization time
 * 1. workerMetricsDAO and workerMetricsManager
 * 2. leaderDecider
 * 3. MigrationAdaptiveLeaseAssignmentModeProvider
 *
 * It updates the following components dynamically:
 * 1. starts/stops LAM
 * 2. starts/stops WorkerMetricStatsReporter
 * 3. updates LeaseAssignmentMode to either DEFAULT_LEASE_COUNT_BASED_ASSIGNMENT or WORKER_UTILIZATION_AWARE_ASSIGNMENT
 * 4. creates GSI (deletion is done by KclMigrationTool)
 * 5. creates WorkerMetricStats table (deletion is done by KclMigrationTool)
 * 6. updates LeaderDecider to either DeterministicShuffleShardSyncLeaderDecider or DynamoDBLockBasedLeaderDecider
 */
@Slf4j
@KinesisClientInternalApi
@ThreadSafe
@Accessors(fluent = true)
public final class DynamicMigrationComponentsInitializer {
    private static final long SCHEDULER_SHUTDOWN_TIMEOUT_SECONDS = 60L;

    @Getter
    private final MetricsFactory metricsFactory;

    @Getter
    private final LeaseRefresher leaseRefresher;

    private final CoordinatorStateDAO coordinatorStateDAO;
    private final ScheduledExecutorService workerMetricsThreadPool;

    @Getter
    private final WorkerMetricStatsDAO workerMetricsDAO;

    private final WorkerMetricStatsManager workerMetricsManager;
    private final ScheduledExecutorService lamThreadPool;
    private final BiFunction lamCreator;
    private final Supplier adaptiveLeaderDeciderCreator;
    private final Supplier deterministicLeaderDeciderCreator;
    private final Supplier ddbLockBasedLeaderDeciderCreator;

    @Getter
    private final String workerIdentifier;

    private final WorkerUtilizationAwareAssignmentConfig workerUtilizationAwareAssignmentConfig;

    @Getter
    private final long workerMetricsExpirySeconds;

    private final MigrationAdaptiveLeaseAssignmentModeProvider leaseModeChangeConsumer;

    @Getter
    private LeaderDecider leaderDecider;

    private LeaseAssignmentManager leaseAssignmentManager;
    private ScheduledFuture workerMetricsReporterFuture;
    private LeaseAssignmentMode currentAssignmentMode;
    private boolean dualMode;
    private boolean initialized;

    @Builder(access = AccessLevel.PACKAGE)
    DynamicMigrationComponentsInitializer(
            final MetricsFactory metricsFactory,
            final LeaseRefresher leaseRefresher,
            final CoordinatorStateDAO coordinatorStateDAO,
            final ScheduledExecutorService workerMetricsThreadPool,
            final WorkerMetricStatsDAO workerMetricsDAO,
            final WorkerMetricStatsManager workerMetricsManager,
            final ScheduledExecutorService lamThreadPool,
            final BiFunction lamCreator,
            final Supplier adaptiveLeaderDeciderCreator,
            final Supplier deterministicLeaderDeciderCreator,
            final Supplier ddbLockBasedLeaderDeciderCreator,
            final String workerIdentifier,
            final WorkerUtilizationAwareAssignmentConfig workerUtilizationAwareAssignmentConfig,
            final MigrationAdaptiveLeaseAssignmentModeProvider leaseAssignmentModeProvider) {
        this.metricsFactory = metricsFactory;
        this.leaseRefresher = leaseRefresher;
        this.coordinatorStateDAO = coordinatorStateDAO;
        this.workerIdentifier = workerIdentifier;
        this.workerUtilizationAwareAssignmentConfig = workerUtilizationAwareAssignmentConfig;
        this.workerMetricsExpirySeconds = Duration.ofMillis(DEFAULT_NO_OF_SKIP_STAT_FOR_DEAD_WORKER_THRESHOLD
                        * workerUtilizationAwareAssignmentConfig.workerMetricsReporterFreqInMillis())
                .getSeconds();
        this.workerMetricsManager = workerMetricsManager;
        this.workerMetricsDAO = workerMetricsDAO;
        this.workerMetricsThreadPool = workerMetricsThreadPool;
        this.lamThreadPool = lamThreadPool;
        this.lamCreator = lamCreator;
        this.adaptiveLeaderDeciderCreator = adaptiveLeaderDeciderCreator;
        this.deterministicLeaderDeciderCreator = deterministicLeaderDeciderCreator;
        this.ddbLockBasedLeaderDeciderCreator = ddbLockBasedLeaderDeciderCreator;
        this.leaseModeChangeConsumer = leaseAssignmentModeProvider;
    }

    public void initialize(final ClientVersion migrationStateMachineStartingClientVersion) throws DependencyException {
        if (initialized) {
            log.info("Already initialized, nothing to do");
            return;
        }

        // always collect metrics so that when we flip to start reporting we will have accurate historical data.
        log.info("Start collection of WorkerMetricStats");
        workerMetricsManager.startManager();
        if (migrationStateMachineStartingClientVersion == ClientVersion.CLIENT_VERSION_3X) {
            initializeComponentsFor3x();
        } else {
            initializeComponentsForMigration(migrationStateMachineStartingClientVersion);
        }
        log.info("Initialized dual mode {} current assignment mode {}", dualMode, currentAssignmentMode);

        log.info("Creating LAM");
        leaseAssignmentManager = lamCreator.apply(lamThreadPool, leaderDecider);
        log.info("Initializing {}", leaseModeChangeConsumer.getClass().getSimpleName());
        leaseModeChangeConsumer.initialize(dualMode, currentAssignmentMode);
        initialized = true;
    }

    private void initializeComponentsFor3x() {
        log.info("Initializing for 3x functionality");
        dualMode = false;
        currentAssignmentMode = WORKER_UTILIZATION_AWARE_ASSIGNMENT;
        log.info("Initializing dualMode {} assignmentMode {}", dualMode, currentAssignmentMode);
        leaderDecider = ddbLockBasedLeaderDeciderCreator.get();
        log.info("Initializing {}", leaderDecider.getClass().getSimpleName());
        leaderDecider.initialize();
    }

    private void initializeComponentsForMigration(final ClientVersion migrationStateMachineStartingClientVersion) {
        log.info("Initializing for migration to 3x");
        dualMode = true;
        final LeaderDecider initialLeaderDecider;
        if (migrationStateMachineStartingClientVersion == ClientVersion.CLIENT_VERSION_3X_WITH_ROLLBACK) {
            currentAssignmentMode = WORKER_UTILIZATION_AWARE_ASSIGNMENT;
            initialLeaderDecider = ddbLockBasedLeaderDeciderCreator.get();
        } else {
            currentAssignmentMode = DEFAULT_LEASE_COUNT_BASED_ASSIGNMENT;
            initialLeaderDecider = deterministicLeaderDeciderCreator.get();
        }
        log.info("Initializing dualMode {} assignmentMode {}", dualMode, currentAssignmentMode);

        final MigrationAdaptiveLeaderDecider adaptiveLeaderDecider = adaptiveLeaderDeciderCreator.get();
        log.info(
                "Initializing MigrationAdaptiveLeaderDecider with {}",
                initialLeaderDecider.getClass().getSimpleName());
        adaptiveLeaderDecider.updateLeaderDecider(initialLeaderDecider);
        this.leaderDecider = adaptiveLeaderDecider;
    }

    void shutdown() {
        log.info("Shutting down components");
        if (initialized) {
            log.info("Stopping LAM, LeaderDecider, workerMetrics reporting and collection");
            leaseAssignmentManager.stop();
            // leader decider is shut down later when scheduler is doing a final shutdown
            // since scheduler still accesses the leader decider while shutting down
            stopWorkerMetricsReporter();
            workerMetricsManager.stopManager();
        }

        // lam does not manage lifecycle of its threadpool to easily stop/start dynamically.
        // once migration code is obsolete (i.e. all 3x functionality is the baseline and no
        // migration is needed), it can be moved inside lam
        log.info("Shutting down lamThreadPool and workerMetrics reporter thread pool");
        lamThreadPool.shutdown();
        workerMetricsThreadPool.shutdown();
        try {
            if (!lamThreadPool.awaitTermination(SCHEDULER_SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
                log.info(
                        "LamThreadPool did not shutdown in {}s, forcefully shutting down",
                        SCHEDULER_SHUTDOWN_TIMEOUT_SECONDS);
                lamThreadPool.shutdownNow();
            }
        } catch (final InterruptedException e) {
            log.warn("Interrupted while waiting for shutdown of LeaseAssignmentManager ThreadPool", e);
            lamThreadPool.shutdownNow();
        }

        try {
            if (!workerMetricsThreadPool.awaitTermination(SCHEDULER_SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
                log.info(
                        "WorkerMetricsThreadPool did not shutdown in {}s, forcefully shutting down",
                        SCHEDULER_SHUTDOWN_TIMEOUT_SECONDS);
                workerMetricsThreadPool.shutdownNow();
            }
        } catch (final InterruptedException e) {
            Thread.currentThread().interrupt();
            log.warn("Interrupted while waiting for shutdown of WorkerMetricStatsManager ThreadPool", e);
            workerMetricsThreadPool.shutdownNow();
        }
    }

    private void startWorkerMetricsReporting() throws DependencyException {
        if (workerMetricsReporterFuture != null) {
            log.info("Worker metrics reporting is already running...");
            return;
        }
        log.info("Initializing WorkerMetricStats");
        this.workerMetricsDAO.initialize();
        log.info("Starting worker metrics reporter");
        // Start with a delay for workerStatsManager to capture some values and start reporting.
        workerMetricsReporterFuture = workerMetricsThreadPool.scheduleAtFixedRate(
                new WorkerMetricStatsReporter(metricsFactory, workerIdentifier, workerMetricsManager, workerMetricsDAO),
                workerUtilizationAwareAssignmentConfig.inMemoryWorkerMetricsCaptureFrequencyMillis() * 2L,
                workerUtilizationAwareAssignmentConfig.workerMetricsReporterFreqInMillis(),
                TimeUnit.MILLISECONDS);
    }

    private void stopWorkerMetricsReporter() {
        log.info("Stopping worker metrics reporter");
        if (workerMetricsReporterFuture != null) {
            workerMetricsReporterFuture.cancel(false);
            workerMetricsReporterFuture = null;
        }
    }

    /**
     * Create LeaseOwnerToLeaseKey GSI for the lease table
     * @param blockingWait  whether to wait for the GSI creation or not, if false, the gsi creation will be initiated
     *                      but this call will not block for its creation
     * @throws DependencyException  If DDB fails unexpectedly when creating the GSI
     */
    private void createGsi(final boolean blockingWait) throws DependencyException {
        log.info("Creating Lease table GSI if it does not exist");
        // KCLv3.0 always starts with GSI available
        leaseRefresher.createLeaseOwnerToLeaseKeyIndexIfNotExists();

        if (blockingWait) {
            log.info("Waiting for Lease table GSI creation");
            final long secondsBetweenPolls = 10L;
            final long timeoutSeconds = 600L;
            final boolean isIndexActive =
                    leaseRefresher.waitUntilLeaseOwnerToLeaseKeyIndexExists(secondsBetweenPolls, timeoutSeconds);

            if (!isIndexActive) {
                throw new DependencyException(
                        new IllegalStateException("Creating LeaseOwnerToLeaseKeyIndex on Lease table timed out"));
            }
        }
    }

    /**
     * Initialize KCL with components and configuration to support upgrade from 2x. This can happen
     * at KCL Worker startup when MigrationStateMachine starts in ClientVersion.CLIENT_VERSION_UPGRADE_FROM_2X.
     * Or Dynamically during roll-forward from ClientVersion.CLIENT_VERSION_2X.
     */
    public synchronized void initializeClientVersionForUpgradeFrom2x(final ClientVersion fromClientVersion)
            throws DependencyException {
        log.info("Initializing KCL components for upgrade from 2x from {}", fromClientVersion);

        createGsi(false);
        startWorkerMetricsReporting();
        // LAM is not started until the dynamic flip to 3xWithRollback
    }

    /**
     * Initialize KCL with components and configuration to run vanilla 3x functionality. This can happen
     * at KCL Worker startup when MigrationStateMachine starts in ClientVersion.CLIENT_VERSION_3X, or dynamically
     * during a new deployment when existing worker are in ClientVersion.CLIENT_VERSION_3X_WITH_ROLLBACK
     */
    public synchronized void initializeClientVersionFor3x(final ClientVersion fromClientVersion)
            throws DependencyException {
        log.info("Initializing KCL components for 3x from {}", fromClientVersion);

        log.info("Initializing LeaseAssignmentManager, DDB-lock-based leader decider, WorkerMetricStats manager"
                + " and creating the Lease table GSI if it does not exist");
        if (fromClientVersion == ClientVersion.CLIENT_VERSION_INIT) {
            // gsi may already exist and be active for migrated application.
            createGsi(true);
            startWorkerMetricsReporting();
            log.info("Starting LAM");
            leaseAssignmentManager.start();
        }
        // nothing to do when transitioning from CLIENT_VERSION_3X_WITH_ROLLBACK.
    }

    /**
     * Initialize KCL with components and configuration to run 2x compatible functionality
     * while allowing roll-forward. This can happen at KCL Worker startup when MigrationStateMachine
     * starts in ClientVersion.CLIENT_VERSION_2X (after a rollback)
     * Or Dynamically during rollback from CLIENT_VERSION_UPGRADE_FROM_2X or CLIENT_VERSION_3X_WITH_ROLLBACK.
     */
    public synchronized void initializeClientVersionFor2x(final ClientVersion fromClientVersion) {
        log.info("Initializing KCL components for rollback to 2x from {}", fromClientVersion);

        if (fromClientVersion != ClientVersion.CLIENT_VERSION_INIT) {
            // dynamic rollback
            stopWorkerMetricsReporter();
            // Migration Tool will delete the lease table LeaseOwner GSI
            // and WorkerMetricStats table
        }

        if (fromClientVersion == ClientVersion.CLIENT_VERSION_3X_WITH_ROLLBACK) {
            // we are rolling back after flip
            currentAssignmentMode = DEFAULT_LEASE_COUNT_BASED_ASSIGNMENT;
            notifyLeaseAssignmentModeChange();
            log.info("Stopping LAM");
            leaseAssignmentManager.stop();
            final LeaderDecider leaderDecider = deterministicLeaderDeciderCreator.get();
            if (this.leaderDecider instanceof MigrationAdaptiveLeaderDecider) {
                log.info(
                        "Updating LeaderDecider to {}", leaderDecider.getClass().getSimpleName());
                ((MigrationAdaptiveLeaderDecider) this.leaderDecider).updateLeaderDecider(leaderDecider);
            } else {
                throw new IllegalStateException(String.format("Unexpected leader decider %s", this.leaderDecider));
            }
        }
    }

    /**
     * Initialize KCL with components and configuration to run vanilla 3x functionality
     * while allowing roll-back to 2x functionality. This can happen at KCL Worker startup
     * when MigrationStateMachine starts in ClientVersion.CLIENT_VERSION_3X_WITH_ROLLBACK (after the flip)
     * Or Dynamically during flip from CLIENT_VERSION_UPGRADE_FROM_2X.
     */
    public synchronized void initializeClientVersionFor3xWithRollback(final ClientVersion fromClientVersion)
            throws DependencyException {
        log.info("Initializing KCL components for 3x with rollback from {}", fromClientVersion);

        if (fromClientVersion == ClientVersion.CLIENT_VERSION_UPGRADE_FROM_2X) {
            // dynamic flip
            currentAssignmentMode = WORKER_UTILIZATION_AWARE_ASSIGNMENT;
            notifyLeaseAssignmentModeChange();
            final LeaderDecider leaderDecider = ddbLockBasedLeaderDeciderCreator.get();
            log.info("Updating LeaderDecider to {}", leaderDecider.getClass().getSimpleName());
            ((MigrationAdaptiveLeaderDecider) this.leaderDecider).updateLeaderDecider(leaderDecider);
        } else {
            startWorkerMetricsReporting();
        }

        log.info("Starting LAM");
        leaseAssignmentManager.start();
    }

    /**
     * Synchronously invoke the consumer to change the lease assignment mode.
     */
    private void notifyLeaseAssignmentModeChange() {
        if (dualMode) {
            log.info("Notifying {} of {}", leaseModeChangeConsumer, currentAssignmentMode);
            if (Objects.nonNull(leaseModeChangeConsumer)) {
                try {
                    leaseModeChangeConsumer.updateLeaseAssignmentMode(currentAssignmentMode);
                } catch (final Exception e) {
                    log.warn("LeaseAssignmentMode change consumer threw exception", e);
                }
            }
        } else {
            throw new IllegalStateException("Unexpected assignment mode change");
        }
    }
}