Extracts ETL events by time

ExtractETLEventsByTime.cpp

#include <Windows.h>
#include <evntcons.h>
#include <relogger.h>
#include <string>
#include <atomic>

constexpr UCHAR DCStart = 3;
constexpr UCHAR DCStop = 4;

constexpr UCHAR LoadVerbose = 143;
constexpr UCHAR UnloadVerbose = 144;
constexpr UCHAR JittingStarted = 145;
constexpr UCHAR ModuleLoad = 152;
constexpr UCHAR ModuleUnload = 153;
constexpr UCHAR AssemblyLoad = 154;
constexpr UCHAR AssemblyUnload = 155;
constexpr UCHAR ILToNativeMap = 190;

constexpr GUID ProcessGuid            = { 0x3d6fa8d0, 0xfe05, 0x11d0, { 0x9d, 0xda, 0x00, 0xc0, 0x4f, 0xd7, 0xba, 0x7c } };
constexpr GUID ImageLoadGuid          = { 0x2cb15d1d, 0x5fc1, 0x11d2, { 0xab, 0xe1, 0x00, 0xa0, 0xc9, 0x11, 0xf5, 0x18 } };
constexpr GUID DotNetRundownGuid      = { 0xa669021c, 0xc450, 0x4609, { 0xa0, 0x35, 0x5a, 0xf5, 0x9a, 0xf4, 0xdf, 0x18 } };
constexpr GUID KernelTraceControlGuid = { 0xb3e675d7, 0x2554, 0x4f18, { 0x83, 0x0b, 0x27, 0x62, 0x73, 0x25, 0x60, 0xde } };
constexpr GUID DotNetRuntimeGuid      = { 0xe13c0d23, 0xccbc, 0x4e12, { 0x93, 0x1b, 0xd9, 0xcc, 0x2e, 0xee, 0x27, 0xe4 } };
constexpr GUID SysConfigExGuid        = { 0x9b79ee91, 0xb5fd, 0x41c0, { 0xa2, 0x43, 0x42, 0x48, 0xe2, 0x66, 0xe9, 0xd0 } };


class TraceEventCallback : public ITraceEventCallback
{
public:
  HRESULT STDMETHODCALLTYPE OnEvent(ITraceEvent *traceEvent, ITraceRelogger *relogger) override
  {
    PEVENT_RECORD eventRecord;
    traceEvent->GetEventRecord(&eventRecord);
    const auto timestamp = eventRecord->EventHeader.TimeStamp.QuadPart;

    const auto offset = (timestamp - this->sessionStartTimeQPC) * 1000.0 / this->qpcFreq;

    // Manifest events, Kernel Trace, SysConfigEx & CLR Rundown Events
    if ((offset >= this->begin && offset <= this->end) || eventRecord->EventHeader.EventDescriptor.Id == 65534 || IsEqualGUID(eventRecord->EventHeader.ProviderId, SysConfigExGuid) || IsEqualGUID(eventRecord->EventHeader.ProviderId, KernelTraceControlGuid) || IsEqualGUID(eventRecord->EventHeader.ProviderId, DotNetRundownGuid))
    {
      goto Inject;
    }

    // Manifest events, Kernel Trace & CLR Rundown Events
    if (IsEqualGUID(eventRecord->EventHeader.ProviderId, DotNetRuntimeGuid))
    {
        switch (eventRecord->EventHeader.EventDescriptor.Opcode)
        {
        case LoadVerbose:
        case UnloadVerbose:
        case JittingStarted:
        case ModuleLoad:
        case ModuleUnload:
        case AssemblyLoad:
        case AssemblyUnload:
        case ILToNativeMap:
            goto Inject;
        default:
            goto Exit;
        }
    }

    // Process/ImageLoad DC Events
    if (IsEqualGUID(eventRecord->EventHeader.ProviderId, ProcessGuid) || IsEqualGUID(eventRecord->EventHeader.ProviderId, ImageLoadGuid))
    {
      goto Inject;
    }

    goto Exit;

  Inject:
    relogger->Inject(traceEvent);

  Exit:

    return S_OK;
  }

  HRESULT STDMETHODCALLTYPE OnBeginProcessTrace(ITraceEvent *header, ITraceRelogger *relogger) override
  {
    return S_OK;
  }

  HRESULT STDMETHODCALLTYPE OnFinalizeProcessTrace(ITraceRelogger *relogger) override
  {
    return S_OK;
  }

  HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, PVOID *ppObj) override
  {
    if (riid == __uuidof(IUnknown) || riid == __uuidof(ITraceEventCallback))
    {
      *ppObj = this;
    }
    else
    {
      *ppObj = nullptr;
      return E_NOINTERFACE;
    }

    this->AddRef();
    return ERROR_SUCCESS;
  }

  ULONG STDMETHODCALLTYPE AddRef() override
  {
    return std::atomic_fetch_add(&this->refCount, 1) + 1;
  }

  ULONG STDMETHODCALLTYPE Release() override
  {
    const int count = std::atomic_fetch_sub(&this->refCount, 1) - 1;
    if (count == 0)
    {
      delete this;
    }

    return count;
  }

  TraceEventCallback(const double begin, const double end, const LONGLONG sessionStartTimeQPC, const LONGLONG qpcFreq) : begin(begin), end(end), sessionStartTimeQPC(sessionStartTimeQPC), qpcFreq(qpcFreq)
  {
  }

  virtual ~TraceEventCallback()
  {
  }

private:
  std::atomic<int> refCount;
  double begin;
  double end;
  LONGLONG sessionStartTimeQPC;
  LONGLONG qpcFreq;
};

extern "C" __declspec(dllexport) HRESULT RelogTraceFile(const PWSTR inputFileName, const PWSTR outputFileName, const double begin, const double end, const LONGLONG sessionStartTimeQPC, const LONGLONG qpcFreq)
{
  HRESULT hr;

  hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
  if (FAILED(hr))
  {
    wprintf(L"Failed to initialize COM: 0x%08x\n", hr);
    return HRESULT_CODE(hr);
  }

  ITraceRelogger *relogger = nullptr;
  TraceEventCallback *callback = nullptr;

  hr = CoCreateInstance(CLSID_TraceRelogger, nullptr, CLSCTX_INPROC_SERVER, __uuidof(ITraceRelogger), reinterpret_cast<LPVOID *>(&relogger));

  if (FAILED(hr))
  {
    wprintf(L"Failed to instantiate ITraceRelogger object: 0x%08x.\n", hr);
    goto cleanup;
  }

  relogger->SetOutputFilename(outputFileName);
  if (FAILED(hr))
  {
    wprintf(L"Failed to set output log file: 0x%08x.\n", hr);
    goto cleanup;
  }

  TRACEHANDLE traceHandle;

  hr = relogger->AddLogfileTraceStream(inputFileName, nullptr, &traceHandle);
  if (FAILED(hr))
  {
    wprintf(L"Failed to set log file input stream: 0x%08x.\n", hr);
    goto cleanup;
  }

  callback = new TraceEventCallback(begin, end, sessionStartTimeQPC, qpcFreq);

  if (callback == nullptr)
  {
    wprintf(L"Failed to allocate callback: 0x%08x.\n", hr);
    goto cleanup;
  }

  hr = relogger->RegisterCallback(callback);
  if (FAILED(hr))
  {
    wprintf(L"Failed to register callback: 0x%08x.\n", hr);
    goto cleanup;
  }

  hr = relogger->ProcessTrace();
  if (FAILED(hr))
  {
    wprintf(L"Failed to process trace: 0x%08x.\n", hr);
  }

cleanup:

  if (relogger != nullptr)
  {
    relogger->Release();
  }

  if (callback != nullptr)
  {
    callback->Release();
  }

  return hr;
}

class StartTimeRetriever
{
public:

  StartTimeRetriever()
  {
    this->sessionStartQPC = 0;
  }

  LONGLONG GetSessionStartQPC() const
  {
    return this->sessionStartQPC;
  }

  VOID WINAPI ProcessEvent(const PEVENT_RECORD eventRecord)
  {
    if (this->sessionStartQPC == 0)
    {
      this->sessionStartQPC = eventRecord->EventHeader.TimeStamp.QuadPart;
    }
  }

private:
  LONGLONG sessionStartQPC;
};

ULONG WINAPI BufferCallback(const PEVENT_TRACE_LOGFILE logfile)
{
  const auto startTimeReceiver = static_cast<StartTimeRetriever*>(logfile->Context);
  return startTimeReceiver->GetSessionStartQPC() == 0;
}

VOID WINAPI ProcessEvent(const PEVENT_RECORD eventRecord)
{
  auto startTimeRetriever = static_cast<StartTimeRetriever*>(eventRecord->UserContext);
  startTimeRetriever->ProcessEvent(eventRecord);
}

extern "C" __declspec(dllexport) HRESULT ExtractByTime(const PWSTR inputFileName, const PWSTR outputFileName, const LONGLONG beginTimeStamp, const LONGLONG endTimeStamp)
{
  EVENT_TRACE_LOGFILE trace;

  ZeroMemory(&trace, sizeof(EVENT_TRACE_LOGFILE));
  trace.LogFileName = inputFileName;
  trace.BufferCallback = static_cast<PEVENT_TRACE_BUFFER_CALLBACK>(BufferCallback);
  trace.EventRecordCallback = static_cast<PEVENT_RECORD_CALLBACK>(ProcessEvent);
  trace.ProcessTraceMode = PROCESS_TRACE_MODE_EVENT_RECORD | PROCESS_TRACE_MODE_RAW_TIMESTAMP;

  auto startTimeReceiver = StartTimeRetriever();
  trace.Context = &startTimeReceiver;

  TRACEHANDLE traceHandle = OpenTrace(&trace);
  if (INVALID_PROCESSTRACE_HANDLE == traceHandle)
  {
    if (INVALID_PROCESSTRACE_HANDLE != traceHandle)
    {
      CloseTrace(traceHandle);
      return HRESULT_FROM_WIN32(GetLastError());
    }
  }

  const DWORD status = ProcessTrace(&traceHandle, 1, nullptr, nullptr);
  if (status != ERROR_SUCCESS && status != ERROR_CANCELLED)
  {
    return HRESULT_FROM_WIN32(status);
  }

  const LONGLONG sessionStartTimeQPC = startTimeReceiver.GetSessionStartQPC(); // Ticks

  const LONGLONG sessionStartTimeStamp = trace.LogfileHeader.StartTime.QuadPart; // FILETIME

  if (INVALID_PROCESSTRACE_HANDLE != traceHandle)
  {
    CloseTrace(traceHandle);
  }

  return RelogTraceFile(inputFileName, outputFileName, (beginTimeStamp - sessionStartTimeStamp) / 10000.0, (endTimeStamp - sessionStartTimeStamp) / 10000.0, sessionStartTimeQPC, trace.LogfileHeader.PerfFreq.QuadPart);
}