Skip to content

Commit

Permalink
fix: comments
Browse files Browse the repository at this point in the history
Signed-off-by: usamoi <usamoi@outlook.com>
  • Loading branch information
usamoi committed Mar 12, 2024
1 parent ebd2478 commit 426d1e3
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 9 deletions.
2 changes: 0 additions & 2 deletions crates/service/src/worker.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
use crate::instance::*;
use crate::version::Version;
use arc_swap::ArcSwap;
use base::index::*;
use base::search::*;
Expand Down Expand Up @@ -31,7 +30,6 @@ impl Worker {
});
let protect = WorkerProtect { startup, indexes };
sync_dir(&path);
Version::write(path.join("metadata"));
Arc::new(Worker {
path,
protect: Mutex::new(protect),
Expand Down
44 changes: 41 additions & 3 deletions src/index/am.rs
Original file line number Diff line number Diff line change
Expand Up @@ -143,12 +143,42 @@ pub unsafe extern "C" fn ambuild(
index_relation,
Some((heap_relation, index_info, result.as_ptr())),
);
make_well_formed(index_relation);
result.into_pg()
}

#[pgrx::pg_guard]
pub unsafe extern "C" fn ambuildempty(index_relation: pgrx::pg_sys::Relation) {
am_build::build(index_relation, None);
pub unsafe extern "C" fn ambuildempty(_index: pgrx::pg_sys::Relation) {}

unsafe fn make_well_formed(index_relation: pgrx::pg_sys::Relation) {
unsafe {
let meta_buffer = pgrx::pg_sys::ReadBuffer(index_relation, 0xFFFFFFFF /* P_NEW */);
pgrx::pg_sys::LockBuffer(meta_buffer, pgrx::pg_sys::BUFFER_LOCK_EXCLUSIVE as _);
assert!(pgrx::pg_sys::BufferGetBlockNumber(meta_buffer) == 0);
let state = pgrx::pg_sys::GenericXLogStart(index_relation);
let meta_page = pgrx::pg_sys::GenericXLogRegisterBuffer(
state,
meta_buffer,
pgrx::pg_sys::GENERIC_XLOG_FULL_IMAGE as _,
);
meta_page.write_bytes(0, pgrx::pg_sys::BLCKSZ as usize);
pgrx::pg_sys::GenericXLogFinish(state);
pgrx::pg_sys::UnlockReleaseBuffer(meta_buffer);
}
}

unsafe fn check_well_formed(index_relation: pgrx::pg_sys::Relation) {
if !test_well_formed(index_relation) {
am_build::build(index_relation, None);
make_well_formed(index_relation);
}
}

unsafe fn test_well_formed(index_relation: pgrx::pg_sys::Relation) -> bool {
pgrx::pg_sys::RelationGetNumberOfBlocksInFork(
index_relation,
pgrx::pg_sys::ForkNumber_MAIN_FORKNUM,
) == 1
}

#[pgrx::pg_guard]
Expand All @@ -162,6 +192,7 @@ pub unsafe extern "C" fn aminsert(
_index_unchanged: bool,
_index_info: *mut pgrx::pg_sys::IndexInfo,
) -> bool {
check_well_formed(index_relation);
let oid = (*index_relation).rd_id;
let id = get_handle(oid);
let vector = from_datum(*values.add(0), *is_null.add(0));
Expand All @@ -177,6 +208,7 @@ pub unsafe extern "C" fn ambeginscan(
n_keys: std::os::raw::c_int,
n_orderbys: std::os::raw::c_int,
) -> pgrx::pg_sys::IndexScanDesc {
check_well_formed(index_relation);
assert!(n_keys == 0);
assert!(n_orderbys == 1);
am_scan::make_scan(index_relation)
Expand Down Expand Up @@ -218,6 +250,9 @@ pub unsafe extern "C" fn ambulkdelete(
callback: pgrx::pg_sys::IndexBulkDeleteCallback,
callback_state: *mut std::os::raw::c_void,
) -> *mut pgrx::pg_sys::IndexBulkDeleteResult {
if !test_well_formed((*info).index) {
pgrx::warning!("The vector index is not initialized.");
}
let oid = (*(*info).index).rd_id;
let id = get_handle(oid);
if let Some(callback) = callback {
Expand All @@ -234,9 +269,12 @@ pub unsafe extern "C" fn ambulkdelete(

#[pgrx::pg_guard]
pub unsafe extern "C" fn amvacuumcleanup(
_info: *mut pgrx::pg_sys::IndexVacuumInfo,
info: *mut pgrx::pg_sys::IndexVacuumInfo,
_stats: *mut pgrx::pg_sys::IndexBulkDeleteResult,
) -> *mut pgrx::pg_sys::IndexBulkDeleteResult {
if !test_well_formed((*info).index) {
pgrx::warning!("The vector index is not initialized.");
}
let result = pgrx::PgBox::<pgrx::pg_sys::IndexBulkDeleteResult>::alloc0();
result.into_pg()
}
8 changes: 4 additions & 4 deletions src/index/hook_maintain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ pub unsafe fn maintain_index_in_object_access(
pgrx::pg_sys::ObjectAccessType_OAT_DROP => {
let search = pgrx::pg_catalog::PgClass::search_reloid(object_id).unwrap();
if let Some(pg_class) = search.get() {
if let Some(()) = check(pg_class) {
if let Some(()) = check_vector_index(pg_class) {
let handle = get_handle(object_id);
let mut t = TRANSACTION.borrow_mut();
match t.index.get(&handle) {
Expand Down Expand Up @@ -110,7 +110,7 @@ pub unsafe fn maintain_index_in_object_access(
}
}

fn check(pg_class: pgrx::pg_catalog::PgClass<'_>) -> Option<()> {
fn check_vector_index(pg_class: pgrx::pg_catalog::PgClass<'_>) -> Option<()> {
if pg_class.relkind() != pgrx::pg_catalog::PgClassRelkind::Index {
return None;
}
Expand All @@ -124,10 +124,10 @@ fn check(pg_class: pgrx::pg_catalog::PgClass<'_>) -> Option<()> {
return None;
}
// probably a vector index, so enter a slow path to ensure it
slow_path(pg_am)
check_vector_index_slow_path(pg_am)
}

fn slow_path(pg_am: pgrx::pg_catalog::PgAm<'_>) -> Option<()> {
fn check_vector_index_slow_path(pg_am: pgrx::pg_catalog::PgAm<'_>) -> Option<()> {
let amhandler = pg_am.amhandler();
let mut flinfo = unsafe {
let mut flinfo = pgrx::pg_sys::FmgrInfo::default();
Expand Down

0 comments on commit 426d1e3

Please sign in to comment.