PresentMon/TraceSession.cpp at master · gjaegy/PresentMon · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
// Code based on:
// http://chabster.blogspot.com/2012/10/realtime-etw-consumer-howto.html
//

#include "TraceSession.hpp"
#include <tchar.h>
#include <cstdio>

static VOID WINAPI EventRecordCallback(_In_ PEVENT_RECORD pEventRecord)
{
    reinterpret_cast<ITraceConsumer *>(pEventRecord->UserContext)->OnEventRecord(pEventRecord);
}

static ULONG WINAPI BufferRecordCallback(_In_ PEVENT_TRACE_LOGFILE Buffer)
{
    return reinterpret_cast<ITraceConsumer *>(Buffer->Context)->ContinueProcessing();
}

TraceSession::TraceSession(LPCTSTR szSessionName)
    : _szSessionName(_tcsdup(szSessionName))
    , _status(0)
    , _pSessionProperties(0)
    , hSession(0)
    , _hTrace(0)
{
}

TraceSession::~TraceSession(void)
{
    free(_szSessionName);
    free(_pSessionProperties);
}

bool TraceSession::Start()
{
    if (!_pSessionProperties) {
        static_assert(sizeof(EVENT_TRACE_PROPERTIES) == 120, "");
        const ULONG buffSize = sizeof(EVENT_TRACE_PROPERTIES) + (ULONG)(_tcslen(_szSessionName) + 1) * sizeof(TCHAR);
        _pSessionProperties = reinterpret_cast<EVENT_TRACE_PROPERTIES *>(malloc(buffSize));
        ZeroMemory(_pSessionProperties, buffSize);
        _pSessionProperties->Wnode.BufferSize = buffSize;
        _pSessionProperties->Wnode.ClientContext = 1;
        _pSessionProperties->LogFileMode = EVENT_TRACE_REAL_TIME_MODE;
        _pSessionProperties->LoggerNameOffset = sizeof(EVENT_TRACE_PROPERTIES);
    }

    // Create the trace session.
    _status = StartTraceW(&hSession, _szSessionName, _pSessionProperties);

    return (_status == ERROR_SUCCESS);
}

bool TraceSession::EnableProvider(const GUID& providerId, UCHAR level, ULONGLONG anyKeyword, ULONGLONG allKeyword)
{
    _status = EnableTraceEx2(hSession, &providerId, EVENT_CONTROL_CODE_ENABLE_PROVIDER, level, anyKeyword, allKeyword, 0, NULL);
    return (_status == ERROR_SUCCESS);
}

bool TraceSession::OpenTrace(ITraceConsumer *pConsumer)
{
    if (!pConsumer)
        return false;

    ZeroMemory(&_logFile, sizeof(EVENT_TRACE_LOGFILE));
    _logFile.LoggerName = _szSessionName;
    _logFile.ProcessTraceMode = PROCESS_TRACE_MODE_REAL_TIME | PROCESS_TRACE_MODE_EVENT_RECORD | PROCESS_TRACE_MODE_RAW_TIMESTAMP;
    _logFile.EventRecordCallback = &EventRecordCallback;
    _logFile.BufferCallback = &BufferRecordCallback;
    _logFile.Context = pConsumer;

    _hTrace = ::OpenTraceW(&_logFile);
    return (_hTrace != 0);
}

bool TraceSession::Process()
{
    _status = ProcessTrace(&_hTrace, 1, NULL, NULL);
    return (_status == ERROR_SUCCESS);
}

bool TraceSession::CloseTrace()
{
    _status = ::CloseTrace(_hTrace);
    return (_status == ERROR_SUCCESS);
}

bool TraceSession::DisableProvider(const GUID& providerId)
{
    _status = EnableTraceEx2(hSession, &providerId, EVENT_CONTROL_CODE_DISABLE_PROVIDER, 0, 0, 0, 0, NULL);
    return (_status == ERROR_SUCCESS);
}

bool TraceSession::Stop()
{
    _status = ControlTraceW(hSession, _szSessionName, _pSessionProperties, EVENT_TRACE_CONTROL_STOP);
    delete _pSessionProperties;
    _pSessionProperties = NULL;

    return (_status == ERROR_SUCCESS);
}

ULONG TraceSession::Status() const
{
    return _status;
}

LONGLONG TraceSession::PerfFreq() const
{
    return _logFile.LogfileHeader.PerfFreq.QuadPart;
}