fix: support multi-threaded RefCount

src/state_transition/cache/effective_balance_increments.zig

@@ -3,11 +3,11 @@ const Allocator = std.mem.Allocator;
 const ssz = @import("consensus_types");
 const preset = @import("preset").preset;
 const BeaconStateAllForks = @import("../types/beacon_state.zig").BeaconStateAllForks;
-const ReferenceCount = @import("../utils/reference_count.zig").ReferenceCount;
+const RefCount = @import("../utils/ref_count.zig").RefCount;
 const EFFECTIVE_BALANCE_INCREMENT = preset.EFFECTIVE_BALANCE_INCREMENT;
 
 pub const EffectiveBalanceIncrements = std.ArrayList(u16);
-pub const EffectiveBalanceIncrementsRc = ReferenceCount(EffectiveBalanceIncrements);
+pub const EffectiveBalanceIncrementsRc = RefCount(EffectiveBalanceIncrements);
 
 pub fn getEffectiveBalanceIncrementsZeroed(allocator: Allocator, len: usize) !EffectiveBalanceIncrements {
     var increments = try EffectiveBalanceIncrements.initCapacity(allocator, len);

src/state_transition/cache/epoch_cache.zig

@@ -47,7 +47,7 @@ const IndexedAttestation = @import("../types/attestation.zig").IndexedAttestatio
 
 const syncPubkeys = @import("./pubkey_cache.zig").syncPubkeys;
 
-const ReferenceCount = @import("../utils/reference_count.zig").ReferenceCount;
+const RefCount = @import("../utils/ref_count.zig").RefCount;
 
 pub const EpochCacheImmutableData = struct {
     config: *const BeaconConfig,
@@ -65,7 +65,7 @@ pub const PROPOSER_WEIGHT_FACTOR = c.PROPOSER_WEIGHT / (c.WEIGHT_DENOMINATOR - c
 /// an EpochCache is shared by multiple CachedBeaconStateAllForks instances
 /// a CachedBeaconStateAllForks should increase the reference count of EpochCache when it is created
 /// and decrease the reference count when it is deinitialized
-pub const EpochCacheRc = ReferenceCount(*EpochCache);
+pub const EpochCacheRc = RefCount(*EpochCache);
 
 pub const EpochCache = struct {
     allocator: Allocator,
@@ -316,16 +316,16 @@ pub const EpochCache = struct {
         // pubkey_to_index and index_to_pubkey are shared across applications, EpochCache does not own this field so should not deinit()
 
         // unref the epoch shufflings
-        self.previous_shuffling.release();
-        self.current_shuffling.release();
-        self.next_shuffling.release();
+        self.previous_shuffling.unref();
+        self.current_shuffling.unref();
+        self.next_shuffling.unref();
 
         // unref the effective balance increments
-        self.effective_balance_increment.release();
+        self.effective_balance_increment.unref();
 
         // unref the sync committee caches
-        self.current_sync_committee_indexed.release();
-        self.next_sync_committee_indexed.release();
+        self.current_sync_committee_indexed.unref();
+        self.next_sync_committee_indexed.unref();
         self.allocator.destroy(self);
     }
 
@@ -341,11 +341,11 @@ pub const EpochCache = struct {
             .proposers = self.proposers,
             .proposer_prev_epoch = self.proposer_prev_epoch,
             // reuse the same instances, increase reference count
-            .previous_shuffling = self.previous_shuffling.acquire(),
-            .current_shuffling = self.current_shuffling.acquire(),
-            .next_shuffling = self.next_shuffling.acquire(),
+            .previous_shuffling = self.previous_shuffling.ref(),
+            .current_shuffling = self.current_shuffling.ref(),
+            .next_shuffling = self.next_shuffling.ref(),
             // reuse the same instances, increase reference count, cloned only when necessary before an epoch transition
-            .effective_balance_increment = self.effective_balance_increment.acquire(),
+            .effective_balance_increment = self.effective_balance_increment.ref(),
             .total_slashings_by_increment = self.total_slashings_by_increment,
             // Basic types (numbers) cloned implicitly
             .sync_participant_reward = self.sync_participant_reward,
@@ -359,8 +359,8 @@ pub const EpochCache = struct {
             .current_target_unslashed_balance_increments = self.current_target_unslashed_balance_increments,
             .previous_target_unslashed_balance_increments = self.previous_target_unslashed_balance_increments,
             // reuse the same instances, increase reference count
-            .current_sync_committee_indexed = self.current_sync_committee_indexed.acquire(),
-            .next_sync_committee_indexed = self.next_sync_committee_indexed.acquire(),
+            .current_sync_committee_indexed = self.current_sync_committee_indexed.ref(),
+            .next_sync_committee_indexed = self.next_sync_committee_indexed.ref(),
             .sync_period = self.sync_period,
             .epoch = self.epoch,
             .next_epoch = self.next_epoch,
@@ -400,7 +400,7 @@ pub const EpochCache = struct {
         const upcoming_epoch = self.next_epoch;
 
         // move current to previous
-        self.previous_shuffling.release();
+        self.previous_shuffling.unref();
         // no need to release current_shuffling and next_shuffling
         self.previous_shuffling = self.current_shuffling;
         self.current_shuffling = self.next_shuffling;
@@ -446,7 +446,7 @@ pub const EpochCache = struct {
         var effective_balance_increment = try EffectiveBalanceIncrements.initCapacity(self.allocator, self.effective_balance_increment.get().items.len);
         try effective_balance_increment.appendSlice(self.effective_balance_increment.get().items);
         // unref the previous effective balance increment
-        self.effective_balance_increment.release();
+        self.effective_balance_increment.unref();
         self.effective_balance_increment = try EffectiveBalanceIncrementsRc.init(self.allocator, effective_balance_increment);
     }
 
@@ -635,9 +635,9 @@ pub const EpochCache = struct {
 
     pub fn rotateSyncCommitteeIndexed(self: *EpochCache, allocator: Allocator, next_sync_committee_indices: []const ValidatorIndex) !void {
         // unref the old instance
-        self.current_sync_committee_indexed.release();
+        self.current_sync_committee_indexed.unref();
         // this is the transfer of reference count
-        // should not do an release() then acquire() here as it may trigger a deinit()
+        // should not do an unref() then ref() here as it may trigger a deinit()
         self.current_sync_committee_indexed = self.next_sync_committee_indexed;
         const next_sync_committee_indexed = try SyncCommitteeCacheAllForks.initValidatorIndices(allocator, next_sync_committee_indices);
         self.next_sync_committee_indexed = try SyncCommitteeCacheRc.init(allocator, next_sync_committee_indexed);
@@ -656,7 +656,7 @@ pub const EpochCache = struct {
         if (index >= effective_balance_increments.items.len) {
             // Clone and extend effectiveBalanceIncrements
             effective_balance_increments = try getEffectiveBalanceIncrementsWithLen(self.allocator, index + 1);
-            self.effective_balance_increment.release();
+            self.effective_balance_increment.unref();
             self.effective_balance_increment = try EffectiveBalanceIncrementsRc.init(allocator, effective_balance_increments);
         }
         self.effective_balance_increment.get().items[index] = @intCast(@divFloor(effective_balance, preset.EFFECTIVE_BALANCE_INCREMENT));

src/state_transition/cache/state_cache.zig

@@ -16,7 +16,7 @@ pub const CachedBeaconStateAllForks = struct {
     /// only a reference to the singleton BeaconConfig
     config: *const BeaconConfig,
     /// only a reference to the shared EpochCache instance
-    /// TODO: before an epoch transition, need to release() epoch_cache before using a new one
+    /// TODO: before an epoch transition, need to unref() epoch_cache before using a new one
     epoch_cache_ref: *EpochCacheRc,
     /// this takes ownership of the state, it is expected to be deinitialized by this struct
     state: *BeaconStateAllForks,
@@ -70,7 +70,7 @@ pub const CachedBeaconStateAllForks = struct {
 
     pub fn deinit(self: *CachedBeaconStateAllForks, allocator: Allocator) void {
         // should not deinit config since we don't take ownership of it, it's singleton across applications
-        self.epoch_cache_ref.release();
+        self.epoch_cache_ref.unref();
         self.state.deinit(allocator);
         self.allocator.destroy(self.state);
     }

src/state_transition/cache/sync_committee_cache.zig

@@ -9,13 +9,13 @@ const BLSPubkey = ssz.primitive.BLSPubkey.Type;
 
 const SyncCommitteeIndices = std.ArrayList(u32);
 const SyncComitteeValidatorIndexMap = std.AutoHashMap(ValidatorIndex, SyncCommitteeIndices);
-const ReferenceCount = @import("../utils/reference_count.zig").ReferenceCount;
+const RefCount = @import("../utils/ref_count.zig").RefCount;
 
-pub const SyncCommitteeCacheRc = ReferenceCount(SyncCommitteeCacheAllForks);
+pub const SyncCommitteeCacheRc = RefCount(SyncCommitteeCacheAllForks);
 
 /// EpochCache is the only consumer of this cache but an instance of SyncCommitteeCacheAllForks is shared across EpochCache instances
 /// no EpochCache instance takes the ownership of SyncCommitteeCacheAllForks instance
-/// instead of that, we count on reference counting to deallocate the memory, see ReferenceCount() utility
+/// instead of that, we count on reference counting to deallocate the memory, see RefCount() utility
 pub const SyncCommitteeCacheAllForks = union(enum) {
     phase0: void,
     altair: *SyncCommitteeCache,

src/state_transition/utils/epoch_shuffling.zig

@@ -8,17 +8,17 @@ const getSeed = @import("./seed.zig").getSeed;
 const c = @import("constants");
 const innerShuffleList = @import("./shuffle.zig").innerShuffleList;
 const Epoch = ssz.primitive.Epoch.Type;
-const ReferenceCount = @import("./reference_count.zig").ReferenceCount;
+const RefCount = @import("./ref_count.zig").RefCount;
 
-pub const EpochShufflingRc = ReferenceCount(*EpochShuffling);
+pub const EpochShufflingRc = RefCount(*EpochShuffling);
 
 const Committee = []const ValidatorIndex;
 const SlotCommittees = []const Committee;
 const EpochCommittees = [preset.SLOTS_PER_EPOCH]SlotCommittees;
 
 /// EpochCache is the only consumer of this cache but an instance of EpochShuffling is shared across EpochCache instances
 /// no EpochCache instance takes the ownership of shuffling
-/// instead of that, we count on reference counting to deallocate the memory, see ReferenceCount() utility
+/// instead of that, we count on reference counting to deallocate the memory, see RefCount() utility
 pub const EpochShuffling = struct {
     allocator: Allocator,
 

src/state_transition/utils/reference_count.zig → src/state_transition/utils/ref_count.zig

@@ -3,54 +3,51 @@ const Allocator = std.mem.Allocator;
 
 /// A reference counted wrapper for a type `T`.
 /// T should be `*Something`, not `*const Something` due to deinit()
-pub fn ReferenceCount(comptime T: type) type {
+pub fn RefCount(comptime T: type) type {
     return struct {
-        // TODO: switch to std.atomic
         allocator: Allocator,
-        _ref_count: usize,
+        _ref_count: std.atomic.Value(u32),
         instance: T,
 
         pub fn init(allocator: Allocator, instance: T) !*@This() {
             const ptr = try allocator.create(@This());
             ptr.* = .{
                 .allocator = allocator,
-                ._ref_count = 1,
+                ._ref_count = std.atomic.Value(u32).init(1),
                 .instance = instance,
             };
             return ptr;
         }
 
-        pub fn deinit(self: *@This()) void {
+        /// private deinit, consumer should call unref() instead
+        fn deinit(self: *@This()) void {
             self.instance.deinit();
             self.allocator.destroy(self);
         }
 
-        pub fn clone(allocator: Allocator, instance: T) !*@This() {
-            const cloned = try instance.clone(allocator);
-            return @This().init(cloned);
-        }
-
+        /// caution: this returns a copy of the instance, make sure this does not have copy overhead
+        /// or use pointer type instead
         pub fn get(self: *@This()) T {
             return self.instance;
         }
 
-        pub fn acquire(self: *@This()) *@This() {
-            self._ref_count += 1;
+        pub fn ref(self: *@This()) *@This() {
+            _ = self._ref_count.fetchAdd(1, .acquire);
             return self;
         }
 
-        pub fn release(self: *@This()) void {
-            self._ref_count -= 1;
-            if (self._ref_count == 0) {
+        pub fn unref(self: *@This()) void {
+            const old_rc = self._ref_count.fetchSub(1, .release);
+            if (old_rc == 1) {
                 self.deinit();
             }
         }
     };
 }
 
-test "ReferenceCount - *std.ArrayList(u32)" {
+test "RefCount - *std.ArrayList(u32)" {
     const allocator = std.testing.allocator;
-    const WrappedArrayList = ReferenceCount(*std.ArrayList(u32));
+    const WrappedArrayList = RefCount(*std.ArrayList(u32));
 
     var array_list = std.ArrayList(u32).init(allocator);
     try array_list.append(1);
@@ -59,29 +56,29 @@ test "ReferenceCount - *std.ArrayList(u32)" {
     // ref_count = 1
     var wrapped_array_list = try WrappedArrayList.init(allocator, &array_list);
     // ref_count = 2
-    _ = wrapped_array_list.acquire();
+    _ = wrapped_array_list.ref();
 
     // ref_count = 1
-    wrapped_array_list.release();
+    wrapped_array_list.unref();
     // ref_count = 0 ===> deinit
-    wrapped_array_list.release();
+    wrapped_array_list.unref();
 
     // the test does not leak any memory because array_list.deinit() is automatically called
 }
 
-test "ReferenceCount - std.ArrayList(u32)" {
+test "RefCount - std.ArrayList(u32)" {
     const allocator = std.testing.allocator;
-    const WrappedArrayList = ReferenceCount(std.ArrayList(u32));
+    const WrappedArrayList = RefCount(std.ArrayList(u32));
 
     // ref_count = 1
     var wrapped_array_list = try WrappedArrayList.init(allocator, std.ArrayList(u32).init(allocator));
     // ref_count = 2
-    _ = wrapped_array_list.acquire();
+    _ = wrapped_array_list.ref();
 
     // ref_count = 1
-    wrapped_array_list.release();
+    wrapped_array_list.unref();
     // ref_count = 0 ===> deinit
-    wrapped_array_list.release();
+    wrapped_array_list.unref();
 
     // the test does not leak any memory because array_list.deinit() is automatically called
 }