From c091ffa6be792b0bd7bb03d604e26e40b2adfda8 Mon Sep 17 00:00:00 2001
From: Timon Vonk <timonv@gmail.com>
Date: Sun, 27 Oct 2024 16:42:34 +0100
Subject: [PATCH] fix(indexing): Use atomics for key generation in memory
 storage (#415)

---
 .../src/persist/memory_storage.rs             | 24 ++++++++++---------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/swiftide-indexing/src/persist/memory_storage.rs b/swiftide-indexing/src/persist/memory_storage.rs
index 222e6fab..c2484922 100644
--- a/swiftide-indexing/src/persist/memory_storage.rs
+++ b/swiftide-indexing/src/persist/memory_storage.rs
@@ -1,4 +1,10 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::HashMap,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
 
 use anyhow::Result;
 use async_trait::async_trait;
@@ -22,13 +28,13 @@ pub struct MemoryStorage {
     data: Arc<RwLock<HashMap<String, Node>>>,
     #[builder(default)]
     batch_size: Option<usize>,
-    #[builder(default = "Arc::new(RwLock::new(0))")]
-    node_count: Arc<RwLock<u64>>,
+    #[builder(default = Arc::new(AtomicUsize::new(0)))]
+    node_count: Arc<AtomicUsize>,
 }
 
 impl MemoryStorage {
-    async fn key(&self) -> String {
-        self.node_count.read().await.to_string()
+    fn key(&self) -> String {
+        self.node_count.fetch_add(1, Ordering::Relaxed).to_string()
     }
 
     /// Retrieve a node by its key
@@ -62,10 +68,8 @@ impl Persist for MemoryStorage {
     ///
     /// If the node does not have an id, a simple counter is used as the key.
     async fn store(&self, node: Node) -> Result<Node> {
-        let key = self.key().await;
-        self.data.write().await.insert(key, node.clone());
+        self.data.write().await.insert(self.key(), node.clone());
 
-        (*self.node_count.write().await) += 1;
         Ok(node)
     }
 
@@ -74,11 +78,9 @@ impl Persist for MemoryStorage {
     /// If a node does not have an id, a simple counter is used as the key.
     async fn batch_store(&self, nodes: Vec<Node>) -> IndexingStream {
         let mut lock = self.data.write().await;
-        let mut last_key = *self.node_count.read().await;
 
         for node in &nodes {
-            lock.insert(last_key.to_string(), node.clone());
-            last_key += 1;
+            lock.insert(self.key(), node.clone());
         }
 
         IndexingStream::iter(nodes.into_iter().map(Ok))