Implement faster pattern of IRQL check + KeAcquireSpinLockAtDpcLevel (#4011)

* Implement faster pattern of IRQL check + KeAcquireSpinLockAtDpcLevel

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

* Fix code annnotation

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

* PR feedback

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

---------

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>
Co-authored-by: Alan Jowett <alan.jowett@microsoft.com>

libs/execution_context/ebpf_maps.c

@@ -1181,7 +1181,6 @@ _insert_into_hot_list(_Inout_ ebpf_core_lru_map_t* map, size_t partition, _Inout
 {
     ebpf_lru_key_state_t key_state = _get_key_state(map, partition, entry);
     ebpf_lock_state_t state = 0;
-    bool is_preemptible = ebpf_is_preemptible();
 
     switch (key_state) {
     case EBPF_LRU_KEY_UNINITIALIZED:
@@ -1194,11 +1193,7 @@ _insert_into_hot_list(_Inout_ ebpf_core_lru_map_t* map, size_t partition, _Inout
         return;
     }
 
-    if (!is_preemptible) {
+    state = ebpf_lock_lock(&map->partitions[partition].lock);
-        ebpf_lock_lock_at_dispatch(&map->partitions[partition].lock);
-    } else {
-        state = ebpf_lock_lock(&map->partitions[partition].lock);
-    }
 
     key_state = _get_key_state(map, partition, entry);
 
@@ -1226,12 +1221,7 @@ _insert_into_hot_list(_Inout_ ebpf_core_lru_map_t* map, size_t partition, _Inout
     }
 
     _merge_hot_into_cold_list_if_needed(map, partition);
-
+    ebpf_lock_unlock(&map->partitions[partition].lock, state);
-    if (!is_preemptible) {
-        ebpf_lock_unlock_at_dispatch(&map->partitions[partition].lock);
-    } else {
-        ebpf_lock_unlock(&map->partitions[partition].lock, state);
-    }
 }
 
 /**

libs/runtime/ebpf_platform.c

@@ -43,26 +43,28 @@ _Requires_lock_not_held_(*lock) _Acquires_lock_(*lock) _IRQL_requires_max_(DISPA
 ebpf_lock_state_t
 ebpf_lock_lock(_Inout_ ebpf_lock_t* lock)
 {
-    return KeAcquireSpinLockRaiseToDpc(lock);
+    KIRQL old_irql = KeGetCurrentIrql();
+
+    if (old_irql < DISPATCH_LEVEL) {
+        old_irql = KeAcquireSpinLockRaiseToDpc(lock);
+    } else {
+        KeAcquireSpinLockAtDpcLevel(lock);
+    }
+    return old_irql;
 }
 
+#pragma warning(push)
+#pragma warning(disable : 28157) // Prefast is not able to understand that the IRQL is lowered if it was raised.
 _Requires_lock_held_(*lock) _Releases_lock_(*lock) _IRQL_requires_(DISPATCH_LEVEL) void ebpf_lock_unlock(
     _Inout_ ebpf_lock_t* lock, _IRQL_restores_ ebpf_lock_state_t state)
 {
-    KeReleaseSpinLock(lock, state);
+    if (state < DISPATCH_LEVEL) {
-}
+        KeReleaseSpinLock(lock, state);
-
+    } else {
-_Requires_lock_not_held_(*lock) _Acquires_lock_(*lock) _IRQL_requires_max_(DISPATCH_LEVEL)
+        KeReleaseSpinLockFromDpcLevel(lock);
-    _IRQL_requires_(DISPATCH_LEVEL) void ebpf_lock_lock_at_dispatch(_Inout_ ebpf_lock_t* lock)
+    }
-{
-    KeAcquireSpinLockAtDpcLevel(lock);
-}
-
-_Requires_lock_held_(*lock) _Releases_lock_(*lock)
-    _IRQL_requires_(DISPATCH_LEVEL) void ebpf_lock_unlock_at_dispatch(_Inout_ ebpf_lock_t* lock)
-{
-    KeReleaseSpinLockFromDpcLevel(lock);
 }
+#pragma warning(pop)
 
 bool
 ebpf_is_preemptible()

libs/runtime/ebpf_platform.h

@@ -212,6 +212,7 @@ extern "C"
      * @param[in, out] lock Pointer to memory location that contains the lock.
      * @returns The previous lock_state required for unlock.
      */
+    EBPF_INLINE_HINT
     _Requires_lock_not_held_(*lock) _Acquires_lock_(*lock) _IRQL_requires_max_(DISPATCH_LEVEL) _IRQL_saves_
         _IRQL_raises_(DISPATCH_LEVEL) ebpf_lock_state_t ebpf_lock_lock(_Inout_ ebpf_lock_t* lock);
 
@@ -220,25 +221,10 @@ extern "C"
      * @param[in, out] lock Pointer to memory location that contains the lock.
      * @param[in] state The state returned from ebpf_lock_lock.
      */
+    EBPF_INLINE_HINT
     _Requires_lock_held_(*lock) _Releases_lock_(*lock) _IRQL_requires_(DISPATCH_LEVEL) void ebpf_lock_unlock(
         _Inout_ ebpf_lock_t* lock, _IRQL_restores_ ebpf_lock_state_t state);
 
-    /**
-     * @brief Acquire exclusive access to the lock.
-     * @param[in, out] lock Pointer to memory location that contains the lock.
-     * @returns The previous lock_state required for unlock.
-     */
-    _Requires_lock_not_held_(*lock) _Acquires_lock_(*lock) _IRQL_requires_(DISPATCH_LEVEL)
-        _IRQL_requires_max_(DISPATCH_LEVEL) void ebpf_lock_lock_at_dispatch(_Inout_ ebpf_lock_t* lock);
-
-    /**
-     * @brief Release exclusive access to the lock.
-     * @param[in, out] lock Pointer to memory location that contains the lock.
-     * @param[in] state The state returned from ebpf_lock_lock.
-     */
-    _Requires_lock_held_(*lock) _Releases_lock_(*lock)
-        _IRQL_requires_(DISPATCH_LEVEL) void ebpf_lock_unlock_at_dispatch(_Inout_ ebpf_lock_t* lock);
-
     /**
      * @brief Raise the IRQL to new_irql.
      *