Sped up importing by TONS. Also made it use zstd instead of bzip2 for chunk compression.

Author: Sweattypalms

Cargo.lock

@@ -47,6 +47,8 @@ uuid = { version = "1.9.1", features = ["v4", "v3", "v5"] }
 include-flate = "0.3.0"
 flate2 = "1.0.31"
 bzip2 = "0.4.4"
+lz4_flex = "0.11.3"
+zstd = "0.13.2"
 
 # OS
 which = "6.0.1"

Cargo.toml

@@ -1,4 +1,6 @@
 use byteorder::LE;
+use zstd::bulk::compress as zstd_compress;
+use zstd::bulk::decompress as zstd_decompress;
 use heed::types::{Bytes, U64};
 use heed::Env;
 use tokio::task::spawn_blocking;
@@ -7,27 +9,27 @@ use tracing::{trace, warn};
 use crate::database::Database;
 use crate::utils::error::Error;
 use crate::utils::hash::hash;
-use crate::world::chunkformat::Chunk;
+use crate::world::chunk_format::Chunk;
 
-use crate::utils::binary_utils::{bzip_compress, bzip_decompress};
+// use crate::utils::binary_utils::{bzip_compress, bzip_decompress};
 use bincode::config::standard;
 use bincode::{decode_from_slice, encode_to_vec};
 
 impl Database {
     /// Fetch chunk from database
     fn get_chunk_from_database(db: &Env, key: &u64) -> Result<Option<Chunk>, Error> {
         // Initialize read transaction and open chunks table
-        let ro_tx = db.read_txn().unwrap();
+        let ro_tx = db.read_txn()?;
         let database = db
-            .open_database::<U64<LE>, Bytes>(&ro_tx, Some("chunks"))
-            .unwrap()
+            .open_database::<U64<LE>, Bytes>(&ro_tx, Some("chunks"))?
             .expect("No table \"chunks\" found. The database should have been initialized");
 
         // Attempt to fetch chunk from table
         if let Ok(data) = database.get(&ro_tx, key) {
             Ok(data.map(|encoded_chunk| {
-                let decompressed =
-                    bzip_decompress(&encoded_chunk).expect("Failed to decompress chunk");
+                // let decompressed =
+                //     bzip_decompress(&encoded_chunk).expect("Failed to decompress chunk");
+                let decompressed = zstd_decompress(&encoded_chunk, 1024*1024*64).expect("Failed to decompress chunk");
                 let chunk: (Chunk, usize) = decode_from_slice(&*decompressed, standard())
                     .expect("Failed to decode chunk from database");
                 chunk.0
@@ -40,15 +42,15 @@ impl Database {
     /// Insert a single chunk into database
     fn insert_chunk_into_database(db: &Env, chunk: &Chunk) -> Result<(), Error> {
         // Initialize write transaction and open chunks table
-        let mut rw_tx = db.write_txn().unwrap();
+        let mut rw_tx = db.write_txn()?;
         let database = db
-            .open_database::<U64<LE>, Bytes>(&rw_tx, Some("chunks"))
-            .unwrap()
+            .open_database::<U64<LE>, Bytes>(&rw_tx, Some("chunks"))?
             .expect("No table \"chunks\" found. The database should have been initialized");
 
         // Encode chunk
         let encoded_chunk = encode_to_vec(chunk, standard()).expect("Failed to encode chunk");
-        let compressed = bzip_compress(&encoded_chunk).expect("Failed to compress chunk");
+        // let compressed = bzip_compress(&encoded_chunk).expect("Failed to compress chunk");
+        let compressed = zstd_compress(&encoded_chunk, 3).expect("Failed to compress chunk");
         let key = hash((chunk.dimension.as_ref().unwrap(), chunk.x_pos, chunk.z_pos));
 
         // Insert chunk
@@ -70,18 +72,18 @@ impl Database {
     /// TODO: Find better name/disambiguation
     fn insert_chunks_into_database(db: &Env, chunks: &[Chunk]) -> Result<(), Error> {
         // Initialize write transaction and open chunks table
-        let mut rw_tx = db.write_txn().unwrap();
+        let mut rw_tx = db.write_txn()?;
         let database = db
-            .open_database::<U64<LE>, Bytes>(&rw_tx, Some("chunks"))
-            .unwrap()
+            .open_database::<U64<LE>, Bytes>(&rw_tx, Some("chunks"))?
             .expect("No table \"chunks\" found. The database should have been initialized");
 
         // Update page
         for chunk in chunks {
             // Encode chunk
             let encoded_chunk = encode_to_vec(chunk, standard()).expect("Failed to encode chunk");
 
-            let compressed = bzip_compress(&encoded_chunk).expect("Failed to compress chunk");
+            // let compressed = bzip_compress(&encoded_chunk).expect("Failed to compress chunk");
+            let compressed = zstd_compress(&encoded_chunk, 3).expect("Failed to compress chunk");
             let key = hash((chunk.dimension.as_ref().unwrap(), chunk.x_pos, chunk.z_pos));
 
             // Insert chunk
@@ -126,7 +128,7 @@ impl Database {
                 warn!("Error getting chunk: {:X}", key,);
             }
         })
-        .await?;
+            .await?;
         Ok(())
     }
 
@@ -310,13 +312,72 @@ impl Database {
     ///
     /// ```
     pub async fn batch_insert(&self, values: Vec<Chunk>) -> Result<(), Error> {
+        /*
+
+        trace!("processing chunks (compressing and encoding)");
+        // Process chunks in parallel
+        let processed_chunks: Vec<(u64, Vec<u8>)> = values
+            .par_iter()
+            .map(|chunk| {
+                let key = hash((
+                    chunk.dimension.as_ref().expect(&format!("Invalid chunk @ ({},{})", chunk.x_pos, chunk.z_pos)),
+                    chunk.x_pos,
+                    chunk.z_pos,
+                ));
+
+                let encoded_chunk = encode_to_vec(chunk, standard())
+                    .expect("Failed to encode chunk");
+                let compressed = zstd_compress(&encoded_chunk, 3)
+                    .expect("Failed to compress chunk.")
+                    ;
+
+                (key, compressed)
+            })
+            .collect();
+        trace!("processed chunks");*/
+
+        // Insert into cache in parallel
+        // TODO: re-enable this?
+        /*values.par_iter().for_each(|chunk| {
+            let key = hash((
+                chunk.dimension.as_ref().expect(&format!("Invalid chunk @ ({},{})", chunk.x_pos, chunk.z_pos)),
+                chunk.x_pos,
+                chunk.z_pos,
+            ));
+
+            // tokio::spawn(self.load_into_cache(key));
+            // if let Err(e) = self.cache.insert(key, chunk.clone()) {
+            //     warn!("Failed to insert chunk into cache: {:?}", e);
+            // }
+        });
+*/
+
+        /*trace!("Inserting chunks into database");
+        // Perform batch insert into LMDB
+        spawn_blocking(move || {
+            let mut rw_tx = db.write_txn()?;
+            let database = db
+                .open_database::<U64<LE>, Bytes>(&rw_tx, Some("chunks"))?
+                .expect("No table \"chunks\" found. The database should have been initialized");
+
+            for (key, compressed) in processed_chunks {
+                database.put(&mut rw_tx, &key, &compressed)?;
+            }
+
+            rw_tx.commit()?;
+            Ok::<_, Error>(())
+        })
+            .await??;
+
+        Ok(())*/
+
         // Clone database pointer
         let db = self.db.clone();
 
         // Calculate all keys
         let keys = values
             .iter()
-            .map(|v| hash((v.dimension.as_ref().unwrap(), v.x_pos, v.z_pos)))
+            .map(|v| hash((v.dimension.as_ref().expect(format!("Invalid chunk @ ({},{})", v.x_pos, v.z_pos).as_str()), v.x_pos, v.z_pos)))
             .collect::<Vec<u64>>();
 
         // WARNING: The previous logic was to first insert in database and then insert in cache using load_into_cache fn.

src/database/chunks.rs

@@ -14,11 +14,10 @@ use tracing::{debug, info, trace};
 use crate::utils::config::get_global_config;
 use crate::utils::error::Error;
 
-use crate::world::chunkformat::Chunk;
+use crate::world::chunk_format::Chunk;
 pub mod chunks;
 
 // MDBX constants
-const LMDB_MAX_TABLE: u32 = 16;
 const LMDB_PAGE_SIZE: usize = 50 * 1024usize.pow(3); // 50GiB
 
 /// Global database structure
@@ -61,7 +60,7 @@ pub async fn start_database() -> Result<Database, Error> {
     let world = get_global_config().world.clone();
     let world_path = root.join("data").join(world);
 
-    debug!("Opening database at {:?}", world_path);
+    debug!("Opening database at {}", world_path.display());
 
     if !fs::try_exists(&world_path).await? {
         fs::create_dir_all(&world_path).await?;
@@ -71,7 +70,7 @@ pub async fn start_database() -> Result<Database, Error> {
     let mut opts = EnvOpenOptions::new();
     opts.max_readers(num_cpus::get() as u32)
         .map_size(LMDB_PAGE_SIZE)
-        .max_dbs(LMDB_MAX_TABLE);
+        .max_dbs(num_cpus::get() as u32 *2);
 
     // Open database (This operation is safe as we assume no other process touched the database)
     let lmdb = unsafe {

src/database/mod.rs

@@ -87,6 +87,8 @@ async fn start_server() -> Result<()> {
     let state = create_state(listener).await?;
 
     if env::args().any(|arg| arg == "--import") {
+        // world::importing::import_regions(state.clone()).await?;
+        rayon::ThreadPoolBuilder::new().num_threads(num_cpus::get() * 2).build_global().expect("Failed to build rayon thread pool");
         world::importing::import_regions(state.clone()).await?;
         exit(0);
     }

src/main.rs

@@ -1,7 +1,7 @@
 use crate::state::GlobalState;
 use crate::utils::encoding::bitset::BitSet;
 use crate::utils::error::Error;
-use crate::world::chunkformat::{Heightmaps};
+use crate::world::chunk_format::{Heightmaps};
 use crate::Result;
 use ferrumc_codec::enc::NetEncode;
 use ferrumc_codec::network_types::varint::VarInt;

src/net/packets/outgoing/chunk_and_light_data.rs

@@ -1,5 +1,6 @@
 use std::convert::Infallible;
 
+use crate::world::importing_v2::ImportingError;
 use config::ConfigError;
 
 #[derive(thiserror::Error, Debug)]
@@ -78,16 +79,23 @@ pub enum Error {
     BitOutputOverflow(usize, usize),
     #[error("Attempted to read more bits than are available: {0} attempted, {1} available")]
     BitReadOverflow(usize, usize),
-    #[error("Attemped to read more bits than are available in the input type: {0} attempted, {1} available")]
+    #[error("Attemped to read more bits than are available in the input type: {0} attempted, {1} available"
+    )]
     BitReadOverflowInput(usize, usize),
-    #[error("Attemped to write more bits than are available in the output type: {0} attempted, {1} available")]
+    #[error("Attemped to write more bits than are available in the output type: {0} attempted, {1} available"
+    )]
     BitWriteOverflow(usize, usize),
     #[error("Codec error")]
     CodecError(#[from] ferrumc_codec::error::CodecError),
     #[error("Conversion error")]
     ConversionError,
     #[error(transparent)]
     CompressionError(std::io::Error),
+
+    #[error(transparent)]
+    ImportingError(#[from] ImportingError),
+    #[error("Database error: {0}")]
+    LmdbError(#[from] heed::Error),
 }
 
 impl From<Infallible> for Error {
@@ -107,3 +115,5 @@ impl From<Error> for std::io::ErrorKind {
         std::io::ErrorKind::Other
     }
 }
+
+

src/utils/error.rs

@@ -1,5 +1,5 @@
 use crate::utils::error::Error;
-use crate::world::chunkformat::{BlockStates, Chunk, Palette, Section};
+use crate::world::chunk_format::{BlockStates, Chunk, Palette, Section};
 use ferrumc_codec::enc::NetEncode;
 use ferrumc_codec::network_types::varint::VarInt;
 use hashbrown::HashMap;