ReplicaFailoverSession.cs

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

using System;
using System.Collections.Generic;
using System.Text;
using System.Threading.Tasks;
using Garnet.client;
using Garnet.common;
using Garnet.server;
using Microsoft.Extensions.Logging;

namespace Garnet.cluster
{
    internal sealed partial class FailoverSession : IDisposable
    {
        /// <summary>
        /// Connection to primary if reachable
        /// </summary>
        GarnetClient primaryClient = null;

        /// <summary>
        /// Send page size for GarnetClient
        /// </summary>
        const int sendPageSize = 1 << 17;

        /// <summary>
        /// Helper method to establish connection towards remote node
        /// </summary>
        /// <param name="nodeId">Id of node to create connection for</param>
        /// <returns></returns>
        private async Task<GarnetClient> CreateConnectionAsync(string nodeId)
        {
            var endpoint = oldConfig.GetEndpointFromNodeId(nodeId);
            var client = new GarnetClient(
                endpoint,
                clusterProvider.serverOptions.TlsOptions?.TlsClientOptions,
                sendPageSize: sendPageSize,
                maxOutstandingTasks: 8,
                authUsername: clusterProvider.ClusterUsername,
                authPassword: clusterProvider.ClusterPassword, epoch: epoch, logger: logger);

            try
            {
                if (!client.IsConnected)
                    await client.ReconnectAsync().WaitAsync(failoverTimeout, cts.Token).ConfigureAwait(false);

                return client;
            }
            catch (Exception ex)
            {
                client?.Dispose();
                logger?.LogError(ex, "ReplicaFailoverSession.CreateConnection");
                return null;
            }
        }

        /// <summary>
        /// Acquire a connection to the node identified by given node-id.
        /// </summary>
        /// <param name="nodeId"></param>
        /// <returns></returns>
        private Task<GarnetClient> GetConnectionAsync(string nodeId)
            => CreateConnectionAsync(nodeId);

        /// <summary>
        /// Send stop writes message to PRIMARY
        /// </summary>
        /// <returns>True on success, false otherwise</returns>
        private async Task<bool> PauseWritesAndWaitForSyncAsync()
        {
            var primaryId = oldConfig.LocalNodePrimaryId;
            var client = await GetConnectionAsync(primaryId).ConfigureAwait(false);
            try
            {
                if (client == null)
                {
                    logger?.LogError("Failed to initialize connection to primary {primaryId}", primaryId);
                    return false;
                }

                // Cache connection for use with next operations
                primaryClient = client;

                // Issue stop writes to the primary
                status = FailoverStatus.ISSUING_PAUSE_WRITES;
                var localIdBytes = Encoding.ASCII.GetBytes(oldConfig.LocalNodeId);

                var resp = await client.ExecuteClusterFailStopWritesAsync(localIdBytes).WaitAsync(failoverTimeout, cts.Token).ConfigureAwait(false);
                var primaryReplicationOffset = AofAddress.FromString(resp);

                // Wait for replica to catch up
                status = FailoverStatus.WAITING_FOR_SYNC;
                while (primaryReplicationOffset.AnyGreater(clusterProvider.replicationManager.ReplicationOffset))
                {
                    // Fail if upper bound time for failover has been reached
                    if (FailoverTimeout)
                    {
                        logger?.LogError("AwaitReplicationSync timed out failoverStart");
                        return false;
                    }
                    await Task.Yield();
                }

                return true;
            }
            catch (Exception ex)
            {
                logger?.LogError(ex, "PauseWritesAndWaitForSync Error");
                return false;
            }
        }

        /// <summary>
        /// Perform series of steps to update local config and take ownership of primary slots.
        /// </summary>
        private async Task<bool> TakeOverAsPrimaryAsync()
        {
            // Take over as primary and inform old primary
            status = FailoverStatus.TAKING_OVER_AS_PRIMARY;
            var acquiredLock = false;

            try
            {
                // Exception injection point for testing: simulates TakeOverAsPrimary failure
                // after PauseWritesAndWaitForSync has already sent failstopwrites to the primary.
                ExceptionInjectionHelper.TriggerException(ExceptionInjectionType.Failover_Fail_TakeOverAsPrimary);

                // Make replica syncing unavailable by setting recovery flag
                if (!clusterProvider.replicationManager.BeginRecovery(RecoveryStatus.ClusterFailover, upgradeLock: false))
                {
                    logger?.LogWarning($"{nameof(TakeOverAsPrimaryAsync)}: {{logMessage}}", Encoding.ASCII.GetString(CmdStrings.RESP_ERR_GENERIC_CANNOT_ACQUIRE_RECOVERY_LOCK));
                    return false;
                }
                acquiredLock = true;
                _ = await clusterProvider.BumpAndWaitForEpochTransitionAsync().ConfigureAwait(false);

                // Take over slots from old primary
                if (!clusterProvider.clusterManager.TryTakeOverForPrimary())
                {
                    logger?.LogWarning($"{nameof(TakeOverAsPrimaryAsync)}: {{logMessage}}", Encoding.ASCII.GetString(CmdStrings.RESP_ERR_GENERIC_CANNOT_TAKEOVER_FROM_PRIMARY));
                    return false;
                }

                // Update replicationIds and replicationOffset2
                clusterProvider.replicationManager.TryUpdateForFailover();

                // Cancel active replication tasks
                clusterProvider.replicationManager.ResetReplicaReplayDriverStore();

                // Update sequence number generator for sharded log if needed
                if (clusterProvider.serverOptions.AofPhysicalSublogCount > 1)
                {
                    clusterProvider.storeWrapper.appendOnlyFile.ResetSequenceNumberGenerator();
                    await clusterProvider.storeWrapper.TaskManager.CancelAsync(TaskType.AdvanceTimeReplicaTask).ConfigureAwait(false);
                }

                // Initialize checkpoint history
                if (!clusterProvider.replicationManager.InitializeCheckpointStore())
                    logger?.LogWarning("Failed acquiring latest memory checkpoint metadata at {method}", nameof(TakeOverAsPrimaryAsync));

                _ = clusterProvider.BumpAndWaitForEpochTransitionAsync().ConfigureAwait(false);

                // Stop advance time task when reconfiguring node to be replica
                if (clusterProvider.storeWrapper.serverOptions.AofPhysicalSublogCount > 1)
                    await clusterProvider.storeWrapper.TaskManager.CancelAsync(TaskType.AdvanceTimeReplicaTask).ConfigureAwait(false);

                // Resume all background maintenance that were possibly shutdown when this node became a replica
                clusterProvider.storeWrapper.StartPrimaryTasks();
            }
            catch (Exception ex)
            {
                logger?.LogError(ex, "{method}", nameof(TakeOverAsPrimaryAsync));
                throw;
            }
            finally
            {
                // Disable recovering as now this node has become a primary or failed in its attempt earlier
                if (acquiredLock)
                    clusterProvider.replicationManager.EndRecovery(RecoveryStatus.NoRecovery, downgradeLock: false);
            }

            return true;
        }

        /// <summary>
        /// Issue gossip and attach request to replica
        /// </summary>
        /// <param name="replicaId">Replica-id to issue gossip and attache request</param>
        /// <param name="configByteArray">Serialized local cluster config data</param>
        /// <returns></returns>
        private async Task BroadcastConfigAndRequestAttachAsync(string replicaId, byte[] configByteArray)
        {
            // Force async
            await Task.Yield();

            var oldPrimaryId = oldConfig.LocalNodePrimaryId;
            var newConfig = clusterProvider.clusterManager.CurrentConfig;
            var client = oldPrimaryId.Equals(replicaId) ? primaryClient : await GetConnectionAsync(replicaId).ConfigureAwait(false);

            try
            {
                if (client == null)
                {
                    logger?.LogError("Failed to initialize connection to replica {replicaId}", replicaId);
                    return;
                }

                // Force send updated config to replica
                var resp = await client.GossipAsync(configByteArray).WaitAsync(failoverTimeout, cts.Token).ConfigureAwait(false);

                try
                {
                    var current = clusterProvider.clusterManager.CurrentConfig;
                    if (resp.Length > 0)
                    {
                        clusterProvider.clusterManager.gossipStats.UpdateGossipBytesRecv(resp.Length);
                        var returnedConfigArray = resp.Span.ToArray();

                        // Validate config version before full deserialization
                        if (!ClusterConfig.TryPeekVersion(returnedConfigArray, out var version) || version != ClusterConfig.ClusterConfigVersion)
                        {
                            logger?.LogWarning("Received failover gossip response with incompatible config version: {version}", version);
                        }
                        else
                        {
                            var other = ClusterConfig.FromByteArray(returnedConfigArray);

                            // Check if gossip is from a node that is known and trusted before merging
                            if (current.IsKnown(other.LocalNodeId))
                                _ = clusterProvider.clusterManager.TryMerge(other);
                            else
                                logger?.LogWarning("Received gossip from unknown node: {node-id}", other.LocalNodeId);
                        }
                    }
                }
                catch (Exception ex)
                {
                    logger?.LogCritical(ex, "IssueAttachReplicas faulted");
                }
                finally
                {
                    resp.Dispose();
                }

                var localAddress = oldConfig.LocalNodeIp;
                var localPort = oldConfig.LocalNodePort;

                // Ask replica to attach and sync
                var replicaOfResp = await client.ReplicaOf(localAddress, localPort).WaitAsync(failoverTimeout, cts.Token).ConfigureAwait(false);

                // Check if response for attach succeeded
                if (!replicaOfResp.Equals("OK"))
                    logger?.LogWarning("IssueAttachReplicas Error: {replicaId} {replicaOfResp}", replicaId, replicaOfResp);
            }
            finally
            {
                client?.Dispose();
            }
        }

        /// <summary>
        /// Issue attach message to remote replicas
        /// </summary>
        /// <returns></returns>
        private async Task IssueAttachReplicasAsync()
        {
            // Get information of local node from newConfig
            var newConfig = clusterProvider.clusterManager.CurrentConfig;
            // Get replica ids for old primary from old configuration
            var oldPrimaryId = oldConfig.LocalNodePrimaryId;
            var replicaIds = newConfig.GetReplicaIds(oldPrimaryId);
            var configByteArray = newConfig.ToByteArray();
            var attachReplicaTasks = new List<Task>();

            // If DEFAULT failover try to make old primary replica of this new primary
            if (option is FailoverOption.DEFAULT)
            {
                replicaIds.Add(oldPrimaryId);
            }

            // Issue gossip and attach request to replicas
            foreach (var replicaId in replicaIds)
            {
                try
                {
                    attachReplicaTasks.Add(BroadcastConfigAndRequestAttachAsync(replicaId, configByteArray));
                }
                catch (Exception ex)
                {
                    logger?.LogError(ex, "IssueAttachReplicas Error");
                }
            }

            // Wait for tasks to complete
            if (attachReplicaTasks.Count > 0)
            {
                try
                {
                    await Task.WhenAll(attachReplicaTasks).ConfigureAwait(false);
                }
                catch (Exception ex)
                {
                    logger?.LogWarning(ex, "WaitingForAttachToComplete Error");
                }
            }
        }

        /// <summary>
        /// Returns true if failstopwrites was confirmed by the primary and the primary's
        /// config was modified (slots given up, role changed to replica). Used to determine
        /// whether the primary needs to be reset on failover failure.
        /// </summary>
        private bool PrimaryNeedsReset()
            => status is FailoverStatus.WAITING_FOR_SYNC or FailoverStatus.TAKING_OVER_AS_PRIMARY;

        /// <summary>
        /// REPLICA main failover task
        /// </summary>
        /// <returns></returns>
        public async Task<bool> BeginAsyncReplicaFailoverAsync()
        {
            // CLUSTER FAILOVER OPTIONS
            // FORCE: Do not await for the primary since it might be unreachable
            // TAKEOVER: Same as force but also do not await for voting from other primaries
            var failoverSucceeded = false;
            try
            {
                // Issue stop writes and on ack wait for replica to catch up
                if (option is FailoverOption.DEFAULT && !await PauseWritesAndWaitForSyncAsync().ConfigureAwait(false))
                {
                    return false;
                }

                // If TAKEOVER option is set skip voting
                if (option is FailoverOption.DEFAULT or FailoverOption.FORCE)
                {
                    //TODO: implement voting
                }

                // Transition to primary role
                if (!await TakeOverAsPrimaryAsync().ConfigureAwait(false))
                {
                    return false;
                }
                failoverSucceeded = true;

                // Attach to old replicas, and old primary if DEFAULT option
                await IssueAttachReplicasAsync().ConfigureAwait(false);

                await clusterProvider.storeWrapper.SuspendReplicaOnlyTasksAsync();
                clusterProvider.storeWrapper.StartPrimaryTasks();

                return true;
            }
            catch (Exception ex)
            {
                logger?.LogWarning(ex, "BeginAsyncReplicaFailover Error");
                return false;
            }
            finally
            {
                // If failstopwrites was confirmed by the primary (status reached WAITING_FOR_SYNC
                // or beyond) but the failover did not succeed, reset the primary back to its
                // original state. Without this, the primary has already given up its slots
                // (via TryStopWrites) but the replica never claimed them, leaving the cluster
                // in an incoherent state where no node owns the slots.
                if (PrimaryNeedsReset() && !failoverSucceeded)
                {
                    try
                    {
                        logger?.LogWarning("Attempting to reset primary after failed failover");
                        if (primaryClient != null)
                        {
                            _ = await primaryClient.ExecuteClusterFailStopWritesAsync(Array.Empty<byte>()).WaitAsync(failoverTimeout, cts.Token).ConfigureAwait(false);
                        }
                    }
                    catch (Exception ex)
                    {
                        logger?.LogError(ex, "Failed to reset primary after failed failover — cluster may be in an incoherent state");
                    }
                }

                primaryClient?.Dispose();
                status = FailoverStatus.NO_FAILOVER;
            }
        }
    }
}