From eee1810861ba01690366c746a58f300b49a1a321 Mon Sep 17 00:00:00 2001
From: Mateusz Guzik <mjguzik@gmail.com>
Date: Sat, 28 Oct 2023 05:36:57 +0000
Subject: [PATCH] Add zfs_refcount_remove_not_last

Signed-off-by: Mateusz Guzik <mjguzik@gmail.com>
---
 include/sys/zfs_refcount.h |  4 ++++
 module/zfs/refcount.c      | 41 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/include/sys/zfs_refcount.h b/include/sys/zfs_refcount.h
index 77965a0aa5..5a1bd400f5 100644
--- a/include/sys/zfs_refcount.h
+++ b/include/sys/zfs_refcount.h
@@ -76,6 +76,8 @@ int zfs_refcount_is_zero(zfs_refcount_t *);
 int64_t zfs_refcount_count(zfs_refcount_t *);
 int64_t zfs_refcount_add(zfs_refcount_t *, const void *);
 int64_t zfs_refcount_remove(zfs_refcount_t *, const void *);
+int zfs_refcount_remove_not_last(zfs_refcount_t *, const void *);
+
 /*
  * Note that (add|remove)_many adds/removes one reference with "number" N,
  * _not_ N references with "number" 1, which is what (add|remove)_few does,
@@ -114,6 +116,8 @@ typedef struct refcount {
 #define	zfs_refcount_count(rc) atomic_load_64(&(rc)->rc_count)
 #define	zfs_refcount_add(rc, holder) atomic_inc_64_nv(&(rc)->rc_count)
 #define	zfs_refcount_remove(rc, holder) atomic_dec_64_nv(&(rc)->rc_count)
+#define	zfs_refcount_remove_not_last(rc, holder) \
+	atomic_dec_not_last_64(&(rc)->rc_count)
 #define	zfs_refcount_add_few(rc, number, holder) \
 	atomic_add_64(&(rc)->rc_count, number)
 #define	zfs_refcount_remove_few(rc, number, holder) \
diff --git a/module/zfs/refcount.c b/module/zfs/refcount.c
index 718bbb34a8..bee8f6f231 100644
--- a/module/zfs/refcount.c
+++ b/module/zfs/refcount.c
@@ -221,6 +221,47 @@ zfs_refcount_remove_few(zfs_refcount_t *rc, uint64_t number, const void *holder)
 		(void) zfs_refcount_remove(rc, holder);
 }
 
+int
+zfs_refcount_remove_not_last(zfs_refcount_t *rc, const void *holder)
+{
+	reference_t *ref, s;
+
+	if (likely(!rc->rc_tracked)) {
+		return (atomic_dec_not_last_64(&(rc)->rc_count));
+	}
+
+	s.ref_holder = holder;
+	s.ref_number = 1;
+	s.ref_search = B_TRUE;
+	mutex_enter(&rc->rc_mtx);
+	ASSERT3U(rc->rc_count, >=, 1);
+	ref = avl_find(&rc->rc_tree, &s, NULL);
+	if (unlikely(ref == NULL)) {
+		panic("No such hold %p on refcount %llx", holder,
+		    (u_longlong_t)(uintptr_t)rc);
+		return (0);
+	}
+	if (rc->rc_count == 1) {
+		mutex_exit(&rc->rc_mtx);
+		return (0);
+	}
+	avl_remove(&rc->rc_tree, ref);
+	if (reference_history > 0) {
+		list_insert_head(&rc->rc_removed, ref);
+		if (rc->rc_removed_count >= reference_history) {
+			ref = list_remove_tail(&rc->rc_removed);
+			kmem_cache_free(reference_cache, ref);
+		} else {
+			rc->rc_removed_count++;
+		}
+	} else {
+		kmem_cache_free(reference_cache, ref);
+	}
+	rc->rc_count -= 1;
+	mutex_exit(&rc->rc_mtx);
+	return (1);
+}
+
 void
 zfs_refcount_transfer(zfs_refcount_t *dst, zfs_refcount_t *src)
 {