/
metrics.go
127 lines (107 loc) · 2.72 KB
/
metrics.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
// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
package core
import (
"time"
"github.com/OWASP/Amass/amass/utils"
)
// MetricsCollector provides Amass services with the ability to track performance.
type MetricsCollector struct {
// The Amass Service collecting the performance metrics
service Service
// The function that will provide the number of DNS names remaining
namesRemaining func() int
// The channel that handles requests for ServiceStats
statsReq chan chan *ServiceStats
// The queue that holds DNS query event times
queries *utils.Queue
// The channel that signals the metrics collector to halt execution
done chan struct{}
}
// NewMetricsCollector returns an initialized MetricsCollector.
func NewMetricsCollector(srv Service) *MetricsCollector {
mc := &MetricsCollector{
service: srv,
statsReq: make(chan chan *ServiceStats, 10),
queries: utils.NewQueue(),
done: make(chan struct{}, 2),
}
go mc.processMetrics()
return mc
}
// Stop halts execution of the metrics collector.
func (mc *MetricsCollector) Stop() {
close(mc.done)
}
// NamesRemainingCallback updates the names remaining callback routine.
func (mc *MetricsCollector) NamesRemainingCallback(nrc func() int) {
mc.namesRemaining = nrc
}
// Stats returns ServiceStats for the metrics collected by this MetricsCollector.
func (mc *MetricsCollector) Stats() *ServiceStats {
c := make(chan *ServiceStats)
mc.statsReq <- c
return <-c
}
func (mc *MetricsCollector) processMetrics() {
var perSec []int
last := time.Now()
t := time.NewTicker(time.Second)
defer t.Stop()
for {
select {
case <-mc.done:
return
case <-t.C:
perSec = append(perSec, mc.eventsPerSec(last, mc.queries))
perSec = trimMetricSlice(perSec)
last = time.Now()
case c := <-mc.statsReq:
var remaining int
if mc.namesRemaining != nil {
remaining = mc.namesRemaining()
}
c <- &ServiceStats{
DNSQueriesPerSec: metricSliceAverage(perSec),
NamesRemaining: remaining,
}
}
}
}
// QueryTime allows a DNS query event time to be posted with the MetricsCollector.
func (mc *MetricsCollector) QueryTime(t time.Time) {
mc.queries.Append(t)
}
func (mc *MetricsCollector) eventsPerSec(last time.Time, q *utils.Queue) int {
var num int
for {
element, ok := q.Next()
if !ok {
break
}
comTime := element.(time.Time)
if comTime.After(last) {
num++
}
}
return num
}
func metricSliceAverage(m []int) int {
var total int
num := len(m)
if num < 10 {
return 0
}
for _, s := range m {
total += s
}
return total / num
}
func trimMetricSlice(m []int) []int {
s := len(m)
if s <= 60 {
return m
}
idx := s - 60
return m[idx:]
}