Parser performance feedback

O365.Security.Native.ETW.Debug.nuspec

@@ -2,7 +2,7 @@
 <package xmlns="http://schemas.microsoft.com/packaging/2011/08/nuspec.xsd">
     <metadata>
         <id>Microsoft.O365.Security.Native.ETW.Debug</id>
-        <version>4.4.8</version>
+        <version>4.4.9</version>
         <title>Microsoft.O365.Security.Native.ETW Debug - managed wrappers for krabsetw</title>
         <authors>Microsoft</authors>
         <owners>Microsoft</owners>

O365.Security.Native.ETW.nuspec

@@ -2,7 +2,7 @@
 <package xmlns="http://schemas.microsoft.com/packaging/2011/08/nuspec.xsd">
     <metadata>
         <id>Microsoft.O365.Security.Native.ETW</id>
-        <version>4.4.8</version>
+        <version>4.4.9</version>
         <title>Microsoft.O365.Security.Native.ETW - managed wrappers for krabsetw</title>
         <authors>Microsoft</authors>
         <owners>Microsoft</owners>

krabs/krabs/parser.hpp

@@ -127,40 +127,17 @@ namespace krabs {
         // A schema contains a collection of properties that are keyed by name.
         // These properties are stored in a blob of bytes that needs to be
         // interpreted according to information that is packaged up in the
-        // schema and that can be retrieved using the Tdh* APIs. This format
-        // requires a linear traversal over the blob, incrementing according to
-        // the contents within it. This is janky, so our strategy is to
-        // minimize this as much as possible via caching.
+        // schema and that can be retrieved using the Tdh* APIs.
+        // We have pre-calculated and cached a name to index map.
 
-        const ULONG totalPropCount = schema_.pSchema_->PropertyCount;
-
-        // Resolve property name to index.
-        ULONG index = totalPropCount; // sentinel = not found
-        if (pPropertyNames_) {
-            // Fast path: use the persistent name to index map shared across
-            // all events of the same type.
-            auto it = pPropertyNames_->find(name);
-            if (it != pPropertyNames_->end()) {
-                index = it->second;
-            }
-        } else {
-            // Fallback: linear scan of property names in the schema.
-            for (ULONG i = 0; i < totalPropCount; ++i) {
-                auto &propInfo = schema_.pSchema_->EventPropertyInfoArray[i];
-                const wchar_t *pName = reinterpret_cast<const wchar_t*>(
-                    reinterpret_cast<const BYTE*>(schema_.pSchema_) +
-                    propInfo.NameOffset);
-                if (name == pName) {
-                    index = i;
-                    break;
-                }
-            }
-        }
-
-        if (index >= totalPropCount) {
+        // Resolve property name to index via map.
+        const auto it = pPropertyNames_->find(name);
+        if (it == pPropertyNames_->end()) { // not found
             return property_info();
         }
 
+        const ULONG index = it->second;
+
         // The first step is to use our cache for the property to see if we've
         // discovered it already.
         if (index < lastPropertyIndex_) {
@@ -172,6 +149,7 @@ namespace krabs {
 
         // accept that last property can be omitted from buffer. this happens if last property
         // is string but empty and the provider strips the null terminator
+        const ULONG totalPropCount = schema_.pSchema_->PropertyCount;
         assert((pBufferIndex_ == pEndBuffer_ ? ((totalPropCount - lastPropertyIndex_) <= 1)
                                              : true)
                && "invariant: if we've exhausted our buffer, then we must've"

krabs/krabs/schema.hpp

@@ -306,6 +306,8 @@ namespace krabs {
         ULONG64 process_start_key() const;
 
     private:
+        schema(const EVENT_RECORD &, std::pair<const PTRACE_EVENT_INFO, const property_name_map *>);
+
         const EVENT_RECORD &record_;
         const TRACE_EVENT_INFO *pSchema_;
         // Persistent name to index map, owned by schema_locator. May be nullptr.
@@ -337,9 +339,13 @@ namespace krabs {
     // ------------------------------------------------------------------------
 
     inline schema::schema(const EVENT_RECORD &record, const krabs::schema_locator &schema_locator)
+        : schema(record, schema_locator.get_schema_and_names(record))
+    { }
+
+    inline schema::schema(const EVENT_RECORD &record, std::pair<const PTRACE_EVENT_INFO, const property_name_map *> p)
         : record_(record)
-        , pSchema_(schema_locator.get_event_schema(record))
-        , pPropertyNames_(schema_locator.get_property_names(pSchema_))
+        , pSchema_(p.first)
+        , pPropertyNames_(p.second)
     { }
 
     inline schema::schema(const EVENT_RECORD &record, const PTRACE_EVENT_INFO pSchema)

krabs/krabs/schema_locator.hpp

@@ -217,19 +217,24 @@ namespace krabs {
 
         /**
          * <summary>
-         * Returns the persistent property name to index map for a schema.
+         * Retrieves the event schema and its persistent property name to index
+         * map in a single cache lookup. Throws on error.
          * The map is built when the schema is first cached.
-         * Returns nullptr if pSchema is null or not in the cache.
          * </summary>
          */
-        const property_name_map* get_property_names(const TRACE_EVENT_INFO* pSchema) const;
+        std::pair<const PTRACE_EVENT_INFO, const property_name_map*>
+        get_schema_and_names(const EVENT_RECORD& record) const;
 
     private:
-        void build_property_names(const TRACE_EVENT_INFO* pSchema) const;
+        struct schema_entry {
+            std::unique_ptr<char[]> buffer;
+            property_name_map name_map;
+        };
 
-        mutable std::unordered_map<schema_key, std::variant<std::unique_ptr<char[]>, TDHSTATUS>> cache_;
-        // Persistent property name to index maps, keyed by schema pointer.
-        mutable std::unordered_map<const TRACE_EVENT_INFO*, property_name_map> property_name_cache_;
+        std::pair<const PTRACE_EVENT_INFO, const property_name_map*>
+        get_schema_and_names_no_throw(const EVENT_RECORD& record, TDHSTATUS& status) const;
+
+        mutable std::unordered_map<schema_key, std::variant<schema_entry, TDHSTATUS>> cache_;
     };
 
     // Implementation
@@ -301,14 +306,26 @@ namespace krabs {
     }
 
     inline const PTRACE_EVENT_INFO schema_locator::get_event_schema(const EVENT_RECORD &record) const
+    {
+        return get_schema_and_names(record).first;
+    }
+
+    inline const PTRACE_EVENT_INFO schema_locator::get_event_schema_no_throw(const EVENT_RECORD& record, TDHSTATUS& status) const
+    {
+        return get_schema_and_names_no_throw(record, status).first;
+    }
+
+    inline std::pair<const PTRACE_EVENT_INFO, const property_name_map*>
+    schema_locator::get_schema_and_names(const EVENT_RECORD& record) const
     {
         TDHSTATUS status = ERROR_SUCCESS;
-        auto buffer = get_event_schema_no_throw(record, status);
+        auto result = get_schema_and_names_no_throw(record, status);
         error_check_common_conditions(status, record);
-        return buffer;
+        return result;
     }
 
-    inline const PTRACE_EVENT_INFO schema_locator::get_event_schema_no_throw(const EVENT_RECORD& record, TDHSTATUS& status) const
+    inline std::pair<const PTRACE_EVENT_INFO, const property_name_map*>
+    schema_locator::get_schema_and_names_no_throw(const EVENT_RECORD& record, TDHSTATUS& status) const
     {
         status = ERROR_SUCCESS;
 
@@ -321,24 +338,33 @@ namespace krabs {
             auto& value = it->second;
             if (std::holds_alternative<TDHSTATUS>(value)) {
                 status = std::get<TDHSTATUS>(value);
-                return nullptr;
+                return {nullptr, nullptr};
             }
-            return (PTRACE_EVENT_INFO)std::get<std::unique_ptr<char[]>>(value).get();
+            auto& entry = std::get<schema_entry>(value);
+            return {(PTRACE_EVENT_INFO)entry.buffer.get(), &entry.name_map};
         }
 
-        // Cache miss. Fetch the schema...
+        // Cache miss. Fetch the schema from TDH.
+        // NB: key's 'internalize_name' gets called by the cctor during emplace.
         auto buffer = get_event_schema_from_tdh_no_throw(record, status);
-        auto returnVal = (PTRACE_EVENT_INFO)buffer.get();
+        auto pSchema = (PTRACE_EVENT_INFO)buffer.get();
 
-        // Add the new instance to the cache.
-        // NB: key's 'internalize_name' gets called by the cctor here.
         if (status == ERROR_SUCCESS) {
-            cache_.emplace(key, std::move(buffer));
-            build_property_names(returnVal);
-        } else
+            schema_entry entry;
+            entry.buffer = std::move(buffer);
+            for (ULONG i = 0; i < pSchema->PropertyCount; ++i) {
+                const wchar_t* pName = reinterpret_cast<const wchar_t*>(
+                    reinterpret_cast<const BYTE*>(pSchema) +
+                    pSchema->EventPropertyInfoArray[i].NameOffset);
+                entry.name_map.emplace(std::wstring_view(pName), i);
+            }
+            auto result = cache_.emplace(key, std::move(entry));
+            auto& inserted = std::get<schema_entry>(result.first->second);
+            return {(PTRACE_EVENT_INFO)inserted.buffer.get(), &inserted.name_map};
+        } else {
             cache_.emplace(key, status);
-
-        return returnVal;
+            return {nullptr, nullptr};
+        }
     }
 
     inline bool schema_locator::has_event_schema(const EVENT_RECORD& record) const
@@ -348,29 +374,6 @@ namespace krabs {
         return status == ERROR_SUCCESS;
     }
 
-    inline void schema_locator::build_property_names(const TRACE_EVENT_INFO* pSchema) const
-    {
-        property_name_map names;
-        for (ULONG i = 0; i < pSchema->PropertyCount; ++i) {
-            const wchar_t* pName = reinterpret_cast<const wchar_t*>(
-                reinterpret_cast<const BYTE*>(pSchema) +
-                pSchema->EventPropertyInfoArray[i].NameOffset);
-            names.emplace(std::wstring_view(pName), i);
-        }
-        property_name_cache_.emplace(pSchema, std::move(names));
-    }
-
-    inline const property_name_map* schema_locator::get_property_names(const TRACE_EVENT_INFO* pSchema) const
-    {
-        if (!pSchema) return nullptr;
-
-        auto it = property_name_cache_.find(pSchema);
-        if (it != property_name_cache_.end()) {
-            return &it->second;
-        }
-        return nullptr;
-    }
-
     inline std::unique_ptr<char[]> get_event_schema_from_tdh(const EVENT_RECORD &record)
     {
         TDHSTATUS status = ERROR_SUCCESS;

krabsetw.nuspec

@@ -2,7 +2,7 @@
 <package xmlns="http://schemas.microsoft.com/packaging/2011/08/nuspec.xsd">
     <metadata>
         <id>Microsoft.O365.Security.Krabsetw</id>
-        <version>4.4.8</version>
+        <version>4.4.9</version>
         <title>Krabs ETW Wrappers</title>
         <authors>Microsoft</authors>
         <owners>Microsoft</owners>