cancellation_token.h

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2024, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#pragma once
 
#include <kcenon/thread/core/error_handling.h>
 
#include <atomic>
#include <functional>
#include <memory>
#include <mutex>
#include <vector>
 
namespace kcenon::thread {
 
class cancellation_token {
private:
    // Private implementation to allow copying/moving the token
    struct token_state {
        std::atomic<bool> is_cancelled{false};
        std::vector<std::function<void()>> callbacks;
        std::mutex callback_mutex;
    };
    
    std::shared_ptr<token_state> state_;
    
    // Private constructor that takes a state
    explicit cancellation_token(std::shared_ptr<token_state> state)
        : state_(std::move(state)) {}
    
public:
    // Default constructor creates a new token
    cancellation_token() : state_(std::make_shared<token_state>()) {}
    
    // Allow copy/move operations
    cancellation_token(const cancellation_token&) = default;
    cancellation_token& operator=(const cancellation_token&) = default;
    cancellation_token(cancellation_token&&) = default;
    cancellation_token& operator=(cancellation_token&&) = default;
    
    static cancellation_token create() {
        return cancellation_token();
    }
    
    static cancellation_token create_linked(std::initializer_list<cancellation_token> tokens) {
        auto new_token = create();
        auto new_state_weak = std::weak_ptr<token_state>(new_token.state_);
        
        for (const auto& token : tokens) {
            auto token_copy = token;
            token_copy.register_callback([new_state_weak]() {
                if (auto state = new_state_weak.lock()) {
                    // Directly set the cancelled flag and invoke callbacks
                    std::vector<std::function<void()>> callbacks_to_invoke;
 
                    {
                        std::lock_guard<std::mutex> lock(state->callback_mutex);
                        bool was_cancelled = state->is_cancelled.exchange(true, std::memory_order_release);
                        if (!was_cancelled) {
                            callbacks_to_invoke = std::move(state->callbacks);
                            state->callbacks.clear();
                        }
                    }
 
                    for (const auto& callback : callbacks_to_invoke) {
                        callback();
                    }
                }
            });
        }
        
        return new_token;
    }
    
    void cancel() {
        std::vector<std::function<void()>> callbacks_to_invoke;
 
        {
            std::lock_guard<std::mutex> lock(state_->callback_mutex);
            bool was_cancelled = state_->is_cancelled.exchange(true, std::memory_order_release);
            if (!was_cancelled) {
                // Move callbacks to local vector to avoid holding lock during invocation
                callbacks_to_invoke = std::move(state_->callbacks);
                state_->callbacks.clear();
            }
        }
 
        // Invoke callbacks outside the lock to prevent deadlock
        for (const auto& callback : callbacks_to_invoke) {
            callback();
        }
    }
 
    [[nodiscard]] bool is_cancelled() const {
        return state_->is_cancelled.load(std::memory_order_acquire);
    }
    
    [[nodiscard]] common::VoidResult check_cancelled() const {
        if (is_cancelled()) {
            return make_error_result(error_code::operation_canceled);
        }
        return common::ok();
    }
    
    void register_callback(std::function<void()> callback) {
        std::unique_lock<std::mutex> lock(state_->callback_mutex);
 
        // Check cancellation state while holding the lock
        if (state_->is_cancelled.load(std::memory_order_acquire)) {
            lock.unlock();
            callback();
            return;
        }
 
        // Add callback while still holding the lock
        state_->callbacks.push_back(std::move(callback));
    }
};
 
} // namespace kcenon::thread