Implement new AtExitEx functionality; make PythonTracer use it to save max ref counts on exit.

Author: tpn

PythonTracer/PythonTracer.c

@@ -389,8 +389,8 @@ InitializePythonTraceContext(
     PVOID UserData
     )
 {
+    BOOL Success;
     PTRACE_STORES TraceStores;
-    PTRACE_STORE EventStore;
     PTRACE_STORE StringBufferStore;
     PTRACE_STORE FunctionTableStore;
     PTRACE_STORE FunctionTableEntryStore;
@@ -401,7 +401,11 @@ InitializePythonTraceContext(
     PPY_TRACE_CALLBACK Callback;
     PYTHON_ALLOCATORS Allocators;
     ULONG NumberOfAllocators = 0;
-    USHORT TraceStoreIndex;
+    TRACE_STORE_ID TraceStoreId;
+
+    //
+    // Validate arguments.
+    //
 
     if (!Context) {
         if (SizeOfContext) {
@@ -440,6 +444,10 @@ InitializePythonTraceContext(
         return FALSE;
     }
 
+    //
+    // Arguments are valid, continue with initialization.
+    //
+
     SecureZeroMemory(Context, sizeof(*Context));
 
     Context->Size = *SizeOfContext;
@@ -456,9 +464,8 @@ InitializePythonTraceContext(
     TraceStores = TraceContext->TraceStores;
 
 #define INIT_STORE_ALLOCATOR(Name)                                       \
-    TraceStoreIndex = TraceStore##Name##Index;                           \
-    Name##Store = &TraceStores->Stores[TraceStore##Name##Index];         \
-    Name##Store->NoRetire = TRUE;                                        \
+    TraceStoreId = TraceStore##Name##Id;                                 \
+    Name##Store = TraceStoreIdToTraceStore(TraceStores, TraceStoreId);   \
     Allocators.##Name##.AllocationRoutine = TraceStoreAllocationRoutine; \
     Allocators.##Name##.AllocationContext = ##Name##Store;               \
     Allocators.##Name##.FreeRoutine = TraceStoreFreeRoutine;             \
@@ -493,13 +500,6 @@ InitializePythonTraceContext(
     INIT_STORE_ALLOCATOR(StringArray);
     INIT_STORE_ALLOCATOR(StringTable);
 
-    EventStore = &TraceStores->Stores[TraceStoreEventIndex];
-
-    if (EventStore->NoRetire) {
-        __debugbreak();
-        return FALSE;
-    }
-
     Allocators.NumberOfAllocators = NumberOfAllocators;
     Allocators.SizeInBytes = sizeof(Allocators);
 
@@ -528,9 +528,104 @@ InitializePythonTraceContext(
     Context->AddModuleName = AddModuleName;
     Context->SetModuleNamesStringTable = SetModuleNamesStringTable;
 
+    //
+    // If we've been configured to track maximum reference counts, register an
+    // extended atexit function that will reflect the maximum values we saw
+    // in the registry on exit (if they exceed previously set maximums).
+    //
+    // (This is also a good test that our SetAtExitEx() machinery is working
+    //  properly, which is why we call the global version instead of going
+    //  through Rtl->RegisterAtExitEx().)
+    //
+    // N.B. Slight quirk: we can't use Context->Flags yet as they won't be set
+    //      correctly until after this call returns to TracedPythonSession, so,
+    //      we have to use to config's flags.
+    //
+
+    if (TraceContext->TracerConfig->Flags.TrackMaxRefCounts) {
+        Success = AtExitEx(SaveMaxRefCountsAtExit,
+                           NULL,
+                           Context,
+                           &Context->AtExitEntry);
+        if (!Success) {
+            return FALSE;
+        }
+    }
+
     return TRUE;
 }
 
+_Use_decl_annotations_
+VOID
+SaveMaxRefCountsAtExit(
+    BOOL IsProcessTerminating,
+    PPYTHON_TRACE_CONTEXT Context
+    )
+{
+    ULONG Result;
+    ULONG Disposition;
+    HKEY RegistryKey;
+    CONST UNICODE_STRING RegistryPath = \
+        RTL_CONSTANT_STRING(L"Software\\PythonTracer");
+
+    //
+    // Validate arguments.
+    //
+
+    if (!ARGUMENT_PRESENT(Context)) {
+        return;
+    }
+
+    //
+    // Arguments are valid, open or create the registry path.
+    //
+
+    Result = RegCreateKeyExW(
+        HKEY_CURRENT_USER,
+        RegistryPath.Buffer,
+        0,          // Reserved
+        NULL,       // Class
+        REG_OPTION_NON_VOLATILE,
+        KEY_ALL_ACCESS,
+        NULL,
+        &RegistryKey,
+        &Disposition
+    );
+
+    if (Result != ERROR_SUCCESS) {
+        OutputDebugStringA("RegCreateKeyExW() failed for: ");
+        OutputDebugStringW(RegistryPath.Buffer);
+        return;
+    }
+
+    //
+    // Write the max ref counts to the registry if they're greater than the
+    // current version.
+    //
+
+#define UPDATE_COUNT(Name)                                      \
+    UPDATE_MAX_REG_QWORD(RegistryKey,                           \
+                         Name,                                  \
+                         (ULONGLONG)Context->##Name##.QuadPart)
+
+    UPDATE_COUNT(MaxNoneRefCount);
+    UPDATE_COUNT(MaxTrueRefCount);
+    UPDATE_COUNT(MaxZeroRefCount);
+    UPDATE_COUNT(MaxFalseRefCount);
+
+    //
+    // Capture MaxDepth as well.
+    //
+
+    UPDATE_COUNT(MaxDepth);
+
+    //
+    // Close the registry key.
+    //
+
+    RegCloseKey(RegistryKey);
+}
+
 _Use_decl_annotations_
 BOOL
 Start(

PythonTracer/PythonTracer.h

@@ -332,6 +332,8 @@ typedef struct _Struct_size_bytes_(Size) _PYTHON_TRACE_CONTEXT {
     PADD_MODULE_NAME AddModuleName;
     PSET_MODULE_NAMES_STRING_TABLE SetModuleNamesStringTable;
 
+    struct _RTL_ATEXIT_ENTRY *AtExitEntry;
+
 } PYTHON_TRACE_CONTEXT, *PPYTHON_TRACE_CONTEXT;
 
 typedef

PythonTracer/PythonTracer.vcxproj

@@ -97,13 +97,15 @@
     </ClCompile>
   </ItemDefinitionGroup>
   <ItemGroup>
+    <ClInclude Include="..\Rtl\AtExitEx.h" />
     <ClInclude Include="PythonTracerConstants.h" />
     <ClInclude Include="stdafx.h" />
     <ClInclude Include="targetver.h" />
     <ClInclude Include="PythonTracer.h" />
     <ClInclude Include="PythonTracerPrivate.h" />
   </ItemGroup>
   <ItemGroup>
+    <ClCompile Include="..\Rtl\AtExitEx.c" />
     <ClCompile Include="..\Rtl\__C_specific_handler.c" />
     <ClCompile Include="dllmain.c">
       <PrecompiledHeader />

PythonTracer/PythonTracer.vcxproj.filters

@@ -30,6 +30,9 @@
     <ClInclude Include="PythonTracerConstants.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\Rtl\AtExitEx.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="stdafx.c">
@@ -62,5 +65,8 @@
     <ClCompile Include="PyTraceEvent4.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\Rtl\AtExitEx.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
 </Project>

PythonTracer/PythonTracerPrivate.h

@@ -15,6 +15,11 @@ Module Name:
 
 --*/
 
+#ifndef _PYTHON_TRACER_INTERNAL_BUILD
+#error PythonTracerPrivate.h being included but _PYTHON_TRACER_INTERNAL_BUILD
+#error not set.
+#endif
+
 #pragma once
 
 #include "stdafx.h"
@@ -23,6 +28,16 @@ Module Name:
 extern "C" {
 #endif
 
+//
+// AtExitEx-related definitions.
+//
+
+ATEXITEX_CALLBACK SaveMaxRefCountsAtExit;
+
+////////////////////////////////////////////////////////////////////////////////
+// Inline functions.
+////////////////////////////////////////////////////////////////////////////////
+
 FORCEINLINE
 BOOL
 IsFunctionOfInterestStringTable(

PythonTracer/stdafx.h

@@ -19,11 +19,15 @@ Module Name:
 #include <Windows.h>
 #include "../Rtl/Rtl.h"
 #include "../Rtl/__C_specific_handler.h"
+#include "../Rtl/AtExitEx.h"
 #include "../Python/Python.h"
 #include "../TraceStore/TraceStore.h"
 #include "../StringTable/StringTable.h"
 #include "PythonTracer.h"
 #include "PythonTracerConstants.h"
+
+#ifdef _PYTHON_TRACER_INTERNAL_BUILD
 #include "PythonTracerPrivate.h"
+#endif
 
 // vim:set ts=8 sw=4 sts=4 tw=80 expandtab                                     :

Rtl/AtExitEx.c

@@ -0,0 +1,44 @@
+/*++
+
+Copyright (c) 2016 Trent Nelson <[email protected]>
+
+Module Name:
+
+    AtExitEx.c
+
+Abstract:
+
+    This module implements the SetAtExitEx() and AtExitEx() stubs.
+
+--*/
+
+#include "stdafx.h"
+
+PATEXITEX AtExitExImpl = NULL;
+
+VOID
+SetAtExitEx(
+    PATEXITEX AtExitEx
+    )
+{
+    AtExitExImpl = AtExitEx;
+}
+
+#pragma warning(push)
+#pragma warning(disable: 4028 4273)
+
+_Use_decl_annotations_
+BOOL
+AtExitEx(
+    PATEXITEX_CALLBACK Callback,
+    PATEXITEX_FLAGS Flags,
+    PVOID Context,
+    struct _RTL_ATEXIT_ENTRY **EntryPointer
+    )
+{
+    return AtExitExImpl(Callback, Flags, Context, EntryPointer);
+}
+
+#pragma warning(pop)
+
+// vim:set ts=8 sw=4 sts=4 tw=80 expandtab                                     :

Rtl/AtExitEx.h

@@ -0,0 +1,66 @@
+/*++
+
+Copyright (c) 2016 Trent Nelson <[email protected]>
+
+Module Name:
+
+    AtExitEx.h
+
+Abstract:
+
+    This module provides the necessary scaffolding to use the AtExitEx()
+    function without needing to have any knowledge of the Rtl library.
+
+    To make use of this functionality, add this file and AtExitEx.c to your
+    project.  For components loaded via the TracedPythonSession component,
+    SetAtExitEx() will be called automatically as soon as the component has
+    been dynamically loaded (via LoadLibrary()).
+
+--*/
+
+#ifdef _NO_RTL
+
+typedef struct _Struct_size_bytes_(sizeof(ULONG)) _ATEXITEX_FLAGS {
+    ULONG SuppressExceptions:1;
+} ATEXITEX_FLAGS, *PATEXITEX_FLAGS;
+C_ASSERT(sizeof(ATEXITEX_FLAGS) == sizeof(ULONG));
+
+typedef
+VOID
+(CALLBACK ATEXITEX_CALLBACK)(
+    _In_ BOOL IsProcessTerminating,
+    _In_opt_ PVOID Context
+    );
+typedef ATEXITEX_CALLBACK *PATEXITEX_CALLBACK;
+
+typedef
+_Success_(return == 0)
+BOOL
+(ATEXITEX)(
+    _In_ PATEXITEX_CALLBACK Callback,
+    _In_opt_ PATEXITEX_FLAGS Flags,
+    _In_opt_ PVOID Context,
+    _Outptr_opt_result_nullonfailure_ struct _RTL_ATEXIT_ENTRY **EntryPointer
+    );
+typedef ATEXITEX *PATEXITEX;
+
+typedef
+VOID
+(SET_ATEXITEX)(
+    _In_ PATEXITEX AtExitEx
+    );
+typedef SET_ATEXITEX *PSET_ATEXITEX;
+
+#endif
+
+extern PATEXITEX AtExitExImpl;
+__declspec(dllexport) SET_ATEXITEX SetAtExitEx;
+
+#pragma warning(push)
+#pragma warning(disable: 4028 4273)
+
+ATEXITEX AtExitEx;
+
+#pragma warning(pop)
+
+// vim:set ts=8 sw=4 sts=4 tw=80 expandtab                                     :

Rtl/Path.c

@@ -610,7 +610,7 @@ Return Value:
     }
 
     AllocSize.LongPart = (
-        (NumberOfChars.LongPart << 1) +
+        (NumberOfChars.LowPart << 1) +
         sizeof(UNICODE_STRING)
     );
     AlignedAllocSize.LongPart = ALIGN_UP_POINTER(AllocSize.LongPart);