From 57e745ab91adbc41b08ae821f1fd5e5e2024349e Mon Sep 17 00:00:00 2001 From: Masahiko Sawada Date: Thu, 24 Oct 2024 17:29:51 -0700 Subject: [PATCH v4 2/4] raidxtree.h: support shared iteration. This commit supports a shared iteration operation on a radix tree with multiple processes. The radix tree must be in shared mode to start a shared itereation. Parallel workers can attach the shared iteration using the iterator handle given by the leader process. Same as the normal interation, it's guarnteed that the shared iteration returns key-values in an ascending order. Author: Reviewed-by: Discussion: https://postgr.es/m/ --- src/include/lib/radixtree.h | 221 +++++++++++++++++++++++++++++++----- 1 file changed, 190 insertions(+), 31 deletions(-) diff --git a/src/include/lib/radixtree.h b/src/include/lib/radixtree.h index 88bf695e3f3..bd5b8eed1bf 100644 --- a/src/include/lib/radixtree.h +++ b/src/include/lib/radixtree.h @@ -177,6 +177,9 @@ #define RT_ATTACH RT_MAKE_NAME(attach) #define RT_DETACH RT_MAKE_NAME(detach) #define RT_GET_HANDLE RT_MAKE_NAME(get_handle) +#define RT_BEGIN_ITERATE_SHARED RT_MAKE_NAME(begin_iterate_shared) +#define RT_ATTACH_ITERATE_SHARED RT_MAKE_NAME(attach_iterate_shared) +#define RT_GET_ITER_HANDLE RT_MAKE_NAME(get_iter_handle) #define RT_LOCK_EXCLUSIVE RT_MAKE_NAME(lock_exclusive) #define RT_LOCK_SHARE RT_MAKE_NAME(lock_share) #define RT_UNLOCK RT_MAKE_NAME(unlock) @@ -236,15 +239,19 @@ #define RT_SHRINK_NODE_16 RT_MAKE_NAME(shrink_child_16) #define RT_SHRINK_NODE_48 RT_MAKE_NAME(shrink_child_48) #define RT_SHRINK_NODE_256 RT_MAKE_NAME(shrink_child_256) +#define RT_INITIALIZE_ITER RT_MAKE_NAME(initialize_iter) #define RT_NODE_ITERATE_NEXT RT_MAKE_NAME(node_iterate_next) #define RT_VERIFY_NODE RT_MAKE_NAME(verify_node) /* type declarations */ #define RT_RADIX_TREE RT_MAKE_NAME(radix_tree) #define RT_RADIX_TREE_CONTROL RT_MAKE_NAME(radix_tree_control) +#define RT_ITER_CONTROL RT_MAKE_NAME(iter_control) #define RT_ITER RT_MAKE_NAME(iter) #ifdef RT_SHMEM #define RT_HANDLE RT_MAKE_NAME(handle) +#define RT_ITER_CONTROL_SHARED RT_MAKE_NAME(iter_control_shared) +#define RT_ITER_HANDLE RT_MAKE_NAME(iter_handle) #endif #define RT_NODE RT_MAKE_NAME(node) #define RT_CHILD_PTR RT_MAKE_NAME(child_ptr) @@ -270,6 +277,7 @@ typedef struct RT_ITER RT_ITER; #ifdef RT_SHMEM typedef dsa_pointer RT_HANDLE; +typedef dsa_pointer RT_ITER_HANDLE; #endif #ifdef RT_SHMEM @@ -687,6 +695,7 @@ typedef struct RT_RADIX_TREE_CONTROL RT_HANDLE handle; uint32 magic; LWLock lock; + int tranche_id; #endif RT_PTR_ALLOC root; @@ -740,11 +749,9 @@ typedef struct RT_NODE_ITER int idx; } RT_NODE_ITER; -/* state for iterating over the whole radix tree */ -struct RT_ITER +/* Contain the iteration state data */ +typedef struct RT_ITER_CONTROL { - RT_RADIX_TREE *tree; - /* * A stack to track iteration for each level. Level 0 is the lowest (or * leaf) level @@ -755,8 +762,36 @@ struct RT_ITER /* The key constructed during iteration */ uint64 key; -}; +} RT_ITER_CONTROL; + +#ifdef RT_SHMEM +/* Contain the shared iteration state data */ +typedef struct RT_ITER_CONTROL_SHARED +{ + /* Actual shared iteration state data */ + RT_ITER_CONTROL common; + + /* protect the control data */ + LWLock lock; + + RT_ITER_HANDLE handle; + pg_atomic_uint32 refcnt; +} RT_ITER_CONTROL_SHARED; +#endif + +/* state for iterating over the whole radix tree */ +struct RT_ITER +{ + RT_RADIX_TREE *tree; + /* pointing to either local memory or DSA */ + RT_ITER_CONTROL *ctl; + +#ifdef RT_SHMEM + /* True if the iterator is for shared iteration */ + bool shared; +#endif +}; /* verification (available only in assert-enabled builds) */ static void RT_VERIFY_NODE(RT_NODE * node); @@ -1848,6 +1883,7 @@ RT_CREATE(MemoryContext ctx) tree->ctl = (RT_RADIX_TREE_CONTROL *) dsa_get_address(dsa, dp); tree->ctl->handle = dp; tree->ctl->magic = RT_RADIX_TREE_MAGIC; + tree->ctl->tranche_id = tranche_id; LWLockInitialize(&tree->ctl->lock, tranche_id); #else tree->ctl = (RT_RADIX_TREE_CONTROL *) palloc0(sizeof(RT_RADIX_TREE_CONTROL)); @@ -1900,6 +1936,9 @@ RT_ATTACH(dsa_area *dsa, RT_HANDLE handle) dsa_pointer control; tree = (RT_RADIX_TREE *) palloc0(sizeof(RT_RADIX_TREE)); + tree->iter_context = AllocSetContextCreate(CurrentMemoryContext, + RT_STR(RT_PREFIX) "_radix_tree iter context", + ALLOCSET_SMALL_SIZES); /* Find the control object in shared memory */ control = handle; @@ -2072,35 +2111,86 @@ RT_FREE(RT_RADIX_TREE * tree) /***************** ITERATION *****************/ +/* Common routine to initialize the given iterator */ +static void +RT_INITIALIZE_ITER(RT_RADIX_TREE * tree, RT_ITER * iter) +{ + RT_CHILD_PTR root; + + iter->tree = tree; + + Assert(RT_PTR_ALLOC_IS_VALID(tree->ctl->root)); + root.alloc = iter->tree->ctl->root; + RT_PTR_SET_LOCAL(tree, &root); + + iter->ctl->top_level = iter->tree->ctl->start_shift / RT_SPAN; + + /* Set the root to start */ + iter->ctl->cur_level = iter->ctl->top_level; + iter->ctl->node_iters[iter->ctl->cur_level].node = root; + iter->ctl->node_iters[iter->ctl->cur_level].idx = 0; +} + /* * Create and return the iterator for the given radix tree. * - * Taking a lock in shared mode during the iteration is the caller's - * responsibility. + * Taking a lock on a radix tree in shared mode during the iteration is the + * caller's responsibility. */ RT_SCOPE RT_ITER * RT_BEGIN_ITERATE(RT_RADIX_TREE * tree) { RT_ITER *iter; - RT_CHILD_PTR root; iter = (RT_ITER *) MemoryContextAllocZero(tree->iter_context, sizeof(RT_ITER)); - iter->tree = tree; + iter->ctl = (RT_ITER_CONTROL *) MemoryContextAllocZero(tree->iter_context, + sizeof(RT_ITER_CONTROL)); - Assert(RT_PTR_ALLOC_IS_VALID(tree->ctl->root)); - root.alloc = iter->tree->ctl->root; - RT_PTR_SET_LOCAL(tree, &root); + RT_INITIALIZE_ITER(tree, iter); - iter->top_level = iter->tree->ctl->start_shift / RT_SPAN; +#ifdef RT_SHMEM + /* we will non-shared iteration on a shared radix tree */ + iter->shared = false; +#endif - /* Set the root to start */ - iter->cur_level = iter->top_level; - iter->node_iters[iter->cur_level].node = root; - iter->node_iters[iter->cur_level].idx = 0; + return iter; +} + +#ifdef RT_SHMEM +/* + * Create and return the shared iterator for the given shard radix tree. + * + * Taking a lock on a radix tree in shared mode during the shared iteration to + * prevent concurrent writes is the caller's responsibility. + */ +RT_SCOPE RT_ITER * +RT_BEGIN_ITERATE_SHARED(RT_RADIX_TREE * tree) +{ + RT_ITER *iter; + RT_ITER_CONTROL_SHARED *ctl_shared; + dsa_pointer dp; + + /* The radix tree must be in shared mode */ + Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC); + + dp = dsa_allocate0(tree->dsa, sizeof(RT_ITER_CONTROL_SHARED)); + ctl_shared = (RT_ITER_CONTROL_SHARED *) dsa_get_address(tree->dsa, dp); + ctl_shared->handle = dp; + LWLockInitialize(&ctl_shared->lock, tree->ctl->tranche_id); + pg_atomic_init_u32(&ctl_shared->refcnt, 1); + + iter = (RT_ITER *) MemoryContextAllocZero(tree->iter_context, + sizeof(RT_ITER)); + + iter->ctl = (RT_ITER_CONTROL *) ctl_shared; + iter->shared = true; + + RT_INITIALIZE_ITER(tree, iter); return iter; } +#endif /* * Scan the inner node and return the next child pointer if one exists, otherwise @@ -2114,12 +2204,18 @@ RT_NODE_ITERATE_NEXT(RT_ITER * iter, int level) RT_CHILD_PTR node; RT_PTR_ALLOC *slot = NULL; + node_iter = &(iter->ctl->node_iters[level]); + node = node_iter->node; + #ifdef RT_SHMEM - Assert(iter->tree->ctl->magic == RT_RADIX_TREE_MAGIC); -#endif - node_iter = &(iter->node_iters[level]); - node = node_iter->node; + /* + * Since the iterator is shared, the local pointer of the node might be + * set by other backends, we need to make sure to use the local pointer. + */ + if (iter->shared) + RT_PTR_SET_LOCAL(iter->tree, &node); +#endif Assert(node.local != NULL); @@ -2192,8 +2288,8 @@ RT_NODE_ITERATE_NEXT(RT_ITER * iter, int level) } /* Update the key */ - iter->key &= ~(((uint64) RT_CHUNK_MASK) << (level * RT_SPAN)); - iter->key |= (((uint64) key_chunk) << (level * RT_SPAN)); + iter->ctl->key &= ~(((uint64) RT_CHUNK_MASK) << (level * RT_SPAN)); + iter->ctl->key |= (((uint64) key_chunk) << (level * RT_SPAN)); return slot; } @@ -2207,18 +2303,29 @@ RT_ITERATE_NEXT(RT_ITER * iter, uint64 *key_p) { RT_PTR_ALLOC *slot = NULL; - while (iter->cur_level <= iter->top_level) +#ifdef RT_SHMEM + /* Prevent the shared iterator from being updated concurrently */ + if (iter->shared) + LWLockAcquire(&((RT_ITER_CONTROL_SHARED *) iter->ctl)->lock, LW_EXCLUSIVE); +#endif + + while (iter->ctl->cur_level <= iter->ctl->top_level) { RT_CHILD_PTR node; - slot = RT_NODE_ITERATE_NEXT(iter, iter->cur_level); + slot = RT_NODE_ITERATE_NEXT(iter, iter->ctl->cur_level); - if (iter->cur_level == 0 && slot != NULL) + if (iter->ctl->cur_level == 0 && slot != NULL) { /* Found a value at the leaf node */ - *key_p = iter->key; + *key_p = iter->ctl->key; node.alloc = *slot; +#ifdef RT_SHMEM + if (iter->shared) + LWLockRelease(&((RT_ITER_CONTROL_SHARED *) iter->ctl)->lock); +#endif + if (RT_CHILDPTR_IS_VALUE(*slot)) return (RT_VALUE_TYPE *) slot; else @@ -2234,17 +2341,23 @@ RT_ITERATE_NEXT(RT_ITER * iter, uint64 *key_p) node.alloc = *slot; RT_PTR_SET_LOCAL(iter->tree, &node); - iter->cur_level--; - iter->node_iters[iter->cur_level].node = node; - iter->node_iters[iter->cur_level].idx = 0; + iter->ctl->cur_level--; + iter->ctl->node_iters[iter->ctl->cur_level].node = node; + iter->ctl->node_iters[iter->ctl->cur_level].idx = 0; } else { /* Not found the child slot, move up the tree */ - iter->cur_level++; + iter->ctl->cur_level++; } + } +#ifdef RT_SHMEM + if (iter->shared) + LWLockRelease(&((RT_ITER_CONTROL_SHARED *) iter->ctl)->lock); +#endif + /* We've visited all nodes, so the iteration finished */ return NULL; } @@ -2255,9 +2368,45 @@ RT_ITERATE_NEXT(RT_ITER * iter, uint64 *key_p) RT_SCOPE void RT_END_ITERATE(RT_ITER * iter) { +#ifdef RT_SHMEM + RT_ITER_CONTROL_SHARED *ctl = (RT_ITER_CONTROL_SHARED *) iter->ctl;; + + if (iter->shared && + pg_atomic_sub_fetch_u32(&ctl->refcnt, 1) == 0) + dsa_free(iter->tree->dsa, ctl->handle); +#endif pfree(iter); } +#ifdef RT_SHMEM +RT_SCOPE RT_ITER_HANDLE +RT_GET_ITER_HANDLE(RT_ITER * iter) +{ + Assert(iter->shared); + return ((RT_ITER_CONTROL_SHARED *) iter->ctl)->handle; + +} + +RT_SCOPE RT_ITER * +RT_ATTACH_ITERATE_SHARED(RT_RADIX_TREE * tree, RT_ITER_HANDLE handle) +{ + RT_ITER *iter; + RT_ITER_CONTROL_SHARED *ctl; + + iter = (RT_ITER *) MemoryContextAllocZero(tree->iter_context, + sizeof(RT_ITER)); + iter->tree = tree; + ctl = (RT_ITER_CONTROL_SHARED *) dsa_get_address(tree->dsa, handle); + iter->ctl = (RT_ITER_CONTROL *) ctl; + iter->shared = true; + + /* For every iterator, increase the refcnt by 1 */ + pg_atomic_add_fetch_u32(&ctl->refcnt, 1); + + return iter; +} +#endif + /***************** DELETION *****************/ #ifdef RT_USE_DELETE @@ -2957,7 +3106,11 @@ RT_DUMP_NODE(RT_NODE * node) #undef RT_PTR_ALLOC #undef RT_INVALID_PTR_ALLOC #undef RT_HANDLE +#undef RT_ITER_HANDLE +#undef RT_ITER_CONTROL +#undef RT_ITER_HANDLE #undef RT_ITER +#undef RT_SHARED_ITER #undef RT_NODE #undef RT_NODE_ITER #undef RT_NODE_KIND_4 @@ -2994,6 +3147,11 @@ RT_DUMP_NODE(RT_NODE * node) #undef RT_LOCK_SHARE #undef RT_UNLOCK #undef RT_GET_HANDLE +#undef RT_BEGIN_ITERATE_SHARED +#undef RT_ATTACH_ITERATE_SHARED +#undef RT_GET_ITER_HANDLE +#undef RT_ATTACH_ITER +#undef RT_GET_ITER_HANDLE #undef RT_FIND #undef RT_SET #undef RT_BEGIN_ITERATE @@ -3050,5 +3208,6 @@ RT_DUMP_NODE(RT_NODE * node) #undef RT_SHRINK_NODE_256 #undef RT_NODE_DELETE #undef RT_NODE_INSERT +#undef RT_INITIALIZE_ITER #undef RT_NODE_ITERATE_NEXT #undef RT_VERIFY_NODE -- 2.43.5