Switch from mod to and for computing bucket index (#2813)

* Switch from mod to and for computing bucket index Signed-off-by: Alan Jowett <[email protected]> * PR feedback Signed-off-by: Alan Jowett <[email protected]> * PR feedback Signed-off-by: Alan Jowett <[email protected]> --------- Signed-off-by: Alan Jowett <[email protected]> Co-authored-by: Alan Jowett <[email protected]>
microsoft · Sep 11, 2023 · ba8f9ec · ba8f9ec
1 parent ea13687
commit ba8f9ec
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 24 deletions.
diff --git a/libs/execution_context/ebpf_maps.c b/libs/execution_context/ebpf_maps.c
@@ -858,7 +858,7 @@ _create_hash_map_internal(
     const ebpf_hash_table_creation_options_t options = {
         .key_size = local_map->ebpf_map_definition.key_size,
         .value_size = local_map->ebpf_map_definition.value_size,
-        .bucket_count = local_map->ebpf_map_definition.max_entries,
+        .minimum_bucket_count = local_map->ebpf_map_definition.max_entries,
         .max_entries = local_map->ebpf_map_definition.max_entries,
         .extract_function = extract_function,
         .supplemental_value_size = supplemental_value_size,

diff --git a/libs/runtime/ebpf_hash_table.c b/libs/runtime/ebpf_hash_table.c
@@ -52,6 +52,7 @@ typedef struct _ebpf_hash_bucket_header_and_lock
 struct _ebpf_hash_table
 {
     size_t bucket_count;            // Count of buckets.
+    size_t bucket_count_mask;       // Mask to use to get bucket index from hash.
     volatile size_t entry_count;    // Count of entries in the hash table.
     size_t max_entry_count;         // Maximum number of entries allowed or EBPF_HASH_TABLE_NO_LIMIT if no maximum.
     uint32_t seed;                  // Seed used for hashing.
@@ -209,14 +210,14 @@ _ebpf_hash_table_compare(_In_ const ebpf_hash_table_t* hash_table, _In_ const ui
 
 /**
  * @brief Given a potentially non-comparable key value, extract the key and
- * compute the hash.
+ * compute the hash and convert it to a bucket index.
  *
  * @param[in] hash_table Hash table the keys belong to.
  * @param[in] key Key to hash.
- * @return Hash of key.
+ * @return Bucket index.
  */
 static uint32_t
-_ebpf_hash_table_compute_hash(_In_ const ebpf_hash_table_t* hash_table, _In_ const uint8_t* key)
+_ebpf_hash_table_compute_bucket_index(_In_ const ebpf_hash_table_t* hash_table, _In_ const uint8_t* key)
 {
     size_t length;
     const uint8_t* data;
@@ -226,7 +227,8 @@ _ebpf_hash_table_compute_hash(_In_ const ebpf_hash_table_t* hash_table, _In_ con
         length = hash_table->key_size * 8;
         data = key;
     }
-    return _ebpf_murmur3_32(data, length, hash_table->seed);
+    uint32_t hash_value = _ebpf_murmur3_32(data, length, hash_table->seed);
+    return hash_value & hash_table->bucket_count_mask;
 }
 
 /**
@@ -466,16 +468,16 @@ _ebpf_hash_table_replace_bucket(
 {
     ebpf_result_t result = EBPF_SUCCESS;
     size_t index;
-    uint32_t hash;
+    uint32_t bucket_index;
     uint8_t* old_data = NULL;
     uint8_t* new_data = NULL;
     ebpf_hash_bucket_header_t* old_bucket = NULL;
     ebpf_hash_bucket_header_t* new_bucket = NULL;
 
-    hash = _ebpf_hash_table_compute_hash(hash_table, key);
+    bucket_index = _ebpf_hash_table_compute_bucket_index(hash_table, key);
 
     // Lock the bucket.
-    ebpf_lock_state_t state = ebpf_lock_lock(&hash_table->buckets[hash % hash_table->bucket_count].lock);
+    ebpf_lock_state_t state = ebpf_lock_lock(&hash_table->buckets[bucket_index].lock);
 
     // Make a copy of the value to insert.
     if (operation != EBPF_HASH_BUCKET_OPERATION_DELETE) {
@@ -497,7 +499,7 @@ _ebpf_hash_table_replace_bucket(
     }
 
     // Find the old bucket.
-    old_bucket = hash_table->buckets[hash % hash_table->bucket_count].header;
+    old_bucket = hash_table->buckets[bucket_index].header;
     size_t old_bucket_count = old_bucket ? old_bucket->count : 0;
 
     // Find the entry in the bucket, if any.
@@ -553,12 +555,12 @@ _ebpf_hash_table_replace_bucket(
 
     // Update the bucket in the hash table.
     // From this point on the new bucket is immutable.
-    hash_table->buckets[hash % hash_table->bucket_count].header = new_bucket;
+    hash_table->buckets[bucket_index].header = new_bucket;
     new_data = NULL;
     new_bucket = NULL;
 
 Done:
-    ebpf_lock_unlock(&hash_table->buckets[hash % hash_table->bucket_count].lock, state);
+    ebpf_lock_unlock(&hash_table->buckets[bucket_index].lock, state);
 
     if (hash_table->notification_callback) {
         if (new_data) {
@@ -589,10 +591,19 @@ ebpf_hash_table_create(_Out_ ebpf_hash_table_t** hash_table, _In_ const ebpf_has
     ebpf_hash_table_t* table = NULL;
     size_t table_size = 0;
     // Select default values for the hash table.
-    size_t bucket_count = options->bucket_count ? options->bucket_count : EBPF_HASH_TABLE_DEFAULT_BUCKET_COUNT;
+    size_t bucket_count =
+        options->minimum_bucket_count ? options->minimum_bucket_count : EBPF_HASH_TABLE_DEFAULT_BUCKET_COUNT;
     void* (*allocate)(size_t size) = options->allocate ? options->allocate : ebpf_epoch_allocate;
     void (*free)(void* memory) = options->free ? options->free : ebpf_epoch_free;
 
+    // Increase bucket_count to next power of 2.
+    unsigned long msb_index;
+    _BitScanReverse64(&msb_index, bucket_count);
+
+    if (bucket_count != (1ull << msb_index)) {
+        bucket_count = 1ull << (msb_index + 1ull);
+    }
+
     retval = ebpf_safe_size_t_multiply(sizeof(ebpf_hash_bucket_header_and_lock_t), bucket_count, &table_size);
     if (retval != EBPF_SUCCESS) {
         goto Done;
@@ -613,6 +624,7 @@ ebpf_hash_table_create(_Out_ ebpf_hash_table_t** hash_table, _In_ const ebpf_has
     table->allocate = allocate;
     table->free = free;
     table->bucket_count = bucket_count;
+    table->bucket_count_mask = bucket_count - 1;
     table->entry_count = 0;
     table->seed = ebpf_random_uint32();
     table->extract = options->extract_function;
@@ -656,7 +668,7 @@ _Must_inspect_result_ ebpf_result_t
 ebpf_hash_table_find(_In_ const ebpf_hash_table_t* hash_table, _In_ const uint8_t* key, _Outptr_ uint8_t** value)
 {
     ebpf_result_t retval;
-    uint32_t hash;
+    uint32_t bucket_index;
     uint8_t* data = NULL;
     size_t index;
     ebpf_hash_bucket_header_t* bucket;
@@ -666,8 +678,8 @@ ebpf_hash_table_find(_In_ const ebpf_hash_table_t* hash_table, _In_ const uint8_
         goto Done;
     }
 
-    hash = _ebpf_hash_table_compute_hash(hash_table, key);
-    bucket = hash_table->buckets[hash % hash_table->bucket_count].header;
+    bucket_index = _ebpf_hash_table_compute_bucket_index(hash_table, key);
+    bucket = hash_table->buckets[bucket_index].header;
     if (!bucket) {
         retval = EBPF_KEY_NOT_FOUND;
         goto Done;
@@ -755,7 +767,7 @@ ebpf_hash_table_next_key_pointer_and_value(
     _Outptr_opt_ uint8_t** value)
 {
     ebpf_result_t result = EBPF_SUCCESS;
-    uint32_t hash;
+    uint32_t starting_bucket_index;
     ebpf_hash_bucket_entry_t* next_entry = NULL;
     size_t bucket_index;
     size_t data_index;
@@ -766,9 +778,10 @@ ebpf_hash_table_next_key_pointer_and_value(
         goto Done;
     }
 
-    hash = (previous_key != NULL) ? _ebpf_hash_table_compute_hash(hash_table, previous_key) : 0;
+    starting_bucket_index =
+        (previous_key != NULL) ? _ebpf_hash_table_compute_bucket_index(hash_table, previous_key) : 0;
 
-    for (bucket_index = hash % hash_table->bucket_count; bucket_index < hash_table->bucket_count; bucket_index++) {
+    for (bucket_index = starting_bucket_index; bucket_index < hash_table->bucket_count; bucket_index++) {
         ebpf_hash_bucket_header_t* bucket = hash_table->buckets[bucket_index].header;
         // Skip empty buckets.
         if (!bucket) {

diff --git a/libs/runtime/ebpf_hash_table.h b/libs/runtime/ebpf_hash_table.h
@@ -58,8 +58,9 @@ extern "C"
         ebpf_hash_table_extract_function extract_function; //< Function to extract key from stored value.
         ebpf_hash_table_allocate allocate; //< Function to allocate memory - defaults to ebpf_epoch_allocate.
         ebpf_hash_table_free free;         //< Function to free memory - defaults to ebpf_epoch_free.
-        size_t bucket_count; //< Number of buckets to use - defaults to EBPF_HASH_TABLE_DEFAULT_BUCKET_COUNT.
-        size_t max_entries;  //< Maximum number of entries in the hash table - defaults to EBPF_HASH_TABLE_NO_LIMIT.
+        size_t minimum_bucket_count;       //< Minimum number of buckets to use - defaults to
+                                           // EBPF_HASH_TABLE_DEFAULT_BUCKET_COUNT.
+        size_t max_entries; //< Maximum number of entries in the hash table - defaults to EBPF_HASH_TABLE_NO_LIMIT.
         size_t supplemental_value_size; //< Size of supplemental value to store in each entry - defaults to 0.
         void* notification_context;     //< Context to pass to notification functions.
         ebpf_hash_table_notification_function

diff --git a/libs/runtime/ebpf_state.c b/libs/runtime/ebpf_state.c
@@ -43,7 +43,7 @@ ebpf_state_initiate()
     const ebpf_hash_table_creation_options_t options = {
         .key_size = sizeof(uint64_t),
         .value_size = sizeof(ebpf_state_entry_t),
-        .bucket_count = ebpf_get_cpu_count(),
+        .minimum_bucket_count = ebpf_get_cpu_count(),
     };
 
     return_value = ebpf_hash_table_create(&_ebpf_state_thread_table, &options);

diff --git a/libs/runtime/unit/platform_unit_test.cpp b/libs/runtime/unit/platform_unit_test.cpp
@@ -148,7 +148,7 @@ TEST_CASE("hash_table_test", "[platform]")
         .value_size = data_1.size(),
         .allocate = ebpf_allocate,
         .free = ebpf_free,
-        .bucket_count = 1,
+        .minimum_bucket_count = 1,
     };
 
     ebpf_hash_table_t* raw_ptr = nullptr;
@@ -292,7 +292,7 @@ TEST_CASE("hash_table_stress_test", "[platform]")
     const ebpf_hash_table_creation_options_t options = {
         .key_size = sizeof(uint32_t),
         .value_size = sizeof(uint64_t),
-        .bucket_count = static_cast<size_t>(worker_threads) * static_cast<size_t>(key_count),
+        .minimum_bucket_count = static_cast<size_t>(worker_threads) * static_cast<size_t>(key_count),
     };
     REQUIRE(ebpf_hash_table_create(&table, &options) == EBPF_SUCCESS);
     auto worker = [table, iterations, key_count, load_factor, &cpu_id]() {

diff --git a/tests/performance/platform.cpp b/tests/performance/platform.cpp
@@ -84,7 +84,7 @@ typedef class _ebpf_hash_table_test_state
         const ebpf_hash_table_creation_options_t options = {
             .key_size = sizeof(uint32_t),
             .value_size = sizeof(uint64_t),
-            .bucket_count = keys.size(),
+            .minimum_bucket_count = keys.size(),
         };
         REQUIRE(ebpf_hash_table_create(&table, &options) == EBPF_SUCCESS);
         for (auto& key : keys) {