forked from DataDog/datadog-agent
/
pdhcounter.go
353 lines (323 loc) · 11.8 KB
/
pdhcounter.go
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016-present Datadog, Inc.
// +build windows
package pdhutil
import (
"fmt"
"github.com/StackVista/stackstate-agent/pkg/util/log"
)
// For testing
var (
pfnMakeCounterSetInstances = makeCounterSetIndexes
pfnPdhOpenQuery = PdhOpenQuery
pfnPdhAddCounter = PdhAddCounter
pfnPdhAddEnglishCounter = PdhAddEnglishCounter
pfnPdhCollectQueryData = PdhCollectQueryData
pfnPdhEnumObjectItems = pdhEnumObjectItems
pfnPdhRemoveCounter = PdhRemoveCounter
pfnPdhLookupPerfNameByIndex = pdhLookupPerfNameByIndex
pfnPdhGetFormattedCounterValueFloat = pdhGetFormattedCounterValueFloat
pfnPdhCloseQuery = PdhCloseQuery
pfnPdhMakeCounterPath = pdhMakeCounterPath
)
// CounterInstanceVerify is a callback function called by GetCounterSet for each
// instance of the counter. Implementation should return true if that instance
// should be included, false otherwise
type CounterInstanceVerify func(string) bool
// PdhCounterSet is the object which represents a pdh counter set.
type PdhCounterSet struct {
className string
query PDH_HQUERY
counterName string
}
// PdhSingleInstanceCounterSet is a specialization for single instance counters
type PdhSingleInstanceCounterSet struct {
PdhCounterSet
singleCounter PDH_HCOUNTER
}
// PdhMultiInstanceCounterSet is a specialization for a multiple instance counter
type PdhMultiInstanceCounterSet struct {
PdhCounterSet
requestedCounterName string
requestedInstances map[string]bool
countermap map[string]PDH_HCOUNTER // map instance name to counter handle
verifyfn CounterInstanceVerify
}
// Initialize initializes a counter set object
func (p *PdhCounterSet) Initialize(className string) error {
// the counter index list may be > 1, but for class name, only take the first
// one. If not present at all, try the english counter name
ndxlist, err := getCounterIndexList(className)
if err != nil {
return err
}
if ndxlist == nil || len(ndxlist) == 0 {
log.Warnf("Didn't find counter index for class %s, attempting english counter", className)
p.className = className
} else {
if len(ndxlist) > 1 {
log.Warnf("Class %s had multiple (%d) indices, using first", className, len(ndxlist))
}
ndx := ndxlist[0]
p.className, err = pfnPdhLookupPerfNameByIndex(ndx)
if err != nil {
return fmt.Errorf("Class name not found: %s", className)
}
log.Debugf("Found class name for %s %s", className, p.className)
}
winerror := pfnPdhOpenQuery(uintptr(0), uintptr(0), &p.query)
if ERROR_SUCCESS != winerror {
err = fmt.Errorf("Failed to open PDH query handle %d", winerror)
return err
}
return nil
}
// GetUnlocalizedCounter wraps the PdhAddEnglishCounter call that takes unlocalized counter names (as opposed to the other functions which use PdhAddCounter)
func GetUnlocalizedCounter(className, counterName, instance string) (PdhSingleInstanceCounterSet, error) {
var p PdhSingleInstanceCounterSet
winerror := pfnPdhOpenQuery(uintptr(0), uintptr(0), &p.query)
if ERROR_SUCCESS != winerror {
return p, fmt.Errorf("Failed to open PDH query handle %d", winerror)
}
path, err := pfnPdhMakeCounterPath("", className, instance, counterName)
if err != nil {
return p, fmt.Errorf("Failed tomake counter path %s: %v", counterName, err)
}
winerror = pfnPdhAddEnglishCounter(p.query, path, uintptr(0), &p.singleCounter)
if ERROR_SUCCESS != winerror {
return p, fmt.Errorf("Failed to add english counter %d", winerror)
}
winerror = pfnPdhCollectQueryData(p.query)
if ERROR_SUCCESS != winerror {
return p, fmt.Errorf("Failed to collect query data %d", winerror)
}
return p, nil
}
// GetSingleInstanceCounter returns a single instance counter object for the given counter class
// TODO: Replace usages of this with GetUnlocalizedCounter using an empty string as instance
func GetSingleInstanceCounter(className, counterName string) (*PdhSingleInstanceCounterSet, error) {
var p PdhSingleInstanceCounterSet
if err := p.Initialize(className); err != nil {
return nil, err
}
// check to make sure this is really a single instance counter
allcounters, instances, _ := pfnPdhEnumObjectItems(p.className)
if len(instances) > 0 {
return nil, fmt.Errorf("Requested counter is not single-instance: %s", p.className)
}
path, err := p.MakeCounterPath("", counterName, "", allcounters)
if err != nil {
log.Warnf("Failed pdhEnumObjectItems %v", err)
return nil, err
}
winerror := pfnPdhAddCounter(p.query, path, uintptr(0), &p.singleCounter)
if ERROR_SUCCESS != winerror {
return nil, fmt.Errorf("Failed to add single counter %d", winerror)
}
// do the initial collect now
pfnPdhCollectQueryData(p.query)
return &p, nil
}
// GetMultiInstanceCounter returns a multi-instance counter object for the given counter class
// TODO: Replace usages of this with a function similar to GetUnlocalizedCounter for multi-instance counters, that uses PdhAddEnglishCounter
func GetMultiInstanceCounter(className, counterName string, requestedInstances *[]string, verifyfn CounterInstanceVerify) (*PdhMultiInstanceCounterSet, error) {
var p PdhMultiInstanceCounterSet
if err := p.Initialize(className); err != nil {
return nil, err
}
p.countermap = make(map[string]PDH_HCOUNTER)
p.verifyfn = verifyfn
p.requestedCounterName = counterName
// check to make sure this is really a single instance counter
_, instances, _ := pfnPdhEnumObjectItems(p.className)
if len(instances) <= 0 {
return nil, fmt.Errorf("Requested counter is a single-instance: %s", p.className)
}
// save the requested instances
if requestedInstances != nil && len(*requestedInstances) > 0 {
p.requestedInstances = make(map[string]bool)
for _, inst := range *requestedInstances {
p.requestedInstances[inst] = true
}
}
if err := p.MakeInstanceList(); err != nil {
return nil, err
}
return &p, nil
}
// MakeInstanceList walks the list of available instances, and adds new
// instances that have appeared since the last check run
func (p *PdhMultiInstanceCounterSet) MakeInstanceList() error {
allcounters, instances, err := pfnPdhEnumObjectItems(p.className)
if err != nil {
return err
}
var instToMake []string
for _, actualInstance := range instances {
// if we have a list of requested instances, walk it, and make sure
// they're here. If not, add them to the list of instances to make
if p.requestedInstances != nil {
// if it's not in the requestedInstances, don't bother
if !p.requestedInstances[actualInstance] {
continue
}
// ok. it was requested. If it's not in our map
// of counters, we have to add it
if p.countermap[actualInstance] == PDH_HCOUNTER(0) {
log.Debugf("Adding requested instance %s", actualInstance)
instToMake = append(instToMake, actualInstance)
}
} else {
// wanted all the instances. Make sure all of the instances
// are present
if p.countermap[actualInstance] == PDH_HCOUNTER(0) {
instToMake = append(instToMake, actualInstance)
}
}
}
added := false
for _, inst := range instToMake {
if p.verifyfn != nil {
if p.verifyfn(inst) == false {
// not interested, moving on
continue
}
}
path, err := p.MakeCounterPath("", p.requestedCounterName, inst, allcounters)
if err != nil {
log.Debugf("Failed tomake counter path %s %s", p.counterName, inst)
continue
}
var hc PDH_HCOUNTER
winerror := pfnPdhAddCounter(p.query, path, uintptr(0), &hc)
if ERROR_SUCCESS != winerror {
log.Debugf("Failed to add counter path %s", path)
continue
}
log.Debugf("Adding missing counter instance %s", inst)
p.countermap[inst] = hc
added = true
}
if added {
// do the initial collect now
pfnPdhCollectQueryData(p.query)
}
return nil
}
//RemoveInvalidInstance removes an instance from the counter that is no longer valid
func (p *PdhMultiInstanceCounterSet) RemoveInvalidInstance(badInstance string) {
hc := p.countermap[badInstance]
if hc != PDH_HCOUNTER(0) {
log.Debugf("Removing non-existent counter instance %s", badInstance)
pfnPdhRemoveCounter(hc)
delete(p.countermap, badInstance)
} else {
log.Debugf("Instance handle not found")
}
}
// MakeCounterPath creates a counter path from the counter instance and
// counter name. Tries all available translated counter indexes from
// the english name
func (p *PdhCounterSet) MakeCounterPath(machine, counterName, instanceName string, counters []string) (string, error) {
/*
When handling non english versions, the counters don't work quite as documented.
This is because strings like "Bytes Sent/sec" might appear multiple times in the
english master, and might not have mappings for each index.
Search each index, and make sure the requested counter name actually appears in
the list of available counters; that's the counter we'll use.
For more information, see README.md.
*/
idxList, err := getCounterIndexList(counterName)
if err != nil {
return "", err
}
for _, ndx := range idxList {
counter, e := pfnPdhLookupPerfNameByIndex(ndx)
if e != nil {
log.Debugf("Counter index %d not found, skipping", ndx)
continue
}
// see if the counter we got back is in the list of counters
if !stringInSlice(counter, counters) {
log.Debugf("counter %s not in counter list", counter)
continue
}
// check to see if we can create the counter
path, err := pfnPdhMakeCounterPath(machine, p.className, instanceName, counter)
if err == nil {
log.Debugf("Successfully created counter path %s", path)
p.counterName = counter
return path, nil
}
// else
log.Debugf("Unable to create path with %s, trying again", counter)
}
// if we get here, was never able to find a counter path or create a valid
// path. Return failure.
log.Warnf("Unable to create counter path for %s %s", counterName, instanceName)
return "", fmt.Errorf("Unable to create counter path %s %s", counterName, instanceName)
}
// GetAllValues returns the data associated with each instance in a query.
func (p *PdhMultiInstanceCounterSet) GetAllValues() (values map[string]float64, err error) {
values = make(map[string]float64)
err = nil
var removeList []string
pfnPdhCollectQueryData(p.query)
for inst, hcounter := range p.countermap {
var retval float64
retval, err = pfnPdhGetFormattedCounterValueFloat(hcounter)
if err != nil {
switch err.(type) {
case *ErrPdhInvalidInstance:
removeList = append(removeList, inst)
log.Debugf("Got invalid instance for %s %s", p.requestedCounterName, inst)
err = nil
continue
default:
log.Debugf("Other Error getting all values %s %s %v", p.requestedCounterName, inst, err)
return
}
}
values[inst] = retval
}
for _, inst := range removeList {
p.RemoveInvalidInstance(inst)
}
// check for newly found instances
p.MakeInstanceList()
return
}
// GetValue returns the data associated with a single-value counter
func (p *PdhSingleInstanceCounterSet) GetValue() (val float64, err error) {
if p.singleCounter == PDH_HCOUNTER(0) {
return 0, fmt.Errorf("Not a single-value counter")
}
pfnPdhCollectQueryData(p.query)
return pfnPdhGetFormattedCounterValueFloat(p.singleCounter)
}
// Close closes the query handle, freeing the underlying windows resources.
func (p *PdhCounterSet) Close() {
PdhCloseQuery(p.query)
}
func getCounterIndexList(cname string) ([]int, error) {
if counterToIndex == nil || len(counterToIndex) == 0 {
if err := pfnMakeCounterSetInstances(); err != nil {
return []int{}, err
}
}
ndxlist, found := counterToIndex[cname]
if !found {
return []int{}, nil
}
return ndxlist, nil
}
func stringInSlice(a string, list []string) bool {
for _, b := range list {
if b == a {
return true
}
}
return false
}