Synchronize world chunk index updates

mosfet/protocol/managers.py

@@ -1,4 +1,5 @@
 import json
+import threading
 import time
 from math import floor
 from munch import Munch
@@ -218,6 +219,7 @@ class WorldState:
     def __init__(self, min_y=-64):
         self.min_y = min_y
         self.min_section = min_y // 16
+        self.lock = threading.RLock()
         self.chunks = {}
         self.index = {}
         self.block_entities = {}
@@ -229,6 +231,7 @@ class WorldState:
         self.player_id = None
 
     def set_dimension(self, dimension, dimension_type=None):
+        with self.lock:
             if dimension:
                 self.dimension = dimension
             if dimension_type:
@@ -238,6 +241,7 @@ class WorldState:
                     self.min_section = min_y // 16
 
     def load_chunk(self, packet):
+        with self.lock:
             # Refresh selectors in case block metadata was reloaded before login.
             from mosfet.info import blocks
             self.indexed_ids = blocks.INDEXED_IDS
@@ -252,7 +256,8 @@ class WorldState:
             for section_y, section in chunk.sections.items():
                 palette = section['blocks']
                 if not self.indexed_ids.intersection(
-                    palette.palette or [palette.get(i) for i in range(palette.size)]):
+                        palette.palette or [palette.get(i)
+                                            for i in range(palette.size)]):
                     continue
                 for index in range(palette.size):
                     state_id = palette.get(index)
@@ -270,20 +275,33 @@ class WorldState:
         self.set_block_at(packet.x, packet.y, packet.z, packet.block_state_id)
 
     def handle_section_update(self, packet):
+        with self.lock:
             for x, y, z, state_id in packet.changes:
                 self.set_block_at(x, y, z, state_id)
 
     def get_chunk(self, x, z):
+        with self.lock:
             return self.chunks.get((x, z))
 
     def get_block_at(self, x, y, z):
-        chunk = self.get_chunk(floor(x / 16), floor(z / 16))
+        with self.lock:
+            chunk = self.chunks.get((floor(x / 16), floor(z / 16)))
             if chunk is None:
                 return None
             return chunk.get_block_at(floor(x), floor(y), floor(z))
 
+    def indexed_positions(self, block_ids):
+        with self.lock:
+            return [position for block_id in block_ids
+                    for position in tuple(self.index.get(block_id, ()))]
+
+    def block_entities_snapshot(self):
+        with self.lock:
+            return dict(self.block_entities)
+
     def set_block_entity(self, position, entity_type, data):
         position = tuple(position)
+        with self.lock:
             if data is None:
                 self.block_entities.pop(position, None)
                 return
@@ -300,7 +318,8 @@ class WorldState:
 
     def set_block_at(self, x, y, z, block_state_id):
         x, y, z = floor(x), floor(y), floor(z)
-        chunk = self.get_chunk(floor(x / 16), floor(z / 16))
+        with self.lock:
+            chunk = self.chunks.get((floor(x / 16), floor(z / 16)))
             if chunk is None:
                 return False
             previous = chunk.set_block_at(x, y, z, block_state_id)
@@ -315,17 +334,21 @@ class WorldState:
             return True
 
     def unload_chunk(self, x, z):
+        with self.lock:
             chunk = self.chunks.pop((x, z), None)
             if chunk is None:
                 return
             x0, z0 = x * 16, z * 16
             for state_id, locations in list(self.index.items()):
-            self.index[state_id] = {
+                remaining = {
-                pos for pos in locations
+                    pos for pos in locations.copy()
-                if not (x0 <= pos[0] < x0 + 16 and z0 <= pos[2] < z0 + 16)
+                    if not (x0 <= pos[0] < x0 + 16 and
+                            z0 <= pos[2] < z0 + 16)
                 }
-            if not self.index[state_id]:
+                if remaining:
-                del self.index[state_id]
+                    self.index[state_id] = remaining
+                else:
+                    self.index.pop(state_id, None)
             self.block_entities = {
                 pos: entity for pos, entity in self.block_entities.items()
                 if not (x0 <= pos[0] < x0 + 16 and
@@ -333,11 +356,13 @@ class WorldState:
             }
 
     def unload_all_chunks(self):
+        with self.lock:
             self.chunks.clear()
             self.index.clear()
             self.block_entities.clear()
 
     def get_loaded_area(self, ignore_empty=False):
+        with self.lock:
             if not self.chunks:
                 return None
             xs = [pos[0] for pos in self.chunks]
@@ -348,6 +373,7 @@ class WorldState:
                     (max(xs), max(ys), max(zs)))
 
     def check_loaded(self, chunk_distance=None):
+        with self.lock:
             if self.center is None:
                 return bool(self.chunks)
             radius = self.radius if chunk_distance is None else chunk_distance
@@ -360,6 +386,7 @@ class WorldState:
 
     @property
     def loaded_chunk_count(self):
+        with self.lock:
             return len(self.chunks)
 
     def __repr__(self):

mosfet/world.py

@@ -60,7 +60,11 @@ class World:
     def find_signs(self, center, distance=16):
         """Return nearby loaded signs as position/side/four-lines records."""
         result = []
-        for position, entity in self.g.chunks.block_entities.items():
+        chunks = self.g.chunks
+        entities = (chunks.block_entities_snapshot()
+                    if hasattr(chunks, 'block_entities_snapshot') else
+                    dict(chunks.block_entities))
+        for position, entity in entities.items():
             if utils.phyp(center, position) > distance:
                 continue
             get_block_at = getattr(self.g.chunks, 'get_block_at', None)
@@ -348,9 +352,13 @@ class World:
                 yield check
 
     def find_blocks_indexed(self, center, block_ids, distance=0):
+        chunks = self.g.chunks
+        if hasattr(chunks, 'indexed_positions'):
+            index = chunks.indexed_positions(block_ids)
+        else:
             index = []
             for bid in block_ids:
-            index.extend(self.g.chunks.index.get(bid, []))
+                index.extend(chunks.index.get(bid, []))
 
         result = []
         for block in index:

tests/test_world_updates.py

@@ -39,6 +39,25 @@ class TestSectionBlockUpdates(unittest.TestCase):
         self.assertEqual(packet.id, 0x56)
         self.assertEqual(packet.changes, [(-173, 70, -311, bed_id)])
 
+    def test_unload_chunk_prunes_index_and_block_entities_safely(self):
+        bed_id = next(iter(blocks.BED_IDS))
+        inside = (1, 64, 1)
+        outside = (17, 64, 1)
+        world = WorldState()
+        world.chunks[(0, 0)] = ChunkColumn(0, 0, [], world.min_section)
+        world.index[bed_id] = {inside, outside}
+        world.block_entities = {
+            inside: {'type': 1, 'data': {}},
+            outside: {'type': 1, 'data': {}},
+        }
+
+        snapshot = world.indexed_positions({bed_id})
+        world.unload_chunk(0, 0)
+
+        self.assertEqual(set(snapshot), {inside, outside})
+        self.assertEqual(world.index[bed_id], {outside})
+        self.assertEqual(set(world.block_entities), {outside})
+
     def test_section_update_replaces_old_block_and_index_entry(self):
         section_x, section_y, section_z = -11, 4, -20
         x, y, z = -173, 70, -311