postgresql_connection_pool.hpp

Go to the documentation of this file.

#pragma once
 
#include "postgresql_connection.hpp"
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <format>
#include <memory>
#include <mutex>
#include <optional>
#include <queue>
#include <string>
#include <string_view>
#include <vector>
 
namespace relx::connection {
 
struct PostgreSQLConnectionPoolConfig {
  PostgreSQLConnectionParams connection_params;
 
  size_t initial_size = 5;
 
  size_t max_size = 10;
 
  std::chrono::milliseconds connection_timeout{5000};
 
  bool validate_connections = true;
 
  std::chrono::milliseconds max_idle_time{60000};
};
 
struct ConnectionPoolError {
  std::string message;
  int error_code = 0;
};
 
inline std::string format_error(const ConnectionPoolError& error) {
  return std::format("Connection pool error: {} (Code: {})", error.message, error.error_code);
}
 
template <typename T>
using ConnectionPoolResult = std::expected<T, ConnectionPoolError>;
 
// Forward declaration
class PostgreSQLConnectionPool;
 
class PostgreSQLConnectionPool : public std::enable_shared_from_this<PostgreSQLConnectionPool> {
private:
  struct PoolEntry {
    std::shared_ptr<PostgreSQLConnection> connection;
    std::chrono::steady_clock::time_point last_used;
  };
 
  explicit PostgreSQLConnectionPool(PostgreSQLConnectionPoolConfig config);
 
public:
  // Forward declaration of the connection wrapper class
  class PooledConnection;
 
 
  static std::shared_ptr<PostgreSQLConnectionPool> create(PostgreSQLConnectionPoolConfig config) {
    // Use new directly instead of make_shared to access the private constructor
    return std::shared_ptr<PostgreSQLConnectionPool>(
        new PostgreSQLConnectionPool(std::move(config)));
  }
 
  ~PostgreSQLConnectionPool();
 
  // Delete copy and move operations
  PostgreSQLConnectionPool(const PostgreSQLConnectionPool&) = delete;
  PostgreSQLConnectionPool& operator=(const PostgreSQLConnectionPool&) = delete;
  PostgreSQLConnectionPool(PostgreSQLConnectionPool&&) = delete;
  PostgreSQLConnectionPool& operator=(PostgreSQLConnectionPool&&) = delete;
 
  [[nodiscard]] ConnectionPoolResult<void> initialize();
 
  [[nodiscard]] ConnectionPoolResult<PooledConnection> get_connection();
 
  size_t active_connections() const;
 
  size_t idle_connections() const;
 
  template <typename Func>
  [[nodiscard]] auto with_connection(Func&& func)
      -> ConnectionPoolResult<std::invoke_result_t<Func, PooledConnection&>> {
    using ResultType = std::invoke_result_t<Func, PooledConnection&>;
 
    auto conn_result = get_connection();
    if (!conn_result) {
      // Convert connection pool error to function result error
      return std::unexpected(conn_result.error());
    }
 
    try {
      // Execute the function with the connection by reference
      if constexpr (std::is_same_v<ResultType, void>) {
        func(*conn_result);
        return {};
      } else {
        return func(*conn_result);
      }
    } catch (...) {
      // Connection will be returned automatically by PooledConnection's destructor
      throw;  // Re-throw the exception
    }
  }
 
  class PooledConnection {
  private:
    std::shared_ptr<PostgreSQLConnection> connection_;
    std::weak_ptr<PostgreSQLConnectionPool> pool_;
 
  public:
    PooledConnection(std::shared_ptr<PostgreSQLConnection> connection,
                     const std::shared_ptr<PostgreSQLConnectionPool>& pool)
        : connection_(std::move(connection)), pool_(pool) {}
 
    ~PooledConnection() {
      if (connection_) {
        // Check if pool still exists
        if (auto pool = pool_.lock()) {
          pool->return_connection(std::move(connection_));
        }
        // If pool no longer exists, connection will simply be destroyed
      }
    }
 
    // Delete copy operations
    PooledConnection(const PooledConnection&) = delete;
    PooledConnection& operator=(const PooledConnection&) = delete;
 
    // Allow move operations
    PooledConnection(PooledConnection&&) = default;
    PooledConnection& operator=(PooledConnection&&) = default;
 
    PostgreSQLConnection* operator->() { return connection_.get(); }
 
    const PostgreSQLConnection* operator->() const { return connection_.get(); }
 
    explicit operator bool() const { return connection_ != nullptr; }
  };
 
private:
  PostgreSQLConnectionPoolConfig config_;
  std::atomic<size_t> active_connections_{0};
  std::atomic<size_t> total_connections_{0};
 
  mutable std::mutex pool_mutex_;
  std::condition_variable conn_available_;
  std::queue<PoolEntry> idle_connections_;
 
  [[nodiscard]] ConnectionPoolResult<std::shared_ptr<PostgreSQLConnection>> get_raw_connection();
 
  void return_connection(std::shared_ptr<PostgreSQLConnection> connection);
 
  ConnectionPoolResult<std::shared_ptr<PostgreSQLConnection>> create_connection();
 
  static bool validate_connection(const std::shared_ptr<PostgreSQLConnection>& connection);
 
  void cleanup_idle_connections();
};
 
}  // namespace relx::connection