Fix synchronization for Monitor
The goroutine running the monitor isn't actually closed. This PR updates the API to allow that Goroutine to properly block program exit. This can be leveraged as we continue to make the configuration non-global.
This commit is contained in:
parent
6848e3ae60
commit
178d984996
@ -4,6 +4,7 @@ import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"runtime/pprof"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"fmt"
|
||||
@ -132,7 +133,8 @@ func ZGrab2Main() {
|
||||
s.Init(flag)
|
||||
zgrab2.RegisterScan(moduleType, s)
|
||||
}
|
||||
monitor := zgrab2.MakeMonitor(1)
|
||||
wg := sync.WaitGroup{}
|
||||
monitor := zgrab2.MakeMonitor(1, &wg)
|
||||
monitor.Callback = func(_ string) {
|
||||
dumpHeapProfile()
|
||||
}
|
||||
@ -141,6 +143,8 @@ func ZGrab2Main() {
|
||||
zgrab2.Process(monitor)
|
||||
end := time.Now()
|
||||
log.Infof("finished grab at %s", end.Format(time.RFC3339))
|
||||
monitor.Stop()
|
||||
wg.Wait()
|
||||
s := Summary{
|
||||
StatusesPerModule: monitor.GetStatuses(),
|
||||
StartTime: start.Format(time.RFC3339),
|
||||
|
13
monitor.go
13
monitor.go
@ -1,5 +1,7 @@
|
||||
package zgrab2
|
||||
|
||||
import "sync"
|
||||
|
||||
// Monitor is a collection of states per scans and a channel to communicate
|
||||
// those scans to the monitor
|
||||
type Monitor struct {
|
||||
@ -34,13 +36,22 @@ func (m *Monitor) GetStatuses() map[string]*State {
|
||||
return m.states
|
||||
}
|
||||
|
||||
// Stop indicates the monitor is done and the internal channel should be closed.
|
||||
// This function does not block, but will allow a call to Wait() on the
|
||||
// WaitGroup passed to MakeMonitor to return.
|
||||
func (m *Monitor) Stop() {
|
||||
close(m.statusesChan)
|
||||
}
|
||||
|
||||
// MakeMonitor returns a Monitor object that can be used to collect and send
|
||||
// the status of a running scan
|
||||
func MakeMonitor(statusChanSize int) *Monitor {
|
||||
func MakeMonitor(statusChanSize int, wg *sync.WaitGroup) *Monitor {
|
||||
m := new(Monitor)
|
||||
m.statusesChan = make(chan moduleStatus, statusChanSize)
|
||||
m.states = make(map[string]*State, 10)
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for s := range m.statusesChan {
|
||||
if m.states[s.name] == nil {
|
||||
m.states[s.name] = new(State)
|
||||
|
Loading…
Reference in New Issue
Block a user