fix memory leak when loading chunks in swarms

2 files changed, 36 insertions, 2 deletions

diff --git a/azalea-world/src/chunk_storage.rs b/azalea-world/src/chunk_storage.rs
index 80854f0a..bbcbeb32 100755
--- a/azalea-world/src/chunk_storage.rs
+++ b/azalea-world/src/chunk_storage.rs
@@ -8,6 +8,7 @@ use azalea_core::position::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlock
 use azalea_nbt::NbtCompound;
 use nohash_hasher::IntMap;
 use parking_lot::RwLock;
+use std::collections::hash_map::Entry;
 use std::str::FromStr;
 use std::{
     collections::HashMap,
@@ -184,7 +185,35 @@ impl PartialChunkStorage {
     /// # Panics
     /// If the chunk is not in the render distance.
     pub fn set(&mut self, pos: &ChunkPos, chunk: Option<Chunk>, chunk_storage: &mut ChunkStorage) {
-        self.set_with_shared_reference(pos, chunk.map(|c| Arc::new(RwLock::new(c))), chunk_storage);
+        let new_chunk;
+
+        if let Some(chunk) = chunk {
+            match chunk_storage.map.entry(*pos) {
+                Entry::Occupied(mut e) => {
+                    if let Some(old_chunk) = e.get_mut().upgrade() {
+                        *old_chunk.write() = chunk;
+                        new_chunk = Some(old_chunk.clone())
+                    } else {
+                        let chunk_lock = Arc::new(RwLock::new(chunk));
+                        e.insert(Arc::downgrade(&chunk_lock));
+                        new_chunk = Some(chunk_lock);
+                    }
+                }
+                Entry::Vacant(e) => {
+                    let chunk_lock = Arc::new(RwLock::new(chunk));
+                    e.insert(Arc::downgrade(&chunk_lock));
+                    new_chunk = Some(chunk_lock);
+                }
+            }
+        } else {
+            // don't remove it from the shared storage, since it'll be removed
+            // automatically if this was the last reference
+
+            new_chunk = None;
+        }
+        if let Some(chunk_mut) = self.limited_get_mut(pos) {
+            *chunk_mut = new_chunk;
+        }
     }
 
     /// Set a chunk in the shared storage and reference it from the limited
diff --git a/azalea/examples/testbot.rs b/azalea/examples/testbot.rs
index 38f8d499..c47aa061 100644
--- a/azalea/examples/testbot.rs
+++ b/azalea/examples/testbot.rs
@@ -56,7 +56,7 @@ async fn main() -> anyhow::Result<()> {
             .add_accounts(accounts.clone())
             .set_handler(handle)
             .set_swarm_handler(swarm_handle)
-            .join_delay(Duration::from_millis(1000))
+            .join_delay(Duration::from_millis(100))
             .start("localhost")
             .await;
         // let e = azalea::ClientBuilder::new()
@@ -154,6 +154,11 @@ async fn handle(mut bot: Client, event: Event, _state: State) -> anyhow::Result<
                 "lag" => {
                     std::thread::sleep(Duration::from_millis(1000));
                 }
+                "quit" => {
+                    bot.disconnect();
+                    tokio::time::sleep(Duration::from_millis(1000)).await;
+                    std::process::exit(0);
+                }
                 "inventory" => {
                     println!("inventory: {:?}", bot.menu());
                 }