diff --git a/discovery/discovery.go b/discovery/discovery.go index 6d6144a..ad91d9b 100644 --- a/discovery/discovery.go +++ b/discovery/discovery.go @@ -32,6 +32,7 @@ func DiscoverZone(probeService *oopcs.ProbeService, requesterID string, dz *disc return } + discoverer.stopWg.Add(1) discoverer.discoverZone(probeService, requesterID, dz) discoverer.stopWg.Wait() @@ -44,6 +45,7 @@ func DiscoverHost(probeService *oopcs.ProbeService, requesterID string, zone *di logging.Logger().Warnf("Discovery: Discovery is running already") return } + discoverer.stopWg.Add(1) discoverer.discoverHost(probeService, requesterID, zone, dh) discoverer.stopWg.Wait() discoverer.complete() @@ -55,6 +57,7 @@ func DiscoverPort(probeService *oopcs.ProbeService, requesterID string, host *di logging.Logger().Warnf("Discovery: Discovery is running already") return } + discoverer.stopWg.Add(1) discoverer.discoverPort(probeService, requesterID, host, dp) discoverer.stopWg.Wait() discoverer.complete() @@ -66,6 +69,7 @@ func DiscoverService(probeService *oopcs.ProbeService, requesterID string, port logging.Logger().Warnf("Discovery: Discovery is running already") return } + discoverer.stopWg.Add(1) discoverer.discoverService(probeService, requesterID, port, ds) discoverer.stopWg.Wait() discoverer.complete() @@ -129,6 +133,9 @@ func (d *discovery) stop() { } func (d *discovery) discoverZone(probeService *oopcs.ProbeService, requesterID string, dz *discoveryM.DiscoveryZone) { + defer func() { + d.stopWg.Done() + }() go taskScan(d, func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) { scanZone(dz, resultChan, errChan, doneChan, stopChan) @@ -144,6 +151,7 @@ func (d *discovery) discoverZone(probeService *oopcs.ProbeService, requesterID s LastScanRange: cr.Last().String(), DiscoveryPort: dz.DiscoveryHost.DiscoveryPort, } + d.stopWg.Add(1) d.discoverHost(probeService, requesterID, z, dh) } }, @@ -151,6 +159,9 @@ func (d *discovery) discoverZone(probeService *oopcs.ProbeService, requesterID s } func (d *discovery) discoverHost(probeService *oopcs.ProbeService, requesterID string, zone *discoveryM.Zone, dh *discoveryM.DiscoveryHost) { + defer func() { + d.stopWg.Done() + }() go taskScan(d, func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) { scanHost(zone, dh, resultChan, errChan, doneChan, stopChan) @@ -160,6 +171,7 @@ func (d *discovery) discoverHost(probeService *oopcs.ProbeService, requesterID s logging.Logger().Debugf("host: %v", h) probeService.Send("DiscoveryService.DiscoveredHost", requesterID, h) if nil != dh.DiscoveryPort { + d.stopWg.Add(1) d.discoverPort(probeService, requesterID, h, dh.DiscoveryPort) } }, @@ -167,6 +179,9 @@ func (d *discovery) discoverHost(probeService *oopcs.ProbeService, requesterID s } func (d *discovery) discoverPort(probeService *oopcs.ProbeService, requesterID string, host *discoveryM.Host, dp *discoveryM.DiscoveryPort) { + defer func() { + d.stopWg.Done() + }() go taskScan(d, func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) { scanPort(host, dp, resultChan, errChan, doneChan, stopChan) @@ -176,6 +191,7 @@ func (d *discovery) discoverPort(probeService *oopcs.ProbeService, requesterID s logging.Logger().Debugf("port: %v", p) probeService.Send("DiscoveryService.DiscoveredPort", requesterID, p) if nil != dp.DiscoveryService { + d.stopWg.Add(1) d.discoverService(probeService, requesterID, p, dp.DiscoveryService) } }, @@ -183,6 +199,9 @@ func (d *discovery) discoverPort(probeService *oopcs.ProbeService, requesterID s } func (d *discovery) discoverService(probeService *oopcs.ProbeService, requesterID string, port *discoveryM.Port, ds *discoveryM.DiscoveryService) { + defer func() { + d.stopWg.Done() + }() go taskScan(d, func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) { scanService(port, ds, resultChan, errChan, doneChan, stopChan) @@ -190,7 +209,6 @@ func (d *discovery) discoverService(probeService *oopcs.ProbeService, requesterI func(result interface{}) { s := result.(*discoveryM.Service) probeService.Send("DiscoveryService.DiscoveredService", requesterID, s) - logging.Logger().Debugf("service: %s(%s)[%s:%s]", s.ServiceName, s.CryptoType, port.Host.IP, port.PortNumber) }, ) }