From: Yuval Lifshitz Date: Tue, 17 Mar 2026 14:12:42 +0000 (+0000) Subject: rgw/s3vectors: skeleton for background process X-Git-Url: http://git-server-git.apps.pok.os.sepia.ceph.com/?a=commitdiff_plain;h=b136444e8655a936d9959e035d924536d2395672;p=ceph.git rgw/s3vectors: skeleton for background process Signed-off-by: Yuval Lifshitz --- diff --git a/src/rgw/CMakeLists.txt b/src/rgw/CMakeLists.txt index 98eb649f7f3..2dcb207ca32 100644 --- a/src/rgw/CMakeLists.txt +++ b/src/rgw/CMakeLists.txt @@ -150,7 +150,8 @@ set(librgw_common_srcs rgw_bucket_sync.cc rgw_rest_restore.cc rgw_rest_s3vector.cc - rgw_s3vector.cc) + rgw_s3vector.cc + rgw_s3vector_background.cc) list(APPEND librgw_common_srcs driver/immutable_config/store.cc diff --git a/src/rgw/rgw_appmain.cc b/src/rgw/rgw_appmain.cc index 56dc4059454..a7858fa471b 100644 --- a/src/rgw/rgw_appmain.cc +++ b/src/rgw/rgw_appmain.cc @@ -77,6 +77,7 @@ #endif #include "rgw_lua_background.h" #include "services/svc_zone.h" +#include "rgw_s3vector_background.h" #ifdef HAVE_SYS_PRCTL_H #include @@ -518,6 +519,7 @@ int rgw::AppMain::init_frontends2(RGWLib* rgwlib) derr << "ERROR: failed to register to service map: " << cpp_strerror(-r) << dendl; /* ignore error */ } + s3vector::init(dpp, env.driver); #ifdef WITH_RADOSGW_RADOS if (env.driver->get_name() == "rados") { @@ -655,6 +657,8 @@ void rgw::AppMain::shutdown(std::function finalize_async_signals) lua_background->shutdown(); } + s3vector::shutdown(); + env.driver->shutdown(); // Do this before closing storage so requests don't try to call into // closed storage. diff --git a/src/rgw/rgw_s3vector.cc b/src/rgw/rgw_s3vector.cc index fe449ada98b..0985115db3a 100644 --- a/src/rgw/rgw_s3vector.cc +++ b/src/rgw/rgw_s3vector.cc @@ -11,6 +11,7 @@ #include #include #include +#include "rgw_s3vector_background.h" #define dout_subsys ceph_subsys_rgw @@ -1075,6 +1076,8 @@ namespace rgw::s3vector { 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); return 0; } diff --git a/src/rgw/rgw_s3vector_background.cc b/src/rgw/rgw_s3vector_background.cc new file mode 100644 index 00000000000..d2dcda0a197 --- /dev/null +++ b/src/rgw/rgw_s3vector_background.cc @@ -0,0 +1,271 @@ +// -*- mode:C++; tab-width:8; c-basic-offset:2; indent-tabs-mode:nil -*- +// vim: ts=8 sw=2 sts=2 expandtab + +#include +#include +#include +#include +#include +#include +#include +#include +#include "common/ceph_time.h" +#include "common/dout.h" +#include +#include +#include "common/async/yield_waiter.h" +#include +#include +#include +#include "rgw_sal.h" + +#define dout_subsys ceph_subsys_rgw + +namespace rgw::s3vector { + +class Manager : public DoutPrefixProvider { + + // use mmap/mprotect to allocate 128k coroutine stacks + auto make_stack_allocator() { + return boost::context::protected_fixedsize_stack{128*1024}; + } + using table_name_t = std::pair; // pair of vector bucket name and index name + using MessageQueue = boost::lockfree::queue>; + using Executor = boost::asio::io_context::executor_type; + bool shutdown = false; + CephContext* const cct; + boost::asio::io_context io_context; + boost::asio::executor_work_guard work_guard; + std::vector workers; + rgw::sal::Driver* const driver; + std::unordered_map> tables; + MessageQueue messages; + + CephContext *get_cct() const override { return cct; } + unsigned get_subsys() const override { return dout_subsys; } + std::ostream& gen_prefix(std::ostream& out) const override { return out << "s3vectors manager: "; } + + class tokens_waiter { + size_t pending_tokens = 0; + DoutPrefixProvider* const dpp; + ceph::async::yield_waiter waiter; + + public: + class token{ + tokens_waiter* tw; + public: + token(const token& other) = delete; + token(token&& other) : tw(other.tw) { + other.tw = nullptr; // mark as moved + } + token& operator=(const token& other) = delete; + token(tokens_waiter* _tw) : tw(_tw) { + ++tw->pending_tokens; + } + + ~token() { + if (!tw) { + return; // already moved + } + --tw->pending_tokens; + if (tw->pending_tokens == 0 && tw->waiter) { + tw->waiter.complete(boost::system::error_code{}); + } + } + }; + + tokens_waiter(DoutPrefixProvider* _dpp) : dpp(_dpp) {} + tokens_waiter(const tokens_waiter& other) = delete; + tokens_waiter& operator=(const tokens_waiter& other) = delete; + + void async_wait(boost::asio::yield_context yield) { + if (pending_tokens == 0) { + return; + } + ldpp_dout(dpp, 20) << "INFO: tokens waiter is waiting on " << + pending_tokens << " tokens" << dendl; + boost::system::error_code ec; + waiter.async_wait(yield[ec]); + ldpp_dout(dpp, 20) << "INFO: tokens waiter finished waiting for all tokens" << dendl; + } + }; + + void async_sleep(boost::asio::yield_context yield, const std::chrono::milliseconds& duration) { + using Clock = ceph::coarse_mono_clock; + using Timer = boost::asio::basic_waitable_timer, Executor>; + Timer timer(io_context); + timer.expires_after(duration); + boost::system::error_code ec; + timer.async_wait(yield[ec]); + if (ec) { + ldpp_dout(this, 1) << "ERROR: async_sleep failed with error: " << ec.message() << dendl; + } + } + + // processing of a specific table + int process_table(const table_name_t& table_name, boost::asio::yield_context yield) { + // TODO: check if processing is needed based on unindexed rows stats and skip if not needed + // TODO: check if processign already started for the table and skip if yes + // TODO: implement actual lancedb table processing logic here + // for PoC just sleep for some time to simulate processing + ldpp_dout(this, 20) << "INFO: started processing table: " << table_name.first << "." << table_name.second << dendl; + async_sleep(yield, std::chrono::milliseconds(1000)); + ldpp_dout(this, 20) << "INFO: done processing table: " << table_name.first << "." << table_name.second << dendl; + return 0; + } + + // process all work items + void process_tables(boost::asio::yield_context yield) { + ldpp_dout(this, 5) << "INFO: start processing tables" << 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); + 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; + } + }); + tokens_waiter tw(this); + for (const auto& table_name : tables_to_process) { + // start processing a table + tokens_waiter::token token(&tw); + boost::asio::spawn(make_strand(io_context), std::allocator_arg, make_stack_allocator(), + [this, table_name](boost::asio::yield_context yield) { + const int rc = process_table(table_name, yield); + if (rc < 0) { + ldpp_dout(this, 1) << "ERROR: failed to process table: " << table_name.first << "." << table_name.second << " with error code: " << rc << dendl; + tables[table_name] = ceph::coarse_real_clock::now() - std::chrono::milliseconds(5000); // set last processed time to past to allow retry on next loop + } + }, [] (std::exception_ptr eptr) { + if (eptr) std::rethrow_exception(eptr); + }); + } + if (!tables_to_process.empty()) { + // 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; + async_sleep(yield, std::chrono::milliseconds(1000)); + } + } + ldpp_dout(this, 5) << "INFO: manager stopped. done processing all tables" << dendl; + } + +public: + + ~Manager() = default; + + void stop() { + ldpp_dout(this, 5) << "INFO: manager received stop signal. shutting down..." << dendl; + shutdown = true; + work_guard.reset(); + for (auto& worker : workers) { + if (worker.joinable()) { + // try graceful shutdown first + auto future = std::async(std::launch::async, [&worker]() {worker.join();}); + if (future.wait_for(std::chrono::milliseconds(1000)) == std::future_status::timeout) { + // force stop if graceful shutdown takes too long + if (!io_context.stopped()) { + ldpp_dout(this, 5) << "INFO: force shutdown of manager" << dendl; + io_context.stop(); + } + worker.join(); + } + } + } + ldpp_dout(this, 5) << "INFO: manager shutdown ended" << dendl; + } + + void init() { + boost::asio::spawn(make_strand(io_context), std::allocator_arg, make_stack_allocator(), + [this](boost::asio::yield_context yield) { + process_tables(yield); + }, [] (std::exception_ptr eptr) { + if (eptr) std::rethrow_exception(eptr); + }); + + // start the worker threads to do the actual queue processing + // TODO: use multiple threads + workers.emplace_back(std::thread([this]() { + ceph_pthread_setname("notif-worker"); + try { + ldpp_dout(this, 10) << "INFO: worker started" << dendl; + io_context.run(); + ldpp_dout(this, 10) << "INFO: worker ended" << dendl; + } catch (const std::exception& err) { + ldpp_dout(this, 1) << "ERROR: worker failed with error: " << err.what() << dendl; + throw err; + } + })); + ldpp_dout(this, 10) << "INfO: started manager" << dendl; + } + + bool notify_index_update(const DoutPrefixProvider* dpp, const std::string& bucket_name, const std::string& index_name) { + if (shutdown) { + ldpp_dout(dpp, 1) << "ERROR: failed to notify s3vectors manager about index update: manager is shutting down" << dendl; + return false; + } + auto message_guard = std::make_unique(bucket_name, index_name); + 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 index update" << dendl; + return true; + } + ldpp_dout(dpp, 1) << "ERROR: failed to notify s3vectors manager about index update: queue is full" << dendl; + return false; + } + + Manager(CephContext* _cct, rgw::sal::Driver* _driver) : + cct(_cct), + work_guard(boost::asio::make_work_guard(io_context)), + driver(_driver), + messages(8192) + {} +}; + +std::unique_ptr s_manager; + +bool init(const DoutPrefixProvider* dpp, rgw::sal::Driver* driver) { + if (s_manager) { + ldpp_dout(dpp, 1) << "ERROR: failed to init s3vectors manager: already exists" << dendl; + return false; + } + s_manager = std::make_unique(dpp->get_cct(), driver); + s_manager->init(); + return true; +} + +void shutdown() { + if (!s_manager) return; + s_manager->stop(); + s_manager.reset(); +} + +bool notify_index_update(const DoutPrefixProvider* dpp, const std::string& bucket_name, const std::string& index_name) { + if (!s_manager) { + ldpp_dout(dpp, 1) << "ERROR: failed to notify s3vectors manager about table update: manager is not initialized" << dendl; + return false; + } + return s_manager->notify_index_update(dpp, bucket_name, index_name); +} + +} // namespace rgw::s3vector + diff --git a/src/rgw/rgw_s3vector_background.h b/src/rgw/rgw_s3vector_background.h new file mode 100644 index 00000000000..2ee8d72a2e9 --- /dev/null +++ b/src/rgw/rgw_s3vector_background.h @@ -0,0 +1,19 @@ +// -*- mode:C++; tab-width:8; c-basic-offset:2; indent-tabs-mode:nil -*- +// vim: ts=8 sw=2 sts=2 expandtab ft=cpp + +#pragma once + +#include + +namespace rgw::sal { + class Driver; +} +class DoutPrefixProvider; + +namespace rgw::s3vector { + bool init(const DoutPrefixProvider* dpp, rgw::sal::Driver* driver); + void shutdown(); + // update whenever new vectors are added to an index + bool notify_index_update(const DoutPrefixProvider* dpp, const std::string& bucket_name, const std::string& index_name); +} + diff --git a/src/test/rgw/s3vectors/s3vector_test.py b/src/test/rgw/s3vectors/s3vector_test.py index acfa9f1641c..cfc508cfe8b 100644 --- a/src/test/rgw/s3vectors/s3vector_test.py +++ b/src/test/rgw/s3vectors/s3vector_test.py @@ -1,16 +1,11 @@ import logging -import json -import tempfile import random -import socket import time import threading import subprocess import os -import stat import string from datetime import datetime, timedelta, timezone -from datetime import datetime, timezone import pytest import boto3 from botocore.config import Config @@ -653,6 +648,49 @@ def test_put_vectors(): _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) +def update_vectors_thread(conn, bucket_name, thread_id): + index_name = 'test-index-'+str(thread_id) + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + dimension = 128 + num_vectors = 10 + vectors = generate_vectors(num_vectors, dimension) + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + time.sleep(0.5) + num_vectors = 20 + vectors = generate_vectors(num_vectors, dimension) + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + time.sleep(5) + num_vectors = 30 + vectors = generate_vectors(num_vectors, dimension) + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + # Get vectors with returnData=True + vector_ids = [f'vec-{i}' for i in range(num_vectors)] + returned_vectors = verify_get_vectors(conn, bucket_name, index_name, vector_ids, expected_dimension=dimension) + + +@pytest.mark.vector_test +def test_update_vectors(): + conn = connection() + bucket_name = gen_bucket_name() + result = conn.create_vector_bucket(vectorBucketName=bucket_name) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + num_indexes = 10 + threads = [] + for i in range(num_indexes): + t = threading.Thread(target=update_vectors_thread, args=(conn, bucket_name, i)) + t.start() + threads.append(t) + for t in threads: + t.join() + + # cleanup + _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) + + @pytest.mark.vector_test def test_put_vectors_dimension_mismatch(): conn = connection()