From 68f047c470b29dc90391151db79c83325133c878 Mon Sep 17 00:00:00 2001
From: Guoteng Rao <3603304+grao1991@users.noreply.github.com>
Date: Tue, 9 Jul 2024 03:46:08 -0700
Subject: [PATCH] [Storage] Implement shard id extraction in a different way.
 (#13940)

---
 Cargo.lock                          | 12 ----------
 Cargo.toml                          |  1 -
 config/Cargo.toml                   |  1 -
 config/src/config/storage_config.rs | 35 +++++++++++++++++++++++------
 4 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 38c71adbaa3ee..531e3cf7f07a6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -820,7 +820,6 @@ dependencies = [
  "get_if_addrs",
  "maplit",
  "num_cpus",
- "number_range",
  "poem-openapi",
  "rand 0.7.3",
  "serde",
@@ -11703,17 +11702,6 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
 
-[[package]]
-name = "number_range"
-version = "0.3.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "60080faccd4ca50ad0b801b2be686136376b13f691f6eac84817e40973b2e1bb"
-dependencies = [
- "anyhow",
- "itertools 0.10.5",
- "num 0.4.1",
-]
-
 [[package]]
 name = "object"
 version = "0.32.2"
diff --git a/Cargo.toml b/Cargo.toml
index 402e12680405c..44846c8fb7d9a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -643,7 +643,6 @@ num_cpus = "1.13.1"
 num-derive = "0.3.3"
 num-integer = "0.1.42"
 num-traits = "0.2.15"
-number_range = "0.3.2"
 once_cell = "1.10.0"
 ordered-float = "3.9.1"
 ouroboros = "0.15.6"
diff --git a/config/Cargo.toml b/config/Cargo.toml
index b36397d17fb5a..1bdf814cd7b5c 100644
--- a/config/Cargo.toml
+++ b/config/Cargo.toml
@@ -28,7 +28,6 @@ cfg-if = { workspace = true }
 get_if_addrs = { workspace = true }
 maplit = { workspace = true }
 num_cpus = { workspace = true }
-number_range = { workspace = true }
 poem-openapi = { workspace = true }
 rand = { workspace = true }
 serde = { workspace = true }
diff --git a/config/src/config/storage_config.rs b/config/src/config/storage_config.rs
index 0234d868489c5..65901a213a278 100644
--- a/config/src/config/storage_config.rs
+++ b/config/src/config/storage_config.rs
@@ -6,16 +6,16 @@ use crate::{
     config::{config_sanitizer::ConfigSanitizer, node_config_loader::NodeType, Error, NodeConfig},
     utils,
 };
-use anyhow::{ensure, Result};
+use anyhow::{bail, ensure, Result};
 use aptos_logger::warn;
 use aptos_types::chain_id::ChainId;
 use arr_macro::arr;
-use number_range::NumberRangeOptions;
 use serde::{Deserialize, Serialize};
 use std::{
     collections::HashMap,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::{Path, PathBuf},
+    str::FromStr,
 };
 
 // Lru cache will consume about 2G RAM based on this default value.
@@ -49,11 +49,7 @@ impl ShardedDbPathConfig {
     pub fn get_shard_paths(&self) -> Result<HashMap<u8, PathBuf>> {
         let mut result = HashMap::new();
         for shard_path in &self.shard_paths {
-            let shard_ids = NumberRangeOptions::<u8>::new()
-                .with_list_sep(',')
-                .with_range_sep('-')
-                .parse(shard_path.shards.as_str())?
-                .collect::<Vec<u8>>();
+            let shard_ids = Self::parse(shard_path.shards.as_str())?;
             let path = &shard_path.path;
             ensure!(
                 path.is_absolute(),
@@ -74,6 +70,31 @@ impl ShardedDbPathConfig {
 
         Ok(result)
     }
+
+    fn parse(path: &str) -> Result<Vec<u8>> {
+        let mut shard_ids = vec![];
+        for p in path.split(',') {
+            let num_or_range: Vec<&str> = p.split('-').collect();
+            match num_or_range.len() {
+                1 => {
+                    let num = u8::from_str(num_or_range[0])?;
+                    ensure!(num < 16);
+                    shard_ids.push(num);
+                },
+                2 => {
+                    let range_start = u8::from_str(num_or_range[0])?;
+                    let range_end = u8::from_str(num_or_range[1])?;
+                    ensure!(range_start <= range_end && range_end < 16);
+                    for num in range_start..=range_end {
+                        shard_ids.push(num);
+                    }
+                },
+                _ => bail!("Invalid path: {path}."),
+            }
+        }
+
+        Ok(shard_ids)
+    }
 }
 
 /// Port selected RocksDB options for tuning underlying rocksdb instance of AptosDB.