postgres_checkpoint.h

Go to the documentation of this file.

 
#pragma once
 
#include <neograph/graph/checkpoint.h>
 
#include <atomic>
#include <condition_variable>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <utility>
#include <vector>
 
// Forward-declare libpq's opaque PGconn. Keeps libpq-fe.h out of
// every translation unit that includes this header.
struct pg_conn;
 
namespace neograph::graph {
 
namespace test_access {
class PostgresCheckpointStoreTestAccess;
}
 
struct PgConn {
    pg_conn* raw = nullptr;
    PgConn()     = default;
    explicit PgConn(pg_conn* p) : raw(p) {}
    ~PgConn();
    PgConn(const PgConn&)            = delete;
    PgConn& operator=(const PgConn&) = delete;
    PgConn(PgConn&&)                 = delete;
    PgConn& operator=(PgConn&&)      = delete;
};
 
class NEOGRAPH_API PostgresCheckpointStore : public CheckpointStore {
public:
    explicit PostgresCheckpointStore(const std::string& conn_str, size_t pool_size = 8);
 
    size_t reconnect_count() const { return reconnect_count_; }
 
    size_t pool_size() const { return pool_.size(); }
 
    ~PostgresCheckpointStore() override;
 
    // Non-copyable, non-movable — mutex would need rebinding on move
    // and connection pool slots are keyed by stable indices.
    PostgresCheckpointStore(const PostgresCheckpointStore&)            = delete;
    PostgresCheckpointStore& operator=(const PostgresCheckpointStore&) = delete;
 
    void                      save(const Checkpoint& cp) override;
    std::optional<Checkpoint> load_latest(const std::string& thread_id) override;
    std::optional<Checkpoint> load_by_id(const std::string& id) override;
    std::vector<Checkpoint>   list(const std::string& thread_id, int limit = 100) override;
    void                      delete_thread(const std::string& thread_id) override;
 
    void                      put_writes(const std::string&  thread_id,
                                         const std::string&  parent_checkpoint_id,
                                         const PendingWrite& write) override;
    std::vector<PendingWrite> get_writes(const std::string& thread_id,
                                         const std::string& parent_checkpoint_id) override;
    void                      clear_writes(const std::string& thread_id,
                                           const std::string& parent_checkpoint_id) override;
 
    // ── Async peers (Sem 4 follow-up) ───────────────────────────────────
    //
    // True async: libpq nonblocking mode + asio::posix::stream_descriptor
    // on PQsocket() + co_await on write/read readiness. The io_context
    // stays free for other coroutines during PG's commit fsync, so N
    // concurrent save_async calls across pool slots actually overlap
    // instead of serialising on the main worker thread.
    //
    // Successful operations use the same storage schema and values as the
    // sync peers. Async connection setup/replacement additionally follows
    // the global deadline policy documented on this class.
 
    asio::awaitable<void>                      save_async(const Checkpoint& cp) override;
    asio::awaitable<std::optional<Checkpoint>> load_latest_async(
        const std::string& thread_id) override;
    asio::awaitable<std::optional<Checkpoint>> load_by_id_async(const std::string& id) override;
    asio::awaitable<std::vector<Checkpoint>>   list_async(const std::string& thread_id,
                                                          int                limit = 100) override;
    asio::awaitable<void> delete_thread_async(const std::string& thread_id) override;
 
    asio::awaitable<void>                      put_writes_async(const std::string&  thread_id,
                                                                const std::string&  parent_checkpoint_id,
                                                                const PendingWrite& write) override;
    asio::awaitable<std::vector<PendingWrite>> get_writes_async(
        const std::string& thread_id, const std::string& parent_checkpoint_id) override;
    asio::awaitable<void> clear_writes_async(const std::string& thread_id,
                                             const std::string& parent_checkpoint_id) override;
 
    bool requires_managed_budget(const std::string& thread_id) override;
    asio::awaitable<bool> requires_managed_budget_async(std::string thread_id) override;
 
    std::shared_ptr<OwnedManagedBudgetLease> acquire_managed_budget_lease(
        const ManagedBudgetLeaseScope& scope, const std::string& expected_checkpoint_id,
        const std::string& expected_checkpoint_commitment) override;
    asio::awaitable<std::shared_ptr<OwnedManagedBudgetLease>> acquire_managed_budget_lease_async(
        ManagedBudgetLeaseScope scope, std::string expected_checkpoint_id,
        std::string expected_checkpoint_commitment) override;
    ManagedBudgetEffectReceipt begin_managed_budget_effect(
        const std::shared_ptr<OwnedManagedBudgetLease>& lease, const std::string& effect_id,
        std::uint64_t exact_claim_amount, const std::string& prepared_request_digest) override;
    asio::awaitable<ManagedBudgetEffectReceipt> begin_managed_budget_effect_async(
        std::shared_ptr<OwnedManagedBudgetLease> lease, std::string effect_id,
        std::uint64_t exact_claim_amount, std::string prepared_request_digest) override;
    void settle_managed_budget_effect(const std::shared_ptr<OwnedManagedBudgetLease>& lease,
        const ManagedBudgetEffectReceipt& effect, sp::runtime::Result genuine_outcome,
        const UsageAccumulator::AuthoritySnapshot& authority) override;
    asio::awaitable<void> settle_managed_budget_effect_async(
        std::shared_ptr<OwnedManagedBudgetLease> lease, ManagedBudgetEffectReceipt effect,
        sp::runtime::Result genuine_outcome, UsageAccumulator::AuthoritySnapshot authority) override;
    void publish_managed_budget_checkpoint(
        const std::shared_ptr<OwnedManagedBudgetLease>& lease, const Checkpoint& checkpoint) override;
    asio::awaitable<void> publish_managed_budget_checkpoint_async(
        std::shared_ptr<OwnedManagedBudgetLease> lease, Checkpoint checkpoint) override;
    void release_managed_budget_lease(const std::shared_ptr<OwnedManagedBudgetLease>& lease) override;
    asio::awaitable<void> release_managed_budget_lease_async(
        std::shared_ptr<OwnedManagedBudgetLease> lease) override;
 
    void drop_schema();
 
    size_t blob_count();
 
private:
    // Constructs only the slot scheduler for deterministic pool tests.
    // The friend test helper never calls a method that dereferences pool_.
    struct PoolOnlyTag {};
    explicit PostgresCheckpointStore(PoolOnlyTag, size_t pool_size);
 
    void ensure_schema();
 
    template <typename Fn>
    auto with_conn(Fn&& fn);
 
    template <typename Fn>
    auto with_conn_async(Fn& fn) -> decltype(fn(std::declval<pg_conn*>()));
 
    size_t                       acquire_slot();
    asio::awaitable<size_t>      acquire_slot_async();
    void                         release_slot(size_t idx);
    void                         rebuild_slot(size_t idx);
    asio::awaitable<void>        rebuild_slot_async(size_t idx);
    size_t                       waiter_count_for_test();
    static asio::awaitable<bool> wait_socket_either_for_test(int fd);
    static void                  set_async_connection_test_seams(int poll_delay_ms, int timeout_ms);
    static int                   async_connection_timeout_ms_for_test(const std::string& conn_str);
 
    std::string conn_str_;
 
    std::vector<std::unique_ptr<PgConn>> pool_;
 
    std::queue<size_t>      free_;
    std::mutex              pool_mutex_;
    std::condition_variable pool_cv_;  // Retained at its original ABI offset.
 
    std::atomic<size_t> reconnect_count_{0};
 
    friend class test_access::PostgresCheckpointStoreTestAccess;
};
 
}  // namespace neograph::graph