From: Adarsh Date: Mon, 4 May 2026 09:16:13 +0000 (+0000) Subject: rgw/s3vectors: add session operations in background X-Git-Url: http://git-server-git.apps.pok.os.sepia.ceph.com/?a=commitdiff_plain;h=85c218c296db1c01da1a93b72e5fd3f3ed6b3667;p=ceph.git rgw/s3vectors: add session operations in background * session is created on vector bucket creation or if any vector level operation is performed and no session for the relevent vector bucket exists * session creation is async and best effors. if operation is done when a session does not exists, it should not fail the operation, but the operation will not benefit for the cahced index and metadata * session is deleted when vector bucket is deleted, or when another operation failed with -ENOENT on the vector bucket (probably deleted by another RGW or concurrently by another request on the same RGW) Signed-off-by: Adarsh --- diff --git a/src/rgw/rgw_rest_s3vector.cc b/src/rgw/rgw_rest_s3vector.cc index f919b7fe7ef..edb086b2ced 100644 --- a/src/rgw/rgw_rest_s3vector.cc +++ b/src/rgw/rgw_rest_s3vector.cc @@ -8,6 +8,7 @@ #include "rgw_process_env.h" #include "common/async/yield_context.h" #include "rgw_arn.h" +#include "rgw_s3vector_background.h" #define dout_context g_ceph_context #define dout_subsys ceph_subsys_rgw @@ -71,6 +72,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if(op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -264,6 +268,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -321,6 +328,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -361,6 +371,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -396,6 +409,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -432,6 +448,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -486,6 +505,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -646,6 +668,9 @@ private: const rgw_bucket bucket_id(s->bucket_tenant, configuration.vector_bucket_name); op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; } } @@ -700,6 +725,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -754,6 +782,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -814,6 +845,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -849,6 +883,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -884,6 +921,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -920,6 +960,9 @@ private: std::unique_ptr bucket; op_ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (op_ret < 0) { + if (op_ret == -ENOENT) { + rgw::s3vector::notify_session_delete(this, bucket_id.name); + } ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } @@ -996,4 +1039,3 @@ RGWOp* RGWHandler_REST_s3Vector::op_post() { return nullptr; } - diff --git a/src/rgw/rgw_s3vector.cc b/src/rgw/rgw_s3vector.cc index 2d67c748b20..23a522fd48b 100644 --- a/src/rgw/rgw_s3vector.cc +++ b/src/rgw/rgw_s3vector.cc @@ -17,6 +17,24 @@ namespace rgw::s3vector { + struct LanceDBSessionConnHandle { + std::shared_ptr session_keepalive; + LanceDBConnection* conn = nullptr; + + explicit operator bool() const { + return conn != nullptr; + } + }; + + struct LanceDBSessionTableHandle { + LanceDBSessionConnHandle conn_handle; + LanceDBTable* table = nullptr; + + explicit operator bool() const { + return table != nullptr; + } + }; + // convert LanceDBError to linux error codes int lancedb_error_to_errno(LanceDBError err) { switch (err) { @@ -66,6 +84,31 @@ namespace rgw::s3vector { return conn; } + LanceDBSessionConnHandle connect_with_session_handle(DoutPrefixProvider* dpp, const std::string& vector_bucket_name){ + const auto dbname = fmt::format("/tmp/lancedb/{}", vector_bucket_name); + //get shared pointer to session for the bucket, if session doesn't exist, fallback to connect w/o session and trigger session creation for future connections + auto session_sp = rgw::s3vector::get_session(dpp, vector_bucket_name); + if (!session_sp) { + rgw::s3vector::notify_session_create(dpp, vector_bucket_name); + return LanceDBSessionConnHandle{ + .conn = connect(dpp, vector_bucket_name) + }; + } + LanceDBConnectBuilder* builder = lancedb_connect(dbname.c_str()); + builder = lancedb_connect_builder_session(builder, session_sp.get()); + LanceDBConnection* conn = lancedb_connect_builder_execute(builder); + if (!conn) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to connect using session to: " << dbname << " falling back to connect without session" << dendl; + return LanceDBSessionConnHandle{ + .conn = connect(dpp, vector_bucket_name) + }; // fallback to connect without session + } + return LanceDBSessionConnHandle{ + .session_keepalive = std::move(session_sp), + .conn = conn + }; + } + LanceDBTable* open_table(DoutPrefixProvider* dpp, const std::string& vector_bucket_name, const std::string& index_name) { LanceDBConnection* conn = connect(dpp, vector_bucket_name); if (!conn) { @@ -80,6 +123,23 @@ namespace rgw::s3vector { return table; } + LanceDBSessionTableHandle open_table_with_session_handle(DoutPrefixProvider* dpp, const std::string& vector_bucket_name, const std::string& index_name) { + auto conn_handle = connect_with_session_handle(dpp, vector_bucket_name); + if (!conn_handle) { + return {}; + } + LanceDBTable* table = lancedb_connection_open_table(conn_handle.conn, index_name.c_str()); + if (!table) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to open index: " << index_name << " in: " << vector_bucket_name << dendl; + lancedb_connection_free(conn_handle.conn); + return {}; + } + return LanceDBSessionTableHandle{ + .conn_handle = std::move(conn_handle), + .table = table + }; + } + // get creation time from the first version of a table // returns 0 on failure uint64_t get_table_creation_time(const LanceDBTable* table, DoutPrefixProvider* dpp) { @@ -737,6 +797,8 @@ namespace rgw::s3vector { lancedb_connection_free(conn); return -EIO; } + ldpp_dout(dpp, 20) << "INFO: deleting in-memory session (if it exists) for bucket: " << configuration.vector_bucket_name << dendl; + rgw::s3vector::notify_session_delete(dpp, configuration.vector_bucket_name); lancedb_connection_free(conn); return 0; } @@ -784,10 +846,11 @@ namespace rgw::s3vector { int create_vector_bucket(const create_vector_bucket_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { log_configuration(dpp, "CreateVectorBucket", configuration); - LanceDBConnection* conn = connect(dpp, configuration.vector_bucket_name); - if (!conn) { + auto conn_handle = connect_with_session_handle(dpp, configuration.vector_bucket_name); + if (!conn_handle) { return -EIO; } + LanceDBConnection* conn = conn_handle.conn; // verify connectivity by listing table names char** table_names; size_t name_count; @@ -946,13 +1009,16 @@ namespace rgw::s3vector { int put_vectors(const put_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { log_configuration(dpp, "PutVectors", configuration); - LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); - if (!table) { + auto table_handle = open_table_with_session_handle(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table_handle) { return -EIO; } + LanceDBTable* table = table_handle.table; + LanceDBConnection* conn = table_handle.conn_handle.conn; if (configuration.vectors.empty()) { ldpp_dout(dpp, 10) << "WARNING: s3vector no vectors provided" << dendl; lancedb_table_free(table); + lancedb_connection_free(conn); return 0; } @@ -960,6 +1026,7 @@ namespace rgw::s3vector { unsigned int dimension = 0; if (int ret = get_vector_dimension(configuration.index_name, table, dpp, dimension); ret < 0) { lancedb_table_free(table); + lancedb_connection_free(conn); return ret; } @@ -973,6 +1040,7 @@ namespace rgw::s3vector { << ". error: " << schema_error_message << dendl; lancedb_free_string(schema_error_message); lancedb_table_free(table); + lancedb_connection_free(conn); return lancedb_error_to_errno(result); } @@ -982,6 +1050,7 @@ namespace rgw::s3vector { << ". error: " << schema.status().ToString() << dendl; lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); lancedb_table_free(table); + lancedb_connection_free(conn); return -EINVAL; } @@ -1022,6 +1091,7 @@ namespace rgw::s3vector { if (num_rows == 0) { ldpp_dout(dpp, 1) << "ERROR: s3vector no valid vectors to insert" << dendl; lancedb_table_free(table); + lancedb_connection_free(conn); return -EINVAL; } @@ -1037,6 +1107,7 @@ namespace rgw::s3vector { ldpp_dout(dpp, 1) << "ERROR: s3vector failed to export record batch to C ABI: " << status.ToString() << dendl; if (c_schema.release) c_schema.release(&c_schema); lancedb_table_free(table); + lancedb_connection_free(conn); return -EINVAL; } @@ -1053,6 +1124,7 @@ namespace rgw::s3vector { if (c_array.release) c_array.release(&c_array); if (c_schema.release) c_schema.release(&c_schema); lancedb_table_free(table); + lancedb_connection_free(conn); return -EINVAL; } @@ -1075,11 +1147,13 @@ namespace rgw::s3vector { ldpp_dout(dpp, 1) << "ERROR: s3vector failed to upsert record batch to index. error: " << error_message << dendl; lancedb_free_string(error_message); lancedb_table_free(table); + lancedb_connection_free(conn); return lancedb_error_to_errno(result); } - lancedb_table_free(table); // we are not failing the operation if we cannot notify the background process on index update notify_index_update(dpp, configuration.vector_bucket_name, configuration.index_name); + lancedb_table_free(table); + lancedb_connection_free(conn); return 0; } @@ -1243,15 +1317,18 @@ namespace rgw::s3vector { int get_vectors(const get_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y, get_vectors_reply_t& reply) { log_configuration(dpp, "GetVectors", configuration); - LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); - if (!table) { + auto table_handle = open_table_with_session_handle(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table_handle) { return -EIO; } + LanceDBTable* table = table_handle.table; + LanceDBConnection* conn = table_handle.conn_handle.conn; LanceDBQuery* query = lancedb_query_new(table); if (!query) { ldpp_dout(dpp, 1) << "ERROR: s3vector failed to create query for index: " << configuration.index_name << dendl; lancedb_table_free(table); + lancedb_connection_free(conn); return -EIO; } @@ -1263,6 +1340,7 @@ namespace rgw::s3vector { lancedb_free_string(error_message); lancedb_query_free(query); lancedb_table_free(table); + lancedb_connection_free(conn); return lancedb_error_to_errno(result); } @@ -1281,6 +1359,7 @@ namespace rgw::s3vector { lancedb_free_string(error_message); lancedb_query_free(query); lancedb_table_free(table); + lancedb_connection_free(conn); return lancedb_error_to_errno(result); } @@ -1289,11 +1368,13 @@ namespace rgw::s3vector { ldpp_dout(dpp, 1) << "ERROR: s3vector failed to execute query on index: " << configuration.index_name << dendl; lancedb_query_free(query); lancedb_table_free(table); + lancedb_connection_free(conn); return -EIO; } auto ret = populate_vectors_from_query(dpp, query_result, reply.vectors, configuration.index_name, configuration.return_data, false, false); lancedb_table_free(table); + lancedb_connection_free(conn); return ret; } @@ -1469,10 +1550,12 @@ namespace rgw::s3vector { int delete_vectors(const delete_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { log_configuration(dpp, "DeleteVectors", configuration); - LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); - if (!table) { + auto table_handle = open_table_with_session_handle(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table_handle) { return -EIO; } + LanceDBTable* table = table_handle.table; + LanceDBConnection* conn = table_handle.conn_handle.conn; // build where filter for keys std::ostringstream oss; const auto keys_size = configuration.keys.size(); @@ -1489,6 +1572,7 @@ namespace rgw::s3vector { lancedb_free_string(error_message); } lancedb_table_free(table); + lancedb_connection_free(conn); return lancedb_error_to_errno(result); } @@ -1544,14 +1628,17 @@ namespace rgw::s3vector { int query_vectors(const query_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y, query_vectors_reply_t& reply) { log_configuration(dpp, "QueryVectors", configuration); - LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); - if (!table) { + auto table_handle = open_table_with_session_handle(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table_handle) { return -EIO; } + LanceDBTable* table = table_handle.table; + LanceDBConnection* conn = table_handle.conn_handle.conn; unsigned int table_dimension; if (int ret = get_vector_dimension(configuration.index_name, table, dpp, table_dimension); ret < 0) { lancedb_table_free(table); + lancedb_connection_free(conn); return ret; } @@ -1559,6 +1646,8 @@ namespace rgw::s3vector { if (table_dimension != query_dimension) { ldpp_dout(dpp, 1) << "ERROR: s3vector query vector dimension (" << query_dimension << ") does not match index: " << configuration.index_name << " vector dimension (" << table_dimension << ")" << dendl; + lancedb_table_free(table); + lancedb_connection_free(conn); return -EINVAL; } @@ -1566,6 +1655,7 @@ namespace rgw::s3vector { if (!query) { ldpp_dout(dpp, 1) << "ERROR: s3vector failed to create vector query for index: " << configuration.index_name << dendl; lancedb_table_free(table); + lancedb_connection_free(conn); return -EIO; } @@ -1577,6 +1667,7 @@ namespace rgw::s3vector { lancedb_free_string(error_message); lancedb_vector_query_free(query); lancedb_table_free(table); + lancedb_connection_free(conn); return lancedb_error_to_errno(result); } @@ -1585,6 +1676,7 @@ namespace rgw::s3vector { lancedb_free_string(error_message); lancedb_vector_query_free(query); lancedb_table_free(table); + lancedb_connection_free(conn); return lancedb_error_to_errno(result); } @@ -1593,6 +1685,7 @@ namespace rgw::s3vector { lancedb_free_string(error_message); lancedb_vector_query_free(query); lancedb_table_free(table); + lancedb_connection_free(conn); return lancedb_error_to_errno(result); } @@ -1601,14 +1694,15 @@ namespace rgw::s3vector { ldpp_dout(dpp, 1) << "ERROR: s3vector failed to execute query on index: " << configuration.index_name << dendl; lancedb_vector_query_free(query); lancedb_table_free(table); + lancedb_connection_free(conn); return -EIO; } int ret = populate_vectors_from_query(dpp, query_result, reply.vectors, configuration.index_name, false, configuration.return_distance, true); reply.distance_metric = get_table_distance_metric(table, dpp); lancedb_table_free(table); + lancedb_connection_free(conn); return ret; } } - diff --git a/src/rgw/rgw_s3vector_background.cc b/src/rgw/rgw_s3vector_background.cc index 63fbfabbf54..d7146eebf0f 100644 --- a/src/rgw/rgw_s3vector_background.cc +++ b/src/rgw/rgw_s3vector_background.cc @@ -18,6 +18,7 @@ #include #include #include "rgw_sal.h" +#include "lancedb.h" #define dout_subsys ceph_subsys_rgw @@ -25,11 +26,14 @@ namespace rgw::s3vector { class Manager : public DoutPrefixProvider { public: + //message_t -> pass in empty index name for session messages (can extend to per table sessions in the future if needed) using table_name_t = std::pair; // pair of vector bucket name and index name struct message_t { enum class Op { UPDATE, - REMOVE + REMOVE, + SESSION_CREATE, + SESSION_DELETE }; message_t(const std::string& bucket_name, const std::string& index_name, Op _type) : table_name(bucket_name, index_name), type(_type) {} @@ -50,6 +54,16 @@ private: boost::asio::executor_work_guard work_guard; std::vector workers; rgw::sal::Driver* const driver; + struct LanceDBSessionDeleter { + void operator()(LanceDBSession* session) const { + if(session) { + lancedb_session_free(session); + } + } + }; + using SessionPtr = std::shared_ptr; + ceph::shared_mutex sessions_mutex = ceph::make_shared_mutex("s3vector::Manager::sessions_mutex"); + std::unordered_map sessions; std::unordered_map> tables; MessageQueue messages; static constexpr auto idle_sleep = std::chrono::milliseconds(1000); // 1s @@ -128,37 +142,78 @@ private: return 0; } - // process all work items - void process_tables(boost::asio::yield_context yield) { - ldpp_dout(this, 5) << "INFO: start processing tables" << dendl; + // process all work items for tables and sessions + void process_messages(boost::asio::yield_context yield) { + ldpp_dout(this, 5) << "INFO: manager started. starting to process messages for background table and session operations" << dendl; while (!shutdown) { std::vector tables_to_process; const auto message_count = messages.consume_all([&tables_to_process, this](auto message) { std::unique_ptr message_guard(message); const auto table_name = std::move(message->table_name); - if (message->type == message_t::Op::REMOVE) { - ldpp_dout(this, 20) << "INFO: received remove message for table: " << table_name.first << "." << table_name.second << dendl; - tables.erase(table_name); - return; - } - auto [it, inserted] = tables.emplace(table_name, ceph::coarse_real_clock::now()); - if (inserted) { - ldpp_dout(this, 20) << "INFO: will try to process new table: " << table_name.first << "." << table_name.second << dendl; - tables_to_process.push_back(table_name); - return; - } - const auto now = ceph::coarse_real_clock::now(); - const auto time_since_last_process = now - it->second; - if (time_since_last_process > std::chrono::milliseconds(5000)) { - ldpp_dout(this, 20) << "INFO: will try to process table: " << table_name.first << "." << table_name.second << - ". " << time_since_last_process << " passed since last processing" << dendl; - it->second = now; - tables_to_process.push_back(table_name); - } else { - ldpp_dout(this, 20) << "INFO: will skip processing table: " << table_name.first << "." << table_name.second << - ". only " << time_since_last_process << " passed since last processing" << dendl; + switch(message->type) { + case message_t::Op::REMOVE: + ldpp_dout(this, 20) << "INFO: received remove message for table: " << table_name.first << "." << table_name.second << dendl; + tables.erase(table_name); + return; + case message_t::Op::UPDATE: + { + ldpp_dout(this, 20) << "INFO: received update message for table: " << table_name.first << "." << table_name.second << dendl; + auto [it, inserted] = tables.emplace(table_name, ceph::coarse_real_clock::now()); + if (inserted) { + ldpp_dout(this, 20) << "INFO: will try to process new table: " << table_name.first << "." << table_name.second << dendl; + tables_to_process.push_back(table_name); + return; + } + const auto now = ceph::coarse_real_clock::now(); + const auto time_since_last_process = now - it->second; + if (time_since_last_process > std::chrono::milliseconds(5000)) { + ldpp_dout(this, 20) << "INFO: will try to process table: " << table_name.first << "." << table_name.second << + ". " << time_since_last_process << " passed since last processing" << dendl; + it->second = now; + tables_to_process.push_back(table_name); + } else { + ldpp_dout(this, 20) << "INFO: will skip processing table: " << table_name.first << "." << table_name.second << + ". only " << time_since_last_process << " passed since last processing" << dendl; + } + return; + } + case message_t::Op::SESSION_CREATE: + { + ldpp_dout(this, 20) << "INFO: received session create message for bucket: " << table_name.first << dendl; + std::unique_lock l(sessions_mutex); + if (sessions.find(table_name.first) == sessions.end()) { + //create session if not exist, otherwise just ignore + //Can define session options in the future if needed, for now just create with default options for cache sizes + LanceDBSession* session = lancedb_session_new(nullptr); + if (session) { + sessions[table_name.first] = SessionPtr(session, LanceDBSessionDeleter()); + ldpp_dout(this, 20) << "INFO: created session for bucket: " << table_name.first << dendl; + } + else { + ldpp_dout(this, 1) << "ERROR: failed to create session for bucket: " << table_name.first << dendl; + } + return; + } + ldpp_dout(this, 20) << "INFO: session already exists for bucket: " << table_name.first << dendl; + return; + } + case message_t::Op::SESSION_DELETE: + { + ldpp_dout(this, 20) << "INFO: received session delete message for bucket: " << table_name.first << dendl; + std::unique_lock l(sessions_mutex); + if (sessions.erase(table_name.first) > 0) { + ldpp_dout(this, 20) << "INFO: deleted session for bucket: " << table_name.first << dendl; + } else { + ldpp_dout(this, 20) << "INFO: session doesn't exist for bucket: " << table_name.first << dendl; + } + return; + } + default: + ldpp_dout(this, 1) << "ERROR: received message with unknown type for bucket: " << table_name.first << " index: " << table_name.second << dendl; + return; } }); + tokens_waiter tw(this); for (const auto& table_name : tables_to_process) { // start processing a table @@ -178,15 +233,16 @@ private: // wait for all pending work to finish tw.async_wait(yield); } + if (message_count == 0) { // if no messages, sleep for a while before checking again - ldpp_dout(this, 20) << "INFO: no tables to process" << dendl; + ldpp_dout(this, 20) << "INFO: no messages to process" << dendl; async_sleep(yield, idle_sleep); } } - ldpp_dout(this, 5) << "INFO: manager stopped. done processing all tables" << dendl; - } - + ldpp_dout(this, 5) << "INFO: manager stopped. done processing all table and session operations" << dendl; + } + public: ~Manager() { @@ -219,7 +275,7 @@ public: void init() { boost::asio::spawn(make_strand(io_context), std::allocator_arg, make_stack_allocator(), [this](boost::asio::yield_context yield) { - process_tables(yield); + process_messages(yield); }, [] (std::exception_ptr eptr) { if (eptr) std::rethrow_exception(eptr); }); @@ -255,6 +311,30 @@ public: return false; } + bool notify_session(const DoutPrefixProvider* dpp, const std::string& bucket_name, message_t::Op op) { + if (shutdown) { + ldpp_dout(dpp, 1) << "ERROR: failed to notify s3vectors manager about session: manager is shutting down" << dendl; + return false; + } + auto message_guard = std::make_unique(bucket_name, "", op); + if (messages.push(message_guard.get())) { + std::ignore = message_guard.release(); // ownership transferred to the queue + ldpp_dout(dpp, 20) << "INFO: notified s3vectors manager about session" << dendl; + return true; + } + ldpp_dout(dpp, 1) << "ERROR: failed to notify s3vectors manager about session: queue is full" << dendl; + return false; + } + + std::shared_ptr get_session(const std::string& bucket_name) { + std::shared_lock l(sessions_mutex); + auto it = sessions.find(bucket_name); + if (it == sessions.end()) { + return nullptr; + } + return it->second; + } + Manager(CephContext* _cct, rgw::sal::Driver* _driver) : cct(_cct), work_guard(boost::asio::make_work_guard(io_context)), @@ -305,5 +385,29 @@ bool notify_index_remove(const DoutPrefixProvider* dpp, const std::string& bucke return s_manager->notify_index(dpp, bucket_name, index_name, Manager::message_t::Op::REMOVE); } +std::shared_ptr get_session(const DoutPrefixProvider* dpp, const std::string& bucket_name) { + if (!s_manager) { + ldpp_dout(dpp, 1) << "ERROR: failed to get LanceDB session for bucket: manager is not initialized" << dendl; + return nullptr; + } + return s_manager->get_session(bucket_name); +} + +bool notify_session_create(const DoutPrefixProvider* dpp, const std::string& bucket_name) { + if (!s_manager) { + ldpp_dout(dpp, 1) << "ERROR: failed to notify s3vectors manager about session creation: manager is not initialized" << dendl; + return false; + } + return s_manager->notify_session(dpp, bucket_name, Manager::message_t::Op::SESSION_CREATE); +} + +bool notify_session_delete(const DoutPrefixProvider* dpp, const std::string& bucket_name) { + if (!s_manager) { + ldpp_dout(dpp, 1) << "ERROR: failed to notify s3vectors manager about session deletion: manager is not initialized" << dendl; + return false; + } + return s_manager->notify_session(dpp, bucket_name, Manager::message_t::Op::SESSION_DELETE); +} + } // namespace rgw::s3vector diff --git a/src/rgw/rgw_s3vector_background.h b/src/rgw/rgw_s3vector_background.h index b6c37c3f7f0..8f3e2b990d1 100644 --- a/src/rgw/rgw_s3vector_background.h +++ b/src/rgw/rgw_s3vector_background.h @@ -3,12 +3,14 @@ #pragma once +#include #include namespace rgw::sal { class Driver; } class DoutPrefixProvider; +struct LanceDBSession; namespace rgw::s3vector { bool init(const DoutPrefixProvider* dpp, rgw::sal::Driver* driver); @@ -19,5 +21,11 @@ namespace rgw::s3vector { bool notify_index_update(const DoutPrefixProvider* dpp, const std::string& bucket_name, const std::string& index_name); // update whenever a index is removed bool notify_index_remove(const DoutPrefixProvider* dpp, const std::string& bucket_name, const std::string& index_name); + // get LanceDB session for a bucket, returns nullptr if session doesn't exist or manager is not initialized + std::shared_ptr get_session(const DoutPrefixProvider* dpp, const std::string& bucket_name); + // notify manager for session creation + bool notify_session_create(const DoutPrefixProvider* dpp, const std::string& bucket_name); + // notify manager for session deletion + bool notify_session_delete(const DoutPrefixProvider* dpp, const std::string& bucket_name); }