Enhance implementation of ShardedMap class

README.md

@@ -4,17 +4,17 @@ zlib-accel is a software shim that intercepts zlib calls and offloads compressio
 The shim allows applications to leverage hardware accelerators transparently without code changes. The only requirement is to preload the shim's shared library (e.g., using LD_PRELOAD).
 
 Two accelerators are supported
-- Intel(R) QuickAssist Technology (QAT)
-- Intel(R) In-Memory Analytics Accelerator (IAA)
+- Intel® QuickAssist Technology (QAT)
+- Intel® In-Memory Analytics Accelerator (IAA)
 
 
 ## Scope/Constraints
 
-This shim is not a general-purpose replacement for zlib, and it is able to offload compression/decompression jobs in certain conditions. Therefore, not all applications can take advantage of the transparent offload, depending on how they use zlib.  It is important to test thoroughly with your specific application and configuration. The use cases we have tested so far are listed in a section below.
+This shim is not a general-purpose replacement for zlib, and it is able to offload compression/decompression jobs in certain conditions. Therefore, not all applications can take advantage of the transparent offload, depending on how they use zlib. It is important to test thoroughly with your specific application and configuration. The use cases we have tested so far are listed in a section below.
 
 In general, the shim is able to offload zlib calls that complete compression/decompression of one deflate stream in one call. "Streaming" compression/decompression (where compression/decompression is done incrementally) are not currently supported. If the shim is not able to offload a job to an accelerator, it will fall back to zlib, ensuring the application still works correctly.
 
-The shim has only been tested on Linux.
+The shim has only been tested on Linux. When building with hardware acceleration enabled (`USE_QAT`, `USE_IAA`, or `USE_IGZIP`), the [Intel® oneTBB](https://www.intel.com/content/www/us/en/developer/tools/oneapi/onetbb.html) performance library is required and can be acquired by installing the [Intel® oneAPI Base Toolkit](https://www.intel.com/content/www/us/en/developer/tools/oneapi/oneapi-toolkit-download.html).
 
 
 ### Hardware Acceleration
@@ -226,6 +226,10 @@ log_file
 - This option applies only if the shim is built with DEBUG_LOG=ON or ENABLE_STATISTICS=ON.
 - If specified, store log messages in the file. If not specified, log messages are printed to stdout.
 
+map_shards
+- Values: 2-65536. Default 64
+- Sets the number of shards in the thread-safe concurrent hash map. Each shard holds an independent map instance.
+- It must be a power of two, so Fibonacci hashing can be used to calculate uniformly distributed shard indexes.
 
 ## Tested Applications/Use Cases
 

common.cmake

@@ -98,4 +98,10 @@ if(USE_QAT)
   link_libraries(qatzip)
 endif()
 
+if(USE_QAT OR USE_IAA OR USE_IGZIP)
+  find_package(TBB REQUIRED COMPONENTS tbb)
+  link_libraries(TBB::tbb)
+  add_compile_definitions(USE_TBB)
+endif()
+
 link_libraries(z)

config/config.cpp

@@ -13,23 +13,24 @@ namespace config {
 
 std::string log_file = "";
 
-// default config values initialization
+// default config values at initialization
 uint32_t configs[CONFIG_MAX] = {
-    1,   /*use_qat_compress*/
-    1,   /*use_qat_uncompress*/
-    0,   /*use_iaa_compress*/
-    0,   /*use_iaa_uncompress*/
-    1,   /*use_zlib_compress*/
-    1,   /*use_zlib_uncompress*/
-    50,  /*iaa_compress_percentage*/
-    50,  /*iaa_uncompress_percentage*/
-    0,   /*iaa_prepend_empty_block*/
-    0,   /*qat_periodical_polling*/
-    1,   /*qat_compression_level*/
-    0,   /*qat_compression_allow_chunking*/
-    0,   /*ignore_zlib_dictionary*/
-    1,   /*log_level*/
-    1000 /*log_stats_samples*/
+    1,    /*use_qat_compress*/
+    1,    /*use_qat_uncompress*/
+    0,    /*use_iaa_compress*/
+    0,    /*use_iaa_uncompress*/
+    1,    /*use_zlib_compress*/
+    1,    /*use_zlib_uncompress*/
+    50,   /*iaa_compress_percentage*/
+    50,   /*iaa_uncompress_percentage*/
+    0,    /*iaa_prepend_empty_block*/
+    0,    /*qat_periodical_polling*/
+    1,    /*qat_compression_level*/
+    0,    /*qat_compression_allow_chunking*/
+    0,    /*ignore_zlib_dictionary*/
+    1,    /*log_level*/
+    1000, /*log_stats_samples*/
+    64    /*map_shards*/
 };
 
 bool LoadConfigFile(std::string& file_content, const char* file_path) {
@@ -50,10 +51,11 @@ bool LoadConfigFile(std::string& file_content, const char* file_path) {
     "iaa_prepend_empty_block",
     "qat_periodical_polling",
     "qat_compression_level",
-	  "qat_compression_allow_chunking",
+    "qat_compression_allow_chunking",
     "ignore_zlib_dictionary",
     "log_level",
-    "log_stats_samples"
+    "log_stats_samples",
+    "map_shards"
   };
   // clang-format on
 
@@ -65,9 +67,10 @@ bool LoadConfigFile(std::string& file_content, const char* file_path) {
   ConfigReader config_reader;
   config_reader.ParseFile(file_path);
 
-  auto trySetConfig = [&](ConfigOption opt, uint32_t max, uint32_t min) {
+  auto trySetConfig = [&](ConfigOption opt, uint32_t max, uint32_t min,
+                          std::function<bool(uint32_t)> validator = nullptr) {
     uint32_t value;
-    if (config_reader.GetValue(config_names[opt], value, max, min)) {
+    if (config_reader.GetValue(config_names[opt], value, max, min, validator)) {
       configs[opt] = value;
     }
   };
@@ -87,7 +90,8 @@ bool LoadConfigFile(std::string& file_content, const char* file_path) {
   trySetConfig(IGNORE_ZLIB_DICTIONARY, 1, 0);
   trySetConfig(LOG_LEVEL, 3, 0);
   trySetConfig(LOG_STATS_SAMPLES, UINT32_MAX, 0);
-
+  trySetConfig(MAP_SHARDS, 65536, 2,
+               [](uint32_t v) { return (v & (v - 1)) == 0; });
   config_reader.GetValue("log_file", log_file);
   file_content.append(config_reader.DumpValues());
 

config/config.h

@@ -25,6 +25,7 @@ enum ConfigOption {
   IGNORE_ZLIB_DICTIONARY,
   LOG_LEVEL,
   LOG_STATS_SAMPLES,
+  MAP_SHARDS,
   CONFIG_MAX
 };
 

config/config_reader.cpp

@@ -16,7 +16,8 @@
 constexpr int CUSTOM_PATH_MAX = 4096;
 
 bool ConfigReader::GetValue(const std::string& tag, uint32_t& value,
-                            uint32_t max_value, uint32_t min_value) {
+                            uint32_t max_value, uint32_t min_value,
+                            std::function<bool(uint32_t)> validator) {
   auto it = config_settings_map.find(tag);
   if (it == config_settings_map.end()) {
     return false;
@@ -41,6 +42,14 @@ bool ConfigReader::GetValue(const std::string& tag, uint32_t& value,
     }
 
     value = static_cast<uint32_t>(temp);
+
+    if (validator && !validator(value)) {
+      Log(LogLevel::LOG_ERROR, "ConfigReader::GetValue Line ", __LINE__,
+          " invalid input value for tag ", tag.c_str(), "\n");
+      value = 0;
+      return false;
+    }
+
     return true;
 
   } catch (const std::exception&) {

config/config_reader.h

@@ -4,6 +4,7 @@
 #pragma once
 
 #include <cstdint>
+#include <functional>
 #include <map>
 #include <string>
 
@@ -19,7 +20,8 @@ class ConfigReader {
  public:
   bool ParseFile(const std::string& file_name);
   bool GetValue(const std::string& tag, uint32_t& value,
-                uint32_t max_value = 100, uint32_t min_value = 0);
+                uint32_t max_value = 100, uint32_t min_value = 0,
+                std::function<bool(uint32_t)> validator = nullptr);
   bool GetValue(const std::string& tag, std::string& value);
 
   std::string DumpValues();

config/default_config

@@ -11,5 +11,6 @@ qat_periodical_polling = 0
 qat_compression_level = 1
 qat_compression_allow_chunking = 0
 ignore_zlib_dictionary = 0
+map_shards = 64
 log_level = 1
 log_file = /tmp/zlib-accel.log

logging.h

@@ -13,8 +13,6 @@
 #include "config/config.h"
 #include "utils.h"
 
-using namespace config;
-
 // Log verbosity levels. 0 = silent; higher values = more verbose.
 // Matches QATzip's QzLogLevel_T convention.
 enum class LogLevel {

sharded_map.h

@@ -4,44 +4,120 @@
 #pragma once
 
 #include <functional>
-#include <shared_mutex>
+#include <memory>
 #include <string>
-#include <unordered_map>
 #include <utility>
 
-static constexpr int SHARDS = 64;
+#ifdef USE_TBB
+// tbb::concurrent_hash_map implementation
+#include <tbb/concurrent_hash_map.h>  // portable: old TBB + oneTBB
+#else
+// stdlib implementation
+#include <mutex>
+#include <shared_mutex>
+#include <unordered_map>
+#endif
+
+#include "config/config.h"
 
 template <typename Key, typename Value>
 class ShardedMap {
  public:
-  auto Get(Key key) -> decltype(std::declval<Value>().get()) {
-    unsigned int shard = GetShard(key);
-    std::shared_lock<std::shared_mutex> lock(shard_mutexes[shard]);
-    auto it = map[shard].find(key);
-    if (it == map[shard].end()) {
+  // Use GetConfig() rather than configs[] directly: configs has hidden
+  // visibility and is not accessible from external translation units under LTO.
+  //
+  // For the three global registry instances, LTO causes the constructor to run
+  // before init_zlib_accel loads the config file, so num_shards_ is
+  // initialised with the default value. InitStreamRegistries() calls Init()
+  // afterwards to re-allocate with the configured shard count.
+  ShardedMap() : num_shards_(config::GetConfig(config::MAP_SHARDS)) {
+    pd_map_arr_ = std::make_unique<PaddedMapType[]>(num_shards_);
+  }
+
+  void Init() {
+    num_shards_ = config::GetConfig(config::MAP_SHARDS);
+    pd_map_arr_ = std::make_unique<PaddedMapType[]>(num_shards_);
+  }
+
+  auto Get(const Key& key) -> decltype(std::declval<Value>().get()) {
+    const auto shard = GetShard(key);
+#ifdef USE_TBB
+    typename MapType::const_accessor acc;
+    if (!pd_map_arr_[shard].map.find(acc, key)) {
+      return nullptr;
+    }
+    return acc->second.get();
+#else
+    std::shared_lock lock(pd_map_arr_[shard].mutex);
+    auto it = pd_map_arr_[shard].map.find(key);
+    if (it == pd_map_arr_[shard].map.end()) {
       return nullptr;
     }
     return it->second.get();
+#endif
   }
 
-  void Set(Key key, Value&& value) {
-    unsigned int shard = GetShard(key);
-    std::unique_lock<std::shared_mutex> lock(shard_mutexes[shard]);
-    map[shard][key] = std::move(value);
+  void Set(const Key& key, Value&& value) {
+    const auto shard = GetShard(key);
+#ifdef USE_TBB
+    typename MapType::accessor acc;
+    pd_map_arr_[shard].map.insert(acc, key);
+    acc->second = std::move(value);
+#else
+    std::unique_lock lock(pd_map_arr_[shard].mutex);
+    pd_map_arr_[shard].map[key] = std::move(value);
+#endif
   }
 
-  void Unset(Key key) {
-    unsigned int shard = GetShard(key);
-    std::unique_lock<std::shared_mutex> lock(shard_mutexes[shard]);
-    auto it = map[shard].find(key);
-    if (it != map[shard].end()) {
-      map[shard].erase(it);
-    }
+  void Unset(const Key& key) {
+    const auto shard = GetShard(key);
+#ifndef USE_TBB
+    std::unique_lock lock(pd_map_arr_[shard].mutex);
+#endif
+    pd_map_arr_[shard].map.erase(key);
   }
 
  private:
-  unsigned int GetShard(Key key) { return std::hash<Key>{}(key) % SHARDS; }
-  std::unordered_map<Key, Value> map[SHARDS];
-  std::shared_mutex shard_mutexes[SHARDS];
-  std::hash<std::string> hash;
+  // Number of bytes in a cache line
+  static constexpr std::size_t CACHE_LINE_SIZE = 64;
+
+  // Fibonacci hash constant: value is (2^64 / phi) to spread bits uniformly,
+  // where phi is golden ratio and 64 is machine word length - constant is
+  // multiplied by hash and then shifted right by top bits to get shard index
+  static constexpr uint64_t FIBONACCI_MULTIPLIER = 11400714819323198485ull;
+
+#ifdef USE_TBB
+  struct HashCompare {
+    static auto hash(const Key& key) -> std::size_t {
+      return std::hash<Key>{}(key);
+    }
+
+    static auto equal(const Key& lhs, const Key& rhs) -> bool {
+      return lhs == rhs;
+    }
+  };
+
+  using MapType = tbb::concurrent_hash_map<Key, Value, HashCompare>;
+#else
+  using MapType = std::unordered_map<Key, Value>;
+#endif
+
+  struct alignas(CACHE_LINE_SIZE) PaddedMapType {
+    MapType map;
+#ifndef USE_TBB
+    mutable std::shared_mutex mutex;
+#endif
+  };
+  std::unique_ptr<PaddedMapType[]> pd_map_arr_;
+  unsigned int num_shards_;
+
+  // Number of map shards must be a power of 2
+  auto GetShard(const Key& key) const -> unsigned int {
+    const auto shard_bits = __builtin_ctz(num_shards_);
+    const auto hash = static_cast<uint64_t>(std::hash<Key>{}(key));
+    const auto dist_hash = hash * FIBONACCI_MULTIPLIER;
+    const auto top_bits = 64 - shard_bits;
+    const auto shard_index = static_cast<unsigned int>(dist_hash >> top_bits);
+    return shard_index;
+  }
 };