Change jit notifications so that they pass the native code address. (#14021)

* Change jit notifications so that they pass the native code address. This fixes !bpmd so that it will set the correct breakpoint on tiered jitted methods.

* code review feedback

* don't handle OnCodeGenerated

8 files changed, 269 insertions, 23 deletions

diff --git a/src/ToolBox/SOS/Strike/strike.cpp b/src/ToolBox/SOS/Strike/strike.cpp
index 84fde91f43..b621023d1f 100644
--- a/src/ToolBox/SOS/Strike/strike.cpp
+++ b/src/ToolBox/SOS/Strike/strike.cpp
@@ -6343,6 +6343,27 @@ public:
         return bNeedUpdates;
     }
 
+    BOOL UpdateKnownCodeAddress(TADDR mod, CLRDATA_ADDRESS bpLocation)
+    {
+        PendingBreakpoint *pCur = m_breakpoints;
+        BOOL bpSet = FALSE;
+
+        while(pCur)
+        {
+            PendingBreakpoint *pNext = pCur->pNext;
+            if (pCur->ModuleMatches(mod))
+            {
+                IssueDebuggerBPCommand(bpLocation);
+                bpSet = TRUE;
+                break;
+            }
+
+            pCur = pNext;
+        }
+
+        return bpSet;
+    }
+
     void RemovePendingForModule(TADDR mod)
     {
         PendingBreakpoint *pCur = m_breakpoints;
@@ -6715,7 +6736,7 @@ BOOL g_stopOnNextCatch = FALSE;
 // According to the latest debuggers these callbacks will not get called
 // unless the user (or an extension, like SOS :-)) had previously enabled
 // clrn with "sxe clrn".
-class CNotification : public IXCLRDataExceptionNotification4
+class CNotification : public IXCLRDataExceptionNotification5
 {
     static int s_condemnedGen;
 
@@ -6741,9 +6762,10 @@ public:
             || IsEqualIID(iid, IID_IXCLRDataExceptionNotification)
             || IsEqualIID(iid, IID_IXCLRDataExceptionNotification2)
             || IsEqualIID(iid, IID_IXCLRDataExceptionNotification3)
-            || IsEqualIID(iid, IID_IXCLRDataExceptionNotification4))
+            || IsEqualIID(iid, IID_IXCLRDataExceptionNotification4)
+            || IsEqualIID(iid, IID_IXCLRDataExceptionNotification5))
         {
-            *ppvObject = static_cast<IXCLRDataExceptionNotification4*>(this);
+            *ppvObject = static_cast<IXCLRDataExceptionNotification5*>(this);
             AddRef();
             return S_OK;
         }
@@ -6769,7 +6791,13 @@ public:
      */
     STDMETHODIMP OnCodeGenerated(IXCLRDataMethodInstance* method)
     {
-        // Some method has been generated, make a breakpoint and remove it.
+        m_dbgStatus = DEBUG_STATUS_GO_HANDLED;
+        return S_OK;
+    }
+
+    STDMETHODIMP OnCodeGenerated2(IXCLRDataMethodInstance* method, CLRDATA_ADDRESS nativeCodeLocation)
+    {
+        // Some method has been generated, make a breakpoint.
         ULONG32 len = mdNameLen;
         LPWSTR szModuleName = (LPWSTR)alloca(mdNameLen * sizeof(WCHAR));
         if (method->GetName(0, mdNameLen, &len, g_mdName) == S_OK)
@@ -6782,12 +6810,11 @@ public:
                 if (pMod->GetName(mdNameLen, &len, szModuleName) == S_OK)
                 {
                     ExtOut("JITTED %S!%S\n", szModuleName, g_mdName);
-
-                    // Add breakpoint, perhaps delete pending breakpoint
+                    
                     DacpGetModuleAddress dgma;
                     if (SUCCEEDED(dgma.Request(pMod)))
                     {
-                        g_bpoints.Update(TO_TADDR(dgma.ModulePtr), FALSE);
+                        g_bpoints.UpdateKnownCodeAddress(TO_TADDR(dgma.ModulePtr), nativeCodeLocation);
                     }
                     else
                     {
diff --git a/src/debug/daccess/daccess.cpp b/src/debug/daccess/daccess.cpp
index 12825e040f..5bfbf6151e 100644
--- a/src/debug/daccess/daccess.cpp
+++ b/src/debug/daccess/daccess.cpp
@@ -4389,6 +4389,7 @@ ClrDataAccess::TranslateExceptionRecordToNotification(
     GcEvtArgs pubGcEvtArgs;
     ULONG32 notifyType = 0;
     DWORD catcherNativeOffset = 0;
+    TADDR nativeCodeLocation = NULL;
 
     DAC_ENTER();
 
@@ -4451,11 +4452,11 @@ ClrDataAccess::TranslateExceptionRecordToNotification(
             break;
         }
 
-        case DACNotify::JIT_NOTIFICATION:
+        case DACNotify::JIT_NOTIFICATION2:
         {
             TADDR methodDescPtr;
 
-            if (DACNotify::ParseJITNotification(exInfo, methodDescPtr))
+            if(DACNotify::ParseJITNotification(exInfo, methodDescPtr, nativeCodeLocation))
             {
                 // Try and find the right appdomain
                 MethodDesc* methodDesc = PTR_MethodDesc(methodDescPtr);
@@ -4488,7 +4489,7 @@ ClrDataAccess::TranslateExceptionRecordToNotification(
             }
             break;
         }
-
+        
         case DACNotify::EXCEPTION_NOTIFICATION:
         {
             TADDR threadPtr;
@@ -4594,6 +4595,13 @@ ClrDataAccess::TranslateExceptionRecordToNotification(
             notify4 = NULL;
         }
 
+        IXCLRDataExceptionNotification5* notify5;
+        if (notify->QueryInterface(__uuidof(IXCLRDataExceptionNotification5),
+                                   (void**)&notify5) != S_OK)
+        {
+            notify5 = NULL;
+        }
+
         switch(notifyType)
         {
         case DACNotify::MODULE_LOAD_NOTIFICATION:
@@ -4604,8 +4612,13 @@ ClrDataAccess::TranslateExceptionRecordToNotification(
             notify->OnModuleUnloaded(pubModule);
             break;
 
-        case DACNotify::JIT_NOTIFICATION:
+        case DACNotify::JIT_NOTIFICATION2:
             notify->OnCodeGenerated(pubMethodInst);
+
+            if (notify5)
+            {
+                notify5->OnCodeGenerated2(pubMethodInst, TO_CDADDR(nativeCodeLocation));
+            }
             break;
 
         case DACNotify::EXCEPTION_NOTIFICATION:
@@ -4647,6 +4660,14 @@ ClrDataAccess::TranslateExceptionRecordToNotification(
         {
             notify3->Release();
         }
+        if (notify4)
+        {
+            notify4->Release();
+        }
+        if (notify5)
+        {
+            notify5->Release();
+        }
     }
 
     if (pubModule)
diff --git a/src/inc/xclrdata.idl b/src/inc/xclrdata.idl
index 952aea4378..fb3f19ade4 100644
--- a/src/inc/xclrdata.idl
+++ b/src/inc/xclrdata.idl
@@ -2536,3 +2536,17 @@ interface IXCLRDataExceptionNotification4 : IXCLRDataExceptionNotification3
      */
     HRESULT ExceptionCatcherEnter([in] IXCLRDataMethodInstance* catchingMethod, DWORD catcherNativeOffset);
 }
+
+[
+    object,
+    local,
+    uuid(e77a39ea-3548-44d9-b171-8569ed1a9423)
+]
+interface IXCLRDataExceptionNotification5 : IXCLRDataExceptionNotification4
+{
+    /*
+     * New code was generated for a method. The given address is the start address of 
+     * the native newly jitted code.
+     */
+    HRESULT OnCodeGenerated2([in] IXCLRDataMethodInstance* method, [in] CLRDATA_ADDRESS nativeCodeLocation);
+}
\ No newline at end of file
diff --git a/src/pal/prebuilt/idl/xclrdata_i.cpp b/src/pal/prebuilt/idl/xclrdata_i.cpp
index 6dd642da4c..349f01108b 100644
--- a/src/pal/prebuilt/idl/xclrdata_i.cpp
+++ b/src/pal/prebuilt/idl/xclrdata_i.cpp
@@ -132,6 +132,9 @@ MIDL_DEFINE_GUID(IID, IID_IXCLRDataExceptionNotification3,0x31201a94,0x4337,0x49
 
 MIDL_DEFINE_GUID(IID, IID_IXCLRDataExceptionNotification4,0xC25E926E,0x5F09,0x4AA2,0xBB,0xAD,0xB7,0xFC,0x7F,0x10,0xCF,0xD7);
 
+
+MIDL_DEFINE_GUID(IID, IID_IXCLRDataExceptionNotification5,0xe77a39ea,0x3548,0x44d9,0xb1,0x71,0x85,0x69,0xed,0x1a,0x94,0x23);
+
 #undef MIDL_DEFINE_GUID
 
 #ifdef __cplusplus
diff --git a/src/pal/prebuilt/inc/xclrdata.h b/src/pal/prebuilt/inc/xclrdata.h
index da7acb34ab..3ce477ec82 100644
--- a/src/pal/prebuilt/inc/xclrdata.h
+++ b/src/pal/prebuilt/inc/xclrdata.h
@@ -7712,6 +7712,185 @@ EXTERN_C const IID IID_IXCLRDataExceptionNotification4;
 #endif 	/* C style interface */
 
 
+#ifndef __IXCLRDataExceptionNotification5_INTERFACE_DEFINED__
+#define __IXCLRDataExceptionNotification5_INTERFACE_DEFINED__
+
+/* interface IXCLRDataExceptionNotification5 */
+/* [uuid][local][object] */ 
+
+
+EXTERN_C const IID IID_IXCLRDataExceptionNotification5;
+
+#if defined(__cplusplus) && !defined(CINTERFACE)
+    
+    MIDL_INTERFACE("e77a39ea-3548-44d9-b171-8569ed1a9423")
+    IXCLRDataExceptionNotification5 : public IXCLRDataExceptionNotification4
+    {
+    public:
+        virtual HRESULT STDMETHODCALLTYPE OnCodeGenerated2( 
+            /* [in] */ IXCLRDataMethodInstance *method,
+            /* [in] */ CLRDATA_ADDRESS nativeCodeLocation) = 0;
+        
+    };
+    
+    
+#else   /* C style interface */
+
+    typedef struct IXCLRDataExceptionNotification5Vtbl
+    {
+        BEGIN_INTERFACE
+        
+        HRESULT ( STDMETHODCALLTYPE *QueryInterface )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ REFIID riid,
+            /* [annotation][iid_is][out] */ 
+            _COM_Outptr_  void **ppvObject);
+        
+        ULONG ( STDMETHODCALLTYPE *AddRef )( 
+            IXCLRDataExceptionNotification5 * This);
+        
+        ULONG ( STDMETHODCALLTYPE *Release )( 
+            IXCLRDataExceptionNotification5 * This);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnCodeGenerated )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataMethodInstance *method);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnCodeDiscarded )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataMethodInstance *method);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnProcessExecution )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ ULONG32 state);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnTaskExecution )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataTask *task,
+            /* [in] */ ULONG32 state);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnModuleLoaded )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataModule *mod);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnModuleUnloaded )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataModule *mod);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnTypeLoaded )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataTypeInstance *typeInst);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnTypeUnloaded )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataTypeInstance *typeInst);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnAppDomainLoaded )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataAppDomain *domain);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnAppDomainUnloaded )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataAppDomain *domain);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnException )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataExceptionState *exception);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnGcEvent )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ GcEvtArgs gcEvtArgs);
+        
+        HRESULT ( STDMETHODCALLTYPE *ExceptionCatcherEnter )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataMethodInstance *catchingMethod,
+            DWORD catcherNativeOffset);
+        
+        HRESULT ( STDMETHODCALLTYPE *OnCodeGenerated2 )( 
+            IXCLRDataExceptionNotification5 * This,
+            /* [in] */ IXCLRDataMethodInstance *method,
+            /* [in] */ CLRDATA_ADDRESS nativeCodeLocation);
+        
+        END_INTERFACE
+    } IXCLRDataExceptionNotification5Vtbl;
+
+    interface IXCLRDataExceptionNotification5
+    {
+        CONST_VTBL struct IXCLRDataExceptionNotification5Vtbl *lpVtbl;
+    };
+
+    
+
+#ifdef COBJMACROS
+
+
+#define IXCLRDataExceptionNotification5_QueryInterface(This,riid,ppvObject) \
+    ( (This)->lpVtbl -> QueryInterface(This,riid,ppvObject) ) 
+
+#define IXCLRDataExceptionNotification5_AddRef(This)    \
+    ( (This)->lpVtbl -> AddRef(This) ) 
+
+#define IXCLRDataExceptionNotification5_Release(This)   \
+    ( (This)->lpVtbl -> Release(This) ) 
+
+
+#define IXCLRDataExceptionNotification5_OnCodeGenerated(This,method)    \
+    ( (This)->lpVtbl -> OnCodeGenerated(This,method) ) 
+
+#define IXCLRDataExceptionNotification5_OnCodeDiscarded(This,method)    \
+    ( (This)->lpVtbl -> OnCodeDiscarded(This,method) ) 
+
+#define IXCLRDataExceptionNotification5_OnProcessExecution(This,state)  \
+    ( (This)->lpVtbl -> OnProcessExecution(This,state) ) 
+
+#define IXCLRDataExceptionNotification5_OnTaskExecution(This,task,state)    \
+    ( (This)->lpVtbl -> OnTaskExecution(This,task,state) ) 
+
+#define IXCLRDataExceptionNotification5_OnModuleLoaded(This,mod)    \
+    ( (This)->lpVtbl -> OnModuleLoaded(This,mod) ) 
+
+#define IXCLRDataExceptionNotification5_OnModuleUnloaded(This,mod)  \
+    ( (This)->lpVtbl -> OnModuleUnloaded(This,mod) ) 
+
+#define IXCLRDataExceptionNotification5_OnTypeLoaded(This,typeInst) \
+    ( (This)->lpVtbl -> OnTypeLoaded(This,typeInst) ) 
+
+#define IXCLRDataExceptionNotification5_OnTypeUnloaded(This,typeInst)   \
+    ( (This)->lpVtbl -> OnTypeUnloaded(This,typeInst) ) 
+
+
+#define IXCLRDataExceptionNotification5_OnAppDomainLoaded(This,domain)  \
+    ( (This)->lpVtbl -> OnAppDomainLoaded(This,domain) ) 
+
+#define IXCLRDataExceptionNotification5_OnAppDomainUnloaded(This,domain)    \
+    ( (This)->lpVtbl -> OnAppDomainUnloaded(This,domain) ) 
+
+#define IXCLRDataExceptionNotification5_OnException(This,exception) \
+    ( (This)->lpVtbl -> OnException(This,exception) ) 
+
+
+#define IXCLRDataExceptionNotification5_OnGcEvent(This,gcEvtArgs)   \
+    ( (This)->lpVtbl -> OnGcEvent(This,gcEvtArgs) ) 
+
+
+#define IXCLRDataExceptionNotification5_ExceptionCatcherEnter(This,catchingMethod,catcherNativeOffset)  \
+    ( (This)->lpVtbl -> ExceptionCatcherEnter(This,catchingMethod,catcherNativeOffset) ) 
+
+
+#define IXCLRDataExceptionNotification5_OnCodeGenerated2(This,method,nativeCodeLocation)    \
+    ( (This)->lpVtbl -> OnCodeGenerated2(This,method,nativeCodeLocation) ) 
+
+#endif /* COBJMACROS */
+
+
+#endif  /* C style interface */
+
+
+
+
+#endif  /* __IXCLRDataExceptionNotification5_INTERFACE_DEFINED__ */
+
+
 
 
 #endif 	/* __IXCLRDataExceptionNotification4_INTERFACE_DEFINED__ */
diff --git a/src/vm/prestub.cpp b/src/vm/prestub.cpp
index f96135008d..d2f24bde29 100644
--- a/src/vm/prestub.cpp
+++ b/src/vm/prestub.cpp
@@ -205,7 +205,7 @@ PCODE MethodDesc::DoBackpatch(MethodTable * pMT, MethodTable *pDispatchingMT, BO
 #pragma optimize("", off)
 #endif
 
-void DACNotifyCompilationFinished(MethodDesc *methodDesc)
+void DACNotifyCompilationFinished(MethodDesc *methodDesc, PCODE pCode)
 {
     CONTRACTL
     {
@@ -231,7 +231,7 @@ void DACNotifyCompilationFinished(MethodDesc *methodDesc)
         if (jnt & CLRDATA_METHNOTIFY_GENERATED)
         {
             // If so, throw an exception!
-            DACNotify::DoJITNotification(methodDesc);
+            DACNotify::DoJITNotification(methodDesc, (TADDR)pCode);
         }
     }
 }
@@ -813,7 +813,7 @@ PCODE MethodDesc::JitCompileCodeLockedEventWrapper(PrepareCodeConfig* pConfig, J
 #endif
     {
         // The notification will only occur if someone has registered for this method.
-        DACNotifyCompilationFinished(this);
+        DACNotifyCompilationFinished(this, pCode);
     }
 
     return pCode;
diff --git a/src/vm/util.cpp b/src/vm/util.cpp
index 2c71289c8d..692b72fc39 100644
--- a/src/vm/util.cpp
+++ b/src/vm/util.cpp
@@ -2849,7 +2849,7 @@ void InitializeClrNotifications()
 #endif // FEATURE_GDBJIT
 
 // called from the runtime
-void DACNotify::DoJITNotification(MethodDesc *MethodDescPtr)
+void DACNotify::DoJITNotification(MethodDesc *MethodDescPtr, TADDR NativeCodeLocation)
 {
     CONTRACTL
     {
@@ -2860,8 +2860,8 @@ void DACNotify::DoJITNotification(MethodDesc *MethodDescPtr)
     }
     CONTRACTL_END;
 
-    TADDR Args[2] = { JIT_NOTIFICATION, (TADDR) MethodDescPtr };
-    DACNotifyExceptionHelper(Args, 2);
+    TADDR Args[3] = { JIT_NOTIFICATION2, (TADDR) MethodDescPtr, NativeCodeLocation };
+    DACNotifyExceptionHelper(Args, 3);
 }
 
 void DACNotify::DoJITPitchingNotification(MethodDesc *MethodDescPtr)
@@ -2986,16 +2986,17 @@ int DACNotify::GetType(TADDR Args[])
     // Type is an enum, and will thus fit into an int.
     return static_cast<int>(Args[0]);
 }
-    
-BOOL DACNotify::ParseJITNotification(TADDR Args[], TADDR& MethodDescPtr)
+
+BOOL DACNotify::ParseJITNotification(TADDR Args[], TADDR& MethodDescPtr, TADDR& NativeCodeLocation)
 {
-    _ASSERTE(Args[0] == JIT_NOTIFICATION);
-    if (Args[0] != JIT_NOTIFICATION)
+    _ASSERTE(Args[0] == JIT_NOTIFICATION2);
+    if (Args[0] != JIT_NOTIFICATION2)
     {
         return FALSE;
     }
 
     MethodDescPtr = Args[1];
+    NativeCodeLocation = Args[2];
 
     return TRUE;
 }
diff --git a/src/vm/util.hpp b/src/vm/util.hpp
index edfd9161e4..ad1261845f 100644
--- a/src/vm/util.hpp
+++ b/src/vm/util.hpp
@@ -1075,10 +1075,11 @@ public:
         EXCEPTION_NOTIFICATION=5,
         GC_NOTIFICATION= 6,
         CATCH_ENTER_NOTIFICATION = 7,
+        JIT_NOTIFICATION2=8,
     };
     
     // called from the runtime
-    static void DoJITNotification(MethodDesc *MethodDescPtr);
+    static void DoJITNotification(MethodDesc *MethodDescPtr, TADDR NativeCodeLocation);
     static void DoJITPitchingNotification(MethodDesc *MethodDescPtr);
     static void DoModuleLoadNotification(Module *Module);
     static void DoModuleUnloadNotification(Module *Module);
@@ -1088,7 +1089,7 @@ public:
 
     // called from the DAC
     static int GetType(TADDR Args[]);
-    static BOOL ParseJITNotification(TADDR Args[], TADDR& MethodDescPtr);
+    static BOOL ParseJITNotification(TADDR Args[], TADDR& MethodDescPtr, TADDR& NativeCodeLocation);
     static BOOL ParseJITPitchingNotification(TADDR Args[], TADDR& MethodDescPtr);
     static BOOL ParseModuleLoadNotification(TADDR Args[], TADDR& ModulePtr);
     static BOOL ParseModuleUnloadNotification(TADDR Args[], TADDR& ModulePtr);