ClickHouse/dbms/include/DB/Storages/MergeTree/ReshardingWorker.h

#pragma once

#include <DB/Storages/MergeTree/ReshardingJob.h>
#include <DB/Storages/AlterCommands.h>
#include <common/logger_useful.h>

#include <zkutil/RWLock.h>
#include <zkutil/SingleBarrier.h>

#include <Poco/Util/LayeredConfiguration.h>

#include <string>
#include <thread>
#include <atomic>
#include <functional>

namespace DB
{

class Context;
class Cluster;
class StorageReplicatedMergeTree;

/// Performer of resharding jobs. It works as a background task within a thread.
/// Its main duties are to keep track of newly submitted resharding jobs and
/// to execute them sequentially.
class ReshardingWorker final
{
    friend class AnomalyMonitor;

public:
    using PartitionList = std::vector<std::string>;

    /// Possible status values of a coordinator or a performer that
    /// has subscribed to a coordinator.
    enum StatusCode
    {
        STATUS_OK = 0,
        STATUS_ERROR,    /// An error occurred on a performer.
        STATUS_ON_HOLD    /// Job is stopped.
    };

public:
    ReshardingWorker(const Poco::Util::AbstractConfiguration & config,
        const std::string & config_name, Context & context_);

    ReshardingWorker(const ReshardingWorker &) = delete;
    ReshardingWorker & operator=(const ReshardingWorker &) = delete;

    ~ReshardingWorker();

    /// Start the job tracker thread.
    void start();

    /// Stop the job tracker thread. If any job is in progress, put it on hold
    /// for future execution.
    void shutdown();

    /// Send a request for resharding on the current performer.
    void submitJob(const ReshardingJob & job);

    /// Is the job tracker thread started?
    bool isStarted() const;

    /// Create a new coordinator for a distributed job. Called from the initiator.
    std::string createCoordinator(const Cluster & cluster);
    /// Register a query into a coordinator.
    void registerQuery(const std::string & coordinator_id, const std::string & query);
    /// Clear all the information related to a given coordinator.
    void deleteCoordinator(const std::string & coordinator_id);

    /// Subscribe the performer, from which this method is called, to a given coordinator.
    UInt64 subscribe(const std::string & coordinator_id, const std::string & query);
    /// Cancel the aforementionned subscription.
    void unsubscribe(const std::string & coordinator_id);
    /// Увеличить количество партиций входящих в одну распределённую задачу. Вызывается с исполнителя.
    void addPartitions(const std::string & coordinator_id, const PartitionList & partition_list);
    /// Rearrange partitions into two categories: coordinated job, uncoordinated job.
    /// Returns an iterator to the beginning of the list of uncoordinated jobs.
    ReshardingWorker::PartitionList::iterator categorizePartitions(const std::string & coordinator_id,
        ReshardingWorker::PartitionList & partition_list);
    /// Get the number of partitions of a distributed job. Called from performers.
    size_t getPartitionCount(const std::string & coordinator_id);
    /// Get the number of performers.
    size_t getNodeCount(const std::string & coordinator_id);
    /// Wait for all the preliminary sanity checks to be completed on all the performers.
    /// Called from performers.
    void waitForCheckCompletion(const std::string & coordinator_id);
    /// Wait for all the required unsubscribe operations to be completed.
    void waitForOptOutCompletion(const std::string & coordinator_id, size_t count);

    /// Set the shard-independent status of a given coordinator.
    void setStatus(const std::string & coordinator_id, StatusCode status, const std::string & msg = "");
    /// Set the status of a shard under a given coordinator.
    void setStatus(const std::string & coordinator_id, const std::string & hostname,
        StatusCode status, const std::string & msg = "");

    zkutil::RWLock createDeletionLock(const std::string & coordinator_id);

    /// Dump the status messages of the coordinator and all the performers.
    std::string dumpCoordinatorState(const std::string & coordinator_id);

    static std::string computeHashFromPart(const std::string & path);

private:
    /// Anomalies that may be detected by the current performer.
    enum AnomalyType
    {
        ANOMALY_NONE = 0,
        ANOMALY_LOCAL_SHUTDOWN,
        ANOMALY_LOCAL_ERROR,
        ANOMALY_REMOTE_NODE_UNAVAILABLE,
        ANOMALY_REMOTE_ERROR
    };

    /// This structure stores various required information for log creation.
    struct TargetShardInfo
    {
        TargetShardInfo(size_t shard_no_, const std::string & part_name_, const std::string & hash_)
            : part_name{part_name_}, hash{hash_}, shard_no{shard_no_}
        {
        }

        std::string part_name;
        std::string hash;
        size_t shard_no;
    };

    using ShardList = std::vector<TargetShardInfo>;

    /// Structure that describes an operation to be performed as part of
    /// a commit procedure.
    class LogRecord final
    {
    public:
        /// Create a new log record.
        LogRecord(zkutil::ZooKeeperPtr zookeeper_);
        /// Open an already existing log record.
        LogRecord(zkutil::ZooKeeperPtr zookeeper_, const std::string & zk_path_);
        /// Append this log record into the specified log onto persistent storage.
        void enqueue(const std::string & log_path);
        /// Update an already existing log record.
        void writeBack();

    public:
        enum Operation
        {
            OP_DROP = 0,
            OP_ATTACH
        };

        enum State
        {
            READY = 0,
            RUNNING,
            DONE
        };

    private:
        /// Serialize this log record.
        std::string toString();

    public:
        zkutil::ZooKeeperPtr zookeeper;
        Operation operation;
        /// Source partition being resharded.
        std::string partition;
        /// Hash of the source partition (for drop operations).
        std::string partition_hash;
        /// Path of this log record on stable storage.
        std::string zk_path;
        /// List of the parts that constitute a sharded partition along with
        /// their respective hash values (for attach operations).
        std::map<std::string, std::string> parts_with_hash;
        /// Target shard index into the list of specified target shards inside
        /// a ALTER TABLE ... RESHARD request.
        size_t shard_no = 0;
        State state;
    };

private:
    /// Keep track of newly submitted jobs. Perform them sequentially.
    void trackAndPerform();

    /// Wake up the tracker thread if it was ever sleeping.
    void wakeUpTrackerThread();

    /// Perform all the jobs that were already pending when the ClickHouse instance
    /// on this node was starting.
    void performPendingJobs();

    /// Sequentially perform jobs which are specified by their corresponding
    /// znodes in ZooKeeper.
    void perform(const Strings & job_nodes);

    /// Perform one job.
    void perform(const std::string & job_descriptor, const std::string & job_name);

    /// Разбить куски входящие в партицию на несколько, согласно ключу шардирования.
    /// Оновременно перегруппировать эти куски по шардам и слить куски в каждой группе.
    /// При завершении этого процесса создаётся новая партиция для каждого шарда.
    void createShardedPartitions();

    /// Upload all the partitions resulting from source partition resharding to their
    /// respective target shards.
    void publishShardedPartitions();

    /// On each target shard attach its corresponding new sharded partition.
    /// If COPY is not specified in the ALTER TABLE ... RESHARD ... query,
    /// drop source partitions on each performer.
    void commit();

    void repairLogRecord(LogRecord & log_record);
    void executeLogRecord(LogRecord & log_record);
    void executeDrop(LogRecord & log_record);
    void executeAttach(LogRecord & log_record);
    bool checkAttachLogRecord(LogRecord & log_record);

    /// Delete temporary data from the local node and ZooKeeper.
    void softCleanup();
    void hardCleanup();

    /// Forcibly abort the job tracker thread if requested.
    void abortTrackingIfRequested();
    /// Forcibly abort the distributed job coordinated by the given coordinator
    /// if requested.
    void abortCoordinatorIfRequested(const std::string & coordinator_id);
    /// Forcibly abort the recovery of the current distributed job if requested.
    void abortRecoveryIfRequested();
    /// Forcibly abort the current job if requested.
    void abortJobIfRequested();

    /// Get the current job-independent status of the coordinator.
    StatusCode getCoordinatorStatus(const std::string & coordinator_id);
    /// Get the status of the current distributed job.
    StatusCode getStatus();

    /// Prepare the current job for execution.
    void initializeJob();
    /// Remove the current job's data from its coordinator.
    /// If there is no job left, also delete the coordinator.
    void finalizeJob();
    /// Wait for all the necesssary sharded partitions uploads to their respective
    /// target shards to be completed.
    void waitForUploadCompletion();
    /// Wait for all the performers to be checked in for leader election.
    void waitForElectionCompletion();
    /// Wait for all the changes required by the current distributed resharding job
    /// to be applied by all the performers on all the target shards.
    void waitForCommitCompletion();

    size_t getPartitionCountUnlocked(const std::string & coordinator_id);

    /// Detect offline performers under a given coordinator.
    bool detectOfflineNodes(const std::string & coordinator_id);
    /// Detect offline performers under the current job.
    bool detectOfflineNodes();

    bool isPublished();
    void markAsPublished();

    bool isLogCreated();
    void markLogAsCreated();

    bool isCommitted();
    void markAsCommitted();

    /// Store onto ZooKeeper information about the target shards. It is used
    /// in order to create the log of operations that apply changes all the
    /// changes required by the resharding operation.
    void storeTargetShardsInfo();

    /// Retrieve from ZooKeeper the aforementioned information about the target shards.
    ShardList getTargetShardsInfo(const std::string & hostname, const std::string & job_name);

    /// Create the log of operations.
    void createLog();

    void electLeader();
    bool isLeader();

    /// Access the global lock handler.
    zkutil::RWLock getGlobalLock();
    /// Acccess the given coordinator lock handler.
    zkutil::RWLock getCoordinatorLock(const std::string & coordinator_id,
        bool usable_in_emergency = false);

    /// The following functions are to design access barrier handlers we use
    /// to synchronize the progression of distributed jobs.
    zkutil::SingleBarrier getCheckBarrier(const std::string & coordinator_id);
    zkutil::SingleBarrier getOptOutBarrier(const std::string & coordinator_id, size_t count);
    zkutil::SingleBarrier getRecoveryBarrier();
    zkutil::SingleBarrier getUploadBarrier();
    zkutil::SingleBarrier getElectionBarrier();
    zkutil::SingleBarrier getCommitBarrier();

    /// Get the ZooKeeper path of a given coordinator.
    std::string getCoordinatorPath(const std::string & coordinator_id) const;
    /// Get the ZooKeeper path of a given job partition.
    std::string getPartitionPath() const;
    /// Get the ZooKeeper path of the job currently running on this performer.
    std::string getLocalJobPath() const;

    /// Common code for softCleanup() and hardCleanup().
    void deleteTemporaryData();
    /// Common code for detectOfflineNodes().
    bool detectOfflineNodesCommon(const std::string & path, const std::string & coordinator_id);
    /// Common code for getStatus() and getCoordinatorStatus().
    StatusCode getStatusCommon(const std::string & path, const std::string & coordinator_id);

    AnomalyType probeForAnomaly();
    void processAnomaly(AnomalyType anomaly_type);

private:
    /// This class is used to spawn a thread which periodically checks for any
    /// anomaly that should lead to abort the job currently being processed.
    /// Jobs perform such checks with a negligible performance overhead by calling
    /// the method getAnomalyType() which merely reads the value of an atomic variable.
    class AnomalyMonitor final
    {
    public:
        AnomalyMonitor(ReshardingWorker & resharding_worker_);
        ~AnomalyMonitor();

        AnomalyMonitor(const AnomalyMonitor &) = delete;
        AnomalyMonitor & operator=(const AnomalyMonitor &) = delete;

        /// Start the thread.
        void start();
        /// Shutdown the thread.
        void shutdown();
        /// Get the type of anomaly that has been registered.
        AnomalyType getAnomalyType() const;

    private:
        void routine();

    private:
        ReshardingWorker & resharding_worker;
        std::thread thread_routine;

        std::atomic<ReshardingWorker::AnomalyType> anomaly_type{ReshardingWorker::ANOMALY_NONE};
        std::atomic<bool> is_started{false};
        std::atomic<bool> must_stop{false};
    };

    /// Guarded exception-safe handling of the class above.
    class ScopedAnomalyMonitor final
    {
    public:
        ScopedAnomalyMonitor(AnomalyMonitor & anomaly_monitor_)
            : anomaly_monitor{anomaly_monitor_}
        {
            anomaly_monitor.start();
        }

        ~ScopedAnomalyMonitor()
        {
            try
            {
                anomaly_monitor.shutdown();
            }
            catch (...)
            {
                tryLogCurrentException(__PRETTY_FUNCTION__);
            }
        }

        ScopedAnomalyMonitor(const ScopedAnomalyMonitor &) = delete;
        ScopedAnomalyMonitor & operator=(const ScopedAnomalyMonitor &) = delete;

    private:
        AnomalyMonitor & anomaly_monitor;
    };

private:
    ReshardingJob current_job;
    std::thread job_tracker;
    AnomalyMonitor anomaly_monitor{*this};

    std::string task_queue_path;
    std::string host_task_queue_path;
    std::string distributed_path;
    std::string distributed_online_path;
    std::string distributed_lock_path;
    std::string coordination_path;

    Context & context;
    Logger * log = &Logger::get("ReshardingWorker");

    zkutil::EventPtr event = std::make_shared<Poco::Event>();

    /// Helper that acquires an alive ZooKeeper session.
    zkutil::GetZooKeeper get_zookeeper;

    std::atomic<bool> is_started{false};
    std::atomic<bool> must_stop{false};
};

using ReshardingWorkerPtr = std::shared_ptr<ReshardingWorker>;

}