Optimize sending tracked entity data (#163)

.gitignore

@@ -8,7 +8,8 @@ Cargo.lock
 /extractor/out
 /extractor/classes
 /extractor/run
-/performance_tests/rust-mc-bot
+/benchmarks/rust-mc-bot
+/velocity
 flamegraph.svg
 perf.data
 perf.data.old

Cargo.toml

@@ -70,9 +70,9 @@ members = [
     "valence_nbt",
     "valence_protocol",
     "packet_inspector",
-    "performance_tests/players"
+    "benchmarks/bench_players"
 ]
-exclude = ["performance_tests/rust-mc-bot"]
+exclude = ["benchmarks/rust-mc-bot"]
 
 [profile.dev.package."*"]
 opt-level = 3

benchmarks/README.md

@@ -3,7 +3,7 @@
 Run the server
 
 ```shell
-cargo r -r -p players
+cargo r -r -p bench_players
 ```
 
 In a separate terminal, start [rust-mc-bot](https://github.com/Eoghanmc22/rust-mc-bot).
@@ -25,7 +25,7 @@ run the server like this:
 
 ```shell
 # You can also try setting the `CARGO_PROFILE_RELEASE_DEBUG` environment variable to `true`.
-cargo flamegraph -p players
+cargo flamegraph -p bench_players
 ```
 
 Run rust-mc-bot as above, and then stop the server after a few seconds. Flamegraph will generate a flamegraph.svg in the

benchmarks/bench_players/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "players"
+name = "bench_players"
 version = "0.1.0"
 edition = "2021"
 

benchmarks/bench_players/src/main.rs

@@ -17,7 +17,6 @@ pub fn main() -> ShutdownResult {
 }
 
 const WITH_PLAYER_ENTITIES: bool = true;
-const MAX_PLAYERS: usize = 10_000;
 
 struct Game;
 
@@ -38,7 +37,7 @@ impl Config for Game {
     type InventoryState = ();
 
     fn max_connections(&self) -> usize {
-        MAX_PLAYERS + 64
+        10_000
     }
 
     fn connection_mode(&self) -> ConnectionMode {
@@ -70,7 +69,7 @@ impl Config for Game {
             online_players: -1,
             max_players: -1,
             player_sample: Default::default(),
-            description: "Test server!".into(),
+            description: "Player Benchmark Server".into(),
             favicon_png: None,
         }
     }
@@ -137,7 +136,10 @@ impl Config for Game {
                         .entities
                         .insert_with_uuid(EntityKind::Player, client.uuid(), ())
                     {
-                        Some((id, _)) => client.state = id,
+                        Some((id, entity)) => {
+                            entity.set_world(world_id);
+                            client.state = id
+                        }
                         None => {
                             client.disconnect("Conflicting UUID");
                             return false;

build/entity.rs

@@ -460,33 +460,27 @@ pub fn build() -> anyhow::Result<TokenStream> {
                 }
             }
 
-            pub(super) fn initial_tracked_data(&self) -> Option<Vec<u8>> {
-                let mut data = Vec::new();
+            pub(super) fn write_initial_tracked_data(&self, buf: &mut Vec<u8>) {
+                debug_assert!(buf.is_empty());
 
                 match self {
-                    #(Self::#concrete_entity_names(e) => e.initial_tracked_data(&mut data),)*
+                    #(Self::#concrete_entity_names(e) => e.initial_tracked_data(buf),)*
                 }
 
-                if data.is_empty() {
-                    None
-                } else {
-                    data.push(0xff);
-                    Some(data)
+                if !buf.is_empty() {
+                    buf.push(0xff);
                 }
             }
 
-            pub(super) fn updated_tracked_data(&self) -> Option<Vec<u8>> {
-                let mut data = Vec::new();
+            pub(super) fn write_updated_tracked_data(&self, buf: &mut Vec<u8>) {
+                debug_assert!(buf.is_empty());
 
                 match self {
-                    #(Self::#concrete_entity_names(e) => e.updated_tracked_data(&mut data),)*
+                    #(Self::#concrete_entity_names(e) => e.updated_tracked_data(buf),)*
                 }
 
-                if data.is_empty() {
-                    None
-                } else {
-                    data.push(0xff);
-                    Some(data)
+                if !buf.is_empty() {
+                    buf.push(0xff);
                 }
             }
 

src/chunk.rs

@@ -561,14 +561,14 @@ impl<C: Config> LoadedChunk<C> {
     }
 
     /// Queues the chunk data packet for this chunk with the given position.
-    pub(crate) fn chunk_data_packet(
+    pub(crate) fn send_chunk_data_packet(
         &self,
         send: &mut PlayPacketSender,
         scratch: &mut Vec<u8>,
         pos: ChunkPos,
         biome_registry_len: usize,
     ) -> anyhow::Result<()> {
-        scratch.clear();
+        debug_assert!(scratch.is_empty());
 
         for sect in self.sections.iter() {
             sect.non_air_count.encode(&mut *scratch)?;

src/client.rs

@@ -201,6 +201,8 @@ pub struct Client<C: Config> {
     /// Ensures that we don't allow more connections to the server until the
     /// client is dropped.
     _permit: OwnedSemaphorePermit,
+    /// General purpose reusable buffer.
+    scratch: Vec<u8>,
     username: Username<String>,
     uuid: Uuid,
     ip: IpAddr,
@@ -297,6 +299,7 @@ impl<C: Config> Client<C> {
             send: Some(send),
             recv,
             _permit: permit,
+            scratch: Vec::new(),
             username: ncd.username,
             uuid: ncd.uuid,
             ip: ncd.ip,
@@ -979,6 +982,8 @@ impl<C: Config> Client<C> {
         player_lists: &PlayerLists<C>,
         inventories: &Inventories<C>,
     ) -> anyhow::Result<()> {
+        debug_assert!(self.scratch.is_empty());
+
         let world = match worlds.get(self.world) {
             Some(world) => world,
             None => bail!("client is in an invalid world and must be disconnected"),
@@ -1135,13 +1140,17 @@ impl<C: Config> Client<C> {
 
         // Load new chunks within the view distance
         {
-            let mut scratch = Vec::new();
             let biome_registry_len = shared.biomes().len();
-
             for pos in chunks_in_view_distance(center, self.view_distance) {
                 if let Some(chunk) = world.chunks.get(pos) {
                     if self.loaded_chunks.insert(pos) {
-                        chunk.chunk_data_packet(send, &mut scratch, pos, biome_registry_len)?;
+                        chunk.send_chunk_data_packet(
+                            send,
+                            &mut self.scratch,
+                            pos,
+                            biome_registry_len,
+                        )?;
+                        self.scratch.clear();
                     }
                 }
             }
@@ -1167,7 +1176,8 @@ impl<C: Config> Client<C> {
                 if self.world == entity.world()
                     && self.position.distance(entity.position()) <= self.view_distance as f64 * 16.0
                 {
-                    let _ = entity.send_updated_tracked_data(send, id);
+                    let _ = entity.send_updated_tracked_data(send, &mut self.scratch, id);
+                    self.scratch.clear();
 
                     let position_delta = entity.position() - entity.old_position();
                     let needs_teleport = position_delta.map(f64::abs).reduce_partial_max() >= 8.0;
@@ -1244,16 +1254,16 @@ impl<C: Config> Client<C> {
         }
 
         // Update the client's own player metadata.
-        let mut data = Vec::new();
-        self.player_data.updated_tracked_data(&mut data);
-
-        if !data.is_empty() {
-            data.push(0xff);
+        self.player_data.updated_tracked_data(&mut self.scratch);
+        if !self.scratch.is_empty() {
+            self.scratch.push(0xff);
 
             send.append_packet(&SetEntityMetadata {
                 entity_id: VarInt(0),
-                metadata: RawBytes(&data),
+                metadata: RawBytes(&self.scratch),
             })?;
+
+            self.scratch.clear();
         }
 
         // Spawn new entities within the view distance.
@@ -1276,9 +1286,10 @@ impl<C: Config> Client<C> {
                         return Some(e);
                     }
 
-                    if let Err(e) = entity.send_initial_tracked_data(send, id) {
+                    if let Err(e) = entity.send_initial_tracked_data(send, &mut self.scratch, id) {
                         return Some(e);
                     }
+                    self.scratch.clear();
 
                     if let Err(e) = send_entity_events(send, id.to_raw(), entity.events()) {
                         return Some(e);

src/entity.rs

@@ -737,13 +737,14 @@ impl<C: Config> Entity<C> {
     pub(crate) fn send_initial_tracked_data(
         &self,
         send: &mut PlayPacketSender,
+        scratch: &mut Vec<u8>,
         this_id: EntityId,
     ) -> anyhow::Result<()> {
-        // TODO: cache metadata buffer?
-        if let Some(metadata) = self.variants.initial_tracked_data() {
+        self.variants.write_initial_tracked_data(scratch);
+        if !scratch.is_empty() {
             send.append_packet(&SetEntityMetadata {
                 entity_id: VarInt(this_id.to_raw()),
-                metadata: RawBytes(&metadata),
+                metadata: RawBytes(scratch),
             })?;
         }
 
@@ -755,13 +756,14 @@ impl<C: Config> Entity<C> {
     pub(crate) fn send_updated_tracked_data(
         &self,
         send: &mut PlayPacketSender,
+        scratch: &mut Vec<u8>,
         this_id: EntityId,
     ) -> anyhow::Result<()> {
-        // TODO: cache metadata buffer?
-        if let Some(metadata) = self.variants.updated_tracked_data() {
+        self.variants.write_updated_tracked_data(scratch);
+        if !scratch.is_empty() {
             send.append_packet(&SetEntityMetadata {
                 entity_id: VarInt(this_id.to_raw()),
-                metadata: RawBytes(&metadata),
+                metadata: RawBytes(scratch),
             })?;
         }