[codex] Prevent concurrent AP and TP modes

include/neug/main/connection.h

@@ -17,8 +17,10 @@
 #include <glog/logging.h>
 
 #include <atomic>
+#include <functional>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "neug/compiler/planner/graph_planner.h"
@@ -83,10 +85,14 @@ class NeugDB;
  */
 class Connection {
  public:
+  using CloseCallback = std::function<void()>;
+
   Connection(GraphSnapshotStore& snapshot_store,
-             std::shared_ptr<QueryProcessor> query_processor)
+             std::shared_ptr<QueryProcessor> query_processor,
+             CloseCallback close_callback = {})
       : snapshot_store_(snapshot_store),
         query_processor_(query_processor),
+        close_callback_(std::move(close_callback)),
         is_closed_(false) {}
   ~Connection() { Close(); }
 
@@ -204,6 +210,7 @@ class Connection {
   GraphSnapshotStore& snapshot_store_;
 
   std::shared_ptr<QueryProcessor> query_processor_;
+  CloseCallback close_callback_;
 
   std::atomic<bool> is_closed_{false};
 };

include/neug/main/connection_manager.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <stddef.h>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <vector>
@@ -39,7 +40,8 @@ class ConnectionManager {
         config_(config) {}
   ~ConnectionManager() { Close(); }
 
-  std::shared_ptr<Connection> CreateConnection();
+  std::shared_ptr<Connection> CreateConnection(
+      std::function<void()> close_callback = {});
 
   /**
    * @brief Close all connections managed by the connection manager.

include/neug/main/neug_db.h

@@ -250,6 +250,8 @@ class NeugDB {
    *
    * @note In READ_ONLY mode, multiple connections can be created.
    * @note In READ_WRITE mode, only one write connection is allowed.
+   * @note Embedded connections cannot be opened while a NeugDBService is
+   * active.
    * @note Connections share the planner instance for efficiency.
    *
    * @throws std::runtime_error if database is not open or closed
@@ -273,8 +275,8 @@ class NeugDB {
 
   /**
    * @brief Remove all connection from the database.
-   * @note This method is used to remove all connection when tp svc created, to
-   * remove the handle from the database.
+   * @note This method is used during database shutdown to close and remove
+   * managed connection handles.
    */
   void CloseAllConnection();
 
@@ -319,14 +321,22 @@ class NeugDB {
    * resets last_ts_ to 0.
    */
   void createCheckpoint(bool reopen = true);
+  void ValidateCanStartTPService() const;
+  void registerTPService();
+  void unregisterTPService() noexcept;
+  void registerAPConnection();
+  void unregisterAPConnection() noexcept;
 
   friend class NeugDBSession;
+  friend class PyDatabase;
   friend class neug::NeugDBService;
 
   timestamp_t last_compaction_ts_;
   timestamp_t last_ts_;
   // Configuration and settings
   std::atomic<bool> closed_;
+  // 0 means idle, -1 means TP service active, >0 counts AP connections.
+  std::atomic<int32_t> mode_state_;
   bool is_pure_memory_;
   int max_thread_num_;
   NeugDBConfig config_;

include/neug/server/neug_db_service.h

@@ -97,11 +97,19 @@ class NeugDBService {
    * @param db Reference to the NeuG database that will handle queries
    *
    * @note The database should be opened and ready before creating the service
+   * @throws std::runtime_error if embedded connections are active
    */
   NeugDBService(neug::NeugDB& db, const ServiceConfig& config = ServiceConfig())
       : db_(db), db_config_(db_.config()), compact_thread_running_(false) {
-    db_.CloseAllConnection();
-    init(config);
+    db_.registerTPService();
+    tp_service_registered_ = true;
+    try {
+      init(config);
+    } catch (...) {
+      db_.unregisterTPService();
+      tp_service_registered_ = false;
+      throw;
+    }
   }
 
   /**
@@ -251,6 +259,7 @@ class NeugDBService {
 
   std::thread compact_thread_;
   bool compact_thread_running_ = false;
+  bool tp_service_registered_ = false;
 
   std::atomic<bool> running_{false};
   std::mutex mtx_;

src/main/connection.cc

@@ -15,6 +15,8 @@
 
 #include "neug/main/connection.h"
 
+#include <exception>
+
 #include "neug/main/neug_db.h"
 #include "neug/main/query_request.h"
 #include "neug/utils/pb_utils.h"
@@ -34,10 +36,26 @@ std::string Connection::GetSchema() const {
 }
 
 void Connection::Close() {
-  if (is_closed_.load(std::memory_order_relaxed)) {
+  if (is_closed_.exchange(true, std::memory_order_acq_rel)) {
     LOG(WARNING) << "Connection is already closed.";
     return;
   }
+  struct CloseCallbackGuard {
+    CloseCallback& callback;
+    ~CloseCallbackGuard() noexcept {
+      if (!callback) {
+        return;
+      }
+      try {
+        callback();
+      } catch (const std::exception& e) {
+        LOG(ERROR) << "Connection close callback failed: " << e.what();
+      } catch (...) {
+        LOG(ERROR) << "Connection close callback failed.";
+      }
+    }
+  } close_callback_guard{close_callback_};
+
   LOG(INFO) << "Closing connection.";
 
   // Clean up all temporary schemas created during this session.
@@ -69,8 +87,6 @@ void Connection::Close() {
   if (!temp_edges.empty() || !temp_vertices.empty()) {
     query_processor_->clear_cache();
   }
-
-  is_closed_.store(true);
 }
 
 result<QueryResult> Connection::Query(const std::string& query_string,

src/main/connection_manager.cc

@@ -16,6 +16,7 @@
 #include "neug/main/connection_manager.h"
 
 #include <glog/logging.h>
+#include <algorithm>
 #include <ostream>
 #include "neug/config.h"
 #include "neug/main/connection.h"
@@ -35,20 +36,30 @@ void ConnectionManager::ConnectionManager::Close() {
   read_only_connections_.clear();
 }
 
-std::shared_ptr<Connection> ConnectionManager::CreateConnection() {
+std::shared_ptr<Connection> ConnectionManager::CreateConnection(
+    std::function<void()> close_callback) {
   std::lock_guard<std::mutex> lock(connection_mutex_);
   if (config_.mode == DBMode::READ_ONLY) {
-    auto conn = std::make_shared<Connection>(snapshot_store_, query_processor_);
+    read_only_connections_.erase(
+        std::remove_if(read_only_connections_.begin(),
+                       read_only_connections_.end(),
+                       [](const auto& conn) { return conn->IsClosed(); }),
+        read_only_connections_.end());
+    auto conn = std::make_shared<Connection>(snapshot_store_, query_processor_,
+                                             close_callback);
     read_only_connections_.push_back(conn);
     return conn;
   } else if (config_.mode == DBMode::READ_WRITE) {
+    if (read_write_connection_ && read_write_connection_->IsClosed()) {
+      read_write_connection_.reset();
+    }
     if (read_write_connection_) {
       LOG(ERROR) << "There is already a read-write connection constructed.";
       THROW_TX_STATE_CONFLICT(
           "There is already a read-write connection constructed.");
     }
-    read_write_connection_ =
-        std::make_shared<Connection>(snapshot_store_, query_processor_);
+    read_write_connection_ = std::make_shared<Connection>(
+        snapshot_store_, query_processor_, close_callback);
     return read_write_connection_;
   } else {
     THROW_RUNTIME_ERROR("Invalid mode.");
@@ -61,6 +72,7 @@ void ConnectionManager::RemoveConnection(std::shared_ptr<Connection> conn) {
     for (auto it = read_only_connections_.begin();
          it != read_only_connections_.end(); ++it) {
       if (*it == conn) {
+        conn->Close();
         read_only_connections_.erase(it);
         VLOG(10) << "Removed a read-only connection.";
         return;
@@ -69,6 +81,7 @@ void ConnectionManager::RemoveConnection(std::shared_ptr<Connection> conn) {
     LOG(ERROR) << "Connection not found in read-only connections.";
   } else if (config_.mode == DBMode::READ_WRITE) {
     if (read_write_connection_ == conn) {
+      conn->Close();
       read_write_connection_.reset();
       VLOG(10) << "Removed the read-write connection.";
       return;

src/main/neug_db.cc

@@ -46,6 +46,18 @@
 
 namespace neug {
 
+namespace {
+
+void ThrowTPServiceModeConflict(int32_t state) {
+  if (state > 0) {
+    THROW_RUNTIME_ERROR(
+        "Cannot start TP service while embedded connections are active.");
+  }
+  THROW_RUNTIME_ERROR("Cannot start TP service while TP service is active.");
+}
+
+}  // namespace
+
 inline std::string allocator_prefix(const std::string& allocator_dir,
                                     int thread_id) {
   return (std::filesystem::path(allocator_dir) /
@@ -79,6 +91,7 @@ NeugDB::NeugDB()
     : last_compaction_ts_(0),
       last_ts_(0),
       closed_(true),
+      mode_state_(0),
       is_pure_memory_(false),
       max_thread_num_(1) {}
 
@@ -178,7 +191,17 @@ void NeugDB::Close() {
 }
 
 std::shared_ptr<Connection> NeugDB::Connect() {
-  return connection_manager_->CreateConnection();
+  if (IsClosed() || !connection_manager_) {
+    THROW_RUNTIME_ERROR("NeugDB instance is not open.");
+  }
+  registerAPConnection();
+  try {
+    return connection_manager_->CreateConnection(
+        [this]() { unregisterAPConnection(); });
+  } catch (...) {
+    unregisterAPConnection();
+    throw;
+  }
 }
 
 void NeugDB::RemoveConnection(std::shared_ptr<Connection> conn) {
@@ -187,6 +210,61 @@ void NeugDB::RemoveConnection(std::shared_ptr<Connection> conn) {
 
 void NeugDB::CloseAllConnection() { connection_manager_->Close(); }
 
+void NeugDB::ValidateCanStartTPService() const {
+  const auto state = mode_state_.load(std::memory_order_acquire);
+  if (state != 0) {
+    ThrowTPServiceModeConflict(state);
+  }
+}
+
+void NeugDB::registerTPService() {
+  int32_t expected = 0;
+  if (mode_state_.compare_exchange_strong(expected, -1,
+                                          std::memory_order_acq_rel,
+                                          std::memory_order_acquire)) {
+    return;
+  }
+  ThrowTPServiceModeConflict(expected);
+}
+
+void NeugDB::unregisterTPService() noexcept {
+  int32_t expected = -1;
+  if (!mode_state_.compare_exchange_strong(expected, 0,
+                                           std::memory_order_acq_rel,
+                                           std::memory_order_acquire)) {
+    LOG(ERROR) << "Unexpected NeugDB mode state when unregistering TP service: "
+               << expected;
+  }
+}
+
+void NeugDB::registerAPConnection() {
+  auto state = mode_state_.load(std::memory_order_acquire);
+  while (true) {
+    if (state < 0) {
+      THROW_RUNTIME_ERROR(
+          "Cannot open embedded connection while TP service is active.");
+    }
+    if (mode_state_.compare_exchange_weak(state, state + 1,
+                                          std::memory_order_acq_rel,
+                                          std::memory_order_acquire)) {
+      return;
+    }
+  }
+}
+
+void NeugDB::unregisterAPConnection() noexcept {
+  auto state = mode_state_.load(std::memory_order_acquire);
+  while (state > 0) {
+    if (mode_state_.compare_exchange_weak(state, state - 1,
+                                          std::memory_order_acq_rel,
+                                          std::memory_order_acquire)) {
+      return;
+    }
+  }
+  LOG(ERROR) << "Unexpected NeugDB mode state when unregistering AP connection: "
+             << state;
+}
+
 void NeugDB::preprocessConfig() {
   if (config_.max_thread_num < 0) {
     THROW_INVALID_ARGUMENT_EXCEPTION(

src/server/neug_db_service.cc

@@ -65,6 +65,10 @@ NeugDBService::~NeugDBService() {
     hdl_mgr_->Stop();
     hdl_mgr_.reset();
   }
+  if (tp_service_registered_) {
+    db_.unregisterTPService();
+    tp_service_registered_ = false;
+  }
 }
 
 const ServiceConfig& NeugDBService::GetServiceConfig() const {