From cc962005ac38e2d7844089934d4910d9918cf5c1 Mon Sep 17 00:00:00 2001
From: Caio Ramos Casimiro <caiorcasimiro@gmail.com>
Date: Tue, 13 Aug 2024 20:40:47 +0100
Subject: [PATCH] feat(shm/kv) add nelts member representing number of entries

Iterating a tree in batches without knowing the number of elements in
it implies an additional traversal through its elements after the last
batch is retrieved. This last traversal returns no elements which signals
the end of the iteration.

Knowing the number of entries in the tree renders the additional traverse
unnecessary, when the last batch is retrieved it can be verified that
the tree has been fully iterated.
---
 src/common/shm/ngx_wa_shm_kv.c | 8 ++++++++
 src/common/shm/ngx_wa_shm_kv.h | 1 +
 2 files changed, 9 insertions(+)

diff --git a/src/common/shm/ngx_wa_shm_kv.c b/src/common/shm/ngx_wa_shm_kv.c
index d989c4e39..a2a1d2704 100644
--- a/src/common/shm/ngx_wa_shm_kv.c
+++ b/src/common/shm/ngx_wa_shm_kv.c
@@ -48,6 +48,8 @@ ngx_wa_shm_kv_init(ngx_wa_shm_t *shm)
     shm->data = kv;
     shm->shpool->log_nomem = 0;
 
+    kv->nelts = 0;
+
     if (shm->eviction == NGX_WA_SHM_EVICTION_LRU) {
         ngx_queue_init(&kv->eviction.lru_queue);
 
@@ -193,6 +195,8 @@ queue_expire(ngx_wa_shm_t *shm, ngx_queue_t *queue, ngx_queue_t *q)
 
     ngx_slab_free_locked(shm->shpool, node);
 
+    kv->nelts--;
+
     return NGX_OK;
 }
 
@@ -287,6 +291,7 @@ ngx_wa_shm_kv_set_locked(ngx_wa_shm_t *shm, ngx_str_t *key,
             ngx_rbtree_delete(&kv->rbtree, &n->key.node);
             ngx_slab_free_locked(shm->shpool, n);
             *written = 1;
+            kv->nelts--;
 
         } else {
             *written = 0;
@@ -343,6 +348,9 @@ ngx_wa_shm_kv_set_locked(ngx_wa_shm_t *shm, ngx_str_t *key,
             n->cas = old->cas;
             ngx_rbtree_delete(&kv->rbtree, &old->key.node);
             ngx_slab_free_locked(shm->shpool, old);
+
+        } else {
+            kv->nelts++;
         }
 
         ngx_memcpy(n->key.str.data, key->data, key->len);
diff --git a/src/common/shm/ngx_wa_shm_kv.h b/src/common/shm/ngx_wa_shm_kv.h
index d3fdd2c18..51a30bdb9 100644
--- a/src/common/shm/ngx_wa_shm_kv.h
+++ b/src/common/shm/ngx_wa_shm_kv.h
@@ -8,6 +8,7 @@
 typedef struct {
     ngx_rbtree_t        rbtree;
     ngx_rbtree_node_t   sentinel;
+    ngx_uint_t          nelts;
     union {
         ngx_queue_t     lru_queue;
         ngx_queue_t     slru_queues[0];