Add rb_concurrent_set_find

This commit is contained in:
Peter Zhu 2025-07-14 11:01:34 -04:00
parent 9ef482bd13
commit f05ee26a1f
2 changed files with 57 additions and 0 deletions

View file

@ -175,6 +175,61 @@ concurrent_set_try_resize(VALUE old_set_obj, VALUE *set_obj_ptr)
}
}
VALUE
rb_concurrent_set_find(VALUE *set_obj_ptr, VALUE key)
{
RUBY_ASSERT(key >= CONCURRENT_SET_SPECIAL_VALUE_COUNT);
VALUE set_obj;
retry:
set_obj = RUBY_ATOMIC_VALUE_LOAD(*set_obj_ptr);
RUBY_ASSERT(set_obj);
struct concurrent_set *set = RTYPEDDATA_GET_DATA(set_obj);
struct concurrent_set_probe probe;
VALUE hash = set->funcs->hash(key);
int idx = concurrent_set_probe_start(&probe, set, hash);
while (true) {
struct concurrent_set_entry *entry = &set->entries[idx];
VALUE curr_key = RUBY_ATOMIC_VALUE_LOAD(entry->key);
switch (curr_key) {
case CONCURRENT_SET_EMPTY:
return 0;
case CONCURRENT_SET_DELETED:
break;
case CONCURRENT_SET_MOVED:
// Wait
RB_VM_LOCKING();
goto retry;
default: {
VALUE curr_hash = RUBY_ATOMIC_VALUE_LOAD(entry->hash);
if ((curr_hash == hash || curr_hash == 0) && set->funcs->cmp(key, curr_key)) {
// We've found a match.
if (UNLIKELY(!RB_SPECIAL_CONST_P(curr_key) && rb_objspace_garbage_object_p(curr_key))) {
// This is a weakref set, so after marking but before sweeping is complete we may find a matching garbage object.
// Skip it and mark it as deleted.
RUBY_ATOMIC_VALUE_CAS(entry->key, curr_key, CONCURRENT_SET_DELETED);
// Fall through and continue our search.
}
else {
RB_GC_GUARD(set_obj);
return curr_key;
}
}
break;
}
}
idx = concurrent_set_probe_next(&probe);
}
}
VALUE
rb_concurrent_set_find_or_insert(VALUE *set_obj_ptr, VALUE key, void *data)
{

View file

@ -1,6 +1,7 @@
#ifndef RUBY_RACTOR_SAFE_TABLE_H
#define RUBY_RACTOR_SAFE_TABLE_H
#include "ruby/atomic.h"
#include "ruby/ruby.h"
typedef VALUE (*rb_concurrent_set_hash_func)(VALUE key);
@ -15,6 +16,7 @@ struct rb_concurrent_set_funcs {
VALUE rb_concurrent_set_new(const struct rb_concurrent_set_funcs *funcs, int capacity);
rb_atomic_t rb_concurrent_set_size(VALUE set_obj);
VALUE rb_concurrent_set_find(VALUE *set_obj_ptr, VALUE key);
VALUE rb_concurrent_set_find_or_insert(VALUE *set_obj_ptr, VALUE key, void *data);
VALUE rb_concurrent_set_delete_by_identity(VALUE set_obj, VALUE key);
void rb_concurrent_set_foreach_with_replace(VALUE set_obj, int (*callback)(VALUE *key, void *data), void *data);